diff --git a/src/server/index.ts b/src/server/index.ts index 054681935..5b15f1e8c 100644 --- a/src/server/index.ts +++ b/src/server/index.ts @@ -185,11 +185,42 @@ import { import { SYSTEM_RESTART_CAPABILITY_VERSION } from "../lib/system-restart-contract"; import { createReadinessGate, type ReadinessGate } from "./readiness"; -const MAX_WS_FRAME_BYTES = 50 * 1024 * 1024; +export const MAX_WS_FRAME_BYTES = 50 * 1024 * 1024; const WEBSOCKET_IDLE_TIMEOUT_SECONDS = 0; const LIVE_SIDEBAND_PENDING_MAX = 32; +const LIVE_SIDEBAND_PENDING_BYTES_MAX = 1024 * 1024; const LIVE_SIDEBAND_CLOSE_FALLBACK_MS = 1_000; +export function exceedsLiveSidebandFrameByteLimit(frameBytes: number): boolean { + return frameBytes > MAX_WS_FRAME_BYTES; +} + +export function exceedsLiveSidebandPendingByteLimit(pendingBytes: number, incomingBytes: number): boolean { + return incomingBytes > LIVE_SIDEBAND_PENDING_BYTES_MAX - pendingBytes; +} + +function webSocketFrameBytes(frame: string | ArrayBuffer | ArrayBufferView | Blob | Buffer): number { + if (typeof frame === "string") return Buffer.byteLength(frame); + if (frame instanceof ArrayBuffer || ArrayBuffer.isView(frame)) return frame.byteLength; + return frame.size; +} + +export type LiveSidebandPendingEnqueueResult = "queued" | "too-many-frames" | "too-many-bytes"; + +export function enqueueLiveSidebandPendingFrame( + data: Pick, + frame: string | Buffer, + frameBytes = webSocketFrameBytes(frame), +): LiveSidebandPendingEnqueueResult { + const pending = data.livePending ?? (data.livePending = []); + if (pending.length >= LIVE_SIDEBAND_PENDING_MAX) return "too-many-frames"; + const pendingBytes = data.livePendingBytes ?? 0; + if (exceedsLiveSidebandPendingByteLimit(pendingBytes, frameBytes)) return "too-many-bytes"; + pending.push(frame); + data.livePendingBytes = pendingBytes + frameBytes; + return "queued"; +} + type LiveSidebandWebSocketFactory = ( url: string, headers: Record, @@ -224,6 +255,7 @@ function finalizeLiveSideband(ws: ServerWebSocket, upstream?: WebSocket) } ws.data.liveUpstream = undefined; ws.data.livePending = undefined; + ws.data.livePendingBytes = undefined; ws.data.cancel = undefined; releaseLiveSidebandAdmission(ws); } @@ -259,6 +291,7 @@ function closeLiveSideband(ws: ServerWebSocket, code = 1000, reason = "" if (ws.data.liveClosing) return; ws.data.liveClosing = true; ws.data.livePending = undefined; + ws.data.livePendingBytes = undefined; ws.data.cancel = undefined; const upstream = ws.data.liveUpstream; // Bun's `WebSocket` type narrows `readyState` to 0|1|2 even though the DOM @@ -314,6 +347,7 @@ function attachLiveSidebandUpstream( ws.data.liveOpened = true; const pending = ws.data.livePending ?? []; ws.data.livePending = undefined; + ws.data.livePendingBytes = undefined; for (const frame of pending) { try { sendUpstreamFrame(upstream, frame); @@ -326,6 +360,10 @@ function attachLiveSidebandUpstream( upstream.addEventListener("message", (event) => { if (ws.data.liveUpstream !== upstream || ws.data.liveClosing) return; try { + if (exceedsLiveSidebandFrameByteLimit(webSocketFrameBytes(event.data))) { + closeLiveSideband(ws, 1009, "message too large"); + return; + } logLiveSidebandFrame("u2c", event.data); if (typeof event.data === "string") ws.send(event.data); else if (event.data instanceof ArrayBuffer) ws.send(event.data); @@ -1255,6 +1293,7 @@ export function startServer(port?: number, deps: StartServerDeps = {}): Server, raw: string | Buffer) { if (ws.data.kind === "live-sideband") { if (ws.data.liveClosing) return; + const rawBytes = webSocketFrameBytes(raw); + if (exceedsLiveSidebandFrameByteLimit(rawBytes)) { + closeLiveSideband(ws, 1009, "message too large"); + return; + } logLiveSidebandFrame("c2u", raw); const upstream = ws.data.liveUpstream; if (!upstream || upstream.readyState === WebSocket.CONNECTING || !ws.data.liveOpened) { - const pending = ws.data.livePending ?? (ws.data.livePending = []); - if (pending.length >= LIVE_SIDEBAND_PENDING_MAX) { + const enqueueResult = enqueueLiveSidebandPendingFrame(ws.data, raw, rawBytes); + if (enqueueResult === "too-many-frames") { closeLiveSideband(ws, 1009, "too many pending frames"); return; } - pending.push(raw); + if (enqueueResult === "too-many-bytes") { + closeLiveSideband(ws, 1009, "too many pending bytes"); + return; + } return; } if (upstream.readyState !== WebSocket.OPEN) { diff --git a/src/server/ws-bridge.ts b/src/server/ws-bridge.ts index da052c1b8..c593861d2 100644 --- a/src/server/ws-bridge.ts +++ b/src/server/ws-bridge.ts @@ -40,6 +40,8 @@ export interface WsData { liveUpstreamUrl?: string; liveUpstreamHeaders?: Record; livePending?: Array; + /** Total encoded bytes retained in livePending while the upstream connects. */ + livePendingBytes?: number; liveOpened?: boolean; /** Once teardown starts, ignore new client frames until the upstream closes. */ liveClosing?: boolean; @@ -467,4 +469,4 @@ export async function readBoundedPrefix( export function looksLikeSse(prefix: Uint8Array): boolean { const text = new TextDecoder().decode(prefix); return /^\s*(event:|data:)/.test(text); -} \ No newline at end of file +} diff --git a/tests/server-live.test.ts b/tests/server-live.test.ts index 329728fdd..fac24ddea 100644 --- a/tests/server-live.test.ts +++ b/tests/server-live.test.ts @@ -14,7 +14,13 @@ import { runStartupReadinessSync, type ReadinessGate, } from "../src/server/readiness"; -import { startServer } from "../src/server"; +import { + enqueueLiveSidebandPendingFrame, + exceedsLiveSidebandFrameByteLimit, + exceedsLiveSidebandPendingByteLimit, + MAX_WS_FRAME_BYTES, + startServer, +} from "../src/server"; import { beginShutdownDrain, isDraining, resetLifecycleDrainStateForTests } from "../src/server/lifecycle"; import type { OcxConfig } from "../src/types"; import { fakeChatGptJwt } from "./helpers/fake-chatgpt-jwt"; @@ -473,7 +479,7 @@ test("a routed pool account's token overrides the caller bearer on the live rela } }); -test("sideband GET /v1/live/{callId} upgrades and relays bidirectionally to ChatGPT backend", async () => { +test("sideband GET /v1/live/{callId} relays the exact frame ceiling bidirectionally", async () => { const seenPaths: string[] = []; const seenUpgradeHeaders: Headers[] = []; const upstream = Bun.serve({ @@ -489,8 +495,9 @@ test("sideband GET /v1/live/{callId} upgrades and relays bidirectionally to Chat return new Response("not found", { status: 404 }); }, websocket: { + maxPayloadLength: MAX_WS_FRAME_BYTES, message(ws, message) { - ws.send(`echo:${typeof message === "string" ? message : message.toString()}`); + ws.send(typeof message === "string" ? `echo:${message}` : `bytes:${message.byteLength}`); }, }, }); @@ -525,13 +532,20 @@ test("sideband GET /v1/live/{callId} upgrades and relays bidirectionally to Chat } as unknown as string[]); await new Promise((resolve, reject) => { - const timer = setTimeout(() => reject(new Error("sideband timeout")), 5_000); + const timer = setTimeout(() => reject(new Error("sideband timeout")), 15_000); + let sawPing = false; client.addEventListener("open", () => { client.send("ping-sideband"); }); client.addEventListener("message", (event) => { try { - expect(String(event.data)).toBe("echo:ping-sideband"); + if (!sawPing) { + expect(String(event.data)).toBe("echo:ping-sideband"); + sawPing = true; + client.send(Buffer.alloc(MAX_WS_FRAME_BYTES)); + return; + } + expect(String(event.data)).toBe(`bytes:${MAX_WS_FRAME_BYTES}`); expect(seenPaths).toContain("/v1/live/rtc_sideband"); expect(seenUpgradeHeaders).toHaveLength(1); expect(seenUpgradeHeaders[0].get("openai-alpha")).toBe("quicksilver=v2"); @@ -555,6 +569,40 @@ test("sideband GET /v1/live/{callId} upgrades and relays bidirectionally to Chat await server.stop(true); await upstream.stop(true); } +}, { timeout: 20_000 }); + +test("sideband byte predicates accept exact limits and reject one byte over", () => { + expect(exceedsLiveSidebandFrameByteLimit(50 * 1024 * 1024)).toBe(false); + expect(exceedsLiveSidebandFrameByteLimit(50 * 1024 * 1024 + 1)).toBe(true); + expect(exceedsLiveSidebandPendingByteLimit(256 * 1024, 768 * 1024)).toBe(false); + expect(exceedsLiveSidebandPendingByteLimit(256 * 1024, 768 * 1024 + 1)).toBe(true); +}); + +test("sideband queue enforces aggregate bytes and frame count without retaining rejected frames", () => { + const byteBounded: { livePending: Array; livePendingBytes: number } = { + livePending: [], + livePendingBytes: 0, + }; + const quarterMiB = Buffer.alloc(256 * 1024); + for (let i = 0; i < 4; i += 1) { + expect(enqueueLiveSidebandPendingFrame(byteBounded, quarterMiB)).toBe("queued"); + } + expect(byteBounded.livePending).toHaveLength(4); + expect(byteBounded.livePendingBytes).toBe(1024 * 1024); + expect(enqueueLiveSidebandPendingFrame(byteBounded, Buffer.from([1]))).toBe("too-many-bytes"); + expect(byteBounded.livePending).toHaveLength(4); + expect(byteBounded.livePendingBytes).toBe(1024 * 1024); + + const countBounded: { livePending: Array; livePendingBytes: number } = { + livePending: [], + livePendingBytes: 0, + }; + for (let i = 0; i < 32; i += 1) { + expect(enqueueLiveSidebandPendingFrame(countBounded, "x")).toBe("queued"); + } + expect(enqueueLiveSidebandPendingFrame(countBounded, "x")).toBe("too-many-frames"); + expect(countBounded.livePending).toHaveLength(32); + expect(countBounded.livePendingBytes).toBe(32); }); test("buildLiveSidebandUpstreamWsUrl maps Frameless and Realtime join shapes", async () => {