diff --git a/packages/client/src/daemon-client.test.ts b/packages/client/src/daemon-client.test.ts index dc264c4f46..1d6da7fb45 100644 --- a/packages/client/src/daemon-client.test.ts +++ b/packages/client/src/daemon-client.test.ts @@ -8607,3 +8607,60 @@ test("RPC deadlines remain active while a preceding frame waits for drain", asyn await first; } }); + +test.each([false, true])( + "event listener errors do not close the transport (chunked=%s)", + async (chunked) => { + const mock = createMockTransport(); + const closed = vi.fn(); + mock.transport.close = closed; + const client = new DaemonClient({ + url: "ws://test", + clientId: "listener-error", + transportFactory: () => mock.transport, + reconnect: { enabled: false }, + }); + clients.push(client); + const connecting = client.connect(); + mock.triggerOpen({ features: { binaryChunkingV1: chunked } }); + await connecting; + const unsubscribe = client.subscribe(() => { + throw new Error("listener failed"); + }); + const text = wrapSessionMessage({ + type: "agent_deleted", + payload: { agentId: "agent-test", requestId: "listener-test" }, + }); + const payload = new TextEncoder().encode(text); + const frame = chunked ? new Uint8Array([0x31, 1, 0, 3, ...payload]) : text; + expect(() => mock.triggerMessage(frame)).toThrow("listener failed"); + expect(closed).not.toHaveBeenCalled(); + unsubscribe(); + const events: string[] = []; + client.subscribe((event) => { + events.push(event.type); + }); + mock.triggerMessage(frame); + expect(events).toEqual(["agent_deleted"]); + await client.sendHostFrame("route", "still connected"); + expect(client.isConnected).toBe(true); + }, +); + +test("malformed chunk metadata still closes the client transport", async () => { + const mock = createMockTransport(); + const closed = vi.fn(); + mock.transport.close = closed; + const client = new DaemonClient({ + url: "ws://test", + clientId: "malformed-chunk", + transportFactory: () => mock.transport, + reconnect: { enabled: false }, + }); + clients.push(client); + const connecting = client.connect(); + mock.triggerOpen({ features: { binaryChunkingV1: true } }); + await connecting; + mock.triggerMessage(new Uint8Array([0x31, 2, 0, 3, 65])); + expect(closed).toHaveBeenCalledWith(4003, "Invalid chunked transport"); +}); diff --git a/packages/client/src/daemon-client.ts b/packages/client/src/daemon-client.ts index 9a2965130b..a83b3b89b8 100644 --- a/packages/client/src/daemon-client.ts +++ b/packages/client/src/daemon-client.ts @@ -1,4 +1,7 @@ -import { ChunkedTransport } from "@getpaseo/protocol/chunked-transport"; +import { + ChunkedTransport, + ChunkedTransportProtocolError, +} from "@getpaseo/protocol/chunked-transport"; import { HOST_ROUTE_OPCODE, encodeHostRouteFrame, @@ -1682,7 +1685,8 @@ export class DaemonClient { if (typeof data === "string" || bytes) this.chunkedTransport!.receive(typeof data === "string" ? data : bytes!); else this.handleTransportMessage(data); - } catch { + } catch (error) { + if (!(error instanceof ChunkedTransportProtocolError)) throw error; transport.close(4003, "Invalid chunked transport"); } }), diff --git a/packages/protocol/src/chunked-transport.test.ts b/packages/protocol/src/chunked-transport.test.ts index 18ea469521..696c0507a0 100644 --- a/packages/protocol/src/chunked-transport.test.ts +++ b/packages/protocol/src/chunked-transport.test.ts @@ -348,3 +348,23 @@ it("a queued pre-write rejection does not poison the next logical send", async ( expect(delivered).toEqual(["sibling"]); expect(failures).toEqual([]); }); + +it.each([false, true])("delivery errors leave framing usable (chunked=%s)", async (chunked) => { + let fail = true; + const delivered: (string | Uint8Array)[] = []; + const transport = new ChunkedTransport({ + frameBytes: 64, + send: () => {}, + deliver: (frame) => { + if (fail) throw new Error("listener failed"); + delivered.push(frame); + }, + }); + transport.enabled = chunked; + const frame = chunked ? new Uint8Array([0x31, 1, 0, 3, 65]) : "A"; + expect(() => transport.receive(frame)).toThrow("listener failed"); + fail = false; + transport.receive(frame); + expect(delivered).toEqual(["A"]); + await transport.send("still connected"); +}); diff --git a/packages/protocol/src/chunked-transport.ts b/packages/protocol/src/chunked-transport.ts index 536c7f1f49..6bfc8aabfc 100644 --- a/packages/protocol/src/chunked-transport.ts +++ b/packages/protocol/src/chunked-transport.ts @@ -1,6 +1,9 @@ import { MAX_WEBSOCKET_MESSAGE_BYTES } from "./websocket-message-limit.js"; type Frame = string | Uint8Array; + +/** A malformed wire frame, distinct from a consumer delivery failure. */ +export class ChunkedTransportProtocolError extends Error {} // Reserved outside the application binary opcode space. Version, kind, flags follow. const OPCODE = 0x31; const HEADER = 4; @@ -119,53 +122,59 @@ export class ChunkedTransport { receive(frame: Frame): void { if (this.closed) return; + let message: Frame | undefined; try { - if (typeof frame === "string" || frame[0] !== OPCODE) { - if (this.assembly) throw new Error("Interrupted chunked message"); - this.io.deliver(frame); - return; - } - if ( - !this.enabled || - frame.length <= HEADER || - frame[1] !== 1 || - frame[2]! > 1 || - (frame[3]! & ~3) !== 0 - ) - throw new Error("Invalid chunk header"); - const start = (frame[3]! & 1) !== 0; - const end = (frame[3]! & 2) !== 0; - if (start === !!this.assembly) throw new Error("Invalid chunk sequence"); - if (start) { - this.assembly = new Uint8Array(0); - this.kind = frame[2]!; - this.size = 0; - } - if (frame[2] !== this.kind) throw new Error("Chunk kind changed"); - this.size += frame.length - HEADER; - if (this.size > MAX_WEBSOCKET_MESSAGE_BYTES) throw new Error("Logical message too large"); - const payload = frame.subarray(HEADER); - if (this.assembly!.byteLength < this.size) { - const capacity = Math.min( - MAX_WEBSOCKET_MESSAGE_BYTES, - Math.max(this.size, this.assembly!.byteLength * 2, 16384), - ); - const grown = new Uint8Array(capacity); - grown.set(this.assembly!.subarray(0, this.size - payload.byteLength)); - this.assembly = grown; - } - this.assembly!.set(payload, this.size - payload.byteLength); - if (!end) return; - const data = this.assembly!.subarray(0, this.size); - this.assembly = undefined; - this.size = 0; - this.io.deliver( - this.kind === 0 ? new TextDecoder("utf-8", { fatal: true }).decode(data) : data.slice(), - ); + message = this.assemble(frame); } catch (error) { this.close(); - throw error; + throw new ChunkedTransportProtocolError( + error instanceof Error ? error.message : String(error), + { cause: error }, + ); } + if (message !== undefined) this.io.deliver(message); + } + + private assemble(frame: Frame): Frame | undefined { + if (typeof frame === "string" || frame[0] !== OPCODE) { + if (this.assembly) throw new Error("Interrupted chunked message"); + return frame; + } + if ( + !this.enabled || + frame.length <= HEADER || + frame[1] !== 1 || + frame[2]! > 1 || + (frame[3]! & ~3) !== 0 + ) + throw new Error("Invalid chunk header"); + const start = (frame[3]! & 1) !== 0; + const end = (frame[3]! & 2) !== 0; + if (start === !!this.assembly) throw new Error("Invalid chunk sequence"); + if (start) { + this.assembly = new Uint8Array(0); + this.kind = frame[2]!; + this.size = 0; + } + if (frame[2] !== this.kind) throw new Error("Chunk kind changed"); + this.size += frame.length - HEADER; + if (this.size > MAX_WEBSOCKET_MESSAGE_BYTES) throw new Error("Logical message too large"); + const payload = frame.subarray(HEADER); + if (this.assembly!.byteLength < this.size) { + const capacity = Math.min( + MAX_WEBSOCKET_MESSAGE_BYTES, + Math.max(this.size, this.assembly!.byteLength * 2, 16384), + ); + const grown = new Uint8Array(capacity); + grown.set(this.assembly!.subarray(0, this.size - payload.byteLength)); + this.assembly = grown; + } + this.assembly!.set(payload, this.size - payload.byteLength); + if (!end) return; + const data = this.assembly!.subarray(0, this.size); + this.assembly = undefined; + this.size = 0; + return this.kind === 0 ? new TextDecoder("utf-8", { fatal: true }).decode(data) : data.slice(); } private fail(error: unknown): void { diff --git a/packages/server/src/server/websocket-server.file-transfer.e2e.test.ts b/packages/server/src/server/websocket-server.file-transfer.e2e.test.ts index b708ebb149..abce664ffd 100644 --- a/packages/server/src/server/websocket-server.file-transfer.e2e.test.ts +++ b/packages/server/src/server/websocket-server.file-transfer.e2e.test.ts @@ -255,3 +255,12 @@ test("an awaited timeline framing overflow closes the actual daemon socket", asy await client.close(); } }, 60000); + +test("malformed chunk metadata closes the daemon socket with a protocol error", async () => { + daemon = await createTestPaseoDaemon(); + const source = await connectSocket(daemon.port, "invalid-framing", { binary_chunking_v1: true }); + sockets.push(source); + const closed = new Promise((resolve) => source.once("close", (code) => resolve(code))); + source.send(new Uint8Array([0x31, 2, 0, 3, 65])); + expect(await closed).toBe(1002); +}); diff --git a/packages/server/src/server/websocket-server.ts b/packages/server/src/server/websocket-server.ts index 7c246a4ba8..7bfc4d5a0f 100644 --- a/packages/server/src/server/websocket-server.ts +++ b/packages/server/src/server/websocket-server.ts @@ -1,4 +1,7 @@ -import { ChunkedTransport } from "@getpaseo/protocol/chunked-transport"; +import { + ChunkedTransport, + ChunkedTransportProtocolError, +} from "@getpaseo/protocol/chunked-transport"; import { type WorkspaceRuntimeStore } from "./workspace-runtime-store/index.js"; import { createAgentMessageDelivery, @@ -2220,7 +2223,8 @@ export class VoiceAssistantWebSocketServer { const data = args[0] as Buffer | ArrayBuffer | Buffer[] | string; try { this.chunkedSocket(ws).receive(typeof data === "string" ? data : bufferFromWsData(data)); - } catch { + } catch (error) { + if (!(error instanceof ChunkedTransportProtocolError)) throw error; ws.close(1002, "Invalid chunked transport"); } });