Merge pull request #4455 from github/shati-patel/fix-cli-extraction

Refactor unzip functionality to limit concurrency and improve error handling in `copyStream`

Author: josefs

extensions/ql-vscode/CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## [UNRELEASED]
 
+- Fix a bug where installing or updating the CodeQL CLI could hang indefinitely while extracting the downloaded archive. Extraction now reports an error if a file cannot be written, and aborts with a clear message if no progress is made within the download timeout (for example due to slow or networked storage, or security software). [#4455](https://github.com/github/vscode-codeql/pull/4455)
 - Remove support for CodeQL CLI versions older than 2.23.9. [#4448](https://github.com/github/vscode-codeql/pull/4448)
 - Added support for selection-based result filtering via a checkbox in the result viewer. When enabled, only results from the currently-viewed file are shown. Additionally, if the editor selection is non-empty, only results within the selection range are shown. [#4362](https://github.com/github/vscode-codeql/pull/4362)
 - Added a new "CodeQL: Go to File in Selected Database" command that allows you to open a file from the source archive of the currently selected database. [#4390](https://github.com/github/vscode-codeql/pull/4390)

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

@@ -560,6 +560,7 @@ class ExtensionSpecificDistributionManager {
               progressCallback,
             )
           : undefined,
+        this.config.downloadTimeout,
       );
     } catch (e) {
       if (e instanceof DOMException && e.name === "AbortError") {

extensions/ql-vscode/src/common/unzip-concurrently.ts

@@ -3,13 +3,22 @@ import type { UnzipProgressCallback } from "./unzip";
 import { unzipToDirectory } from "./unzip";
 import PQueue from "p-queue";
 
+/**
+ * Maximum number of files to extract concurrently. We cap this rather than
+ * using the full core count so that we don't open an excessive number of
+ * simultaneous file writes, which can overwhelm slower or contended storage
+ * during extraction.
+ */
+const MAX_UNZIP_CONCURRENCY = 4;
+
 export async function unzipToDirectoryConcurrently(
   archivePath: string,
   destinationPath: string,
   progress?: UnzipProgressCallback,
+  timeoutSeconds?: number,
 ): Promise<void> {
   const queue = new PQueue({
-    concurrency: availableParallelism(),
+    concurrency: Math.min(availableParallelism(), MAX_UNZIP_CONCURRENCY),
   });
 
   return unzipToDirectory(
@@ -19,5 +28,6 @@ export async function unzipToDirectoryConcurrently(
     async (tasks) => {
       await queue.addAll(tasks);
     },
+    timeoutSeconds,
   );
 }

extensions/ql-vscode/src/common/unzip.ts

@@ -2,10 +2,23 @@ import type { Entry as ZipEntry, Options as ZipOptions, ZipFile } from "yauzl";
 import { open } from "yauzl";
 import type { Readable } from "stream";
 import { Transform } from "stream";
+import { pipeline } from "stream/promises";
 import { dirname, join } from "path";
 import type { WriteStream } from "fs";
 import { createWriteStream, ensureDir } from "fs-extra";
 import { asError } from "./helpers-pure";
+import { createTimeoutSignal } from "./fetch-stream";
+
+/**
+ * Default idle timeout (in seconds) used when a caller does not specify one. If
+ * no bytes are extracted and no files complete within this window, the
+ * extraction is aborted. This is a safety net that guards against a single
+ * stalled write hanging the whole operation forever. 30 seconds of zero
+ * extraction progress is well beyond any healthy pause, so this won't trigger
+ * false positives. (The CLI downloader passes the configurable
+ * `codeQL.cli.downloadTimeout` instead of relying on this default.)
+ */
+const DEFAULT_UNZIP_IDLE_TIMEOUT_SECONDS = 30;
 
 // We can't use promisify because it picks up the wrong overload.
 export function openZip(
@@ -88,31 +101,23 @@ export async function openZipBuffer(
   });
 }
 
-async function copyStream(
+export async function copyStream(
   readable: Readable,
   writeStream: WriteStream,
   bytesExtractedCallback?: (bytesExtracted: number) => void,
+  signal?: AbortSignal,
 ): Promise<void> {
-  return new Promise((resolve, reject) => {
-    readable.on("error", (err) => {
-      reject(err);
-    });
-    readable.on("end", () => {
-      resolve();
-    });
-
-    readable
-      .pipe(
-        new Transform({
-          transform(chunk, _encoding, callback) {
-            bytesExtractedCallback?.(chunk.length);
-            this.push(chunk);
-            callback();
-          },
-        }),
-      )
-      .pipe(writeStream);
-  });
+  await pipeline(
+    readable,
+    new Transform({
+      transform(chunk, _encoding, callback) {
+        bytesExtractedCallback?.(chunk.length);
+        callback(null, chunk);
+      },
+    }),
+    writeStream,
+    { signal },
+  );
 }
 
 type UnzipProgress = {
@@ -139,6 +144,7 @@ async function unzipFile(
   entry: ZipEntry,
   rootDestinationPath: string,
   bytesExtractedCallback?: (bytesExtracted: number) => void,
+  signal?: AbortSignal,
 ): Promise<number> {
   const path = join(rootDestinationPath, entry.fileName);
 
@@ -164,7 +170,7 @@ async function unzipFile(
       mode,
     });
 
-    await copyStream(readable, writeStream, bytesExtractedCallback);
+    await copyStream(readable, writeStream, bytesExtractedCallback, signal);
 
     return entry.uncompressedSize;
   }
@@ -185,6 +191,7 @@ export async function unzipToDirectory(
   destinationPath: string,
   progress: UnzipProgressCallback | undefined,
   taskRunner: (tasks: Array<() => Promise<void>>) => Promise<void>,
+  timeoutSeconds: number = DEFAULT_UNZIP_IDLE_TIMEOUT_SECONDS,
 ): Promise<void> {
   const zipFile = await openZip(archivePath, {
     autoClose: false,
@@ -211,28 +218,53 @@ export async function unzipToDirectory(
 
     reportProgress();
 
-    await taskRunner(
-      entries.map((entry) => async () => {
-        let entryBytesExtracted = 0;
-
-        const totalEntryBytesExtracted = await unzipFile(
-          zipFile,
-          entry,
-          destinationPath,
-          (thisBytesExtracted) => {
-            entryBytesExtracted += thisBytesExtracted;
-            bytesExtracted += thisBytesExtracted;
-            reportProgress();
-          },
+    // Abort extraction if no progress is made for `timeoutSeconds`. `pipeline`
+    // only rejects on a stream error, so without this a single write that
+    // blocks without erroring (e.g. slow/networked storage or security
+    // software holding a file) would hang the extraction indefinitely.
+    const { signal, onData, dispose } = createTimeoutSignal(timeoutSeconds);
+
+    try {
+      await taskRunner(
+        entries.map((entry) => async () => {
+          let entryBytesExtracted = 0;
+
+          const totalEntryBytesExtracted = await unzipFile(
+            zipFile,
+            entry,
+            destinationPath,
+            (thisBytesExtracted) => {
+              // Reset the idle timeout: we are making progress.
+              onData();
+              entryBytesExtracted += thisBytesExtracted;
+              bytesExtracted += thisBytesExtracted;
+              reportProgress();
+            },
+            signal,
+          );
+
+          // Should be 0, but just in case.
+          bytesExtracted += -entryBytesExtracted + totalEntryBytesExtracted;
+
+          // Reset the idle timeout on completion too, so extracting many empty
+          // files or directories doesn't trip the timeout.
+          onData();
+          filesExtracted++;
+          reportProgress();
+        }),
+      );
+    } catch (e) {
+      if (signal.aborted) {
+        throw new Error(
+          `Timed out while extracting archive: no progress was made for ${timeoutSeconds} seconds. ` +
+            "This can happen if a file cannot be written, for example because of slow or networked storage, " +
+            "or security software holding a file open.",
         );
-
-        // Should be 0, but just in case.
-        bytesExtracted += -entryBytesExtracted + totalEntryBytesExtracted;
-
-        filesExtracted++;
-        reportProgress();
-      }),
-    );
+      }
+      throw e;
+    } finally {
+      dispose();
+    }
   } finally {
     zipFile.close();
   }
@@ -251,6 +283,7 @@ export async function unzipToDirectorySequentially(
   archivePath: string,
   destinationPath: string,
   progress?: UnzipProgressCallback,
+  timeoutSeconds?: number,
 ): Promise<void> {
   return unzipToDirectory(
     archivePath,
@@ -261,5 +294,6 @@ export async function unzipToDirectorySequentially(
         await task();
       }
     },
+    timeoutSeconds,
   );
 }

extensions/ql-vscode/test/unit-tests/common/unzip.test.ts

@@ -1,10 +1,13 @@
 import { createHash } from "crypto";
 import { open } from "fs/promises";
 import { join, relative, resolve, sep } from "path";
+import { Readable } from "stream";
+import { createWriteStream } from "fs";
 import { chmod, pathExists, readdir } from "fs-extra";
 import type { DirectoryResult } from "tmp-promise";
 import { dir } from "tmp-promise";
 import {
+  copyStream,
   excludeDirectories,
   openZip,
   openZipBuffer,
@@ -13,6 +16,7 @@ import {
 } from "../../../src/common/unzip";
 import { walkDirectory } from "../../../src/common/files";
 import { unzipToDirectoryConcurrently } from "../../../src/common/unzip-concurrently";
+import { createTimeoutSignal } from "../../../src/common/fetch-stream";
 
 const zipPath = resolve(__dirname, "../data/unzip/test-zip.zip");
 
@@ -276,3 +280,104 @@ async function computeHash(contents: Buffer) {
 
   return hash.digest("hex");
 }
+
+describe("copyStream error handling", () => {
+  let tmpDir: DirectoryResult;
+
+  beforeEach(async () => {
+    tmpDir = await dir({
+      unsafeCleanup: true,
+    });
+  });
+
+  afterEach(async () => {
+    await tmpDir?.cleanup();
+  });
+
+  it("rejects when the write stream errors mid-extraction", async () => {
+    // Use a real zip to trigger unzip, but make the destination read-only
+    // so the write stream fails. This verifies the promise rejects rather
+    // than hanging indefinitely.
+    const destPath = join(tmpDir.path, "output");
+
+    // Extract once to create the directory structure
+    await unzipToDirectorySequentially(zipPath, destPath);
+
+    // Make a file read-only so re-extraction will fail on write
+    const targetFile = join(destPath, "directory", "file.txt");
+    await chmod(targetFile, 0o000);
+
+    // On Windows, chmod doesn't prevent writes, so skip assertion there
+    if (process.platform === "win32") {
+      await chmod(targetFile, 0o644);
+      return;
+    }
+
+    // Re-extract — should reject with a write error, not hang
+    await expect(
+      unzipToDirectorySequentially(zipPath, destPath),
+    ).rejects.toThrow();
+
+    // Restore permissions for cleanup
+    await chmod(targetFile, 0o644);
+  });
+
+  it("rejects when the write stream is destroyed mid-copy", async () => {
+    // A destroyed write stream should cause `copyStream` to reject, not hang.
+    const destFile = join(tmpDir.path, "output.bin");
+    const writeStream = createWriteStream(destFile);
+
+    // A readable that emits a chunk and then destroys the write stream, to
+    // simulate a mid-copy write failure.
+    let pushed = false;
+    const readable = new Readable({
+      read() {
+        if (pushed) {
+          return;
+        }
+        pushed = true;
+        this.push(Buffer.alloc(1024, "x"));
+        setImmediate(() => {
+          writeStream.destroy(new Error("simulated write failure"));
+        });
+      },
+    });
+
+    await expect(copyStream(readable, writeStream)).rejects.toThrow(
+      "simulated write failure",
+    );
+  });
+
+  it("rejects (rather than hanging) when the idle timeout aborts a stalled copy", async () => {
+    // A readable that emits one chunk and then stalls forever: it never pushes
+    // more data, never ends, and never errors. Without the abort signal this
+    // would hang indefinitely (the original bug). `copyStream` must honour the
+    // signal and reject once the idle timeout fires.
+    let pushed = false;
+    const stalled = new Readable({
+      read() {
+        if (!pushed) {
+          pushed = true;
+          this.push(Buffer.alloc(1024, "x"));
+        }
+        // Then never push again and never call push(null) -> stalled.
+      },
+    });
+
+    const destFile = join(tmpDir.path, "stalled-output.bin");
+    const writeStream = createWriteStream(destFile);
+
+    // Short idle timeout so the test is fast; well within Jest's default limit.
+    const { signal, dispose } = createTimeoutSignal(0.05);
+
+    try {
+      await expect(
+        copyStream(stalled, writeStream, undefined, signal),
+      ).rejects.toThrow();
+      expect(signal.aborted).toBe(true);
+    } finally {
+      dispose();
+      stalled.destroy();
+    }
+  });
+});