Merge pull request #3373 from github/koesie10/download-timeout

Add timeouts to downloading databases and the CLI

Author: koesie10

extensions/ql-vscode/CHANGELOG.md

@@ -3,6 +3,7 @@
 ## [UNRELEASED]
 
 - Remove support for CodeQL CLI versions older than 2.13.5. [#3371](https://github.com/github/vscode-codeql/pull/3371)
+- Add a timeout to downloading databases and the CodeQL CLI. These can be changed using the `codeQL.addingDatabases.downloadTimeout` and `codeQL.cli.downloadTimeout` settings respectively. [#3373](https://github.com/github/vscode-codeql/pull/3373)
 
 ## 1.12.2 - 14 February 2024
 

extensions/ql-vscode/package.json

@@ -178,6 +178,11 @@
             "type": "string",
             "default": "",
             "markdownDescription": "Path to the CodeQL executable that should be used by the CodeQL extension. The executable is named `codeql` on Linux/Mac and `codeql.exe` on Windows. If empty, the extension will look for a CodeQL executable on your shell PATH, or if CodeQL is not on your PATH, download and manage its own CodeQL executable (note: if you later introduce CodeQL on your PATH, the extension will prefer a CodeQL executable it has downloaded itself)."
+          },
+          "codeQL.cli.downloadTimeout": {
+            "type": "integer",
+            "default": 10,
+            "description": "Download timeout in seconds for downloading the CLI distribution."
           }
         }
       },
@@ -376,6 +381,11 @@
         "title": "Adding databases",
         "order": 6,
         "properties": {
+          "codeQL.addingDatabases.downloadTimeout": {
+            "type": "integer",
+            "default": 10,
+            "description": "Download timeout in seconds for adding a CodeQL database."
+          },
           "codeQL.addingDatabases.allowHttp": {
             "type": "boolean",
             "default": false,

extensions/ql-vscode/src/codeql-cli/distribution.ts

@@ -1,3 +1,4 @@
+import type { WriteStream } from "fs";
 import { createWriteStream, mkdtemp, pathExists, remove } from "fs-extra";
 import { tmpdir } from "os";
 import { delimiter, dirname, join } from "path";
@@ -26,6 +27,8 @@ import { unzipToDirectoryConcurrently } from "../common/unzip-concurrently";
 import { reportUnzipProgress } from "../common/vscode/unzip-progress";
 import type { Release } from "./distribution/release";
 import { ReleasesApiConsumer } from "./distribution/releases-api-consumer";
+import { createTimeoutSignal } from "../common/fetch-stream";
+import { AbortError } from "node-fetch";
 
 /**
  * distribution.ts
@@ -384,15 +387,25 @@ class ExtensionSpecificDistributionManager {
       );
     }
 
-    const assetStream =
-      await this.createReleasesApiConsumer().streamBinaryContentOfAsset(
-        assets[0],
-      );
+    const {
+      signal,
+      onData,
+      dispose: disposeTimeout,
+    } = createTimeoutSignal(this.config.downloadTimeout);
+
     const tmpDirectory = await mkdtemp(join(tmpdir(), "vscode-codeql"));
 
+    let archiveFile: WriteStream | undefined = undefined;
+
     try {
+      const assetStream =
+        await this.createReleasesApiConsumer().streamBinaryContentOfAsset(
+          assets[0],
+          signal,
+        );
+
       const archivePath = join(tmpDirectory, "distributionDownload.zip");
-      const archiveFile = createWriteStream(archivePath);
+      archiveFile = createWriteStream(archivePath);
 
       const contentLength = assetStream.headers.get("content-length");
       const totalNumBytes = contentLength
@@ -405,12 +418,23 @@ class ExtensionSpecificDistributionManager {
         progressCallback,
       );
 
-      await new Promise((resolve, reject) =>
+      assetStream.body.on("data", onData);
+
+      await new Promise((resolve, reject) => {
+        if (!archiveFile) {
+          throw new Error("Invariant violation: archiveFile not set");
+        }
+
         assetStream.body
           .pipe(archiveFile)
           .on("finish", resolve)
-          .on("error", reject),
-      );
+          .on("error", reject);
+
+        // If an error occurs on the body, we also want to reject the promise (e.g. during a timeout error).
+        assetStream.body.on("error", reject);
+      });
+
+      disposeTimeout();
 
       await this.bumpDistributionFolderIndex();
 
@@ -427,7 +451,19 @@ class ExtensionSpecificDistributionManager {
             )
           : undefined,
       );
+    } catch (e) {
+      if (e instanceof AbortError) {
+        const thrownError = new AbortError("The download timed out.");
+        thrownError.stack = e.stack;
+        throw thrownError;
+      }
+
+      throw e;
     } finally {
+      disposeTimeout();
+
+      archiveFile?.close();
+
       await remove(tmpDirectory);
     }
   }

extensions/ql-vscode/src/codeql-cli/distribution/releases-api-consumer.ts

@@ -90,21 +90,28 @@ export class ReleasesApiConsumer {
 
   public async streamBinaryContentOfAsset(
     asset: ReleaseAsset,
+    signal?: AbortSignal,
   ): Promise<Response> {
     const apiPath = `/repos/${this.repositoryNwo}/releases/assets/${asset.id}`;
 
-    return await this.makeApiCall(apiPath, {
-      accept: "application/octet-stream",
-    });
+    return await this.makeApiCall(
+      apiPath,
+      {
+        accept: "application/octet-stream",
+      },
+      signal,
+    );
   }
 
   protected async makeApiCall(
     apiPath: string,
     additionalHeaders: { [key: string]: string } = {},
+    signal?: AbortSignal,
   ): Promise<Response> {
     const response = await this.makeRawRequest(
       ReleasesApiConsumer.apiBase + apiPath,
       Object.assign({}, this.defaultHeaders, additionalHeaders),
+      signal,
     );
 
     if (!response.ok) {
@@ -129,11 +136,13 @@ export class ReleasesApiConsumer {
   private async makeRawRequest(
     requestUrl: string,
     headers: { [key: string]: string },
+    signal?: AbortSignal,
     redirectCount = 0,
   ): Promise<Response> {
     const response = await fetch(requestUrl, {
       headers,
       redirect: "manual",
+      signal,
     });
 
     const redirectUrl = response.headers.get("location");
@@ -153,7 +162,12 @@ export class ReleasesApiConsumer {
         // mechanism is provided.
         delete headers["authorization"];
       }
-      return await this.makeRawRequest(redirectUrl, headers, redirectCount + 1);
+      return await this.makeRawRequest(
+        redirectUrl,
+        headers,
+        signal,
+        redirectCount + 1,
+      );
     }
 
     return response;

extensions/ql-vscode/src/common/fetch-stream.ts

@@ -0,0 +1,36 @@
+import { clearTimeout } from "node:timers";
+
+export function createTimeoutSignal(timeoutSeconds: number): {
+  signal: AbortSignal;
+  onData: () => void;
+  dispose: () => void;
+} {
+  const timeout = timeoutSeconds * 1000;
+
+  const abortController = new AbortController();
+
+  let timeoutId: NodeJS.Timeout;
+
+  // If we don't get any data within the timeout, abort the download
+  timeoutId = setTimeout(() => {
+    abortController.abort();
+  }, timeout);
+
+  // If we receive any data within the timeout, reset the timeout
+  const onData = () => {
+    clearTimeout(timeoutId);
+    timeoutId = setTimeout(() => {
+      abortController.abort();
+    }, timeout);
+  };
+
+  const dispose = () => {
+    clearTimeout(timeoutId);
+  };
+
+  return {
+    signal: abortController.signal,
+    onData,
+    dispose,
+  };
+}

extensions/ql-vscode/src/config.ts

@@ -93,6 +93,10 @@ const PERSONAL_ACCESS_TOKEN_SETTING = new Setting(
   "personalAccessToken",
   DISTRIBUTION_SETTING,
 );
+const CLI_DOWNLOAD_TIMEOUT_SETTING = new Setting(
+  "downloadTimeout",
+  DISTRIBUTION_SETTING,
+);
 const CLI_CHANNEL_SETTING = new Setting("channel", DISTRIBUTION_SETTING);
 
 // Query History configuration
@@ -118,6 +122,7 @@ export interface DistributionConfig {
   updateCustomCodeQlPath: (newPath: string | undefined) => Promise<void>;
   includePrerelease: boolean;
   personalAccessToken?: string;
+  downloadTimeout: number;
   channel: CLIChannel;
   onDidChangeConfiguration?: Event<void>;
 }
@@ -272,6 +277,10 @@ export class DistributionConfigListener
     return PERSONAL_ACCESS_TOKEN_SETTING.getValue() || undefined;
   }
 
+  public get downloadTimeout(): number {
+    return CLI_DOWNLOAD_TIMEOUT_SETTING.getValue() || 10;
+  }
+
   public async updateCustomCodeQlPath(newPath: string | undefined) {
     await CUSTOM_CODEQL_PATH_SETTING.updateValue(
       newPath,
@@ -644,8 +653,16 @@ const DEPRECATED_ALLOW_HTTP_SETTING = new Setting(
 
 const ADDING_DATABASES_SETTING = new Setting("addingDatabases", ROOT_SETTING);
 
+const DOWNLOAD_TIMEOUT_SETTING = new Setting(
+  "downloadTimeout",
+  ADDING_DATABASES_SETTING,
+);
 const ALLOW_HTTP_SETTING = new Setting("allowHttp", ADDING_DATABASES_SETTING);
 
+export function downloadTimeout(): number {
+  return DOWNLOAD_TIMEOUT_SETTING.getValue<number>() || 10;
+}
+
 export function allowHttp(): boolean {
   return (
     ALLOW_HTTP_SETTING.getValue<boolean>() ||

extensions/ql-vscode/src/databases/database-fetcher.ts

@@ -1,5 +1,5 @@
 import type { Response } from "node-fetch";
-import fetch from "node-fetch";
+import fetch, { AbortError } from "node-fetch";
 import { zip } from "zip-a-folder";
 import type { InputBoxOptions } from "vscode";
 import { Uri, window } from "vscode";
@@ -28,11 +28,16 @@ import {
 } from "../common/github-url-identifier-helper";
 import type { Credentials } from "../common/authentication";
 import type { AppCommandManager } from "../common/commands";
-import { addDatabaseSourceToWorkspace, allowHttp } from "../config";
+import {
+  addDatabaseSourceToWorkspace,
+  allowHttp,
+  downloadTimeout,
+} from "../config";
 import { showAndLogInformationMessage } from "../common/logging";
 import { AppOctokit } from "../common/octokit";
 import { getLanguageDisplayName } from "../common/query-language";
 import type { DatabaseOrigin } from "./local-databases/database-origin";
+import { createTimeoutSignal } from "../common/fetch-stream";
 
 /**
  * Prompts a user to fetch a database from a remote location. Database is assumed to be an archive file.
@@ -478,10 +483,33 @@ async function fetchAndUnzip(
     step: 1,
   });
 
-  const response = await checkForFailingResponse(
-    await fetch(databaseUrl, { headers: requestHeaders }),
-    "Error downloading database",
-  );
+  const {
+    signal,
+    onData,
+    dispose: disposeTimeout,
+  } = createTimeoutSignal(downloadTimeout());
+
+  let response: Response;
+  try {
+    response = await checkForFailingResponse(
+      await fetch(databaseUrl, {
+        headers: requestHeaders,
+        signal,
+      }),
+      "Error downloading database",
+    );
+  } catch (e) {
+    disposeTimeout();
+
+    if (e instanceof AbortError) {
+      const thrownError = new AbortError("The request timed out.");
+      thrownError.stack = e.stack;
+      throw thrownError;
+    }
+
+    throw e;
+  }
+
   const archiveFileStream = createWriteStream(archivePath);
 
   const contentLength = response.headers.get("content-length");
@@ -493,12 +521,34 @@ async function fetchAndUnzip(
     progress,
   );
 
-  await new Promise((resolve, reject) =>
-    response.body
-      .pipe(archiveFileStream)
-      .on("finish", resolve)
-      .on("error", reject),
-  );
+  response.body.on("data", onData);
+
+  try {
+    await new Promise((resolve, reject) => {
+      response.body
+        .pipe(archiveFileStream)
+        .on("finish", resolve)
+        .on("error", reject);
+
+      // If an error occurs on the body, we also want to reject the promise (e.g. during a timeout error).
+      response.body.on("error", reject);
+    });
+  } catch (e) {
+    // Close and remove the file if an error occurs
+    archiveFileStream.close(() => {
+      void remove(archivePath);
+    });
+
+    if (e instanceof AbortError) {
+      const thrownError = new AbortError("The download timed out.");
+      thrownError.stack = e.stack;
+      throw thrownError;
+    }
+
+    throw e;
+  } finally {
+    disposeTimeout();
+  }
 
   await readAndUnzip(
     Uri.file(archivePath).toString(true),