diff --git a/src/server/images.ts b/src/server/images.ts index a35ca84e6..28229df3e 100644 --- a/src/server/images.ts +++ b/src/server/images.ts @@ -52,6 +52,45 @@ const IMAGES_UPSTREAM_TIMEOUT_MS = 300_000; */ const IMAGES_RESPONSE_MAX_BYTES = 100 * 1024 * 1024; +async function readImageResponseBody(response: Response): Promise | undefined> { + const contentLength = response.headers.get("content-length"); + if (contentLength !== null) { + const declaredBytes = Number(contentLength); + if (Number.isFinite(declaredBytes) && declaredBytes > IMAGES_RESPONSE_MAX_BYTES) { + try { void response.body?.cancel().catch(() => undefined); } catch { /* best-effort upstream abort */ } + return undefined; + } + } + + const reader = response.body?.getReader(); + if (!reader) return new Uint8Array(); + let payload = new Uint8Array(Math.min(IMAGES_RESPONSE_MAX_BYTES, 64 * 1024)); + let length = 0; + try { + for (;;) { + const { done, value } = await reader.read(); + if (done) return payload.subarray(0, length); + if (!value || value.byteLength === 0) continue; + if (value.byteLength > IMAGES_RESPONSE_MAX_BYTES - length) { + try { void reader.cancel().catch(() => undefined); } catch { /* best-effort upstream abort */ } + return undefined; + } + if (length + value.byteLength > payload.byteLength) { + const grown = new Uint8Array(Math.min( + IMAGES_RESPONSE_MAX_BYTES, + Math.max(payload.byteLength * 2, length + value.byteLength), + )); + grown.set(payload.subarray(0, length)); + payload = grown; + } + payload.set(value, length); + length += value.byteLength; + } + } finally { + reader.releaseLock(); + } +} + const CCA_IMAGE_MODEL = "gemini-3.1-flash-image"; /** @@ -449,12 +488,11 @@ export async function handleImages( body: JSON.stringify(body), signal: linkedSignal.signal, }); - // Buffer rather than stream: the payload is one JSON document (base64 image, typically a few - // MB), and buffering keeps the timeout window covering the whole exchange. Cap the size to - // prevent an oversized response from exhausting process memory. - const payload = await upstreamResponse.arrayBuffer(); - if (payload.byteLength > IMAGES_RESPONSE_MAX_BYTES) { - return formatErrorResponse(502, "upstream_error", `image ${endpoint} response too large (${payload.byteLength} bytes)`); + // The payload is one JSON document, so buffer it under a strict byte ceiling. The reader + // cancels as soon as the response exceeds the ceiling rather than allocating the whole body. + const payload = await readImageResponseBody(upstreamResponse); + if (!payload) { + return formatErrorResponse(502, "upstream_error", `image ${endpoint} response too large (exceeded ${IMAGES_RESPONSE_MAX_BYTES} bytes)`); } forward?.recordOutcome?.(upstreamResponse.status); const relayHeaders: Record = {}; diff --git a/tests/server-images.test.ts b/tests/server-images.test.ts index 437c76c25..2fcc87d55 100644 --- a/tests/server-images.test.ts +++ b/tests/server-images.test.ts @@ -290,6 +290,45 @@ test("falls back to a keyed openai-responses provider when no forward provider e } }); +test("rejects an oversized image relay response before buffering its body", async () => { + let bodyCanceled = false; + globalThis.fetch = (async (input: RequestInfo | URL) => { + const requestUrl = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; + if (requestUrl === "https://api.openai.com/v1/images/generations") { + return new Response(new ReadableStream({ + cancel() { bodyCanceled = true; }, + }), { + status: 200, + headers: { + "content-type": "application/json", + "content-length": String(100 * 1024 * 1024 + 1), + }, + }); + } + return originalFetch(input); + }) as typeof fetch; + + const config = { + port: 0, + defaultProvider: "openai-apikey", + openaiProviderTierVersion: 2, + providers: { + openai: disabledOpenAiProvider, + "openai-apikey": keyedProvider(), + }, + } as OcxConfig; + const { handleImages } = await import("../src/server/images"); + const response = await handleImages(new Request("http://localhost/v1/images/generations", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ prompt: "a cat", model: "gpt-image-2" }), + }), config, "generations", { model: "", provider: "" } as never); + + expect(response.status).toBe(502); + expect(await response.text()).toContain("response too large"); + expect(bodyCanceled).toBe(true); +}); + test("an explicit custom Images provider uses its configured endpoint, key, and headers", async () => { const captured: CapturedRequest[] = []; const upstream = Bun.serve({