diff --git a/scripts/check-memory-fd-repro.mjs b/scripts/check-memory-fd-repro.mjs index cb22c28ace2..1293a5e1baf 100644 --- a/scripts/check-memory-fd-repro.mjs +++ b/scripts/check-memory-fd-repro.mjs @@ -25,6 +25,7 @@ const ISSUE_MEMORY_FILE_COUNT = ISSUE_FILE_COUNTS.reduce((sum, [, count]) => sum const DEFAULT_FILE_COUNT = 512; const DEFAULT_MAX_WORKSPACE_REG_FDS = process.platform === "darwin" ? 8 : 64; export const GATEWAY_READY_OUTPUT_MAX_CHARS = 128 * 1024; +export const MEMORY_SEARCH_RESPONSE_MAX_BYTES = 256 * 1024; const SKIP_GATEWAY_ENV = { NODE_ENV: "test", @@ -431,6 +432,55 @@ export async function stopGatewayWithRuntime({ } } +function responseBodyTooLargeError(label, maxBytes) { + return new Error(`${label} response body exceeded ${maxBytes} bytes`); +} + +export async function readBoundedResponseText(response, label, maxBytes) { + const contentLength = Number(response.headers.get("content-length") ?? ""); + if (Number.isSafeInteger(contentLength) && contentLength > maxBytes) { + await response.body?.cancel().catch(() => undefined); + throw responseBodyTooLargeError(label, maxBytes); + } + + if (!response.body) { + return ""; + } + + const reader = response.body.getReader(); + const decoder = new TextDecoder(); + const chunks = []; + let totalBytes = 0; + let canceled = false; + + try { + for (;;) { + const { done, value } = await reader.read(); + if (done) { + const tail = decoder.decode(); + if (tail) { + chunks.push(tail); + } + break; + } + + totalBytes += value.byteLength; + if (totalBytes > maxBytes) { + canceled = true; + await reader.cancel().catch(() => undefined); + throw responseBodyTooLargeError(label, maxBytes); + } + chunks.push(decoder.decode(value, { stream: true })); + } + } finally { + if (!canceled) { + reader.releaseLock(); + } + } + + return chunks.join(""); +} + async function invokeMemorySearch({ port, token, timeoutMs }) { const controller = new AbortController(); const timer = setTimeout(() => controller.abort(), timeoutMs); @@ -453,7 +503,11 @@ async function invokeMemorySearch({ port, token, timeoutMs }) { }), signal: controller.signal, }); - const text = await res.text(); + const text = await readBoundedResponseText( + res, + "memory_search", + MEMORY_SEARCH_RESPONSE_MAX_BYTES, + ); return { ok: res.ok, status: res.status, diff --git a/test/scripts/check-memory-fd-repro.test.ts b/test/scripts/check-memory-fd-repro.test.ts index 2a65a2a9537..20d351fc7e0 100644 --- a/test/scripts/check-memory-fd-repro.test.ts +++ b/test/scripts/check-memory-fd-repro.test.ts @@ -2,8 +2,10 @@ import { EventEmitter } from "node:events"; import { describe, expect, it, vi } from "vitest"; import { GATEWAY_READY_OUTPUT_MAX_CHARS, + MEMORY_SEARCH_RESPONSE_MAX_BYTES, hasChildExited, parseArgs, + readBoundedResponseText, readNumber, readPositiveNumber, stopGatewayWithRuntime, @@ -178,4 +180,44 @@ describe("check-memory-fd-repro", () => { expect(state.readySeen).toBe(true); expect(state.tail).toBe("w output"); }); + + it("reads memory_search response bodies under the byte cap", async () => { + await expect( + readBoundedResponseText( + new Response("ok"), + "memory_search", + MEMORY_SEARCH_RESPONSE_MAX_BYTES, + ), + ).resolves.toBe("ok"); + }); + + it("rejects oversized memory_search response bodies from content-length", async () => { + const response = new Response("ignored", { + headers: { "content-length": String(MEMORY_SEARCH_RESPONSE_MAX_BYTES + 1) }, + }); + + await expect( + readBoundedResponseText(response, "memory_search", MEMORY_SEARCH_RESPONSE_MAX_BYTES), + ).rejects.toThrow( + `memory_search response body exceeded ${MEMORY_SEARCH_RESPONSE_MAX_BYTES} bytes`, + ); + }); + + it("stops reading memory_search response streams after the byte cap", async () => { + const chunk = new Uint8Array(MEMORY_SEARCH_RESPONSE_MAX_BYTES + 1); + const response = new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(chunk); + controller.close(); + }, + }), + ); + + await expect( + readBoundedResponseText(response, "memory_search", MEMORY_SEARCH_RESPONSE_MAX_BYTES), + ).rejects.toThrow( + `memory_search response body exceeded ${MEMORY_SEARCH_RESPONSE_MAX_BYTES} bytes`, + ); + }); });