perf: Fix streaming file outputs passthrough (#13153)

* perf: Fix streaming file outputs passthrough correctly

* test

server/routes/api/documents/documents.test.ts

@@ -39,8 +39,13 @@ import {
   buildGroup,
   buildAdmin,
   buildTemplate,
+  buildAttachment,
 } from "@server/test/factories";
-import { getTestServer, withAPIContext } from "@server/test/support";
+import {
+  getTestServer,
+  readZipResponse,
+  withAPIContext,
+} from "@server/test/support";
 
 const server = getTestServer();
 
@@ -604,6 +609,10 @@ describe("#documents.info", () => {
 });
 
 describe("#documents.export", () => {
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
   it("should return published document", async () => {
     const user = await buildUser();
     const document = await buildDocument({
@@ -688,6 +697,45 @@ describe("#documents.export", () => {
     expect(body.data).toEqual(await DocumentHelper.toMarkdown(document));
   });
 
+  it("should stream a zip when the document has attachments", async () => {
+    const user = await buildUser();
+    const attachment = await buildAttachment({
+      teamId: user.teamId,
+      userId: user.id,
+    });
+    const document = await buildDocument({
+      title: "Export Test",
+      userId: user.id,
+      teamId: user.teamId,
+      text: `![image](${attachment.redirectUrl})`,
+    });
+    vi.spyOn(FileStorage, "getFileBuffer").mockResolvedValue(
+      Buffer.from("image-data")
+    );
+
+    const res = await server.post("/api/documents.export", user, {
+      body: {
+        id: document.id,
+      },
+      headers: {
+        accept: "text/markdown",
+      },
+    });
+
+    expect(res.status).toEqual(200);
+    expect(res.headers.get("content-type")).toEqual("application/zip");
+    expect(res.headers.get("content-disposition")).toContain(
+      `filename="export-test.zip"`
+    );
+
+    const location = `attachments/${attachment.id}.png`;
+    const entries = await readZipResponse(res);
+    expect(Object.keys(entries).sort()).toEqual([location, "export-test.md"]);
+    expect(entries[location]).toEqual("image-data");
+    expect(entries["export-test.md"]).toContain(location);
+    expect(entries["export-test.md"]).not.toContain(attachment.redirectUrl);
+  });
+
   it("should require authorization without token", async () => {
     const document = await buildDocument();
     const res = await server.post("/api/documents.export", {

server/routes/api/documents/documents.ts

@@ -958,7 +958,7 @@ router.post(
       return;
     }
 
-    await streamZipResponse(ctx, `${fileName}.zip`, async (zip) => {
+    streamZipResponse(ctx, `${fileName}.zip`, (zip) => {
       for (const { attachment, buffer } of externalAttachments) {
         const location = path.join(
           "attachments",

server/routes/api/revisions/revisions.test.ts

@@ -1,12 +1,14 @@
 import { createContext } from "@server/context";
 import { UserMembership, Revision } from "@server/models";
+import FileStorage from "@server/storage/files";
 import {
   buildAdmin,
+  buildAttachment,
   buildCollection,
   buildDocument,
   buildUser,
 } from "@server/test/factories";
-import { getTestServer } from "@server/test/support";
+import { getTestServer, readZipResponse } from "@server/test/support";
 
 const server = getTestServer();
 
@@ -245,6 +247,53 @@ describe("#revisions.list", () => {
 });
 
 describe("#revisions.export", () => {
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  it("should stream a zip when the revision has attachments", async () => {
+    const user = await buildUser();
+    const attachment = await buildAttachment({
+      teamId: user.teamId,
+      userId: user.id,
+    });
+    const document = await buildDocument({
+      title: "Export Test",
+      userId: user.id,
+      teamId: user.teamId,
+      text: `![image](${attachment.redirectUrl})`,
+    });
+    const revision = await Revision.createFromDocument(
+      createContext({ user }),
+      document
+    );
+    vi.spyOn(FileStorage, "getFileBuffer").mockResolvedValue(
+      Buffer.from("image-data")
+    );
+
+    const res = await server.post("/api/revisions.export", user, {
+      body: {
+        id: revision.id,
+      },
+      headers: {
+        accept: "text/markdown",
+      },
+    });
+
+    expect(res.status).toEqual(200);
+    expect(res.headers.get("content-type")).toEqual("application/zip");
+    expect(res.headers.get("content-disposition")).toContain(
+      `filename="export-test.zip"`
+    );
+
+    const location = `attachments/${attachment.id}.png`;
+    const entries = await readZipResponse(res);
+    expect(Object.keys(entries).sort()).toEqual([location, "export-test.md"]);
+    expect(entries[location]).toEqual("image-data");
+    expect(entries["export-test.md"]).toContain(location);
+    expect(entries["export-test.md"]).not.toContain(attachment.redirectUrl);
+  });
+
   it("should return revision as markdown by default", async () => {
     const user = await buildUser();
     const document = await buildDocument({

server/routes/api/revisions/revisions.ts

@@ -195,7 +195,7 @@ router.post(
       return;
     }
 
-    await streamZipResponse(ctx, `${fileName}.zip`, async (zip) => {
+    streamZipResponse(ctx, `${fileName}.zip`, async (zip) => {
       for (const attachment of attachments) {
         const location = path.join(
           "attachments",

server/test/support.ts

@@ -1,3 +1,7 @@
+import fs from "node:fs/promises";
+import os from "node:os";
+import path from "node:path";
+import { randomUUID } from "node:crypto";
 import { faker } from "@faker-js/faker";
 import type { Transaction } from "sequelize";
 import { afterEach, beforeEach, vi } from "vitest";
@@ -11,6 +15,7 @@ import webService from "@server/services/web";
 import { sequelize } from "@server/storage/database";
 import type { APIContext } from "@server/types";
 import { AuthenticationType } from "@server/types";
+import ZipHelper from "@server/utils/ZipHelper";
 import TestServer from "./TestServer";
 
 export function getTestServer() {
@@ -78,6 +83,31 @@ export function withAPIContext<T>(
   });
 }
 
+/**
+ * Read a zip archive returned from an API endpoint into memory.
+ *
+ * @param res The response to read the archive from.
+ * @returns a map of entry filename to the entry contents as a string.
+ */
+export async function readZipResponse(
+  res: Awaited<ReturnType<TestServer["post"]>>
+): Promise<Record<string, string>> {
+  const filePath = path.join(os.tmpdir(), `test-${randomUUID()}.zip`);
+  await fs.writeFile(filePath, Buffer.from(await res.arrayBuffer()));
+
+  try {
+    const entries: Record<string, string> = {};
+    await ZipHelper.walk(filePath, async (entry) => {
+      entries[entry.fileName] = (
+        await entry.readBuffer(1024 * 1024)
+      ).toString();
+    });
+    return entries;
+  } finally {
+    await fs.rm(filePath, { force: true });
+  }
+}
+
 /**
  * Helper function to convert an object to form-urlencoded string.
  * Useful for testing OAuth endpoints that expect application/x-www-form-urlencoded content type.

server/utils/koa.ts

@@ -36,13 +36,12 @@ export const getFileFromRequest = (
  * @param fileName The filename to advertise in the Content-Disposition header.
  * @param build Callback that adds entries to the provided ZipFile.
  */
-export const streamZipResponse = async (
+export const streamZipResponse = (
   ctx: Context,
   fileName: string,
   build: (zip: ZipFile) => void | Promise<void>
-): Promise<void> => {
+): void => {
   const zip = new ZipFile();
-  await build(zip);
 
   ctx.set("Content-Type", "application/zip");
   ctx.set(
@@ -50,10 +49,20 @@ export const streamZipResponse = async (
     contentDisposition(fileName, { type: "attachment" })
   );
 
-  zip.outputStream.on("error", (err) => {
+  const handleError = (err: Error) => {
     ctx.app.emit("error", err, ctx);
     ctx.res.destroy(err);
-  });
+  };
+
+  zip.outputStream.on("error", handleError);
   ctx.body = zip.outputStream;
-  zip.end();
+
+  void (async () => {
+    try {
+      await build(zip);
+      zip.end();
+    } catch (err) {
+      handleError(err instanceof Error ? err : new Error(String(err)));
+    }
+  })();
 };

shared/editor/nodes/Image.tsx

@@ -33,7 +33,6 @@ const imageSizeRegex = /\s=(\d+)?x(\d+)?$/;
 // parse so the embed type survives the API markdown round-trip.
 const SOURCE_TOKEN_PREFIX = "source=";
 
-
 type TitleAttributes = {
   layoutClass?: string;
   title?: string;
@@ -520,8 +519,7 @@ export default class Image extends SimpleImage {
     }
 
     if (titleParts.length > 0 || size) {
-      markdown +=
-        ' "' + state.esc(titleParts.join(" "), false) + size + '"';
+      markdown += ' "' + state.esc(titleParts.join(" "), false) + size + '"';
     }
     markdown += ")";
     state.write(markdown);