anonymous_github/test/large-notebook-download.test.js

const { expect } = require("chai");
const { Readable, PassThrough } = require("stream");
const { once } = require("events");
const http = require("node:http");
const express = require("express");
const compression = require("compression");
const { randomBytes } = require("node:crypto");
const { gunzipSync } = require("node:zlib");
const archiver = require("archiver");
const { Parse } = require("unzip-stream");
require("ts-node/register/transpile-only");
const got = require("got");
const GitHubDownload = require("../src/core/source/GitHubDownload").default;
const config = require("../src/config").default;
const { AnonymizeTransformer, ContentAnonimizer } = require("../src/core/anonymize-utils");
const { streamAnonymizedZip } = require("../src/core/zipStream");

function notebook() {
  return JSON.stringify({ nbformat: 4, nbformat_minor: 0, metadata: {}, cells: [{
    cell_type: "code", metadata: {}, execution_count: 1,
    source: ["# Alice's experiment"],
    outputs: [{ output_type: "display_data", metadata: {}, data: {
      "image/png": "AbCd".repeat(12 * 1024 * 1024),
      "text/plain": ["Alice completed the experiment"],
    } }],
  }] });
}

async function collect(stream) {
  const chunks = [];
  for await (const chunk of stream) chunks.push(chunk);
  return Buffer.concat(chunks);
}

async function zip(entries) {
  const archive = archiver("zip");
  const result = collect(archive);
  for (const [name, content] of Object.entries(entries)) archive.append(content, { name });
  await archive.finalize();
  return result;
}

async function unzip(buffer) {
  const parser = Parse();
  const entries = {};
  const pending = [];
  parser.on("entry", entry => pending.push(collect(entry).then(data => { entries[entry.path] = data; })));
  const done = once(parser, "finish");
  Readable.from([buffer]).pipe(parser);
  await done;
  await Promise.all(pending);
  return entries;
}

describe("large notebook downloads", function () {
  this.timeout(15000);

  it("redacts names before and after large outputs within the anonymization deadline", async function () {
    const input = notebook();
    const transformer = new AnonymizeTransformer({ filePath: "example.ipynb", terms: ["Alice"] });
    const output = collect(transformer);
    Readable.from([Buffer.from(input)]).pipe(transformer);
    const result = JSON.parse((await output).toString());
    expect(result.cells[0].source[0]).to.equal("# XXXX-1's experiment");
    expect(result.cells[0].outputs[0].data["text/plain"][0]).to.equal("XXXX-1 completed the experiment");
    expect(result.cells[0].outputs[0].data["image/png"]).to.equal("AbCd".repeat(12 * 1024 * 1024));
  });

  it("matches RE2 boundaries and case folding for nearby literal candidates", function () {
    for (const term of ["k", "s", "Alice", "a-a", "Σ", "研究", "@Alice", "😀", "---", "@@", " "]) {
      const input = ["", " ", "é", "_", "😀", "𐐀"].flatMap(left =>
        ["", " ", "é", "_", "😀", "𐐀"].map(right => `${left}${term} ${term.toUpperCase()} K ſ${right}`)
      ).join(" ") + term.repeat(10);
      const reference = new ContentAnonimizer({ terms: [term] });
      for (const compiled of reference.compiledTerms) delete compiled.literalPrefilter;
      expect(new ContentAnonimizer({ terms: [term] }).anonymize(input)).to.equal(reference.anonymize(input));
    }
  });

  describe("ZIP streaming", function () {
    let originalStream, originalUrl, originalLimit, upstream;
    beforeEach(function () {
      originalStream = got.stream;
      originalUrl = GitHubDownload.prototype.getZipUrl;
      originalLimit = config.MAX_FILE_SIZE;
      GitHubDownload.prototype.getZipUrl = async () => ({ url: "https://example.invalid/source.zip" });
      got.stream = () => upstream;
    });
    afterEach(async function () {
      got.stream = originalStream;
      GitHubDownload.prototype.getZipUrl = originalUrl;
      config.MAX_FILE_SIZE = originalLimit;
      upstream?.destroy();
      await require("../src/core/cache-coordination").closeCacheRedis();
    });
    const options = {
      repoId: "fixture", organization: "owner", repoName: "repo", commit: "HEAD",
      getToken: () => "test", anonymizerOptions: { filePath: "", terms: ["Alice"] },
    };

    for (const compressed of [false, true]) {
      it(`finishes ZIP responses through Express compression with gzip ${compressed}`, async function () {
        const text = "Alice Σ 研究\n".repeat(25000), binary = randomBytes(128 * 1024);
        const input = await zip({ "root/empty.txt": "", "root/large.txt": text, "root/image.png": binary, "root/after.txt": "Alice" });
        upstream = Readable.from((async function* () {
          for (let offset = 0; offset < input.length; offset += 1021) {
            yield input.subarray(offset, offset + 1021);
          }
        })());
        const app = express();
        app.use(compression({ threshold: 0, highWaterMark: 1024, filter: () => compressed }));
        const drainEmitters = new Set();
        let drainSubscriptions = 0;
        app.get("/archive", async (_req, res) => {
          const on = res.on;
          res.on = function (event, listener) {
            const emitter = on.call(this, event, listener);
            if (event === "drain") { drainSubscriptions++; drainEmitters.add(emitter); }
            return emitter;
          };
          res.type("application/zip");
          await streamAnonymizedZip(options, res);
        });
        const server = await new Promise(resolve => {
          const listening = http.createServer({ highWaterMark: 1024 }, app).listen(0, "127.0.0.1", () => resolve(listening));
        });
        let client;
        try {
          const response = await new Promise((resolve, reject) => {
            client = http.get(`http://127.0.0.1:${server.address().port}/archive`, { headers: { "Accept-Encoding": "gzip" } }, resolve);
            client.setTimeout(8000, () => client.destroy(new Error("ZIP response stalled")));
            client.on("error", reject);
          });
          const bytes = await collect(response);
          expect(response.statusCode).to.equal(200);
          expect(response.headers["content-encoding"]).to.equal(compressed ? "gzip" : undefined);
          const entries = await unzip(compressed ? gunzipSync(bytes) : bytes);
          expect(Object.keys(entries)).to.have.members(["empty.txt", "large.txt", "image.png", "after.txt"]);
          expect(entries["empty.txt"]).to.have.length(0);
          expect(entries["large.txt"].toString()).to.equal(new ContentAnonimizer(options.anonymizerOptions).anonymize(text));
          expect(entries["image.png"]).to.deep.equal(binary);
          expect(entries["after.txt"].toString()).to.equal("XXXX-1");
          expect(drainSubscriptions).to.be.greaterThan(0);
          for (const emitter of drainEmitters) expect(emitter.listenerCount("drain")).to.equal(0);
        } finally {
          client?.destroy(); server.closeAllConnections();
          await new Promise(resolve => server.close(resolve));
        }
      });
    }

    it("completes an archive containing a large notebook and subsequent entries", async function () {
      upstream = Readable.from([await zip({ "root/example.ipynb": notebook(), "root/after.txt": "Alice" })]);
      const response = new PassThrough();
      const output = collect(response);
      await streamAnonymizedZip(options, response);
      const entries = await unzip(await output);
      expect(Object.keys(entries)).to.have.members(["example.ipynb", "after.txt"]);
      expect(entries["after.txt"].toString()).to.equal("XXXX-1");
      expect(JSON.parse(entries["example.ipynb"].toString()).cells[0].source[0]).to.equal("# XXXX-1's experiment");
    });

    it("completes large safe regex entries and continues to the next ZIP entry", async function () {
      const payload = "123,456,78.91234,-45.123456\n".repeat(280000);
      upstream = Readable.from([await zip({ "root/large.csv": `Alice\n${payload}Alice`, "root/after.txt": "Alice" })]);
      const response = new PassThrough(), output = collect(response);
      await streamAnonymizedZip({ ...options, anonymizerOptions: {
        terms: ["Alice", ...Array.from({ length: 8 }, (_, i) => `Researcher${i}[0-9]+`)],
      } }, response);
      const entries = await unzip(await output);
      expect(entries["large.csv"].toString()).to.equal(`XXXX-1\n${payload}XXXX-1`);
      expect(entries["after.txt"].toString()).to.equal("XXXX-1");
      expect(response.writableFinished).to.equal(true);
    });

    it("aborts the response when an entry cannot be anonymized", async function () {
      upstream = Readable.from([await zip({ "root/large.txt": "Alice".repeat(1000) })]);
      config.MAX_FILE_SIZE = 100;
      const response = new PassThrough();
      response.resume();
      const error = once(response, "error");
      await streamAnonymizedZip(options, response);
      expect((await error)[0].message).to.contain("Text file exceeded");
      expect(response.destroyed).to.equal(true);
      expect(response.writableFinished).to.equal(false);
      expect(upstream.destroyed).to.equal(true);
    });

    it("stops downloading when the client disconnects", async function () {
      upstream = new PassThrough();
      const response = new PassThrough();
      await streamAnonymizedZip(options, response);
      const closed = once(response, "close");
      response.destroy();
      await closed;
      expect(upstream.destroyed).to.equal(true);
    });
  });
});