From c2185ff3aaeee48d9328794fcc597de21f583b09 Mon Sep 17 00:00:00 2001 From: Marius Wichtner Date: Tue, 1 Sep 2026 15:30:06 +0200 Subject: [PATCH] fix(gateway): decode bounded Zstandard inference requests --- .../api/openrouter/[...path]/route.test.ts | 91 ++++++++ .../src/app/api/openrouter/[...path]/route.ts | 17 +- .../src/lib/ai-gateway/request-body.test.ts | 210 ++++++++++++++++++ apps/web/src/lib/ai-gateway/request-body.ts | 116 ++++++++++ 4 files changed, 433 insertions(+), 1 deletion(-) create mode 100644 apps/web/src/lib/ai-gateway/request-body.test.ts create mode 100644 apps/web/src/lib/ai-gateway/request-body.ts diff --git a/apps/web/src/app/api/openrouter/[...path]/route.test.ts b/apps/web/src/app/api/openrouter/[...path]/route.test.ts index 88d28ed773..6df40fd20f 100644 --- a/apps/web/src/app/api/openrouter/[...path]/route.test.ts +++ b/apps/web/src/app/api/openrouter/[...path]/route.test.ts @@ -1,4 +1,6 @@ import { describe, it, expect, beforeEach, afterEach } from '@jest/globals'; +import { NextRequest } from 'next/server'; +import { zstdCompressSync } from 'node:zlib'; import type { User } from '@kilocode/db/schema'; import { getUserFromAuth } from '@/lib/user/server'; import { getBalanceAndOrgSettings } from '@/lib/organizations/organization-usage'; @@ -294,6 +296,95 @@ describe('POST /api/openrouter/v1/chat/completions rules-engine actions', () => jest.useRealTimers(); }); + it.each(['openrouter', 'gateway'])( + 'decodes zstd JSON before blank-model validation through /api/%s', + async alias => { + const { POST } = await import('@/app/api/gateway/[...path]/route'); + const response = await POST( + new NextRequest(`http://localhost/api/${alias}/chat/completions`, { + method: 'POST', + headers: { 'content-type': 'application/json', 'content-encoding': 'zstd' }, + body: zstdCompressSync(JSON.stringify(makeBody(''))), + }) + ); + + expect(response.status).toBe(404); + expect(await response.json()).toMatchObject({ + error: 'Model not found', + error_type: 'model_not_found', + }); + expect(mockedGetProvider).not.toHaveBeenCalled(); + expect(mockedUpstreamRequest).not.toHaveBeenCalled(); + } + ); + + it.each([ + { encoding: 'zstd', status: 400, error: 'Invalid zstd request body.' }, + { + encoding: 'zstd, identity', + status: 415, + error: 'Unsupported Content-Encoding. Use identity or zstd.', + }, + ])( + 'returns a stable $status error for invalid $encoding bodies', + async ({ encoding, status, error }) => { + const { POST } = await import('./route'); + const response = await POST( + new NextRequest('http://localhost/api/openrouter/chat/completions', { + method: 'POST', + headers: { 'content-type': 'application/json', 'content-encoding': encoding }, + body: JSON.stringify(makeBody('')), + }) + ); + + expect(response.status).toBe(status); + expect(await response.json()).toEqual({ + error, + error_type: 'invalid_request', + message: error, + }); + expect(mockedGetUserFromAuth).not.toHaveBeenCalled(); + expect(mockedUpstreamRequest).not.toHaveBeenCalled(); + } + ); + + it('forwards decoded JSON without compressed transport headers', async () => { + const actual = jest.requireActual<{ upstreamRequest: typeof upstreamRequest }>( + '@/lib/ai-gateway/providers/upstream-request' + ); + mockedUpstreamRequest.mockImplementation(actual.upstreamRequest); + const fetch = jest + .spyOn(global, 'fetch') + .mockResolvedValue(upstreamJsonResponse({ choices: [] })); + try { + const { POST } = await import('./route'); + const body = zstdCompressSync(JSON.stringify(makeBody())); + const response = await POST( + new NextRequest('http://localhost/api/openrouter/chat/completions', { + method: 'POST', + headers: { + 'content-type': 'application/json', + 'content-encoding': 'zstd', + 'content-length': String(body.length), + 'x-forwarded-for': '127.0.0.1', + }, + body, + }) + ); + + expect(response.status).toBe(200); + expect(fetch).toHaveBeenCalledTimes(1); + const init = fetch.mock.calls.at(0)?.[1]; + const headers = new Headers(init?.headers); + expect(headers.get('content-type')).toBe('application/json'); + expect(headers.has('content-encoding')).toBe(false); + expect(headers.has('content-length')).toBe(false); + expect(JSON.parse(String(init?.body))).toMatchObject(makeBody()); + } finally { + fetch.mockRestore(); + } + }); + it('rejects providerOptions and directs clients to provider', async () => { const { POST } = await import('./route'); const response = await POST( diff --git a/apps/web/src/app/api/openrouter/[...path]/route.ts b/apps/web/src/app/api/openrouter/[...path]/route.ts index 326544a859..22039975e2 100644 --- a/apps/web/src/app/api/openrouter/[...path]/route.ts +++ b/apps/web/src/app/api/openrouter/[...path]/route.ts @@ -3,6 +3,7 @@ import { type NextRequest } from 'next/server'; import { stripRequiredPrefix, toMicrodollars } from '@/lib/utils'; import { extractPromptInfo } from '@/lib/ai-gateway/extractPromptInfo'; import { determineFallbackFeature } from '@/lib/ai-gateway/determineFallbackFeature'; +import { readGatewayRequestBody } from '@/lib/ai-gateway/request-body'; import { validateFeatureHeader, FEATURE_HEADER, @@ -187,7 +188,21 @@ export async function POST(request: NextRequest): Promise({ + start(controller) { + for (let offset = 0; offset < buffer.length; offset += 7) { + controller.enqueue(buffer.subarray(offset, offset + 7)); + } + controller.close(); + }, + }); +} + +describe('readGatewayRequestBody', () => { + it.each([undefined, 'identity', ' IDENTITY\t', 'zstd', '\tZsTd '])( + 'preserves exact image and encrypted reasoning JSON with encoding %s', + async encoding => { + const body = encoding?.trim().toLowerCase() === 'zstd' ? frame : bytes; + await expect(readGatewayRequestBody(request(chunked(body), encoding))).resolves.toEqual({ + text, + }); + } + ); + + it('matches Request.text UTF-8 and BOM handling', async () => { + const body = Buffer.concat([Buffer.from([0xef, 0xbb, 0xbf]), bytes, Buffer.from([0xff])]); + const expected = await request(body).text(); + await expect(readGatewayRequestBody(request(chunked(body)))).resolves.toEqual({ + text: expected, + }); + await expect(readGatewayRequestBody(request(zstdCompressSync(body), 'zstd'))).resolves.toEqual({ + text: expected, + }); + }); + + it('preserves empty identity bodies', async () => { + await expect(readGatewayRequestBody(request(null))).resolves.toEqual({ text: '' }); + }); + + it.each(['', 'gzip', 'br', 'unknown', 'zstd, identity', 'identity, zstd', 'zstd, zstd'])( + 'rejects unsupported or stacked encoding %s', + async encoding => { + const input = request(frame, encoding); + await expect(readGatewayRequestBody(input)).resolves.toEqual({ + status: 415, + error: 'Unsupported Content-Encoding. Use identity or zstd.', + }); + expect(input.bodyUsed).toBe(false); + } + ); + + it.each([0, 4, 6, frame.length - 1])('rejects a frame truncated to %i bytes', async length => { + await expect( + readGatewayRequestBody(request(chunked(frame.subarray(0, length)), 'zstd')) + ).resolves.toEqual({ status: 400, error: 'Invalid zstd request body.' }); + }); + + it.each([bytes, Buffer.concat([frame, Buffer.from([0])])])( + 'rejects invalid frame data', + async body => { + await expect(readGatewayRequestBody(request(body, 'zstd'))).resolves.toEqual({ + status: 400, + error: 'Invalid zstd request body.', + }); + } + ); + + it('rejects a complete frame with a corrupt checksum', async () => { + const corrupt = Buffer.from(frame); + corrupt.writeUInt8(corrupt.readUInt8(corrupt.length - 1) ^ 0xff, corrupt.length - 1); + await expect(readGatewayRequestBody(request(corrupt, 'zstd'))).resolves.toEqual({ + status: 400, + error: 'Invalid zstd request body.', + }); + }); + + it('rejects concatenated and skippable frames instead of silently dropping bytes', async () => { + const skip = Buffer.from([0x50, 0x2a, 0x4d, 0x18, 3, 0, 0, 0, 1, 2, 3]); + for (const body of [Buffer.concat([frame, frame]), Buffer.concat([frame, skip]), skip]) { + await expect(readGatewayRequestBody(request(body, 'zstd'))).resolves.toEqual({ + status: 400, + error: 'Invalid zstd request body.', + }); + } + }); + + it('accepts frames without a declared content size', async () => { + const body = zstdCompressSync(bytes, { params: { [constants.ZSTD_c_contentSizeFlag]: 0 } }); + await expect(readGatewayRequestBody(request(body, 'zstd'))).resolves.toEqual({ text }); + }); + + it.each(['identity', 'zstd'])( + 'enforces actual input and output byte limits for %s', + async encoding => { + const body = encoding === 'zstd' ? frame : bytes; + await expect( + readGatewayRequestBody(request(chunked(body), encoding), { + input: body.length, + output: bytes.length, + }) + ).resolves.toEqual({ text }); + const input = request(chunked(body), encoding); + input.headers.set('content-length', '1'); + await expect( + readGatewayRequestBody(input, { input: body.length - 1, output: bytes.length }) + ).resolves.toEqual({ + status: 413, + error: 'Request body exceeds the gateway resource limit.', + }); + await expect( + readGatewayRequestBody(request(body, encoding), { output: bytes.length - 1 }) + ).resolves.toEqual({ + status: 413, + error: 'Request body exceeds the gateway resource limit.', + }); + } + ); + + it('enforces the native decoder window limit before output exceeds its limit', async () => { + const body = zstdCompressSync(Buffer.alloc(2048, 65), { + params: { [constants.ZSTD_c_contentSizeFlag]: 0, [constants.ZSTD_c_windowLog]: 12 }, + }); + await expect( + readGatewayRequestBody(request(body, 'zstd'), { window: 10, output: 4096 }) + ).resolves.toEqual({ status: 413, error: 'Request body exceeds the gateway resource limit.' }); + }); + + it.each(['identity', 'zstd'])('cancels a stalled %s body read', async encoding => { + const controller = new AbortController(); + const started = Promise.withResolvers(); + const cancelled = Promise.withResolvers(); + const body = new ReadableStream({ + start(stream) { + stream.enqueue((encoding === 'zstd' ? frame : bytes).subarray(0, 4)); + }, + pull() { + started.resolve(); + }, + cancel(reason) { + cancelled.resolve(reason); + }, + }); + const result = readGatewayRequestBody(request(body, encoding, controller.signal)); + await started.promise; + controller.abort(); + await expect(result).resolves.toEqual({ + status: 499, + error: 'Request cancelled while reading request body.', + }); + await expect(cancelled.promise).resolves.toMatchObject({ name: 'AbortError' }); + }); + + it('does not decode an already cancelled request', async () => { + const controller = new AbortController(); + controller.abort(); + await expect( + readGatewayRequestBody(request(frame, 'zstd', controller.signal)) + ).resolves.toMatchObject({ status: 499 }); + }); +}); diff --git a/apps/web/src/lib/ai-gateway/request-body.ts b/apps/web/src/lib/ai-gateway/request-body.ts new file mode 100644 index 0000000000..e2075d671d --- /dev/null +++ b/apps/web/src/lib/ai-gateway/request-body.ts @@ -0,0 +1,116 @@ +import { Readable, Transform, Writable } from 'node:stream'; +import { pipeline } from 'node:stream/promises'; +import type { ReadableStream } from 'node:stream/web'; +import { constants, createZstdDecompress } from 'node:zlib'; + +const MAX_BODY_BYTES = 32 * 1024 * 1024; + +function complete(buffer: Buffer): boolean { + if (buffer.length < 5 || buffer.readUInt32LE(0) !== 0xfd2fb528) return false; + const descriptor = buffer.readUInt8(4); + const single = (descriptor & 0x20) !== 0; + const dictionary = descriptor & 3; + const size = descriptor >>> 6; + let offset = + 5 + + (single ? 0 : 1) + + (dictionary === 3 ? 4 : dictionary) + + (size === 0 ? (single ? 1 : 0) : 2 ** size); + + for (;;) { + if (buffer.length - offset < 3) return false; + const block = buffer.readUIntLE(offset, 3); + const type = (block >>> 1) & 3; + if (type === 3) return false; + offset += 3 + (type === 1 ? 1 : block >>> 3); + if (offset > buffer.length) return false; + if (block & 1) break; + } + return offset + (descriptor & 4 ? 4 : 0) === buffer.length; +} + +export async function readGatewayRequestBody( + request: Request, + limits: { input?: number; output?: number; window?: number } = {} +): Promise<{ text: string } | { status: 400 | 413 | 415 | 499; error: string }> { + const encoding = request.headers.get('content-encoding')?.trim().toLowerCase() ?? 'identity'; + if (encoding !== 'identity' && encoding !== 'zstd') { + return { status: 415, error: 'Unsupported Content-Encoding. Use identity or zstd.' }; + } + + const exceeded = new Error('Request body exceeds the gateway resource limit.'); + const chunks: Buffer[] = []; + const utf8 = new TextDecoder(); + let received = 0; + let decoded = 0; + let text = ''; + const input = new Transform({ + transform(chunk: Buffer, _encoding, callback) { + received += chunk.length; + if (received > (limits.input ?? MAX_BODY_BYTES)) { + callback(exceeded); + return; + } + if (encoding === 'zstd') chunks.push(chunk); + callback(null, chunk); + }, + flush(callback) { + if (encoding === 'zstd' && !complete(Buffer.concat(chunks, received))) { + callback(new Error('Incomplete zstd request body.')); + return; + } + chunks.length = 0; + callback(); + }, + }); + const output = new Writable({ + write(chunk: Buffer, _encoding, callback) { + decoded += chunk.length; + if (decoded > (limits.output ?? MAX_BODY_BYTES)) { + callback(exceeded); + return; + } + text += utf8.decode(chunk, { stream: true }); + callback(); + }, + }); + + try { + const source = request.body + ? Readable.fromWeb(request.body as ReadableStream) + : Readable.from([]); + await pipeline( + [ + source, + input, + ...(encoding === 'zstd' + ? [ + createZstdDecompress({ + params: { [constants.ZSTD_d_windowLogMax]: limits.window ?? 25 }, + }), + ] + : []), + output, + ], + { signal: request.signal } + ); + return { text: text + utf8.decode() }; + } catch (error) { + if (request.signal.aborted) { + return { status: 499, error: 'Request cancelled while reading request body.' }; + } + if ( + error === exceeded || + (typeof error === 'object' && + error !== null && + 'code' in error && + error.code === 'ZSTD_error_frameParameter_windowTooLarge') + ) { + return { status: 413, error: exceeded.message }; + } + return { + status: 400, + error: encoding === 'zstd' ? 'Invalid zstd request body.' : 'Could not read request body.', + }; + } +}