Break out download code

The main thing here is to pass 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.

Otherwise the structure remains the same.  Some variables are renamed
due to being in a new context and some strings messages are simplified.

Author: code-asher

src/storage.ts

@@ -1,3 +1,4 @@
+import type { AxiosInstance, AxiosRequestConfig } from "axios";
 import { Api } from "coder/site/src/api/api";
 import { createWriteStream } from "fs";
 import fs from "fs/promises";
@@ -202,122 +203,22 @@ export class Storage {
 		const etag = stat !== undefined ? await cli.eTag(binPath) : "";
 		this.output.info("Using ETag", etag);
 
-		// Make the download request.
-		const controller = new AbortController();
-		const resp = await restClient.getAxiosInstance().get(binSource, {
-			signal: controller.signal,
-			baseURL: baseUrl,
-			responseType: "stream",
-			headers: {
-				"Accept-Encoding": "gzip",
-				"If-None-Match": `"${etag}"`,
-			},
-			decompress: true,
-			// Ignore all errors so we can catch a 404!
-			validateStatus: () => true,
+		// Download the binary to a temporary file.
+		await fs.mkdir(path.dirname(binPath), { recursive: true });
+		const tempFile =
+			binPath + ".temp-" + Math.random().toString(36).substring(8);
+		const writeStream = createWriteStream(tempFile, {
+			autoClose: true,
+			mode: 0o755,
+		});
+		const client = restClient.getAxiosInstance();
+		const status = await this.download(client, binSource, writeStream, {
+			"Accept-Encoding": "gzip",
+			"If-None-Match": `"${etag}"`,
 		});
-		this.output.info("Got status code", resp.status);
 
-		switch (resp.status) {
+		switch (status) {
 			case 200: {
-				const rawContentLength = resp.headers["content-length"];
-				const contentLength = Number.parseInt(rawContentLength);
-				if (Number.isNaN(contentLength)) {
-					this.output.warn(
-						"Got invalid or missing content length",
-						rawContentLength,
-					);
-				} else {
-					this.output.info("Got content length", prettyBytes(contentLength));
-				}
-
-				// Download to a temporary file.
-				await fs.mkdir(path.dirname(binPath), { recursive: true });
-				const tempFile =
-					binPath + ".temp-" + Math.random().toString(36).substring(8);
-
-				// Track how many bytes were written.
-				let written = 0;
-
-				const completed = await vscode.window.withProgress<boolean>(
-					{
-						location: vscode.ProgressLocation.Notification,
-						title: `Downloading ${buildInfo.version} from ${baseUrl} to ${binPath}`,
-						cancellable: true,
-					},
-					async (progress, token) => {
-						const readStream = resp.data as IncomingMessage;
-						let cancelled = false;
-						token.onCancellationRequested(() => {
-							controller.abort();
-							readStream.destroy();
-							cancelled = true;
-						});
-
-						// Reverse proxies might not always send a content length.
-						const contentLengthPretty = Number.isNaN(contentLength)
-							? "unknown"
-							: prettyBytes(contentLength);
-
-						// Pipe data received from the request to the temp file.
-						const writeStream = createWriteStream(tempFile, {
-							autoClose: true,
-							mode: 0o755,
-						});
-						readStream.on("data", (buffer: Buffer) => {
-							writeStream.write(buffer, () => {
-								written += buffer.byteLength;
-								progress.report({
-									message: `${prettyBytes(written)} / ${contentLengthPretty}`,
-									increment: Number.isNaN(contentLength)
-										? undefined
-										: (buffer.byteLength / contentLength) * 100,
-								});
-							});
-						});
-
-						// Wait for the stream to end or error.
-						return new Promise<boolean>((resolve, reject) => {
-							writeStream.on("error", (error) => {
-								readStream.destroy();
-								reject(
-									new Error(
-										`Unable to download binary: ${errToStr(error, "no reason given")}`,
-									),
-								);
-							});
-							readStream.on("error", (error) => {
-								writeStream.close();
-								reject(
-									new Error(
-										`Unable to download binary: ${errToStr(error, "no reason given")}`,
-									),
-								);
-							});
-							readStream.on("close", () => {
-								writeStream.close();
-								if (cancelled) {
-									resolve(false);
-								} else {
-									resolve(true);
-								}
-							});
-						});
-					},
-				);
-
-				// False means the user canceled, although in practice it appears we
-				// would not get this far because VS Code already throws on cancelation.
-				if (!completed) {
-					this.output.warn("User aborted download");
-					throw new Error("User aborted download");
-				}
-
-				this.output.info(
-					`Downloaded ${prettyBytes(written)} to`,
-					path.basename(tempFile),
-				);
-
 				// Move the old binary to a backup location first, just in case.  And,
 				// on Linux at least, you cannot write onto a binary that is in use so
 				// moving first works around that (delete would also work).
@@ -389,7 +290,7 @@ export class Storage {
 						}
 						const params = new URLSearchParams({
 							title: `Failed to download binary on \`${cli.goos()}-${cli.goarch()}\``,
-							body: `Received status code \`${resp.status}\` when downloading the binary.`,
+							body: `Received status code \`${status}\` when downloading the binary.`,
 						});
 						const uri = vscode.Uri.parse(
 							`https://github.com/coder/vscode-coder/issues/new?` +
@@ -402,6 +303,121 @@ export class Storage {
 		}
 	}
 
+	/**
+	 * Download the source to the provided stream with a progress dialog.  Return
+	 * the status code or throw if the user aborts or there is an error.
+	 */
+	private async download(
+		client: AxiosInstance,
+		source: string,
+		writeStream: WriteStream,
+		headers?: AxiosRequestConfig["headers"],
+	): Promise<number> {
+		const baseUrl = client.defaults.baseURL;
+
+		const controller = new AbortController();
+		const resp = await client.get(source, {
+			signal: controller.signal,
+			baseURL: baseUrl,
+			responseType: "stream",
+			headers,
+			decompress: true,
+			// Ignore all errors so we can catch a 404!
+			validateStatus: () => true,
+		});
+		this.output.info("Got status code", resp.status);
+
+		if (resp.status === 200) {
+			const rawContentLength = resp.headers["content-length"];
+			const contentLength = Number.parseInt(rawContentLength);
+			if (Number.isNaN(contentLength)) {
+				this.output.warn(
+					"Got invalid or missing content length",
+					rawContentLength,
+				);
+			} else {
+				this.output.info("Got content length", prettyBytes(contentLength));
+			}
+
+			// Track how many bytes were written.
+			let written = 0;
+
+			const completed = await vscode.window.withProgress<boolean>(
+				{
+					location: vscode.ProgressLocation.Notification,
+					title: `Downloading ${baseUrl}`,
+					cancellable: true,
+				},
+				async (progress, token) => {
+					const readStream = resp.data as IncomingMessage;
+					let cancelled = false;
+					token.onCancellationRequested(() => {
+						controller.abort();
+						readStream.destroy();
+						cancelled = true;
+					});
+
+					// Reverse proxies might not always send a content length.
+					const contentLengthPretty = Number.isNaN(contentLength)
+						? "unknown"
+						: prettyBytes(contentLength);
+
+					// Pipe data received from the request to the stream.
+					readStream.on("data", (buffer: Buffer) => {
+						writeStream.write(buffer, () => {
+							written += buffer.byteLength;
+							progress.report({
+								message: `${prettyBytes(written)} / ${contentLengthPretty}`,
+								increment: Number.isNaN(contentLength)
+									? undefined
+									: (buffer.byteLength / contentLength) * 100,
+							});
+						});
+					});
+
+					// Wait for the stream to end or error.
+					return new Promise<boolean>((resolve, reject) => {
+						writeStream.on("error", (error) => {
+							readStream.destroy();
+							reject(
+								new Error(
+									`Unable to download binary: ${errToStr(error, "no reason given")}`,
+								),
+							);
+						});
+						readStream.on("error", (error) => {
+							writeStream.close();
+							reject(
+								new Error(
+									`Unable to download binary: ${errToStr(error, "no reason given")}`,
+								),
+							);
+						});
+						readStream.on("close", () => {
+							writeStream.close();
+							if (cancelled) {
+								resolve(false);
+							} else {
+								resolve(true);
+							}
+						});
+					});
+				},
+			);
+
+			// False means the user canceled, although in practice it appears we
+			// would not get this far because VS Code already throws on cancelation.
+			if (!completed) {
+				this.output.warn("User aborted download");
+				throw new Error("Download aborted");
+			}
+
+			this.output.info(`Downloaded ${prettyBytes(written)}`);
+		}
+
+		return resp.status;
+	}
+
 	/**
 	 * Return the directory for a deployment with the provided label to where its
 	 * binary is cached.