Skip to content

Commit 8beb761

Browse files
code-ashercode-asher
committed
Break out download code
No code changes aside from some variable names and passing in an axios client instead of the sdk client (since this does not need to make API calls and we will need to pass a separate client without headers when downloading external signatures).
1 parent 22d24da commit 8beb761

File tree

1 file changed

+130
-114
lines changed

1 file changed

+130
-114
lines changed

‎src/storage.ts

Lines changed: 130 additions & 114 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import type { AxiosInstance, AxiosRequestConfig } from "axios";
12
import { Api } from "coder/site/src/api/api";
23
import { createWriteStream } from "fs";
34
import fs from "fs/promises";
@@ -202,122 +203,22 @@ export class Storage {
202203
const etag = stat !== undefined ? await cli.eTag(binPath) : "";
203204
this.output.info("Using ETag", etag);
204205

205-
// Make the download request.
206-
const controller = new AbortController();
207-
const resp = await restClient.getAxiosInstance().get(binSource, {
208-
signal: controller.signal,
209-
baseURL: baseUrl,
210-
responseType: "stream",
211-
headers: {
212-
"Accept-Encoding": "gzip",
213-
"If-None-Match": `"${etag}"`,
214-
},
215-
decompress: true,
216-
// Ignore all errors so we can catch a 404!
217-
validateStatus: () => true,
206+
// Download the binary to a temporary file.
207+
await fs.mkdir(path.dirname(binPath), { recursive: true });
208+
const tempFile =
209+
binPath + ".temp-" + Math.random().toString(36).substring(8);
210+
const writeStream = createWriteStream(tempFile, {
211+
autoClose: true,
212+
mode: 0o755,
213+
});
214+
const client = restClient.getAxiosInstance();
215+
const status = await this.download(client, binSource, writeStream, {
216+
"Accept-Encoding": "gzip",
217+
"If-None-Match": `"${etag}"`,
218218
});
219-
this.output.info("Got status code", resp.status);
220219

221-
switch (resp.status) {
220+
switch (status) {
222221
case 200: {
223-
const rawContentLength = resp.headers["content-length"];
224-
const contentLength = Number.parseInt(rawContentLength);
225-
if (Number.isNaN(contentLength)) {
226-
this.output.warn(
227-
"Got invalid or missing content length",
228-
rawContentLength,
229-
);
230-
} else {
231-
this.output.info("Got content length", prettyBytes(contentLength));
232-
}
233-
234-
// Download to a temporary file.
235-
await fs.mkdir(path.dirname(binPath), { recursive: true });
236-
const tempFile =
237-
binPath + ".temp-" + Math.random().toString(36).substring(8);
238-
239-
// Track how many bytes were written.
240-
let written = 0;
241-
242-
const completed = await vscode.window.withProgress<boolean>(
243-
{
244-
location: vscode.ProgressLocation.Notification,
245-
title: `Downloading ${buildInfo.version} from ${baseUrl} to ${binPath}`,
246-
cancellable: true,
247-
},
248-
async (progress, token) => {
249-
const readStream = resp.data as IncomingMessage;
250-
let cancelled = false;
251-
token.onCancellationRequested(() => {
252-
controller.abort();
253-
readStream.destroy();
254-
cancelled = true;
255-
});
256-
257-
// Reverse proxies might not always send a content length.
258-
const contentLengthPretty = Number.isNaN(contentLength)
259-
? "unknown"
260-
: prettyBytes(contentLength);
261-
262-
// Pipe data received from the request to the temp file.
263-
const writeStream = createWriteStream(tempFile, {
264-
autoClose: true,
265-
mode: 0o755,
266-
});
267-
readStream.on("data", (buffer: Buffer) => {
268-
writeStream.write(buffer, () => {
269-
written += buffer.byteLength;
270-
progress.report({
271-
message: `${prettyBytes(written)} / ${contentLengthPretty}`,
272-
increment: Number.isNaN(contentLength)
273-
? undefined
274-
: (buffer.byteLength / contentLength) * 100,
275-
});
276-
});
277-
});
278-
279-
// Wait for the stream to end or error.
280-
return new Promise<boolean>((resolve, reject) => {
281-
writeStream.on("error", (error) => {
282-
readStream.destroy();
283-
reject(
284-
new Error(
285-
`Unable to download binary: ${errToStr(error, "no reason given")}`,
286-
),
287-
);
288-
});
289-
readStream.on("error", (error) => {
290-
writeStream.close();
291-
reject(
292-
new Error(
293-
`Unable to download binary: ${errToStr(error, "no reason given")}`,
294-
),
295-
);
296-
});
297-
readStream.on("close", () => {
298-
writeStream.close();
299-
if (cancelled) {
300-
resolve(false);
301-
} else {
302-
resolve(true);
303-
}
304-
});
305-
});
306-
},
307-
);
308-
309-
// False means the user canceled, although in practice it appears we
310-
// would not get this far because VS Code already throws on cancelation.
311-
if (!completed) {
312-
this.output.warn("User aborted download");
313-
throw new Error("User aborted download");
314-
}
315-
316-
this.output.info(
317-
`Downloaded ${prettyBytes(written)} to`,
318-
path.basename(tempFile),
319-
);
320-
321222
// Move the old binary to a backup location first, just in case. And,
322223
// on Linux at least, you cannot write onto a binary that is in use so
323224
// moving first works around that (delete would also work).
@@ -389,7 +290,7 @@ export class Storage {
389290
}
390291
const params = new URLSearchParams({
391292
title: `Failed to download binary on \`${cli.goos()}-${cli.goarch()}\``,
392-
body: `Received status code \`${resp.status}\` when downloading the binary.`,
293+
body: `Received status code \`${status}\` when downloading the binary.`,
393294
});
394295
const uri = vscode.Uri.parse(
395296
`https://github.com/coder/vscode-coder/issues/new?` +
@@ -402,6 +303,121 @@ export class Storage {
402303
}
403304
}
404305

306+
/**
307+
* Download the source to the provided stream with a progress dialog. Return
308+
* the status code or throw if the user aborts or there is an error.
309+
*/
310+
private async download(
311+
client: AxiosInstance,
312+
source: string,
313+
writeStream: WriteStream,
314+
headers?: AxiosRequestConfig["headers"],
315+
): Promise<number> {
316+
const baseUrl = client.defaults.baseURL;
317+
318+
const controller = new AbortController();
319+
const resp = await client.get(source, {
320+
signal: controller.signal,
321+
baseURL: baseUrl,
322+
responseType: "stream",
323+
headers,
324+
decompress: true,
325+
// Ignore all errors so we can catch a 404!
326+
validateStatus: () => true,
327+
});
328+
this.output.info("Got status code", resp.status);
329+
330+
if (resp.status === 200) {
331+
const rawContentLength = resp.headers["content-length"];
332+
const contentLength = Number.parseInt(rawContentLength);
333+
if (Number.isNaN(contentLength)) {
334+
this.output.warn(
335+
"Got invalid or missing content length",
336+
rawContentLength,
337+
);
338+
} else {
339+
this.output.info("Got content length", prettyBytes(contentLength));
340+
}
341+
342+
// Track how many bytes were written.
343+
let written = 0;
344+
345+
const completed = await vscode.window.withProgress<boolean>(
346+
{
347+
location: vscode.ProgressLocation.Notification,
348+
title: `Downloading ${baseUrl}`,
349+
cancellable: true,
350+
},
351+
async (progress, token) => {
352+
const readStream = resp.data as IncomingMessage;
353+
let cancelled = false;
354+
token.onCancellationRequested(() => {
355+
controller.abort();
356+
readStream.destroy();
357+
cancelled = true;
358+
});
359+
360+
// Reverse proxies might not always send a content length.
361+
const contentLengthPretty = Number.isNaN(contentLength)
362+
? "unknown"
363+
: prettyBytes(contentLength);
364+
365+
// Pipe data received from the request to the stream.
366+
readStream.on("data", (buffer: Buffer) => {
367+
writeStream.write(buffer, () => {
368+
written += buffer.byteLength;
369+
progress.report({
370+
message: `${prettyBytes(written)} / ${contentLengthPretty}`,
371+
increment: Number.isNaN(contentLength)
372+
? undefined
373+
: (buffer.byteLength / contentLength) * 100,
374+
});
375+
});
376+
});
377+
378+
// Wait for the stream to end or error.
379+
return new Promise<boolean>((resolve, reject) => {
380+
writeStream.on("error", (error) => {
381+
readStream.destroy();
382+
reject(
383+
new Error(
384+
`Unable to download binary: ${errToStr(error, "no reason given")}`,
385+
),
386+
);
387+
});
388+
readStream.on("error", (error) => {
389+
writeStream.close();
390+
reject(
391+
new Error(
392+
`Unable to download binary: ${errToStr(error, "no reason given")}`,
393+
),
394+
);
395+
});
396+
readStream.on("close", () => {
397+
writeStream.close();
398+
if (cancelled) {
399+
resolve(false);
400+
} else {
401+
resolve(true);
402+
}
403+
});
404+
});
405+
},
406+
);
407+
408+
// False means the user canceled, although in practice it appears we
409+
// would not get this far because VS Code already throws on cancelation.
410+
if (!completed) {
411+
this.output.warn("User aborted download");
412+
throw new Error("Download aborted");
413+
}
414+
415+
this.output.info(`Downloaded ${prettyBytes(written)}`);
416+
}
417+
418+
return resp.status;
419+
}
420+
405421
/**
406422
* Return the directory for a deployment with the provided label to where its
407423
* binary is cached.

0 commit comments

Comments
 (0)