feat: add inbound media understanding

Co-authored-by: Tristan Manchester <tmanchester96@gmail.com>

src/media/fetch.test.ts

@@ -0,0 +1,47 @@
+import { describe, expect, it } from "vitest";
+
+import { fetchRemoteMedia } from "./fetch.js";
+
+function makeStream(chunks: Uint8Array[]) {
+  return new ReadableStream<Uint8Array>({
+    start(controller) {
+      for (const chunk of chunks) {
+        controller.enqueue(chunk);
+      }
+      controller.close();
+    },
+  });
+}
+
+describe("fetchRemoteMedia", () => {
+  it("rejects when content-length exceeds maxBytes", async () => {
+    const fetchImpl = async () =>
+      new Response(makeStream([new Uint8Array([1, 2, 3, 4, 5])]), {
+        status: 200,
+        headers: { "content-length": "5" },
+      });
+
+    await expect(
+      fetchRemoteMedia({
+        url: "https://example.com/file.bin",
+        fetchImpl,
+        maxBytes: 4,
+      }),
+    ).rejects.toThrow("exceeds maxBytes");
+  });
+
+  it("rejects when streamed payload exceeds maxBytes", async () => {
+    const fetchImpl = async () =>
+      new Response(makeStream([new Uint8Array([1, 2, 3]), new Uint8Array([4, 5, 6])]), {
+        status: 200,
+      });
+
+    await expect(
+      fetchRemoteMedia({
+        url: "https://example.com/file.bin",
+        fetchImpl,
+        maxBytes: 4,
+      }),
+    ).rejects.toThrow("exceeds maxBytes");
+  });
+});

src/media/fetch.ts

@@ -14,6 +14,7 @@ type FetchMediaOptions = {
   url: string;
   fetchImpl?: FetchLike;
   filePathHint?: string;
+  maxBytes?: number;
 };
 
 function stripQuotes(value: string): string {
@@ -51,7 +52,7 @@ async function readErrorBodySnippet(res: Response, maxChars = 200): Promise<stri
 }
 
 export async function fetchRemoteMedia(options: FetchMediaOptions): Promise<FetchMediaResult> {
-  const { url, fetchImpl, filePathHint } = options;
+  const { url, fetchImpl, filePathHint, maxBytes } = options;
   const fetcher: FetchLike | undefined = fetchImpl ?? globalThis.fetch;
   if (!fetcher) {
     throw new Error("fetch is not available");
@@ -77,7 +78,19 @@ export async function fetchRemoteMedia(options: FetchMediaOptions): Promise<Fetc
     throw new Error(`Failed to fetch media from ${url}${redirected}: ${detail}`);
   }
 
-  const buffer = Buffer.from(await res.arrayBuffer());
+  const contentLength = res.headers.get("content-length");
+  if (maxBytes && contentLength) {
+    const length = Number(contentLength);
+    if (Number.isFinite(length) && length > maxBytes) {
+      throw new Error(
+        `Failed to fetch media from ${url}: content length ${length} exceeds maxBytes ${maxBytes}`,
+      );
+    }
+  }
+
+  const buffer = maxBytes
+    ? await readResponseWithLimit(res, maxBytes)
+    : Buffer.from(await res.arrayBuffer());
   let fileNameFromUrl: string | undefined;
   try {
     const parsed = new URL(url);
@@ -109,3 +122,47 @@ export async function fetchRemoteMedia(options: FetchMediaOptions): Promise<Fetc
     fileName,
   };
 }
+
+async function readResponseWithLimit(res: Response, maxBytes: number): Promise<Buffer> {
+  const body = res.body;
+  if (!body || typeof body.getReader !== "function") {
+    const fallback = Buffer.from(await res.arrayBuffer());
+    if (fallback.length > maxBytes) {
+      throw new Error(
+        `Failed to fetch media from ${res.url || "response"}: payload exceeds maxBytes ${maxBytes}`,
+      );
+    }
+    return fallback;
+  }
+
+  const reader = body.getReader();
+  const chunks: Uint8Array[] = [];
+  let total = 0;
+  try {
+    while (true) {
+      const { done, value } = await reader.read();
+      if (done) break;
+      if (value?.length) {
+        total += value.length;
+        if (total > maxBytes) {
+          try {
+            await reader.cancel();
+          } catch {}
+          throw new Error(
+            `Failed to fetch media from ${res.url || "response"}: payload exceeds maxBytes ${maxBytes}`,
+          );
+        }
+        chunks.push(value);
+      }
+    }
+  } finally {
+    try {
+      reader.releaseLock();
+    } catch {}
+  }
+
+  return Buffer.concat(
+    chunks.map((chunk) => Buffer.from(chunk)),
+    total,
+  );
+}