Merge pull request #931 from tdurieux/fix/bounded-zip-finalization

fix: bound ZIP directory memory for slow downloads

src/core/bounded-zip.ts

@@ -0,0 +1,109 @@
+import archiver = require("archiver");
+import { once } from "events";
+import { createReadStream, closeSync, mkdtempSync, openSync, rmSync, writeSync, utimesSync, ReadStream } from "fs";
+import { setInterval, clearInterval } from "timers";
+import { tmpdir } from "os";
+import { join } from "path";
+
+// archiver 7 / zip-stream 6 / compress-commons 6 retain all directory entries
+// and synchronously emit millions of buffers at finalize(). Keep their ZIP
+// encoding (including ZIP64), but spool directory records as entries finish.
+interface ZipEngine extends NodeJS.WritableStream {
+  offset: number;
+  _entries: { length: number; push(entry: unknown): number };
+  _archive: { centralOffset: number; centralLength: number; processing: boolean; finish: boolean; finished: boolean };
+  _finish(): void;
+  _writeCentralFileHeader(entry: unknown): void;
+  _writeCentralDirectoryZip64(): void;
+  _writeCentralDirectoryEnd(): void;
+  isZip64(): boolean;
+  emit(event: string, error: Error): boolean;
+}
+
+export function boundedZip() {
+  const archive = archiver("zip", {});
+  const engine = (archive as unknown as { _module: { engine: ZipEngine } })._module.engine;
+  if (!Array.isArray(engine._entries) || typeof engine._writeCentralFileHeader !== "function" ||
+      typeof engine._writeCentralDirectoryZip64 !== "function" || typeof engine._writeCentralDirectoryEnd !== "function") {
+    throw new Error("unsupported_zip_engine");
+  }
+  const cancellation = new AbortController();
+  let directory: string | undefined;
+  let fd: number | undefined;
+  let reader: ReadStream | undefined;
+  let heartbeat: NodeJS.Timeout | undefined;
+  let count = 0;
+  let finalizing = false;
+  const cleanup = () => {
+    clearInterval(heartbeat); heartbeat = undefined;
+    reader?.destroy();
+    reader = undefined;
+    if (fd !== undefined) { closeSync(fd); fd = undefined; }
+    if (directory) { rmSync(directory, { recursive: true, force: true }); directory = undefined; }
+  };
+  const stop = () => { cancellation.abort(); cleanup(); };
+  const abort = archive.abort;
+  archive.abort = function () { stop(); return abort.call(this); };
+  archive.once("error", stop);
+  archive.once("close", stop);
+  archive.once("end", cleanup);
+
+  engine._entries = {
+    get length() { return count; },
+    push(entry) {
+      try {
+        if (cancellation.signal.aborted) return count;
+        if (fd === undefined) {
+          directory = mkdtempSync(join(tmpdir(), "anonymous-zip-directory-"));
+          fd = openSync(join(directory, "central"), "wx", 0o600);
+          // A paused download can outlive the stale-spool cleanup threshold.
+          heartbeat = setInterval(() => {
+            if (directory) {
+              const now = new Date();
+              try { utimesSync(join(directory, "central"), now, now); } catch { /* writes/reads report I/O errors */ }
+            }
+          }, 60_000);
+          heartbeat.unref();
+        }
+        // A directory record is bounded by ZIP's 16-bit name/extra/comment
+        // lengths. Serialize only this record, without touching output offsets.
+        const chunks: Buffer[] = [];
+        const write = engine.write;
+        const offset = engine.offset;
+        try {
+          engine.write = (chunk: Uint8Array | string) => { chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk)); return true; };
+          engine._writeCentralFileHeader(entry);
+        } finally { engine.write = write; engine.offset = offset; }
+        const record = Buffer.concat(chunks);
+        let written = 0;
+        while (written < record.length) written += writeSync(fd, record, written, record.length - written);
+        return ++count;
+      } catch (error) { engine.emit("error", error as Error); return count; }
+    },
+  };
+  engine._finish = () => {
+    if (finalizing || cancellation.signal.aborted) return;
+    finalizing = true;
+    engine._archive.finish = true;
+    engine._archive.centralOffset = engine.offset;
+    void (async () => {
+      if (fd !== undefined) { closeSync(fd); fd = undefined; }
+      if (directory) {
+        reader = createReadStream(join(directory, "central"), { highWaterMark: 64 * 1024 });
+        for await (const chunk of reader) {
+          cancellation.signal.throwIfAborted();
+          if (!engine.write(chunk)) await once(engine, "drain", { signal: cancellation.signal });
+        }
+      }
+      cancellation.signal.throwIfAborted();
+      engine._archive.centralLength = engine.offset - engine._archive.centralOffset;
+      if (engine.isZip64()) engine._writeCentralDirectoryZip64();
+      engine._writeCentralDirectoryEnd();
+      engine._archive.processing = false;
+      engine._archive.finished = true;
+      engine.end();
+      cleanup();
+    })().catch(error => { if (!cancellation.signal.aborted) engine.emit("error", error); cleanup(); });
+  };
+  return archive;
+}

src/core/storage/FileSystem.ts

@@ -1,3 +1,4 @@
+import { boundedZip } from "../bounded-zip";
 import { appendArchiveEntry, zipEntryOptions } from "../archive-entry";
 import config from "../../config";
 import * as fs from "fs";
@@ -261,7 +262,7 @@ export default class FileSystem extends StorageBase {
     }
   ) {
     this.assertSafePath(dir);
-    const archive = archiver(opt?.format || "zip", {});
+    const archive = opt?.format === "tar" ? archiver("tar", {}) : boundedZip();
     let active: Readable | undefined;
     let stopped = false;
     archive.once("close", () => { stopped = true; active?.destroy(); });

src/core/storage/S3.ts

@@ -1,3 +1,4 @@
+import { boundedZip } from "../bounded-zip";
 import { appendArchiveEntry, zipEntryOptions } from "../archive-entry";
 import {
   GetObjectCommand,
@@ -371,7 +372,7 @@ export default class S3Storage extends StorageBase {
   ) {
     if (!config.S3_BUCKET) throw new Error("S3_BUCKET not set");
     this.assertSafePath(dir);
-    const archive = archiver(opt?.format || "zip", {});
+    const archive = opt?.format === "tar" ? archiver("tar", {}) : boundedZip();
     if (dir && dir[dir.length - 1] != "/") dir = dir + "/";
 
     let active: Readable | undefined;

src/core/temporary-storage.ts

@@ -13,10 +13,11 @@ let cleanup: Promise<void> | undefined;
 export async function removeStaleTextSpools(directory = tmpdir(), now = Date.now()) {
   const entries = await fs.opendir(directory);
   for await (const entry of entries) {
-    if (!entry.isDirectory() || !entry.name.startsWith("anonymous-text-")) continue;
+    if (!entry.isDirectory() || (!entry.name.startsWith("anonymous-text-") && !entry.name.startsWith("anonymous-zip-directory-"))) continue;
     const path = join(directory, entry.name);
     try {
-      if (now - (await fs.stat(path)).mtimeMs > 3600_000) await fs.rm(path, { recursive: true, force: true });
+      const agePath = entry.name.startsWith("anonymous-zip-directory-") ? join(path, "central") : path;
+      if (now - (await fs.stat(agePath)).mtimeMs > 3600_000) await fs.rm(path, { recursive: true, force: true });
     } catch { /* another process can finish cleanup first */ }
   }
 }

src/core/zipStream.ts

@@ -3,7 +3,7 @@ import config from "../config";
 import { MAX_SUBMODULE_ENTRIES, SubmoduleMount, safeRepositoryPath, validateMounts } from "./source/github-submodules";
 import { Readable, Transform, Writable } from "stream";
 import { Parse } from "unzip-stream";
-import archiver = require("archiver");
+import { boundedZip } from "./bounded-zip";
 
 import GitHubDownload from "./source/GitHubDownload";
 import { classifyGitHubMissError } from "./source/GitHubBase";
@@ -118,7 +118,7 @@ export async function streamAnonymizedZip(
     });
   }
   let downloadStream: Readable;
-  const archive = archiver("zip", {});
+  const archive = boundedZip();
   let parser: Transform;
   let childEntries = 0;
   const activeStreams = new Set<Readable>();

test/bounded-zip.test.js

@@ -0,0 +1,96 @@
+const { expect } = require("chai");
+const { setImmediate, setTimeout } = require("timers");
+const { once } = require("events");
+const { Writable } = require("stream");
+const archiver = require("archiver");
+const fs = require("fs");
+const os = require("os");
+const { execFileSync } = require("child_process");
+require("ts-node/register/transpile-only");
+const { boundedZip } = require("../src/core/bounded-zip");
+
+async function bytes(archive, entries) {
+  const chunks = [];
+  archive.on("data", chunk => chunks.push(chunk));
+  const end = once(archive, "end");
+  for (const entry of entries) archive.append(entry.body, { date: new Date("2020-01-01T00:00:00Z"), ...entry.options });
+  await archive.finalize(); await end;
+  return Buffer.concat(chunks);
+}
+
+describe("bounded ZIP central directory", function () {
+  this.timeout(20000);
+  it("preserves archive bytes, compression, Unicode names, modes, and empty files", async () => {
+    const entries = [
+      { body: "private text\n".repeat(100), options: { name: "folder/é中.txt", mode: 0o100755 } },
+      { body: Buffer.from([0, 1, 255]), options: { name: "binary.dat", store: true } },
+      { body: "", options: { name: "empty.txt" } },
+    ];
+    expect(await bytes(boundedZip(), entries)).to.deep.equal(await bytes(archiver("zip"), entries));
+    expect(await bytes(boundedZip(), [])).to.deep.equal(await bytes(archiver("zip"), []));
+  });
+
+  it("preserves ZIP64 size and offset encoding beyond 4 GiB", async () => {
+    const Entry = require("compress-commons/lib/archivers/zip/zip-archive-entry");
+    async function directory(archive) {
+      const entry = new Entry("large.bin");
+      entry.setMethod(0); entry.setTime(new Date("2020-01-01"));
+      entry.setSize(2 ** 32 + 123); entry.setCompressedSize(2 ** 32 + 123);
+      entry._offsets = { file: 2 ** 32 + 456 };
+      archive._module.engine._entries.push(entry);
+      archive._module.engine.offset = 2 ** 33;
+      return bytes(archive, []);
+    }
+    expect(await directory(boundedZip())).to.deep.equal(await directory(archiver("zip")));
+  });
+
+  it("preserves ZIP64 counts and drains 180,636 records under a 128 MiB heap", function () {
+    this.timeout(60000);
+    const output = execFileSync(process.execPath, ["--max-old-space-size=128", "test/fixtures/zip-directory-memory.js"],
+      { cwd: process.cwd(), timeout: 55000, encoding: "utf8" });
+    const result = JSON.parse(output.trim().split("\n").pop());
+    expect(result.entries).to.equal(180636);
+    expect(result.zip64Count).to.equal(180636);
+    expect(result.centralRecords).to.equal(180636);
+    expect(result.heapMiB).to.be.lessThan(128);
+  });
+
+  it("reports spool write failures through the stream error handler", () => {
+    const archive = boundedZip();
+    let failure;
+    archive.on("error", error => { failure = error; });
+    const Entry = require("compress-commons/lib/archivers/zip/zip-archive-entry");
+    const entry = new Entry("file.txt"); entry.setMethod(0); entry.setTime(new Date()); entry._offsets = { file: 0 };
+    const write = fs.writeSync;
+    try {
+      fs.writeSync = () => { throw Object.assign(new Error("disk full"), { code: "ENOSPC" }); };
+      expect(() => archive._module.engine._entries.push(entry)).not.to.throw();
+      expect(failure?.code).to.equal("ENOSPC");
+    } finally { fs.writeSync = write; archive.abort(); }
+  });
+
+  it("removes its spool when aborted while waiting for a slow consumer", async () => {
+    const existing = new Set(fs.readdirSync(os.tmpdir()));
+    const archive = boundedZip();
+    archive.on("error", () => {});
+    const sink = new Writable({ highWaterMark: 1, write(_chunk, _encoding, _callback) {} });
+    archive.pipe(sink);
+    archive.append("small file", { name: "test.txt" });
+    await once(archive, "entry");
+    const created = fs.readdirSync(os.tmpdir()).filter(n => n.startsWith("anonymous-zip-directory-") && !existing.has(n));
+    expect(created).to.have.length(1);
+    const engine = archive._module.engine;
+    const Entry = require("compress-commons/lib/archivers/zip/zip-archive-entry");
+    for (let i = 0; i < 50000; i++) {
+      const entry = new Entry(`${"directory/".repeat(32)}file-${i}.txt`); entry.setMethod(0); entry.setTime(new Date()); entry._offsets = { file: 0 };
+      engine._entries.push(entry);
+    }
+    try {
+      void archive.finalize().catch(() => {});
+      await new Promise(resolve => setTimeout(resolve, 50));
+      expect(engine._archive.finished).to.equal(false);
+    } finally { archive.abort(); sink.destroy(); }
+    await new Promise(resolve => setImmediate(resolve));
+    expect(fs.existsSync(`${os.tmpdir()}/${created[0]}`)).to.equal(false);
+  });
+});

test/fixtures/zip-directory-memory.js

@@ -0,0 +1,43 @@
+// Exercise real library directory encoding without spending minutes compressing
+// file bodies: these are the same completed entries held at finalization.
+const { setTimeout } = require("timers");
+const { Writable } = require("stream");
+const { once } = require("events");
+const Entry = require("compress-commons/lib/archivers/zip/zip-archive-entry");
+const { boundedZip } = require("../../build/core/bounded-zip");
+(async () => {
+  const archive = boundedZip();
+  const engine = archive._module.engine;
+  let centralRecords = 0, zip64Count = 0, pending = Buffer.alloc(0);
+  const sink = new Writable({ highWaterMark: 1024, write(chunk, _encoding, done) {
+    pending = Buffer.concat([pending, chunk]);
+    while (pending.length >= 4) {
+      const signature = pending.readUInt32LE(0);
+      let length;
+      if (signature === 0x02014b50) {
+        if (pending.length < 46) break;
+        length = 46 + pending.readUInt16LE(28) + pending.readUInt16LE(30) + pending.readUInt16LE(32);
+        if (pending.length < length) break;
+        centralRecords++;
+      } else if (signature === 0x06064b50) {
+        if (pending.length < 56) break;
+        length = 56; zip64Count = Number(pending.readBigUInt64LE(32));
+      } else if (signature === 0x07064b50) length = 20;
+      else if (signature === 0x06054b50) length = 22;
+      else throw new Error("invalid ZIP directory signature");
+      if (pending.length < length) break;
+      pending = pending.subarray(length);
+    }
+    setTimeout(done, 1);
+  } });
+  archive.pipe(sink);
+  const finished = once(sink, "finish");
+  const entries = 180636;
+  for (let i = 0; i < entries; i++) {
+    const entry = new Entry(`file-${i}.txt`); entry.setMethod(0); entry.setTime(new Date("2020-01-01")); entry._offsets = { file: 0 };
+    engine._entries.push(entry);
+  }
+  await archive.finalize(); await finished;
+  if (pending.length) throw new Error("incomplete directory");
+  console.log(JSON.stringify({ entries, centralRecords, zip64Count, heapMiB: process.memoryUsage().heapUsed / 1048576 }));
+})().catch(error => { console.error(error); process.exitCode = 1; });

test/performance-worker.test.js

@@ -104,8 +104,13 @@ describe("large anonymization resources", function () {
       await Promise.all([old, current, other].map(path => fs.mkdir(path)));
       await fs.writeFile(join(old, "input"), "private");
       await fs.utimes(old, new Date(0), new Date(0));
+      const oldZip = join(root, "anonymous-zip-directory-old"), liveZip = join(root, "anonymous-zip-directory-live");
+      await Promise.all([oldZip, liveZip].map(path => fs.mkdir(path)));
+      await Promise.all([oldZip, liveZip].map(path => fs.writeFile(join(path, "central"), "directory records")));
+      await fs.utimes(join(oldZip, "central"), new Date(0), new Date(0));
+      await fs.utimes(liveZip, new Date(0), new Date(0));
       await removeStaleTextSpools(root);
-      expect((await fs.readdir(root)).sort()).to.deep.equal(["anonymous-text-current", "other"]);
+      expect((await fs.readdir(root)).sort()).to.deep.equal(["anonymous-text-current", "anonymous-zip-directory-live", "other"]);
     } finally { await fs.rm(root, { force: true, recursive: true }); }
   });
   it("bounds an unresponsive Redis command and disconnects its socket", async () => {