From 0351dff52b0e40f8c6012b8a2321b7341ab07fda Mon Sep 17 00:00:00 2001 From: Ben Vinegar Date: Sat, 29 Aug 2026 10:34:05 -0400 Subject: [PATCH 1/5] fix(session): bind broker state to peer ownership --- .changeset/secure-session-peer-ownership.md | 2 + .../src/brokerState.test.ts | 123 +++++++++++---- .../session-broker-core/src/brokerState.ts | 109 ++++++++----- packages/session-broker/src/broker.test.ts | 8 +- packages/session-broker/src/broker.ts | 57 ++++--- .../session-broker/src/connection.test.ts | 146 ++++++++++++++++++ packages/session-broker/src/connection.ts | 64 +++++--- packages/session-broker/src/daemon.test.ts | 116 +++++++++++++- packages/session-broker/src/daemon.ts | 43 ++++-- src/session/broker/brokerServer.test.ts | 4 +- src/session/broker/brokerServer.ts | 7 +- src/session/broker/state.ts | 6 +- test/helpers/review-session-harness.ts | 3 +- 13 files changed, 545 insertions(+), 143 deletions(-) create mode 100644 .changeset/secure-session-peer-ownership.md diff --git a/.changeset/secure-session-peer-ownership.md b/.changeset/secure-session-peer-ownership.md new file mode 100644 index 000000000..a845151cc --- /dev/null +++ b/.changeset/secure-session-peer-ownership.md @@ -0,0 +1,2 @@ +--- +--- diff --git a/packages/session-broker-core/src/brokerState.test.ts b/packages/session-broker-core/src/brokerState.test.ts index abf198181..5d8a732ef 100644 --- a/packages/session-broker-core/src/brokerState.test.ts +++ b/packages/session-broker-core/src/brokerState.test.ts @@ -324,7 +324,7 @@ describe("session broker state", () => { createSnapshot(), ); - expect(accepted).toBe(false); + expect(accepted).toBe("invalid"); expect(state.listSessions()).toEqual([]); }); @@ -336,7 +336,7 @@ describe("session broker state", () => { state.registerSession(socket, createRegistration(), createSnapshot()); - const result = state.updateSnapshot("session-1", { + const result = state.updateSnapshot(socket, "session-1", { selectedIndex: "oops", }); @@ -344,14 +344,12 @@ describe("session broker state", () => { expect(state.getSession({ sessionId: "session-1" }).snapshot.state.selectedIndex).toBe(0); }); - test("reports missing sessions separately from invalid snapshot payloads", () => { + test("rejects snapshot and heartbeat assertions from an unregistered peer", () => { const state = createState(); + const socket = { send() {} }; - expect( - state.updateSnapshot("missing-session", { - selectedIndex: 0, - }), - ).toBe("not-found"); + expect(state.updateSnapshot(socket, "missing-session", createSnapshot())).toBe("not-owner"); + expect(state.markSessionSeen(socket, "missing-session")).toBe("not-owner"); }); test("routes one opaque broker command to the live session and resolves the async result", async () => { @@ -397,7 +395,7 @@ describe("session broker state", () => { annotationId: "annotation-1", }; - state.handleCommandResult({ + state.handleCommandResult(socket, { requestId: outgoing.requestId, ok: true, result, @@ -406,40 +404,67 @@ describe("session broker state", () => { await expect(pending).resolves.toEqual(result); }); - test("rejects in-flight commands when the session disconnects", async () => { + test("rejects cross-peer mutation and leaves the owner's command pending", async () => { const state = createState(); - const socket = { - send() {}, + const ownerSent: string[] = []; + const owner = { + send(data: string) { + ownerSent.push(data); + }, }; + const other = { send() {} }; + + state.registerSession(owner, createRegistration(), createSnapshot()); + state.registerSession( + other, + createRegistration({ sessionId: "session-2", cwd: "/other", repoRoot: "/other" }), + createSnapshot(), + ); + + expect( + state.updateSnapshot( + other, + "session-1", + createSnapshot({ updatedAt: "2026-03-22T00:00:01.000Z", selectedIndex: 1 }), + ), + ).toBe("not-owner"); + expect(state.markSessionSeen(other, "session-1")).toBe("not-owner"); + expect(state.getSession({ sessionId: "session-1" }).snapshot.state.selectedIndex).toBe(0); - state.registerSession(socket, createRegistration(), createSnapshot()); const pending = state.dispatchCommand<{ kind: "annotated"; annotationId: string }, "annotate">({ - selector: { - sessionId: "session-1", - }, + selector: { sessionId: "session-1" }, command: "annotate", - input: { - filePath: "src/example.ts", - summary: "Review note", - }, + input: { filePath: "src/example.ts", summary: "Review note" }, timeoutMessage: "Timed out waiting for the session to apply the note.", }); + const outgoing = JSON.parse(ownerSent[0]!) as { requestId: string }; - state.unregisterSocket(socket); + expect( + state.handleCommandResult(other, { + requestId: outgoing.requestId, + ok: true, + result: { kind: "annotated", annotationId: "forged" }, + }), + ).toBe("not-owner"); + expect(state.getPendingCommandCount()).toBe(1); - await expect(pending).rejects.toThrow("disconnected"); + expect( + state.handleCommandResult(owner, { + requestId: outgoing.requestId, + ok: true, + result: { kind: "annotated", annotationId: "owned" }, + }), + ).toBe("handled"); + await expect(pending).resolves.toEqual({ kind: "annotated", annotationId: "owned" }); }); - test("rejects in-flight commands when a session reconnects on a new socket", async () => { + test("rejects in-flight commands when the session disconnects", async () => { const state = createState(); - const originalSocket = { - send() {}, - }; - const replacementSocket = { + const socket = { send() {}, }; - state.registerSession(originalSocket, createRegistration(), createSnapshot()); + state.registerSession(socket, createRegistration(), createSnapshot()); const pending = state.dispatchCommand<{ kind: "annotated"; annotationId: string }, "annotate">({ selector: { sessionId: "session-1", @@ -452,13 +477,43 @@ describe("session broker state", () => { timeoutMessage: "Timed out waiting for the session to apply the note.", }); - state.registerSession( - replacementSocket, - createRegistration(), - createSnapshot({ updatedAt: "2026-03-22T00:00:01.000Z" }), + state.unregisterSocket(socket); + + await expect(pending).rejects.toThrow("disconnected"); + }); + + test("rejects a second live peer and allows reconnect only after the owner closes", () => { + const state = createState(); + const originalSocket = { send() {} }; + const replacementSocket = { send() {} }; + + expect(state.registerSession(originalSocket, createRegistration(), createSnapshot())).toBe( + "registered", ); + expect( + state.registerSession( + replacementSocket, + createRegistration(), + createSnapshot({ updatedAt: "2026-03-22T00:00:01.000Z" }), + ), + ).toBe("already-connected"); + expect(state.listSessions()[0]?.snapshot.updatedAt).toBe("2026-03-22T00:00:00.000Z"); - await expect(pending).rejects.toThrow("reconnected before the command completed"); + state.unregisterSocket(originalSocket); + + expect( + state.registerSession( + replacementSocket, + createRegistration(), + createSnapshot({ updatedAt: "2026-03-22T00:00:01.000Z" }), + ), + ).toBe("registered"); + expect(state.listSessions()[0]?.snapshot.updatedAt).toBe("2026-03-22T00:00:01.000Z"); + expect(state.updateSnapshot(originalSocket, "session-1", createSnapshot())).toBe("not-owner"); + expect(state.markSessionSeen(originalSocket, "session-1")).toBe("not-owner"); + + // A delayed close callback from the retired transport cannot unregister its replacement. + state.unregisterSocket(originalSocket); expect(state.listSessions()).toHaveLength(1); }); @@ -537,7 +592,7 @@ describe("session broker state", () => { }), ).toBe(0); - state.markSessionSeen("session-1"); + expect(state.markSessionSeen(socket, "session-1")).toBe("seen"); expect( state.pruneStaleSessions({ diff --git a/packages/session-broker-core/src/brokerState.ts b/packages/session-broker-core/src/brokerState.ts index 43c104f71..99b94fae1 100644 --- a/packages/session-broker-core/src/brokerState.ts +++ b/packages/session-broker-core/src/brokerState.ts @@ -9,6 +9,7 @@ import type { interface PendingCommand { sessionId: string; + socket: DaemonSessionSocket; resolve: (result: Result) => void; reject: (error: Error) => void; timeout: ReturnType; @@ -58,7 +59,10 @@ export interface SessionBrokerViewAdapter< listComments: (session: ListedSession, filter: { filePath?: string }) => SessionCommentSummary[]; } -export type UpdateSnapshotResult = "updated" | "invalid" | "not-found"; +export type RegisterSessionResult = "registered" | "invalid" | "already-connected"; +export type UpdateSnapshotResult = "updated" | "invalid" | "not-owner"; +export type MarkSessionSeenResult = "seen" | "not-owner"; +export type HandleCommandResult = "handled" | "not-found" | "not-owner"; export type SessionTargetSelector = SessionTargetInput; @@ -203,7 +207,11 @@ export class SessionBrokerState< return this.pendingCommands.size; } - registerSession(socket: DaemonSessionSocket, registrationInput: unknown, snapshotInput: unknown) { + registerSession( + socket: DaemonSessionSocket, + registrationInput: unknown, + snapshotInput: unknown, + ): RegisterSessionResult { const registration = this.view.parseRegistration(registrationInput); const snapshot = this.view.parseSnapshot(snapshotInput); if (!registration || !snapshot) { @@ -217,7 +225,14 @@ export class SessionBrokerState< ); } - return false; + return "invalid"; + } + + const existing = this.sessions.get(registration.sessionId); + if (existing && existing.socket !== socket) { + // Reconnect proof is not available yet, so an unauthenticated peer cannot supersede a live + // owner. Once the owner closes, normal unregister cleanup makes this ID available again. + return "already-connected"; } const previousSessionId = this.sessionIdsBySocket.get(socket); @@ -225,17 +240,6 @@ export class SessionBrokerState< this.unregisterSocket(socket); } - const existing = this.sessions.get(registration.sessionId); - if (existing && existing.socket !== socket) { - this.sessionIdsBySocket.delete(existing.socket); - // A reconnect on a new socket supersedes the old transport immediately. Reject in-flight - // commands so callers do not wait on a connection that can never answer. - this.rejectPendingCommandsForSession( - registration.sessionId, - new Error("Session reconnected before the command completed."), - ); - } - const now = new Date().toISOString(); this.sessions.set(registration.sessionId, { registration, @@ -245,13 +249,22 @@ export class SessionBrokerState< lastSeenAt: now, }); this.sessionIdsBySocket.set(socket, registration.sessionId); - return true; + return "registered"; } - updateSnapshot(sessionId: string, snapshotInput: unknown): UpdateSnapshotResult { - const entry = this.sessions.get(sessionId); - if (!entry) { - return "not-found"; + updateSnapshot( + socket: DaemonSessionSocket, + sessionIdAssertion: string, + snapshotInput: unknown, + ): UpdateSnapshotResult { + const ownedSessionId = this.sessionIdsBySocket.get(socket); + if (!ownedSessionId || ownedSessionId !== sessionIdAssertion) { + return "not-owner"; + } + + const entry = this.sessions.get(ownedSessionId); + if (!entry || entry.socket !== socket) { + return "not-owner"; } const snapshot = this.view.parseSnapshot(snapshotInput); @@ -259,7 +272,7 @@ export class SessionBrokerState< return "invalid"; } - this.sessions.set(sessionId, { + this.sessions.set(ownedSessionId, { ...entry, snapshot, lastSeenAt: new Date().toISOString(), @@ -267,16 +280,22 @@ export class SessionBrokerState< return "updated"; } - markSessionSeen(sessionId: string) { - const entry = this.sessions.get(sessionId); - if (!entry) { - return; + markSessionSeen(socket: DaemonSessionSocket, sessionIdAssertion: string): MarkSessionSeenResult { + const ownedSessionId = this.sessionIdsBySocket.get(socket); + if (!ownedSessionId || ownedSessionId !== sessionIdAssertion) { + return "not-owner"; + } + + const entry = this.sessions.get(ownedSessionId); + if (!entry || entry.socket !== socket) { + return "not-owner"; } - this.sessions.set(sessionId, { + this.sessions.set(ownedSessionId, { ...entry, lastSeenAt: new Date().toISOString(), }); + return "seen"; } unregisterSocket(socket: DaemonSessionSocket) { @@ -346,21 +365,21 @@ export class SessionBrokerState< // Record the pending request before sending so synchronous transport failures and later close // events can both resolve the same command bookkeeping path. - this.pendingCommands.set(requestId, { - sessionId: session.sessionId, - resolve: (result) => resolve(result as ResultType), - reject, - timeout, - }); - const entry = this.sessions.get(session.sessionId); if (!entry) { clearTimeout(timeout); - this.pendingCommands.delete(requestId); reject(new Error("The targeted session is no longer connected.")); return; } + this.pendingCommands.set(requestId, { + sessionId: session.sessionId, + socket: entry.socket, + resolve: (result) => resolve(result as ResultType), + reject, + timeout, + }); + try { const message = { type: "command", @@ -382,15 +401,22 @@ export class SessionBrokerState< }); } - handleCommandResult(message: { - requestId: string; - ok: boolean; - result?: CommandResult; - error?: string; - }) { + handleCommandResult( + socket: DaemonSessionSocket, + message: { + requestId: string; + ok: boolean; + result?: CommandResult; + error?: string; + }, + ): HandleCommandResult { const pending = this.pendingCommands.get(message.requestId); if (!pending) { - return; + return "not-found"; + } + + if (pending.socket !== socket) { + return "not-owner"; } clearTimeout(pending.timeout); @@ -398,10 +424,11 @@ export class SessionBrokerState< if (message.ok) { pending.resolve(message.result as CommandResult); - return; + return "handled"; } pending.reject(new Error(message.error ?? "The session failed to handle the command.")); + return "handled"; } shutdown(error = new Error("The session broker daemon shut down.")) { diff --git a/packages/session-broker/src/broker.test.ts b/packages/session-broker/src/broker.test.ts index 424ae90cb..0e637e18d 100644 --- a/packages/session-broker/src/broker.test.ts +++ b/packages/session-broker/src/broker.test.ts @@ -103,7 +103,9 @@ describe("session broker wrapper", () => { const broker = createBroker(); const connection = { send() {} }; - expect(broker.registerSession(connection, createRegistration(), createSnapshot())).toBe(true); + expect(broker.registerSession(connection, createRegistration(), createSnapshot())).toBe( + "registered", + ); expect(broker.listSessions()).toEqual([ { @@ -132,7 +134,7 @@ describe("session broker wrapper", () => { }, createSnapshot(), ), - ).toBe(false); + ).toBe("invalid"); expect(broker.listSessions()).toEqual([]); }); @@ -157,7 +159,7 @@ describe("session broker wrapper", () => { const outgoing = JSON.parse(sent[0]!) as { requestId: string; command: string }; expect(outgoing.command).toBe("annotate"); - broker.handleCommandResult({ + broker.handleCommandResult(connection, { requestId: outgoing.requestId, ok: true, result: { ok: true }, diff --git a/packages/session-broker/src/broker.ts b/packages/session-broker/src/broker.ts index 37c4201f4..37a4b6300 100644 --- a/packages/session-broker/src/broker.ts +++ b/packages/session-broker/src/broker.ts @@ -1,5 +1,8 @@ import { SessionBrokerState, + type HandleCommandResult, + type MarkSessionSeenResult, + type RegisterSessionResult, type SessionBrokerEntry, type SessionRegistration, type SessionServerMessage, @@ -50,9 +53,13 @@ export interface SessionBrokerController< connection: SessionBrokerPeer, registrationInput: unknown, snapshotInput: unknown, - ): boolean; - updateSnapshot(sessionId: string, snapshotInput: unknown): UpdateSnapshotResult; - markSessionSeen(sessionId: string): void; + ): RegisterSessionResult; + updateSnapshot( + connection: SessionBrokerPeer, + sessionIdAssertion: string, + snapshotInput: unknown, + ): UpdateSnapshotResult; + markSessionSeen(connection: SessionBrokerPeer, sessionIdAssertion: string): MarkSessionSeenResult; unregisterConnection(connection: SessionBrokerPeer): void; pruneStaleSessions(options: { ttlMs: number; now?: number }): number; dispatchCommand(options: { @@ -62,12 +69,15 @@ export interface SessionBrokerController< timeoutMessage: string; timeoutMs?: number; }): Promise; - handleCommandResult(message: { - requestId: string; - ok: boolean; - result?: CommandResult; - error?: string; - }): void; + handleCommandResult( + connection: SessionBrokerPeer, + message: { + requestId: string; + ok: boolean; + result?: CommandResult; + error?: string; + }, + ): HandleCommandResult; shutdown(error?: Error): void; } @@ -150,12 +160,16 @@ export class SessionBroker< return this.state.registerSession(connection, registrationInput, snapshotInput); } - updateSnapshot(sessionId: string, snapshotInput: unknown): UpdateSnapshotResult { - return this.state.updateSnapshot(sessionId, snapshotInput); + updateSnapshot( + connection: SessionBrokerPeer, + sessionIdAssertion: string, + snapshotInput: unknown, + ): UpdateSnapshotResult { + return this.state.updateSnapshot(connection, sessionIdAssertion, snapshotInput); } - markSessionSeen(sessionId: string) { - this.state.markSessionSeen(sessionId); + markSessionSeen(connection: SessionBrokerPeer, sessionIdAssertion: string) { + return this.state.markSessionSeen(connection, sessionIdAssertion); } unregisterConnection(connection: SessionBrokerPeer) { @@ -201,13 +215,16 @@ export class SessionBroker< }); } - handleCommandResult(message: { - requestId: string; - ok: boolean; - result?: CommandResult; - error?: string; - }) { - this.state.handleCommandResult(message); + handleCommandResult( + connection: SessionBrokerPeer, + message: { + requestId: string; + ok: boolean; + result?: CommandResult; + error?: string; + }, + ) { + return this.state.handleCommandResult(connection, message); } shutdown(error = new Error("The session broker shut down.")) { diff --git a/packages/session-broker/src/connection.test.ts b/packages/session-broker/src/connection.test.ts index 54ad11d33..0cf19653d 100644 --- a/packages/session-broker/src/connection.test.ts +++ b/packages/session-broker/src/connection.test.ts @@ -175,6 +175,152 @@ describe("session broker connection", () => { expect(resultMessage).toMatchObject({ type: "command-result", ok: true }); }); + test("does not migrate a late command result onto a replacement socket", async () => { + const sockets: TestSocket[] = []; + let resolveCommand!: (result: { ok: true }) => void; + const commandResult = new Promise<{ ok: true }>((resolve) => { + resolveCommand = resolve; + }); + const connection = createSessionBrokerConnection< + TestSessionInfo, + TestSessionState, + TestSocket, + TestServerMessage, + { ok: true } + >({ + url: "ws://broker.test/session", + createSocket: () => { + const socket = new TestSocket(); + sockets.push(socket); + return socket; + }, + registration: createRegistration(), + snapshot: createSnapshot(), + bridge: { dispatchCommand: () => commandResult }, + reconnectDelayMs: 1, + }); + + connection.start(); + sockets[0]!.emitOpen(); + sockets[0]!.emitMessage( + JSON.stringify({ + type: "command", + requestId: "request-1", + command: "annotate", + input: { summary: "Review note" }, + }), + ); + sockets[0]!.emitClose(); + await Bun.sleep(5); + sockets[1]!.emitOpen(); + + resolveCommand({ ok: true }); + await Bun.sleep(0); + + expect(sockets[0]!.sent.map((message) => JSON.parse(message).type)).toEqual(["register"]); + expect(sockets[1]!.sent.map((message) => JSON.parse(message).type)).toEqual(["register"]); + connection.stop(); + }); + + test("discards queued commands when their source socket disconnects", async () => { + const sockets: TestSocket[] = []; + const dispatched: string[] = []; + const connection = createSessionBrokerConnection< + TestSessionInfo, + TestSessionState, + TestSocket, + TestServerMessage, + { ok: true } + >({ + url: "ws://broker.test/session", + createSocket: () => { + const socket = new TestSocket(); + sockets.push(socket); + return socket; + }, + registration: createRegistration(), + snapshot: createSnapshot(), + reconnectDelayMs: 1, + }); + + connection.start(); + sockets[0]!.emitOpen(); + sockets[0]!.emitMessage( + JSON.stringify({ + type: "command", + requestId: "request-old", + command: "annotate", + input: { summary: "Old review note" }, + }), + ); + sockets[0]!.emitClose(); + await Bun.sleep(5); + sockets[1]!.emitOpen(); + + connection.setBridge({ + dispatchCommand: async (message) => { + dispatched.push(message.requestId); + return { ok: true }; + }, + }); + await Bun.sleep(0); + + expect(dispatched).toEqual([]); + expect(sockets[1]!.sent.map((message) => JSON.parse(message).type)).toEqual(["register"]); + connection.stop(); + }); + + test("stops replaying a queued batch when its source socket disconnects", async () => { + const socket = new TestSocket(); + const dispatched: string[] = []; + let resolveFirst!: (result: { ok: true }) => void; + const firstResult = new Promise<{ ok: true }>((resolve) => { + resolveFirst = resolve; + }); + const connection = createSessionBrokerConnection< + TestSessionInfo, + TestSessionState, + TestSocket, + TestServerMessage, + { ok: true } + >({ + url: "ws://broker.test/session", + createSocket: () => socket, + registration: createRegistration(), + snapshot: createSnapshot(), + reconnectDelayMs: 1_000, + }); + + connection.start(); + socket.emitOpen(); + for (const requestId of ["request-1", "request-2"]) { + socket.emitMessage( + JSON.stringify({ + type: "command", + requestId, + command: "annotate", + input: { summary: "Review note" }, + }), + ); + } + + connection.setBridge({ + dispatchCommand: (message) => { + dispatched.push(message.requestId); + return message.requestId === "request-1" + ? firstResult + : Promise.resolve({ ok: true as const }); + }, + }); + await Bun.sleep(0); + socket.emitClose(); + resolveFirst({ ok: true }); + await Bun.sleep(0); + + expect(dispatched).toEqual(["request-1"]); + connection.stop(); + }); + test("reconnects after socket close unless a close directive disables it", async () => { const sockets: TestSocket[] = []; const warnings: string[] = []; diff --git a/packages/session-broker/src/connection.ts b/packages/session-broker/src/connection.ts index f9cee9966..4eab9db2e 100644 --- a/packages/session-broker/src/connection.ts +++ b/packages/session-broker/src/connection.ts @@ -53,7 +53,7 @@ export class SessionBrokerConnection< > { private socket: Socket | null = null; private bridge: SessionBrokerConnectionBridge | null; - private queuedMessages: ServerMessage[] = []; + private queuedMessages: Array<{ socket: Socket; message: ServerMessage }> = []; private reconnectTimer: ReturnType | null = null; private heartbeatTimer: ReturnType | null = null; private stopped = false; @@ -135,14 +135,13 @@ export class SessionBrokerConnection< socket.onopen = () => { this.startHeartbeat(); - // Always register again on a fresh socket so the broker can replace any stale connection for - // the same session id before later snapshots or commands arrive. - this.send({ + // Register on every fresh socket after the prior close retired its broker-side ownership. + this.sendToSocket(socket, { type: "register", registration: this.registration, snapshot: this.snapshot, }); - void this.flushQueuedMessages(); + void this.flushQueuedMessages(socket); }; socket.onmessage = (event) => { @@ -157,15 +156,16 @@ export class SessionBrokerConnection< return; } - void this.handleServerMessage(parsed); + void this.handleServerMessage(socket, parsed); }; socket.onclose = (event) => { if (this.socket === socket) { this.socket = null; + this.stopHeartbeat(); } - this.stopHeartbeat(); + this.queuedMessages = this.queuedMessages.filter((queued) => queued.socket !== socket); if (this.stopped) { return; } @@ -225,34 +225,43 @@ export class SessionBrokerConnection< } private send(message: SessionClientMessage) { + if (!this.socket) { + return; + } + + this.sendToSocket(this.socket, message); + } + + /** Send a response only through the still-active socket that received its command. */ + private sendToSocket(socket: Socket, message: SessionClientMessage) { if ( - !this.socket || - this.socket.readyState !== (this.options.openState ?? DEFAULT_SOCKET_OPEN_STATE) + this.socket !== socket || + socket.readyState !== (this.options.openState ?? DEFAULT_SOCKET_OPEN_STATE) ) { return; } - this.socket.send(JSON.stringify(message)); + socket.send(JSON.stringify(message)); } - private async handleServerMessage(message: ServerMessage) { + private async handleServerMessage(socket: Socket, message: ServerMessage) { if (!this.bridge) { - // Sessions may connect before the host app has finished wiring its command bridge. Queue - // broker commands so startup races do not drop user-triggered actions. - this.queuedMessages.push(message); + // Sessions may connect before the host app has finished wiring its command bridge. Bind each + // queued command to its source so a reconnect cannot inherit work from a disconnected socket. + this.queuedMessages.push({ socket, message }); return; } try { const result = await this.bridge.dispatchCommand(message); - this.send({ + this.sendToSocket(socket, { type: "command-result", requestId: message.requestId, ok: true, result, }); } catch (error) { - this.send({ + this.sendToSocket(socket, { type: "command-result", requestId: message.requestId, ok: false, @@ -261,18 +270,25 @@ export class SessionBrokerConnection< } } - private async flushQueuedMessages() { - if (!this.bridge || this.queuedMessages.length === 0) { + private async flushQueuedMessages(socket = this.socket) { + if (!this.bridge || !socket || this.queuedMessages.length === 0) { return; } - // Snapshot the queue up front so commands dispatched while we replay are handled in a later - // pass and the original broker ordering stays intact. - const queued = [...this.queuedMessages]; - this.queuedMessages = []; + // Snapshot only this transport's queue so commands cannot cross a disconnect. Commands received + // while replay runs stay in the queue for a later pass and preserve their original ordering. + const queued = this.queuedMessages.filter((entry) => entry.socket === socket); + this.queuedMessages = this.queuedMessages.filter((entry) => entry.socket !== socket); + + for (const entry of queued) { + if ( + this.socket !== entry.socket || + entry.socket.readyState !== (this.options.openState ?? DEFAULT_SOCKET_OPEN_STATE) + ) { + break; + } - for (const message of queued) { - await this.handleServerMessage(message); + await this.handleServerMessage(entry.socket, entry.message); } } } diff --git a/packages/session-broker/src/daemon.test.ts b/packages/session-broker/src/daemon.test.ts index bbee4e8c8..cca6a166f 100644 --- a/packages/session-broker/src/daemon.test.ts +++ b/packages/session-broker/src/daemon.test.ts @@ -271,7 +271,7 @@ describe("session broker daemon", () => { daemon.shutdown(); }); - test("closes incompatible snapshot updates with a specific reason", () => { + test("closes snapshot assertions from unregistered peers", () => { const daemon = createSessionBrokerDaemon({ broker: createBroker(), capabilities: { version: 1 }, @@ -290,11 +290,123 @@ describe("session broker daemon", () => { expect(session.closed).toEqual({ code: 1008, - reason: "Session not registered with broker.", + reason: "Session ownership rejected.", }); daemon.shutdown(); }); + test("rejects duplicate live registration without retiring the owner", () => { + const daemon = createSessionBrokerDaemon({ + broker: createBroker(), + capabilities: { version: 1 }, + }); + const owner = createConnection(); + const duplicate = createConnection(); + const register = JSON.stringify({ + type: "register", + registration: createRegistration(), + snapshot: createSnapshot(), + }); + + daemon.handleConnectionMessage(owner.connection, register); + daemon.handleConnectionMessage(duplicate.connection, register); + + expect(duplicate.closed).toEqual({ code: 1008, reason: "Session registration rejected." }); + daemon.handleConnectionClose(duplicate.connection); + expect(daemon.listSessions()).toHaveLength(1); + expect(daemon.listSessions()[0]).toMatchObject({ sessionId: "session-1" }); + daemon.shutdown(); + }); + + test("rejects cross-peer snapshot, heartbeat, and result authority", async () => { + const daemon = createSessionBrokerDaemon({ + broker: createBroker(), + capabilities: { version: 1 }, + exposeHttpApi: true, + }); + const owner = createConnection(); + const snapshotPeer = createConnection(); + const heartbeatPeer = createConnection(); + const resultPeer = createConnection(); + + for (const [session, sessionId] of [ + [owner, "session-1"], + [snapshotPeer, "session-2"], + [heartbeatPeer, "session-3"], + [resultPeer, "session-4"], + ] as const) { + daemon.handleConnectionMessage( + session.connection, + JSON.stringify({ + type: "register", + registration: createRegistration({ sessionId, cwd: `/${sessionId}` }), + snapshot: createSnapshot(), + }), + ); + } + + daemon.handleConnectionMessage( + snapshotPeer.connection, + JSON.stringify({ + type: "snapshot", + sessionId: "session-1", + snapshot: createSnapshot({ updatedAt: "2026-04-15T00:00:01.000Z", selectedIndex: 1 }), + }), + ); + expect(snapshotPeer.closed).toEqual({ code: 1008, reason: "Session ownership rejected." }); + expect(daemon.getSession({ sessionId: "session-1" })).toMatchObject({ + snapshot: { state: { selectedIndex: 0 } }, + }); + + const ownerSeenAt = daemon.getSession({ sessionId: "session-1" }).lastSeenAt; + daemon.handleConnectionMessage( + heartbeatPeer.connection, + JSON.stringify({ type: "heartbeat", sessionId: "session-1" }), + ); + expect(heartbeatPeer.closed).toEqual({ code: 1008, reason: "Session ownership rejected." }); + expect(daemon.getSession({ sessionId: "session-1" }).lastSeenAt).toBe(ownerSeenAt); + + const pendingResponse = daemon.handleRequest( + new Request("http://broker.test/broker", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + action: "dispatch", + selector: { sessionId: "session-1" }, + command: "annotate", + input: { summary: "Review note" }, + }), + }), + ); + await Bun.sleep(0); + const outgoing = JSON.parse(owner.sent.at(-1)!) as { requestId: string }; + + daemon.handleConnectionMessage( + resultPeer.connection, + JSON.stringify({ + type: "command-result", + requestId: outgoing.requestId, + ok: true, + result: { applied: "forged" }, + }), + ); + expect(resultPeer.closed).toEqual({ code: 1008, reason: "Command ownership rejected." }); + expect(daemon.getHealth().pendingCommands).toBe(1); + + daemon.handleConnectionMessage( + owner.connection, + JSON.stringify({ + type: "command-result", + requestId: outgoing.requestId, + ok: true, + result: { applied: true }, + }), + ); + const response = await pendingResponse; + await expect(response?.json()).resolves.toEqual({ result: { applied: true } }); + daemon.shutdown(); + }); + test("requests shutdown after the idle timeout when no sessions remain", async () => { const daemon = createSessionBrokerDaemon({ broker: createBroker(), diff --git a/packages/session-broker/src/daemon.ts b/packages/session-broker/src/daemon.ts index 3e1e38d16..c03fa5e52 100644 --- a/packages/session-broker/src/daemon.ts +++ b/packages/session-broker/src/daemon.ts @@ -194,13 +194,23 @@ export class SessionBrokerDaemon< switch (parsed.type) { case "register": { - if (!this.broker.registerSession(connection, parsed.registration, parsed.snapshot)) { + const registrationResult = this.broker.registerSession( + connection, + parsed.registration, + parsed.snapshot, + ); + if (registrationResult === "invalid") { // Close immediately when the registration payload is incompatible so the session does not // stay connected under stale assumptions after an upgrade. connection.close?.(INCOMPATIBLE_PAYLOAD_CLOSE_CODE, "Incompatible session registration."); return; } + if (registrationResult === "already-connected") { + connection.close?.(INCOMPATIBLE_PAYLOAD_CLOSE_CODE, "Session registration rejected."); + return; + } + this.noteActivity(); break; } @@ -211,12 +221,13 @@ export class SessionBrokerDaemon< // Snapshot updates are only valid after registration. Closing missing or invalid sessions // keeps the broker state single-sourced instead of guessing how to recover. - const updateResult = this.broker.updateSnapshot(parsed.sessionId, parsed.snapshot); - if (updateResult === "not-found") { - connection.close?.( - INCOMPATIBLE_PAYLOAD_CLOSE_CODE, - "Session not registered with broker.", - ); + const updateResult = this.broker.updateSnapshot( + connection, + parsed.sessionId, + parsed.snapshot, + ); + if (updateResult === "not-owner") { + connection.close?.(INCOMPATIBLE_PAYLOAD_CLOSE_CODE, "Session ownership rejected."); return; } @@ -233,7 +244,12 @@ export class SessionBrokerDaemon< return; } - this.broker.markSessionSeen(parsed.sessionId); + const seenResult = this.broker.markSessionSeen(connection, parsed.sessionId); + if (seenResult === "not-owner") { + connection.close?.(INCOMPATIBLE_PAYLOAD_CLOSE_CODE, "Session ownership rejected."); + return; + } + this.noteActivity(); break; } @@ -242,13 +258,20 @@ export class SessionBrokerDaemon< return; } - this.broker.handleCommandResult({ + const result = this.broker.handleCommandResult(connection, { requestId: parsed.requestId, ok: parsed.ok, result: parsed.result as CommandResult | undefined, error: typeof parsed.error === "string" ? parsed.error : undefined, }); - this.noteActivity(); + if (result === "not-owner") { + connection.close?.(INCOMPATIBLE_PAYLOAD_CLOSE_CODE, "Command ownership rejected."); + return; + } + + if (result === "handled") { + this.noteActivity(); + } break; } } diff --git a/src/session/broker/brokerServer.test.ts b/src/session/broker/brokerServer.test.ts index 8614ec473..7bdd791af 100644 --- a/src/session/broker/brokerServer.test.ts +++ b/src/session/broker/brokerServer.test.ts @@ -425,7 +425,7 @@ describe("Hunk session daemon server", () => { } }); - test("closes snapshots for missing sessions with a specific not-registered reason", async () => { + test("closes snapshot assertions from peers that do not own the session", async () => { // Bun's Windows WebSocket client does not reliably surface this immediate server close. // The daemon-core test covers the close code/reason without the flaky transport layer. if (platform() === "win32") { @@ -455,7 +455,7 @@ describe("Hunk session daemon server", () => { await expect(closed).resolves.toEqual({ code: 1008, - reason: "Session not registered with broker.", + reason: "Session ownership rejected.", }); } finally { socket.close(); diff --git a/src/session/broker/brokerServer.ts b/src/session/broker/brokerServer.ts index 7d1340b2a..dd44cbcb7 100644 --- a/src/session/broker/brokerServer.ts +++ b/src/session/broker/brokerServer.ts @@ -466,15 +466,16 @@ function createHunkBrokerController( getPendingCommandCount: () => state.getPendingCommandCount(), registerSession: (connection, registrationInput, snapshotInput) => state.registerSession(connection, registrationInput, snapshotInput), - updateSnapshot: (sessionId, snapshotInput) => state.updateSnapshot(sessionId, snapshotInput), - markSessionSeen: (sessionId) => state.markSessionSeen(sessionId), + updateSnapshot: (connection, sessionId, snapshotInput) => + state.updateSnapshot(connection, sessionId, snapshotInput), + markSessionSeen: (connection, sessionId) => state.markSessionSeen(connection, sessionId), unregisterConnection: (connection) => state.unregisterSocket(connection), pruneStaleSessions: (options) => state.pruneStaleSessions(options), dispatchCommand: (options) => state.dispatchCommand( options as Parameters[0], ), - handleCommandResult: (message) => state.handleCommandResult(message), + handleCommandResult: (connection, message) => state.handleCommandResult(connection, message), shutdown: (error) => state.shutdown(error), }; } diff --git a/src/session/broker/state.ts b/src/session/broker/state.ts index eb02f4313..a4478acdf 100644 --- a/src/session/broker/state.ts +++ b/src/session/broker/state.ts @@ -202,7 +202,7 @@ export class HunkSessionBrokerState extends SessionBrokerState< ) { const registered = super.registerSession(socket, registrationInput, snapshotInput); this.reconcileMirroredSessions(); - if (!registered) { + if (registered !== "registered") { return registered; } @@ -223,8 +223,8 @@ export class HunkSessionBrokerState extends SessionBrokerState< return registered; } - override updateSnapshot(sessionId: string, snapshotInput: unknown) { - const result = super.updateSnapshot(sessionId, snapshotInput); + override updateSnapshot(socket: HunkBrokerConnection, sessionId: string, snapshotInput: unknown) { + const result = super.updateSnapshot(socket, sessionId, snapshotInput); if (result === "updated") { // A snapshot carries no catalog, so only a further revision of the generation // already mirrored can be adopted from one; a generation change waits for the diff --git a/test/helpers/review-session-harness.ts b/test/helpers/review-session-harness.ts index 61c2c8994..f20a2f6ec 100644 --- a/test/helpers/review-session-harness.ts +++ b/test/helpers/review-session-harness.ts @@ -137,7 +137,7 @@ export function connectReviewSession(files: DiffFile[], options: ReviewSessionHa }, } : result; - state.handleCommandResult({ + state.handleCommandResult(socket, { requestId: message.requestId, ok: true, result: delivered as typeof result, @@ -166,6 +166,7 @@ export function connectReviewSession(files: DiffFile[], options: ReviewSessionHa const publication = producer.getPublication(); const snapshot = createInitialSessionSnapshot(bootstrap, publication); state.updateSnapshot( + socket, sessionId, JSON.parse( JSON.stringify({ From 0cb1a1055e5218aa9a5808f45e741190307f968c Mon Sep 17 00:00:00 2001 From: Ben Vinegar Date: Sat, 29 Aug 2026 11:33:15 -0400 Subject: [PATCH 2/5] feat(session): add signed broker authorization --- .../secure-session-broker-authentication.md | 2 + packages/session-broker-bun/src/serve.test.ts | 44 +- packages/session-broker-core/src/auth.test.ts | 232 +++++ packages/session-broker-core/src/auth.ts | 439 ++++++++ .../session-broker-core/src/brokerState.ts | 7 + .../src/canonicalJson.test.ts | 29 + .../session-broker-core/src/canonicalJson.ts | 72 ++ packages/session-broker-core/src/index.ts | 2 + .../session-broker-core/src/limits.test.ts | 16 +- packages/session-broker-core/src/limits.ts | 80 +- packages/session-broker-core/src/types.ts | 1 + .../session-broker-node/src/serve.test.ts | 44 +- packages/session-broker/README.md | 43 +- .../session-broker/src/authentication.test.ts | 569 +++++++++++ packages/session-broker/src/authentication.ts | 953 ++++++++++++++++++ packages/session-broker/src/broker.ts | 6 + packages/session-broker/src/crypto.test.ts | 83 ++ packages/session-broker/src/crypto.ts | 62 ++ packages/session-broker/src/daemon.test.ts | 260 ++++- packages/session-broker/src/daemon.ts | 329 +++++- packages/session-broker/src/index.ts | 2 + packages/session-broker/src/types.ts | 52 +- 22 files changed, 3208 insertions(+), 119 deletions(-) create mode 100644 .changeset/secure-session-broker-authentication.md create mode 100644 packages/session-broker-core/src/auth.test.ts create mode 100644 packages/session-broker-core/src/auth.ts create mode 100644 packages/session-broker-core/src/canonicalJson.test.ts create mode 100644 packages/session-broker-core/src/canonicalJson.ts create mode 100644 packages/session-broker/src/authentication.test.ts create mode 100644 packages/session-broker/src/authentication.ts create mode 100644 packages/session-broker/src/crypto.test.ts create mode 100644 packages/session-broker/src/crypto.ts diff --git a/.changeset/secure-session-broker-authentication.md b/.changeset/secure-session-broker-authentication.md new file mode 100644 index 000000000..a845151cc --- /dev/null +++ b/.changeset/secure-session-broker-authentication.md @@ -0,0 +1,2 @@ +--- +--- diff --git a/packages/session-broker-bun/src/serve.test.ts b/packages/session-broker-bun/src/serve.test.ts index 19eadf84a..58086272d 100644 --- a/packages/session-broker-bun/src/serve.test.ts +++ b/packages/session-broker-bun/src/serve.test.ts @@ -124,8 +124,10 @@ async function waitForSessionCount(port: number, count: number) { return null; } - const payload = (await response.json()) as { sessions: { sessionId: string }[] }; - return payload.sessions.length === count ? payload : null; + const payload = (await response.json()) as { + body: { sessions: { sessionId: string }[] }; + }; + return payload.body.sessions.length === count ? payload : null; }); } @@ -143,6 +145,34 @@ describe("session broker bun adapter", () => { broker, capabilities: { version: 1 }, exposeHttpApi: true, + appId: "test.app", + appRevision: 1, + callerAuthenticator: { + authenticate: async () => ({ + principal: { + kind: "caller" as const, + appId: "test.app", + principalId: "test-caller", + keyId: "test-key", + grantId: "test-grant", + operations: ["list", "get"] as const, + commands: [], + }, + requestId: "request-1", + assertActive() {}, + signResponse: async ({ httpStatus, appContract }) => ({ + generation: "generation-1", + brokerRevision: 1 as const, + ...(appContract ? { appContract } : {}), + requestId: "request-1", + httpStatus, + bodyDigest: "test-digest", + daemonKeyId: "daemon-key-1", + daemonSignature: "test-signature", + }), + }), + }, + authorizer: async () => true, }); const port = await reserveLoopbackPort(); const server = serveSessionBrokerDaemon({ @@ -152,7 +182,7 @@ describe("session broker bun adapter", () => { }); try { - await expect(readHealth(port)).resolves.toMatchObject({ ok: true, sessions: 0 }); + await expect(readHealth(port)).resolves.toMatchObject({ ok: true }); const socket = new WebSocket(`ws://127.0.0.1:${port}/session`); await new Promise((resolve, reject) => { @@ -195,9 +225,11 @@ describe("session broker bun adapter", () => { }); expect(response.status).toBe(200); await expect(response.json()).resolves.toMatchObject({ - session: { - registration: { sessionId: "session-1" }, - snapshot: { state: { selectedIndex: 0 } }, + body: { + session: { + registration: { sessionId: "session-1" }, + snapshot: { state: { selectedIndex: 0 } }, + }, }, }); diff --git a/packages/session-broker-core/src/auth.test.ts b/packages/session-broker-core/src/auth.test.ts new file mode 100644 index 000000000..7e4fb3f70 --- /dev/null +++ b/packages/session-broker-core/src/auth.test.ts @@ -0,0 +1,232 @@ +import { describe, expect, test } from "bun:test"; +import { + CallerSequenceAllocator, + CallerSequenceReplayWindow, + MAX_CALLER_SEQUENCE, + buildBrokerChallengeTranscript, + buildBrokerHelloAckTranscript, + buildBrokerResponseTranscript, + buildCallerRequestTranscript, + callerPrincipalAllows, + freezeBrokerGrant, + isGrantActive, + isGrantNarrowing, + parseCallerSequence, + principalFromGrant, + producerPrincipalAllows, + type CallerGrant, + type ProducerGrant, +} from "./auth"; + +const callerGrant: CallerGrant = { + kind: "caller", + appId: "dev.example", + principalId: "caller-1", + keyId: "caller-key-1", + grantId: "caller-grant-1", + algorithm: "Ed25519", + issuedAt: 1_000, + expiresAt: 2_000, + revocationId: "revoke-1", + mayDelegate: false, + operations: ["get", "dispatch"], + commands: [{ name: "review", version: 1 }], +}; + +describe("session broker authentication core", () => { + test("builds golden canonical domain-separated transcripts", () => { + expect( + new TextDecoder().decode( + buildBrokerChallengeTranscript({ + role: "caller", + appId: "dev.example", + generation: "generation-1", + endpoint: "http://127.0.0.1:47657/broker", + keyId: "caller-key-1", + grantId: "caller-grant-1", + initiatorNonce: "nonce-a", + responderNonce: "nonce-b", + proposal: { brokerRevision: 1, appRevision: 7, features: ["z", "a"] }, + }), + ), + ).toBe( + '{"appId":"dev.example","domain":"dev.hunk.session-broker.v1/caller-hello","endpoint":"http://127.0.0.1:47657/broker","generation":"generation-1","grantId":"caller-grant-1","initiatorNonce":"nonce-a","keyId":"caller-key-1","proposal":{"appRevision":7,"brokerRevision":1,"features":["a","z"]},"responderNonce":"nonce-b"}', + ); + expect( + new TextDecoder().decode( + buildBrokerHelloAckTranscript({ + role: "producer", + appId: "dev.example", + generation: "generation-1", + keyId: "producer-key-1", + grantId: "producer-grant-1", + helloTranscriptHash: "hello-hash", + selection: { brokerRevision: 1, appRevision: 7, features: [] }, + connectionId: "connection-1", + }), + ), + ).toBe( + '{"appId":"dev.example","connectionId":"connection-1","domain":"dev.hunk.session-broker.v1/producer-hello-ack","generation":"generation-1","grantId":"producer-grant-1","helloTranscriptHash":"hello-hash","keyId":"producer-key-1","selection":{"appRevision":7,"brokerRevision":1,"features":[]}}', + ); + expect( + new TextDecoder().decode( + buildBrokerResponseTranscript({ + appId: "dev.example", + generation: "generation-1", + brokerRevision: 1, + appContract: { appRevision: 7, features: [] }, + requestId: "request-1", + httpStatus: 200, + bodyDigest: "body-hash", + }), + ), + ).toBe( + '{"appContract":{"appRevision":7,"features":[]},"appId":"dev.example","bodyDigest":"body-hash","brokerRevision":1,"domain":"dev.hunk.session-broker.v1/caller-response","generation":"generation-1","httpStatus":200,"requestId":"request-1"}', + ); + expect( + new TextDecoder().decode( + buildCallerRequestTranscript({ + appId: "dev.example", + generation: "generation-1", + callerSessionId: "caller-session-1", + keyId: "caller-key-1", + grantId: "caller-grant-1", + helloTranscriptHash: "hello-hash", + method: "post", + target: "/broker?a=1&b=2", + bodyDigest: "body-hash", + requestId: "request-1", + sequence: "1", + }), + ), + ).toBe( + '{"appId":"dev.example","bodyDigest":"body-hash","callerSessionId":"caller-session-1","domain":"dev.hunk.session-broker.v1/caller-request","generation":"generation-1","grantId":"caller-grant-1","helloTranscriptHash":"hello-hash","keyId":"caller-key-1","method":"POST","requestId":"request-1","sequence":"1","target":"/broker?a=1&b=2"}', + ); + }); + + test("enforces immutable grants, expiry, revocation, and command scope separation", () => { + const grant = freezeBrokerGrant(callerGrant); + const principal = principalFromGrant(grant); + expect(Object.isFrozen(grant)).toBe(true); + expect(Object.isFrozen(grant.commands)).toBe(true); + expect(isGrantActive(grant, { appId: "dev.example", now: 1_500 })).toBe(true); + expect(isGrantActive(grant, { appId: "wrong.app", now: 1_500 })).toBe(false); + expect(isGrantActive(grant, { appId: "dev.example", now: 2_000 })).toBe(false); + expect( + isGrantNarrowing( + { ...grant, mayDelegate: true }, + { + ...grant, + keyId: "delegated-key", + grantId: "delegated-grant", + operations: ["get"], + commands: [], + expiresAt: 1_900, + }, + ), + ).toBe(true); + expect( + isGrantNarrowing( + { ...grant, mayDelegate: true }, + { ...grant, operations: ["list"], expiresAt: 1_900 }, + ), + ).toBe(false); + expect( + isGrantActive(grant, { + appId: "dev.example", + now: 1_500, + isRevoked: (id) => id === "revoke-1", + }), + ).toBe(false); + expect( + callerPrincipalAllows(principal, { + appId: "dev.example", + operation: "dispatch", + command: "review", + commandVersion: 1, + }), + ).toBe(true); + expect( + callerPrincipalAllows(principal, { + appId: "dev.example", + operation: "dispatch", + command: "review", + commandVersion: 2, + }), + ).toBe(false); + expect(callerPrincipalAllows(principal, { appId: "dev.example", operation: "list" })).toBe( + false, + ); + + const producer = principalFromGrant( + freezeBrokerGrant({ + kind: "producer", + appId: "dev.example", + principalId: "producer-1", + keyId: "producer-key-1", + grantId: "producer-grant-1", + algorithm: "Ed25519", + issuedAt: 1_000, + expiresAt: 2_000, + revocationId: "producer-revocation-1", + mayDelegate: false, + sessionId: "session-1", + operations: ["register"], + }), + ); + expect( + producerPrincipalAllows(producer, { + appId: "dev.example", + operation: "register", + sessionId: "session-1", + }), + ).toBe(true); + expect( + producerPrincipalAllows(producer, { + appId: "dev.example", + operation: "reconnect", + sessionId: "session-1", + }), + ).toBe(false); + }); +}); + +describe("caller uint64 replay window", () => { + test("rejects zero, non-canonical values, overflow, duplicate, and old sequences", () => { + expect(parseCallerSequence(MAX_CALLER_SEQUENCE.toString())).toBe(MAX_CALLER_SEQUENCE); + expect(parseCallerSequence("18446744073709551616")).toBeNull(); + expect(parseCallerSequence("01")).toBeNull(); + const replay = new CallerSequenceReplayWindow(); + expect(replay.admit("0")).toBe("zero"); + expect(replay.admit("1")).toBe("accepted"); + expect(replay.admit("1")).toBe("duplicate"); + expect(replay.admit("65")).toBe("accepted"); + expect(replay.admit("1")).toBe("too-old"); + }); + + test("accepts out-of-order values and exact +64 while rejecting +65", () => { + const replay = new CallerSequenceReplayWindow(); + expect(replay.admit("4")).toBe("accepted"); + expect(replay.admit("2")).toBe("accepted"); + expect(replay.admit("3")).toBe("accepted"); + expect(replay.admit("68")).toBe("accepted"); + expect(replay.admit("133")).toBe("too-far-ahead"); + expect(replay.admit("132")).toBe("accepted"); + }); + + test("handles uint64 exhaustion without addition overflow or wrapping", () => { + const allocator = new CallerSequenceAllocator(MAX_CALLER_SEQUENCE); + expect(allocator.allocate()).toBe(MAX_CALLER_SEQUENCE.toString()); + expect(allocator.allocate()).toBeNull(); + + const replay = new CallerSequenceReplayWindow(); + expect(replay.admit((MAX_CALLER_SEQUENCE - 64n).toString())).toBe("too-far-ahead"); + const nearMax = new CallerSequenceReplayWindow({ + highest: MAX_CALLER_SEQUENCE - 64n, + bitmap: 1n, + }); + expect(nearMax.admit(MAX_CALLER_SEQUENCE.toString())).toBe("accepted"); + expect(nearMax.admit(MAX_CALLER_SEQUENCE.toString())).toBe("duplicate"); + expect(nearMax.admit("18446744073709551616")).toBe("invalid"); + }); +}); diff --git a/packages/session-broker-core/src/auth.ts b/packages/session-broker-core/src/auth.ts new file mode 100644 index 000000000..e75f3f038 --- /dev/null +++ b/packages/session-broker-core/src/auth.ts @@ -0,0 +1,439 @@ +import { canonicalJsonBytes, type CanonicalJsonValue } from "./canonicalJson"; + +export const SESSION_BROKER_PROTOCOL_REVISION = 1 as const; +export const SESSION_BROKER_SIGNATURE_ALGORITHM = "Ed25519" as const; +export const SESSION_BROKER_AUTH_DOMAIN = "dev.hunk.session-broker.v1" as const; +export const MAX_CALLER_SEQUENCE = 18_446_744_073_709_551_615n; +export const MAX_BROKER_IDENTIFIER_LENGTH = 128; +export const MAX_BROKER_COMMAND_SCOPES = 256; +const REPLAY_BITMAP_MASK = (1n << 64n) - 1n; +const APP_ID_PATTERN = /^[a-z0-9](?:[a-z0-9._-]{0,126}[a-z0-9])?$/; +const IDENTIFIER_PATTERN = /^[A-Za-z0-9](?:[A-Za-z0-9._~-]{0,126}[A-Za-z0-9])?$/; + +export type ProducerOperation = "register" | "reconnect"; +export type CallerOperation = + | "list" + | "get" + | "dispatch" + | "diagnostics" + | "shutdown" + | "capability:issue"; + +export interface BrokerCommandScope { + readonly name: string; + readonly version: number; +} + +interface BrokerGrantBase { + readonly appId: string; + readonly principalId: string; + readonly keyId: string; + readonly grantId: string; + readonly algorithm: typeof SESSION_BROKER_SIGNATURE_ALGORITHM; + readonly issuedAt: number; + readonly expiresAt: number; + readonly revocationId: string; + readonly mayDelegate: boolean; +} + +export interface ProducerGrant extends BrokerGrantBase { + readonly kind: "producer"; + readonly sessionId?: string; + readonly operations: readonly ProducerOperation[]; +} + +export interface CallerGrant extends BrokerGrantBase { + readonly kind: "caller"; + readonly sessionId?: string; + readonly operations: readonly CallerOperation[]; + readonly commands: readonly BrokerCommandScope[]; +} + +export type BrokerGrant = ProducerGrant | CallerGrant; + +export interface ProducerPrincipal { + readonly kind: "producer"; + readonly appId: string; + readonly principalId: string; + readonly keyId: string; + readonly grantId: string; + readonly sessionId?: string; + readonly scopes: readonly ProducerOperation[]; +} + +export interface CallerPrincipal { + readonly kind: "caller"; + readonly appId: string; + readonly principalId: string; + readonly keyId: string; + readonly grantId: string; + readonly sessionId?: string; + readonly operations: readonly CallerOperation[]; + readonly commands: readonly BrokerCommandScope[]; +} + +export type BrokerPrincipal = ProducerPrincipal | CallerPrincipal; + +export interface BrokerHelloProposal { + readonly brokerRevision: typeof SESSION_BROKER_PROTOCOL_REVISION; + readonly appRevision: number; + readonly features: readonly string[]; +} + +export interface BrokerAppContract { + readonly appRevision: number; + readonly features: readonly string[]; +} + +export interface BrokerChallengeTranscriptInput { + readonly role: "producer" | "caller"; + readonly appId: string; + readonly generation: string; + readonly endpoint: string; + readonly keyId: string; + readonly grantId: string; + readonly initiatorNonce: string; + readonly responderNonce: string; + readonly proposal: BrokerHelloProposal; +} + +interface BrokerHelloAckTranscriptBase { + readonly appId: string; + readonly generation: string; + readonly keyId: string; + readonly grantId: string; + readonly helloTranscriptHash: string; + readonly selection: BrokerHelloProposal; +} + +export type BrokerHelloAckTranscriptInput = BrokerHelloAckTranscriptBase & + ( + | { readonly role: "producer"; readonly connectionId: string } + | { + readonly role: "caller"; + readonly callerSessionId: string; + readonly initialSequence: string; + } + ); + +export interface CallerRequestTranscriptInput { + readonly appId: string; + readonly generation: string; + readonly callerSessionId: string; + readonly keyId: string; + readonly grantId: string; + readonly helloTranscriptHash: string; + readonly method: string; + readonly target: string; + readonly bodyDigest: string; + readonly requestId: string; + readonly sequence: string; +} + +export interface BrokerResponseTranscriptInput { + readonly appId: string; + readonly generation: string; + readonly brokerRevision: typeof SESSION_BROKER_PROTOCOL_REVISION; + readonly requestId: string; + readonly httpStatus: number; + readonly bodyDigest: string; + readonly appContract?: BrokerAppContract; +} + +/** Return whether a value follows the immutable application identity grammar. */ +export function isValidBrokerAppId(value: unknown): value is string { + return typeof value === "string" && APP_ID_PATTERN.test(value); +} + +/** Return whether a value follows the bounded opaque broker identifier grammar. */ +export function isValidBrokerIdentifier(value: unknown): value is string { + return typeof value === "string" && IDENTIFIER_PATTERN.test(value); +} + +/** Return whether a value is a supported positive integer protocol or command revision. */ +export function isValidBrokerRevision(value: unknown): value is number { + return Number.isSafeInteger(value) && (value as number) > 0; +} +export function freezeBrokerGrant(grant: Grant): Grant { + const commands = + grant.kind === "caller" + ? grant.commands.map((scope) => Object.freeze({ ...scope })) + : undefined; + return Object.freeze({ + ...grant, + operations: Object.freeze([...grant.operations]), + ...(commands ? { commands: Object.freeze(commands) } : {}), + }) as unknown as Grant; +} + +/** Project one immutable grant into the redacted principal used by authorization hooks. */ +export function principalFromGrant(grant: ProducerGrant): ProducerPrincipal; +export function principalFromGrant(grant: CallerGrant): CallerPrincipal; +export function principalFromGrant(grant: BrokerGrant): BrokerPrincipal { + if (grant.kind === "producer") { + return Object.freeze({ + kind: "producer", + appId: grant.appId, + principalId: grant.principalId, + keyId: grant.keyId, + grantId: grant.grantId, + ...(grant.sessionId !== undefined ? { sessionId: grant.sessionId } : {}), + scopes: Object.freeze([...grant.operations]), + }); + } + + return Object.freeze({ + kind: "caller", + appId: grant.appId, + principalId: grant.principalId, + keyId: grant.keyId, + grantId: grant.grantId, + ...(grant.sessionId !== undefined ? { sessionId: grant.sessionId } : {}), + operations: Object.freeze([...grant.operations]), + commands: Object.freeze(grant.commands.map((scope) => Object.freeze({ ...scope }))), + }); +} + +/** Return whether a grant is currently valid, app-scoped, and not revoked. */ +export function isGrantActive( + grant: BrokerGrant, + facts: { appId: string; now: number; isRevoked?: (revocationId: string) => boolean }, +): boolean { + return ( + grant.appId === facts.appId && + grant.issuedAt <= facts.now && + facts.now < grant.expiresAt && + !facts.isRevoked?.(grant.revocationId) + ); +} + +/** Return whether a delegated grant is an immutable subset of its parent authority. */ +export function isGrantNarrowing(parent: BrokerGrant, child: BrokerGrant): boolean { + if ( + !parent.mayDelegate || + parent.kind !== child.kind || + parent.appId !== child.appId || + child.issuedAt < parent.issuedAt || + child.expiresAt > parent.expiresAt || + (parent.sessionId !== undefined && child.sessionId !== parent.sessionId) || + child.operations.some( + (operation) => !(parent.operations as readonly string[]).includes(operation), + ) + ) { + return false; + } + if (parent.kind === "producer" || child.kind === "producer") { + return parent.kind === child.kind; + } + return child.commands.every((childScope) => + parent.commands.some( + (parentScope) => + parentScope.name === childScope.name && parentScope.version === childScope.version, + ), + ); +} + +/** Check producer operation and session authority without an allow-by-default path. */ +export function producerPrincipalAllows( + principal: ProducerPrincipal, + facts: { appId: string; operation: ProducerOperation; sessionId?: string }, +): boolean { + return ( + principal.appId === facts.appId && + principal.scopes.includes(facts.operation) && + (principal.sessionId === undefined || principal.sessionId === facts.sessionId) + ); +} + +/** Check caller operation, session, and command authority without an allow-by-default path. */ +export function callerPrincipalAllows( + principal: CallerPrincipal, + facts: { + appId: string; + operation: CallerOperation; + sessionId?: string; + command?: string; + commandVersion?: number; + }, +): boolean { + if (principal.appId !== facts.appId || !principal.operations.includes(facts.operation)) { + return false; + } + if (principal.sessionId !== undefined && principal.sessionId !== facts.sessionId) { + return false; + } + if (facts.operation !== "dispatch") { + return true; + } + if (!facts.command) { + return false; + } + return principal.commands.some( + (scope) => scope.name === facts.command && scope.version === facts.commandVersion, + ); +} + +/** Build the deterministic, domain-separated producer or caller hello transcript. */ +export function buildBrokerChallengeTranscript(input: BrokerChallengeTranscriptInput): Uint8Array { + return canonicalJsonBytes({ + appId: input.appId, + domain: `${SESSION_BROKER_AUTH_DOMAIN}/${input.role}-hello`, + endpoint: input.endpoint, + generation: input.generation, + grantId: input.grantId, + initiatorNonce: input.initiatorNonce, + keyId: input.keyId, + proposal: { + appRevision: input.proposal.appRevision, + brokerRevision: input.proposal.brokerRevision, + features: [...input.proposal.features].sort(), + }, + responderNonce: input.responderNonce, + }); +} + +/** Build the signed hello acknowledgement binding identity, selection, and connection/session. */ +export function buildBrokerHelloAckTranscript(input: BrokerHelloAckTranscriptInput): Uint8Array { + const binding: Record = + input.role === "producer" + ? { connectionId: input.connectionId } + : { + callerSessionId: input.callerSessionId, + initialSequence: input.initialSequence, + }; + return canonicalJsonBytes({ + appId: input.appId, + ...binding, + domain: `${SESSION_BROKER_AUTH_DOMAIN}/${input.role}-hello-ack`, + generation: input.generation, + grantId: input.grantId, + helloTranscriptHash: input.helloTranscriptHash, + keyId: input.keyId, + selection: { + appRevision: input.selection.appRevision, + brokerRevision: input.selection.brokerRevision, + features: [...input.selection.features].sort(), + }, + }); +} + +/** Build the deterministic, domain-separated signed HTTP caller request transcript. */ +export function buildCallerRequestTranscript(input: CallerRequestTranscriptInput): Uint8Array { + return canonicalJsonBytes({ + appId: input.appId, + bodyDigest: input.bodyDigest, + callerSessionId: input.callerSessionId, + domain: `${SESSION_BROKER_AUTH_DOMAIN}/caller-request`, + generation: input.generation, + grantId: input.grantId, + helloTranscriptHash: input.helloTranscriptHash, + keyId: input.keyId, + method: input.method.toUpperCase(), + requestId: input.requestId, + sequence: input.sequence, + target: input.target, + } satisfies CanonicalJsonValue); +} + +/** Build the signed response transcript binding status and the structured body digest. */ +export function buildBrokerResponseTranscript(input: BrokerResponseTranscriptInput): Uint8Array { + return canonicalJsonBytes({ + appId: input.appId, + ...(input.appContract + ? { + appContract: { + appRevision: input.appContract.appRevision, + features: [...input.appContract.features].sort(), + }, + } + : {}), + bodyDigest: input.bodyDigest, + brokerRevision: input.brokerRevision, + domain: `${SESSION_BROKER_AUTH_DOMAIN}/caller-response`, + generation: input.generation, + httpStatus: input.httpStatus, + requestId: input.requestId, + }); +} + +/** Parse the canonical decimal uint64 representation used by caller replay admission. */ +export function parseCallerSequence(value: string): bigint | null { + if (!/^(?:0|[1-9][0-9]{0,19})$/.test(value)) { + return null; + } + const sequence = BigInt(value); + return sequence <= MAX_CALLER_SEQUENCE ? sequence : null; +} + +/** Allocate canonical caller sequences monotonically and stop rather than wrapping at uint64 max. */ +export class CallerSequenceAllocator { + private next: bigint; + + constructor(initialNext = 1n) { + if (initialNext < 1n || initialNext > MAX_CALLER_SEQUENCE) { + throw new RangeError("Invalid initial caller sequence."); + } + this.next = initialNext; + } + + allocate(): string | null { + if (this.next > MAX_CALLER_SEQUENCE) return null; + const allocated = this.next; + this.next += 1n; + return allocated.toString(); + } +} + +export type CallerSequenceAdmission = + | "accepted" + | "zero" + | "duplicate" + | "too-old" + | "too-far-ahead" + | "invalid"; + +/** Atomically admit canonical caller sequences through the contract's 64-value replay bitmap. */ +export class CallerSequenceReplayWindow { + private highest: bigint; + private bitmap: bigint; + + constructor(initial: { highest: bigint; bitmap: bigint } = { highest: 0n, bitmap: 0n }) { + if ( + initial.highest < 0n || + initial.highest > MAX_CALLER_SEQUENCE || + initial.bitmap < 0n || + initial.bitmap > REPLAY_BITMAP_MASK + ) { + throw new RangeError("Invalid caller replay window state."); + } + this.highest = initial.highest; + this.bitmap = initial.bitmap; + } + + admit(canonicalSequence: string): CallerSequenceAdmission { + const sequence = parseCallerSequence(canonicalSequence); + if (sequence === null) return "invalid"; + if (sequence === 0n) return "zero"; + + if (sequence <= this.highest) { + const distance = this.highest - sequence; + if (distance >= 64n) return "too-old"; + const bit = 1n << distance; + if ((this.bitmap & bit) !== 0n) return "duplicate"; + this.bitmap = (this.bitmap | bit) & REPLAY_BITMAP_MASK; + return "accepted"; + } + + const delta = sequence - this.highest; + if (delta > 64n) return "too-far-ahead"; + this.bitmap = delta === 64n ? 0n : (this.bitmap << delta) & REPLAY_BITMAP_MASK; + this.highest = sequence; + this.bitmap |= 1n; + return "accepted"; + } + + /** Return a read-only diagnostic snapshot containing no credential material. */ + snapshot() { + return { highest: this.highest.toString(), bitmap: this.bitmap.toString(16) } as const; + } +} diff --git a/packages/session-broker-core/src/brokerState.ts b/packages/session-broker-core/src/brokerState.ts index 99b94fae1..cb97dadb3 100644 --- a/packages/session-broker-core/src/brokerState.ts +++ b/packages/session-broker-core/src/brokerState.ts @@ -1,4 +1,5 @@ import { randomUUID } from "node:crypto"; +import { isValidBrokerRevision } from "./auth"; import { matchesSessionSelector, repoSelectorDistance, type SelectableSession } from "./selectors"; import type { SessionRegistration, @@ -343,16 +344,21 @@ export class SessionBrokerState< dispatchCommand({ selector, command, + commandVersion = 1, input, timeoutMessage, timeoutMs = 15_000, }: { selector: SessionTargetInput; command: CommandName; + commandVersion?: number; input: Extract["input"]; timeoutMessage: string; timeoutMs?: number; }) { + if (!isValidBrokerRevision(commandVersion)) { + throw new TypeError("Command version must be a positive safe integer."); + } const session = resolveSessionTarget(this.listSessions(), selector); const requestId = randomUUID(); @@ -385,6 +391,7 @@ export class SessionBrokerState< type: "command", requestId, command, + commandVersion, input, } as Extract; diff --git a/packages/session-broker-core/src/canonicalJson.test.ts b/packages/session-broker-core/src/canonicalJson.test.ts new file mode 100644 index 000000000..f5a6a2974 --- /dev/null +++ b/packages/session-broker-core/src/canonicalJson.test.ts @@ -0,0 +1,29 @@ +import { describe, expect, test } from "bun:test"; +import { canonicalizeJson } from "./canonicalJson"; + +describe("RFC 8785 canonical JSON", () => { + test("matches the RFC primitive and UTF-16 property ordering examples", () => { + expect( + canonicalizeJson({ + numbers: [333333333.3333333, 1e30, 4.5, 0.002, 1e-27], + string: '€$\u000f\nA\'B"\\"/', + literals: [null, true, false], + }), + ).toBe( + '{"literals":[null,true,false],"numbers":[333333333.3333333,1e+30,4.5,0.002,1e-27],"string":"€$\\u000f\\nA\'B\\\"\\\\\\\"/"}', + ); + expect(canonicalizeJson({ "\u20ac": "Euro", "\r": "CR", "1": "one" })).toBe( + '{"\\r":"CR","1":"one","€":"Euro"}', + ); + }); + + test("rejects values JSON cannot represent deterministically", () => { + const sparse: unknown[] = []; + sparse.length = 1; + expect(() => canonicalizeJson(Number.NaN)).toThrow("non-finite"); + expect(() => canonicalizeJson(Number.POSITIVE_INFINITY)).toThrow("non-finite"); + expect(() => canonicalizeJson("\ud800")).toThrow("surrogates"); + expect(() => canonicalizeJson(sparse as never)).toThrow("sparse"); + expect(() => canonicalizeJson(new Date() as never)).toThrow("plain JSON objects"); + }); +}); diff --git a/packages/session-broker-core/src/canonicalJson.ts b/packages/session-broker-core/src/canonicalJson.ts new file mode 100644 index 000000000..e7a8e2b4c --- /dev/null +++ b/packages/session-broker-core/src/canonicalJson.ts @@ -0,0 +1,72 @@ +export type CanonicalJsonPrimitive = null | boolean | number | string; +export type CanonicalJsonValue = + | CanonicalJsonPrimitive + | readonly CanonicalJsonValue[] + | { readonly [key: string]: CanonicalJsonValue }; + +function assertUnicodeScalarString(value: string): void { + for (let index = 0; index < value.length; index += 1) { + const unit = value.charCodeAt(index); + if (unit >= 0xd800 && unit <= 0xdbff) { + const trailing = value.charCodeAt(index + 1); + if (!(trailing >= 0xdc00 && trailing <= 0xdfff)) { + throw new TypeError("Canonical JSON rejects lone Unicode surrogates."); + } + index += 1; + } else if (unit >= 0xdc00 && unit <= 0xdfff) { + throw new TypeError("Canonical JSON rejects lone Unicode surrogates."); + } + } +} + +/** Serialize one JSON value using RFC 8785 key ordering and ECMAScript primitive encoding. */ +export function canonicalizeJson(value: CanonicalJsonValue): string { + if (value === null || typeof value === "boolean") { + return JSON.stringify(value); + } + + if (typeof value === "string") { + assertUnicodeScalarString(value); + return JSON.stringify(value); + } + + if (typeof value === "number") { + if (!Number.isFinite(value)) { + throw new TypeError("Canonical JSON does not support non-finite numbers."); + } + return JSON.stringify(value); + } + + if (Array.isArray(value)) { + const entries: string[] = []; + for (let index = 0; index < value.length; index += 1) { + if (!Object.hasOwn(value, index)) { + throw new TypeError("Canonical JSON rejects sparse arrays."); + } + entries.push(canonicalizeJson(value[index]!)); + } + return `[${entries.join(",")}]`; + } + + if (typeof value === "object") { + const prototype = Object.getPrototypeOf(value); + if (prototype !== Object.prototype && prototype !== null) { + throw new TypeError("Canonical JSON requires plain JSON objects."); + } + const record = value as Record; + const entries = Object.keys(record) + .sort() + .map((key) => { + assertUnicodeScalarString(key); + return `${JSON.stringify(key)}:${canonicalizeJson(record[key]!)}`; + }); + return `{${entries.join(",")}}`; + } + + throw new TypeError("Canonical JSON supports JSON values only."); +} + +/** Encode canonical JSON as the exact UTF-8 bytes covered by broker signatures. */ +export function canonicalJsonBytes(value: CanonicalJsonValue): Uint8Array { + return new TextEncoder().encode(canonicalizeJson(value)); +} diff --git a/packages/session-broker-core/src/index.ts b/packages/session-broker-core/src/index.ts index c75d74644..c7f918b59 100644 --- a/packages/session-broker-core/src/index.ts +++ b/packages/session-broker-core/src/index.ts @@ -1,4 +1,6 @@ export * from "./types"; +export * from "./canonicalJson"; +export * from "./auth"; export * from "./brokerWire"; export * from "./limits"; export * from "./brokerState"; diff --git a/packages/session-broker-core/src/limits.test.ts b/packages/session-broker-core/src/limits.test.ts index ce9d38a17..fef455aa3 100644 --- a/packages/session-broker-core/src/limits.test.ts +++ b/packages/session-broker-core/src/limits.test.ts @@ -1,5 +1,10 @@ import { describe, expect, test } from "bun:test"; -import { PayloadTooLargeError, readRequestTextWithLimit, utf8ByteLength } from "./limits"; +import { + PayloadTooLargeError, + readRequestBytesWithLimit, + readRequestTextWithLimit, + utf8ByteLength, +} from "./limits"; /** Build a streaming request body so the read path runs without a Content-Length header. */ function streamingRequest(byteLength: number, chunkSize = 64 * 1024) { @@ -59,6 +64,15 @@ describe("readRequestTextWithLimit", () => { ); }); + test("returns exact bytes and rejects malformed UTF-8 only during strict text decoding", async () => { + const bytes = new Uint8Array([0x7b, 0xc0, 0xaf, 0x7d]); + const byteRequest = new Request("http://broker.test/api", { method: "POST", body: bytes }); + await expect(readRequestBytesWithLimit(byteRequest, 1024)).resolves.toEqual(bytes); + + const textRequest = new Request("http://broker.test/api", { method: "POST", body: bytes }); + await expect(readRequestTextWithLimit(textRequest, 1024)).rejects.toBeInstanceOf(TypeError); + }); + test("treats a missing body as an empty string", async () => { const request = new Request("http://broker.test/api", { method: "GET" }); diff --git a/packages/session-broker-core/src/limits.ts b/packages/session-broker-core/src/limits.ts index 47781063d..275f9f04d 100644 --- a/packages/session-broker-core/src/limits.ts +++ b/packages/session-broker-core/src/limits.ts @@ -38,6 +38,7 @@ export class PayloadTooLargeError extends Error { // Reused across every websocket message, HTTP body, and patch check to avoid a per-call alloc. const sharedTextEncoder = new TextEncoder(); +const fatalTextDecoder = new TextDecoder("utf-8", { fatal: true }); /** UTF-8 byte length of a string without allocating a Buffer in non-Node runtimes. */ export function utf8ByteLength(value: string): number { @@ -45,80 +46,57 @@ export function utf8ByteLength(value: string): number { } /** - * Read one request body as text while enforcing a hard byte ceiling. + * Read one request body as exact bytes while enforcing a hard byte ceiling. * * The Content-Length header is rejected early when it already declares an oversized body, and the * stream is aborted mid-read so a missing or lying Content-Length cannot force the daemon to * buffer an unbounded body before the cap is noticed. */ -export async function readRequestTextWithLimit( +export async function readRequestBytesWithLimit( request: Request, maxBytes: number, -): Promise { +): Promise { const declared = request.headers.get("content-length"); - if (declared) { - const length = Number.parseInt(declared, 10); - if (Number.isInteger(length) && length > maxBytes) { - throw new PayloadTooLargeError(maxBytes); - } + if (declared && /^(?:0|[1-9][0-9]*)$/.test(declared) && Number(declared) > maxBytes) { + throw new PayloadTooLargeError(maxBytes); } const body = request.body; - if (!body) { - // Some runtimes do not expose a streaming body; the Content-Length guard above still bounds - // well-behaved clients, and the post-read check bounds the rest. - const text = await request.text(); - if (utf8ByteLength(text) > maxBytes) { - throw new PayloadTooLargeError(maxBytes); - } - - return text; - } + if (!body) return new Uint8Array(); const reader = body.getReader(); const chunks: Uint8Array[] = []; let total = 0; - for (;;) { - let done: boolean; - let value: Uint8Array | undefined; - try { - const result = await reader.read(); - done = result.done; - value = result.value; - } catch (error) { - reader.releaseLock(); - throw error; - } - - if (done) { - break; - } - - if (!value) { - continue; + try { + for (;;) { + const { done, value } = await reader.read(); + if (done) break; + if (!value) continue; + + total += value.byteLength; + if (total > maxBytes) { + await reader.cancel().catch(() => {}); + throw new PayloadTooLargeError(maxBytes); + } + chunks.push(value); } - - total += value.byteLength; - if (total > maxBytes) { - // Stop pulling from the stream immediately so the body cannot grow past the cap. - await reader.cancel().catch(() => {}); - // cancel() does not release the lock per the Streams spec; release it explicitly so the - // over-limit path matches the normal-exit path instead of waiting for GC. - reader.releaseLock(); - throw new PayloadTooLargeError(maxBytes); - } - - chunks.push(value); + } finally { + reader.releaseLock(); } - reader.releaseLock(); - const merged = new Uint8Array(total); let offset = 0; for (const chunk of chunks) { merged.set(chunk, offset); offset += chunk.byteLength; } + return merged; +} - return new TextDecoder().decode(merged); +/** Read and strictly decode one bounded request body as UTF-8 text. */ +export async function readRequestTextWithLimit( + request: Request, + maxBytes: number, +): Promise { + return fatalTextDecoder.decode(await readRequestBytesWithLimit(request, maxBytes)); } diff --git a/packages/session-broker-core/src/types.ts b/packages/session-broker-core/src/types.ts index af5ce74ed..092e36d59 100644 --- a/packages/session-broker-core/src/types.ts +++ b/packages/session-broker-core/src/types.ts @@ -71,5 +71,6 @@ export type SessionServerMessage { broker, capabilities: { version: 1 }, exposeHttpApi: true, + appId: "test.app", + appRevision: 1, + callerAuthenticator: { + authenticate: async () => ({ + principal: { + kind: "caller" as const, + appId: "test.app", + principalId: "test-caller", + keyId: "test-key", + grantId: "test-grant", + operations: ["list", "get"] as const, + commands: [], + }, + requestId: "request-1", + assertActive() {}, + signResponse: async ({ httpStatus, appContract }) => ({ + generation: "generation-1", + brokerRevision: 1 as const, + ...(appContract ? { appContract } : {}), + requestId: "request-1", + httpStatus, + bodyDigest: "test-digest", + daemonKeyId: "daemon-key-1", + daemonSignature: "test-signature", + }), + }), + }, + authorizer: async () => true, }); const port = await reserveLoopbackPort(); const server = await serveSessionBrokerDaemon({ @@ -120,7 +148,7 @@ describe("session broker node adapter", () => { try { const health = await fetch(`http://127.0.0.1:${port}/health`); - await expect(health.json()).resolves.toMatchObject({ ok: true, sessions: 0 }); + await expect(health.json()).resolves.toMatchObject({ ok: true }); const socket = new WebSocket(`ws://127.0.0.1:${port}/session`); await new Promise((resolve, reject) => { @@ -165,8 +193,10 @@ describe("session broker node adapter", () => { return null; } - const payload = (await response.json()) as { sessions: { sessionId: string }[] }; - return payload.sessions.length === 1 ? payload : null; + const payload = (await response.json()) as { + body: { sessions: { sessionId: string }[] }; + }; + return payload.body.sessions.length === 1 ? payload : null; }); const response = await fetch(`http://127.0.0.1:${port}/broker`, { @@ -175,9 +205,11 @@ describe("session broker node adapter", () => { body: JSON.stringify({ action: "get", selector: { sessionId: "session-1" } }), }); await expect(response.json()).resolves.toMatchObject({ - session: { - registration: { sessionId: "session-1" }, - snapshot: { state: { selectedIndex: 0 } }, + body: { + session: { + registration: { sessionId: "session-1" }, + snapshot: { state: { selectedIndex: 0 } }, + }, }, }); diff --git a/packages/session-broker/README.md b/packages/session-broker/README.md index 4ef89f6c6..0c82d7680 100644 --- a/packages/session-broker/README.md +++ b/packages/session-broker/README.md @@ -6,10 +6,12 @@ The implementation and release contract for turning these internal workspaces in per-application SDK lives in [`docs/session-broker-sdk.md`](https://github.com/modem-dev/hunk/blob/main/docs/session-broker-sdk.md). Current package APIs predate that contract and do not yet satisfy every security, compatibility, -supervision, or packaging gate it defines. **Both the generic WebSocket producer path and optional -raw HTTP control path are currently unauthenticated and are internal-only. Do not deploy them as a -security boundary.** Hunk currently adds loopback Host/Origin checks around its custom session -routes; only its separate browser-review routes add a per-session capability. +supervision, or packaging gate it defines. The package now provides signed producer/caller hello, +short-lived caller sessions, replay admission, and default-deny raw HTTP authorization primitives. +**Hunk credential discovery and automatic producer/caller activation intentionally remain deferred +to the later Hunk runtime-adapter change (PR 5 of this stack).** Until that composition lands, the +legacy Hunk WebSocket and custom session routes remain internal-only; no bearer fallback is +available. Hunk's separate browser-review capabilities remain independent. This is the **main broker package** in the workspace. It owns the reusable broker behavior without committing to Bun or Node server APIs. @@ -117,21 +119,13 @@ At this point the daemon can: - process websocket register/snapshot/heartbeat/result messages - prune stale sessions and request idle shutdown -The raw HTTP broker API is opt-in. **The current internal API has no generic authentication or -Host/Origin enforcement. Do not deploy this option, even on loopback, until the security contract -in `docs/session-broker-sdk.md` is implemented.** It exists only for controlled package tests -today: - -```ts -const daemon = createSessionBrokerDaemon({ - broker, - capabilities: { - version: 1, - name: "example-broker", - }, - exposeHttpApi: true, -}); -``` +The raw HTTP broker API is opt-in and fails closed: `exposeHttpApi: true` exposes no control route +unless an explicit immutable `appId`, a singleton `appRevision`, a `callerAuthenticator`, and an app +`authorizer` are supplied. The included +`SessionBrokerAuthenticator` implements Ed25519 challenge/proof and signed requests; applications +inject app-scoped grants, public verifiers, daemon signing identity, revocation policy, and their +own default-deny authorization hook. It performs no filesystem, environment, coordinator, or Hunk +credential discovery. ### 3. Serve it through a runtime adapter @@ -193,7 +187,8 @@ The helper owns: ## Raw broker API The daemon always serves `GET /health`. Its raw capability/control API is intentionally small and -disabled by default. When `exposeHttpApi: true` is set, it additionally serves: +disabled by default. When `exposeHttpApi: true` is set together with an explicit `appId`, singleton +`appRevision`, caller authenticator, and authorizer, it additionally serves: - `GET /broker/capabilities` - `POST /broker` @@ -203,10 +198,14 @@ Request body shapes: ```ts { action: "list" } { action: "get", selector: { sessionId: "..." } } -{ action: "dispatch", selector: { sessionId: "..." }, command: "...", input: {...} } +{ action: "dispatch", selector: { sessionId: "..." }, command: "...", commandVersion: 1, input: {...} } ``` -Responses return raw session records or command results. +An omitted `commandVersion` is validated and deliberately defaults to revision `1` for current +internal callers. Authentication covers the exact bounded HTTP body bytes before strict UTF-8 and +JSON decoding. Authenticated responses use `{ body, authentication }`; the signed authentication +record binds daemon generation, broker revision, target application contract when applicable, +request ID, HTTP status, and the canonical structured-body digest. ## Hunk-specific layering diff --git a/packages/session-broker/src/authentication.test.ts b/packages/session-broker/src/authentication.test.ts new file mode 100644 index 000000000..0c42ab250 --- /dev/null +++ b/packages/session-broker/src/authentication.test.ts @@ -0,0 +1,569 @@ +import { describe, expect, test } from "bun:test"; +import { + buildBrokerResponseTranscript, + buildCallerRequestTranscript, + type CallerGrant, + type ProducerGrant, +} from "@hunk/session-broker-core"; +import { + SessionBrokerAuthenticationError, + SessionBrokerAuthenticator, + canonicalHttpTarget, + challengeTranscriptForClient, + type SessionBrokerHelloChallengeRequest, +} from "./authentication"; +import { decodeBase64Url, encodeBase64Url, webSessionBrokerCrypto } from "./crypto"; +import type { SessionBrokerCrypto } from "./crypto"; + +async function keyPair() { + return crypto.subtle.generateKey("Ed25519", false, ["sign", "verify"]); +} + +function callerGrant(overrides: Partial = {}): CallerGrant { + return { + kind: "caller", + appId: "dev.example", + principalId: "caller-1", + keyId: "caller-key-1", + grantId: "caller-grant-1", + algorithm: "Ed25519", + issuedAt: 1_000, + expiresAt: 10_000, + revocationId: "caller-revocation-1", + mayDelegate: false, + operations: ["list", "get", "dispatch"], + commands: [{ name: "review", version: 1 }], + ...overrides, + }; +} + +function producerGrant(): ProducerGrant { + return { + kind: "producer", + appId: "dev.example", + principalId: "producer-1", + keyId: "producer-key-1", + grantId: "producer-grant-1", + algorithm: "Ed25519", + issuedAt: 1_000, + expiresAt: 10_000, + revocationId: "producer-revocation-1", + mayDelegate: false, + operations: ["register"], + }; +} + +async function setup( + options: { + revoked?: () => boolean; + maxChallenges?: number; + maxChallengeTranscriptBytes?: number; + maxCallerSessions?: number; + callerSessionTtlMs?: number; + crypto?: SessionBrokerCrypto; + } = {}, +) { + let now = 2_000; + const daemon = await keyPair(); + const caller = await keyPair(); + const producer = await keyPair(); + const authenticator = new SessionBrokerAuthenticator({ + appId: "dev.example", + appRevision: 1, + generation: "generation-1", + daemonIdentity: { keyId: "daemon-key-1", privateKey: daemon.privateKey }, + credentials: [ + { grant: callerGrant(), publicKey: caller.publicKey }, + { grant: producerGrant(), publicKey: producer.publicKey }, + ], + now: () => now, + isRevoked: options.revoked, + maxChallenges: options.maxChallenges, + maxChallengeTranscriptBytes: options.maxChallengeTranscriptBytes, + maxCallerSessions: options.maxCallerSessions, + callerSessionTtlMs: options.callerSessionTtlMs, + crypto: options.crypto, + }); + return { + authenticator, + daemon, + caller, + producer, + setNow(value: number) { + now = value; + }, + }; +} + +function challengeRequest( + role: "caller" | "producer" = "caller", +): SessionBrokerHelloChallengeRequest { + return { + role, + appId: "dev.example", + endpoint: "http://127.0.0.1:47657/broker", + keyId: `${role}-key-1`, + grantId: `${role}-grant-1`, + initiatorNonce: "initiator-nonce", + proposal: { brokerRevision: 1, appRevision: 1, features: [] }, + }; +} + +async function openCallerSession(setupResult: Awaited>) { + const request = challengeRequest(); + const challenge = await setupResult.authenticator.issueChallenge(request, request.endpoint); + const transcript = challengeTranscriptForClient(request, challenge, "generation-1"); + expect( + await webSessionBrokerCrypto.verify( + setupResult.daemon.publicKey, + Uint8Array.from( + atob( + challenge.daemonSignature + .replaceAll("-", "+") + .replaceAll("_", "/") + .padEnd(Math.ceil(challenge.daemonSignature.length / 4) * 4, "="), + ), + (character) => character.charCodeAt(0), + ), + transcript, + ), + ).toBe(true); + const signature = encodeBase64Url( + await webSessionBrokerCrypto.sign(setupResult.caller.privateKey, transcript), + ); + const session = await setupResult.authenticator.completeCallerHello({ + challengeId: challenge.challengeId, + signature, + }); + return { session, transcript }; +} + +async function signedRequest( + setupResult: Awaited>, + caller: Awaited>, + sequence: string, + bodyText = '{"action":"list"}', +) { + const body = new TextEncoder().encode(bodyText); + const url = new URL("http://broker.test/broker?z=2&a=hello%20world"); + const bodyDigest = encodeBase64Url(await webSessionBrokerCrypto.sha256(body)); + const helloTranscriptHash = encodeBase64Url( + await webSessionBrokerCrypto.sha256(caller.transcript), + ); + const transcript = buildCallerRequestTranscript({ + appId: "dev.example", + generation: "generation-1", + callerSessionId: caller.session.callerSessionId, + keyId: "caller-key-1", + grantId: "caller-grant-1", + helloTranscriptHash, + method: "POST", + target: canonicalHttpTarget(url), + bodyDigest, + requestId: `request-${sequence}`, + sequence, + }); + const signature = encodeBase64Url( + await webSessionBrokerCrypto.sign(setupResult.caller.privateKey, transcript), + ); + const request = new Request(url, { + method: "POST", + headers: { + "x-session-broker-caller-session": caller.session.callerSessionId, + "x-session-broker-request-id": `request-${sequence}`, + "x-session-broker-sequence": sequence, + "x-session-broker-signature": signature, + }, + }); + return { request, body }; +} + +describe("session broker signed authentication", () => { + test("canonicalizes encoded paths and sorted duplicate query values", () => { + expect( + canonicalHttpTarget(new URL("http://broker.test/review/%7euser?z=2&a=hello+world&a=%2F")), + ).toBe("/review/~user?a=%2F&a=hello%20world&z=2"); + }); + + test("verifies daemon identity before accepting caller and producer proofs", async () => { + const values = await setup(); + const caller = await openCallerSession(values); + expect(caller.session).toMatchObject({ initialSequence: "1", brokerRevision: 1 }); + + const request = challengeRequest("producer"); + const challenge = await values.authenticator.issueChallenge(request, request.endpoint); + const transcript = challengeTranscriptForClient(request, challenge, "generation-1"); + const signature = encodeBase64Url( + await webSessionBrokerCrypto.sign(values.producer.privateKey, transcript), + ); + await expect( + values.authenticator.completeProducerHello( + { challengeId: challenge.challengeId, signature }, + "connection-1", + ), + ).resolves.toMatchObject({ + connectionId: "connection-1", + daemonKeyId: "daemon-key-1", + principal: { kind: "producer", scopes: ["register"] }, + }); + }); + + test("rejects missing, wrong, expired, revoked, and reused credentials with redacted errors", async () => { + const values = await setup(); + await expect( + values.authenticator.issueChallenge( + { ...challengeRequest(), keyId: "wrong-key" }, + challengeRequest().endpoint, + ), + ).rejects.toMatchObject({ code: "invalid-credential" }); + await expect( + values.authenticator.authenticate({ + request: new Request("http://broker.test/broker"), + body: new Uint8Array(), + }), + ).rejects.toMatchObject({ code: "authentication-required" }); + + const request = challengeRequest(); + const challenge = await values.authenticator.issueChallenge(request, request.endpoint); + const transcript = challengeTranscriptForClient(request, challenge, "generation-1"); + const signature = encodeBase64Url( + await webSessionBrokerCrypto.sign(values.caller.privateKey, transcript), + ); + await values.authenticator.completeCallerHello({ + challengeId: challenge.challengeId, + signature, + }); + await expect( + values.authenticator.completeCallerHello({ challengeId: challenge.challengeId, signature }), + ).rejects.toMatchObject({ code: "challenge-used" }); + + values.setNow(11_000); + await expect( + values.authenticator.issueChallenge(challengeRequest(), challengeRequest().endpoint), + ).rejects.toMatchObject({ + code: "credential-expired", + }); + + const revoked = await setup({ revoked: () => true }); + await expect( + revoked.authenticator.issueChallenge(challengeRequest(), challengeRequest().endpoint), + ).rejects.toMatchObject({ + code: "credential-revoked", + }); + + const error = new SessionBrokerAuthenticationError("invalid-signature"); + expect(JSON.stringify({ message: error.message, code: error.code })).not.toContain( + "signature-", + ); + expect(error.message).toBe("Session broker authentication failed."); + }); + + test("expires challenges and short-lived caller sessions", async () => { + const challengeValues = await setup(); + const request = challengeRequest(); + const challenge = await challengeValues.authenticator.issueChallenge(request, request.endpoint); + challengeValues.setNow(challenge.expiresAt); + await expect( + challengeValues.authenticator.completeCallerHello({ + challengeId: challenge.challengeId, + signature: "not-used-after-expiry", + }), + ).rejects.toMatchObject({ code: "challenge-expired" }); + + const sessionValues = await setup({ callerSessionTtlMs: 500 }); + const caller = await openCallerSession(sessionValues); + sessionValues.setNow(caller.session.expiresAt); + const signed = await signedRequest(sessionValues, caller, "1"); + await expect(sessionValues.authenticator.authenticate(signed)).rejects.toMatchObject({ + code: "caller-session-expired", + }); + }); + + test("binds method, canonical target, body digest, request ID, and replay sequence", async () => { + const values = await setup(); + const caller = await openCallerSession(values); + const first = await signedRequest(values, caller, "1"); + await expect(values.authenticator.authenticate(first)).resolves.toMatchObject({ + principal: { principalId: "caller-1" }, + }); + await expect(values.authenticator.authenticate(first)).rejects.toMatchObject({ + code: "replay-rejected", + }); + + const tampered = await signedRequest(values, caller, "2"); + await expect( + values.authenticator.authenticate({ + request: tampered.request, + body: new TextEncoder().encode('{"action":"get"}'), + }), + ).rejects.toMatchObject({ code: "invalid-signature" }); + await expect(values.authenticator.authenticate(tampered)).resolves.toMatchObject({ + principal: { principalId: "caller-1" }, + }); + }); + + test("rechecks revocation after asynchronous signature verification", async () => { + let revoked = false; + let revokeAfterVerify = false; + const cryptoWithRevocation: SessionBrokerCrypto = { + ...webSessionBrokerCrypto, + async verify(publicKey, signature, value) { + const verified = await webSessionBrokerCrypto.verify(publicKey, signature, value); + if (revokeAfterVerify) revoked = true; + return verified; + }, + }; + const values = await setup({ revoked: () => revoked, crypto: cryptoWithRevocation }); + const caller = await openCallerSession(values); + const signed = await signedRequest(values, caller, "1"); + + revokeAfterVerify = true; + await expect(values.authenticator.authenticate(signed)).rejects.toMatchObject({ + code: "credential-revoked", + }); + }); + + test("rejects non-HTTP request targets even when path and query are signable", async () => { + const values = await setup(); + const caller = await openCallerSession(values); + const signed = await signedRequest(values, caller, "1"); + const request = new Request("ftp://broker.test/broker?z=2&a=hello%20world", { + method: "POST", + headers: signed.request.headers, + }); + + await expect( + values.authenticator.authenticate({ request, body: signed.body }), + ).rejects.toMatchObject({ code: "invalid-signature" }); + }); + + test("signs response envelopes and rejects transcript tampering or the wrong generation", async () => { + const values = await setup(); + const caller = await openCallerSession(values); + const signed = await signedRequest(values, caller, "1"); + const authenticated = await values.authenticator.authenticate(signed); + const response = await authenticated.signResponse({ + httpStatus: 200, + body: { sessions: [] }, + appContract: { appRevision: 1, features: [] }, + }); + const signature = decodeBase64Url(response.daemonSignature)!; + const transcriptInput = { + appId: "dev.example", + generation: response.generation, + brokerRevision: 1 as const, + appContract: { appRevision: 1, features: [] }, + requestId: response.requestId, + httpStatus: response.httpStatus, + bodyDigest: response.bodyDigest, + }; + expect( + await webSessionBrokerCrypto.verify( + values.daemon.publicKey, + signature, + buildBrokerResponseTranscript(transcriptInput), + ), + ).toBe(true); + expect( + await webSessionBrokerCrypto.verify( + values.daemon.publicKey, + signature, + buildBrokerResponseTranscript({ ...transcriptInput, generation: "generation-2" }), + ), + ).toBe(false); + expect( + await webSessionBrokerCrypto.verify( + values.daemon.publicKey, + signature, + buildBrokerResponseTranscript({ ...transcriptInput, httpStatus: 201 }), + ), + ).toBe(false); + expect( + await webSessionBrokerCrypto.verify( + values.daemon.publicKey, + signature, + buildBrokerResponseTranscript({ ...transcriptInput, bodyDigest: "tampered" }), + ), + ).toBe(false); + }); + + test("uses bounded collision-safe IDs with a deterministic custom random source", async () => { + const randomValues = [1, 2, 1, 3]; + const deterministicCrypto: SessionBrokerCrypto = { + ...webSessionBrokerCrypto, + randomBytes(length) { + return new Uint8Array(length).fill(randomValues.shift() ?? 9); + }, + }; + const values = await setup({ crypto: deterministicCrypto, maxChallenges: 2 }); + const first = await values.authenticator.issueChallenge( + challengeRequest(), + challengeRequest().endpoint, + ); + const second = await values.authenticator.issueChallenge( + challengeRequest(), + challengeRequest().endpoint, + ); + expect(first.challengeId).not.toBe(second.challengeId); + + const callerRandomValues = [1, 2, 3, 1, 4, 3, 5]; + const callerCrypto: SessionBrokerCrypto = { + ...webSessionBrokerCrypto, + randomBytes(length) { + return new Uint8Array(length).fill(callerRandomValues.shift() ?? 9); + }, + }; + const callerValues = await setup({ crypto: callerCrypto, maxCallerSessions: 2 }); + const firstCaller = await openCallerSession(callerValues); + const secondCaller = await openCallerSession(callerValues); + expect(firstCaller.session.callerSessionId).not.toBe(secondCaller.session.callerSessionId); + }); + + test("validates and snapshots mutable startup authority before accepting traffic", async () => { + const daemon = await keyPair(); + const caller = await keyPair(); + let now = 2_000; + const grant = callerGrant(); + const options = { + appId: "dev.example", + appRevision: 1, + generation: "generation-1", + daemonIdentity: { keyId: "daemon-key-1", privateKey: daemon.privateKey }, + credentials: [{ grant, publicKey: caller.publicKey }], + now: () => now, + isRevoked: () => false, + maxChallenges: 1, + }; + const authenticator = new SessionBrokerAuthenticator(options); + (grant.operations as CallerGrant["operations"] & string[]).splice(0, grant.operations.length); + options.generation = "generation-2"; + options.maxChallenges = 0; + options.isRevoked = () => true; + + const request = challengeRequest(); + await expect(authenticator.issueChallenge(request, request.endpoint)).resolves.toMatchObject({ + daemonKeyId: "daemon-key-1", + }); + now = 18_000; + await expect(authenticator.issueChallenge(request, request.endpoint)).rejects.toMatchObject({ + code: "credential-expired", + }); + + expect( + () => + new SessionBrokerAuthenticator({ + ...options, + generation: "generation-1", + credentials: [ + { grant: callerGrant({ algorithm: "RSA" as "Ed25519" }), publicKey: caller.publicKey }, + ], + }), + ).toThrow("algorithm"); + expect( + () => + new SessionBrokerAuthenticator({ + ...options, + generation: "generation-1", + credentials: [ + { grant: callerGrant({ issuedAt: Number.NaN }), publicKey: caller.publicKey }, + ], + }), + ).toThrow("timestamps"); + expect( + () => + new SessionBrokerAuthenticator({ + ...options, + appId: "Wrong.App", + generation: "generation-1", + credentials: [], + }), + ).toThrow("appId"); + expect( + () => + new SessionBrokerAuthenticator({ + ...options, + generation: "generation-1", + credentials: [ + { + grant: callerGrant({ appId: "other.app", operations: ["unknown" as "list"] }), + publicKey: caller.publicKey, + }, + ], + }), + ).toThrow("configured appId"); + expect( + () => + new SessionBrokerAuthenticator({ + ...options, + generation: "generation-1", + credentials: [ + { + grant: callerGrant({ operations: ["unknown" as "list"] }), + publicKey: caller.publicKey, + }, + ], + }), + ).toThrow("recognized"); + expect( + () => + new SessionBrokerAuthenticator({ + ...options, + generation: "generation-1", + credentials: [ + { + grant: callerGrant({ + commands: [ + { name: "review", version: 1 }, + { name: "review", version: 1 }, + ], + }), + publicKey: caller.publicKey, + }, + ], + }), + ).toThrow("unique"); + }); + + test("rejects malformed fixed hello proposals, identifiers, and endpoints", async () => { + const values = await setup(); + const valid = challengeRequest(); + for (const malformed of [ + { ...valid, initiatorNonce: "bad nonce" }, + { ...valid, endpoint: "http://user@127.0.0.1/broker" }, + { ...valid, proposal: { ...valid.proposal, appRevision: 2 } }, + { ...valid, proposal: { ...valid.proposal, features: ["unexpected.feature"] } }, + ]) { + await expect( + values.authenticator.issueChallenge(malformed, malformed.endpoint), + ).rejects.toMatchObject({ + code: "invalid-credential", + }); + } + }); + + test("rejects malformed percent encodings and UTF-8 in canonical HTTP targets", () => { + for (const target of ["/%", "/%GG", "/%C0%AF", "/broker?q=%ED%A0%80"]) { + expect(() => canonicalHttpTarget(new URL(`http://broker.test${target}`))).toThrow(); + } + }); + + test("bounds pending challenge counts and retained transcript bytes", async () => { + const values = await setup({ maxChallenges: 1 }); + await values.authenticator.issueChallenge(challengeRequest(), challengeRequest().endpoint); + await expect( + values.authenticator.issueChallenge(challengeRequest(), challengeRequest().endpoint), + ).rejects.toMatchObject({ + code: "authentication-capacity", + }); + + const byteBound = await setup({ maxChallengeTranscriptBytes: 128 }); + await expect( + byteBound.authenticator.issueChallenge(challengeRequest(), challengeRequest().endpoint), + ).rejects.toMatchObject({ code: "authentication-capacity" }); + + const noCallerCapacity = await setup({ maxCallerSessions: 0 }); + await expect(openCallerSession(noCallerCapacity)).rejects.toMatchObject({ + code: "authentication-capacity", + }); + }); +}); diff --git a/packages/session-broker/src/authentication.ts b/packages/session-broker/src/authentication.ts new file mode 100644 index 000000000..c9399fa26 --- /dev/null +++ b/packages/session-broker/src/authentication.ts @@ -0,0 +1,953 @@ +import { + CallerSequenceReplayWindow, + MAX_BROKER_COMMAND_SCOPES, + SESSION_BROKER_PROTOCOL_REVISION, + SESSION_BROKER_SIGNATURE_ALGORITHM, + buildBrokerChallengeTranscript, + buildBrokerHelloAckTranscript, + buildBrokerResponseTranscript, + buildCallerRequestTranscript, + canonicalJsonBytes, + freezeBrokerGrant, + isGrantActive, + isValidBrokerAppId, + isValidBrokerIdentifier, + isValidBrokerRevision, + principalFromGrant, + type BrokerAppContract, + type BrokerChallengeTranscriptInput, + type BrokerGrant, + type BrokerHelloProposal, + type CallerGrant, + type CallerOperation, + type CallerPrincipal, + type CanonicalJsonValue, + type ProducerGrant, + type ProducerOperation, + type ProducerPrincipal, +} from "@hunk/session-broker-core"; +import { + decodeBase64Url, + encodeBase64Url, + webSessionBrokerCrypto, + type SessionBrokerCrypto, +} from "./crypto"; + +const DEFAULT_CHALLENGE_TTL_MS = 15_000; +const DEFAULT_CALLER_SESSION_TTL_MS = 5 * 60_000; +const DEFAULT_MAX_CHALLENGES = 128; +const DEFAULT_MAX_CHALLENGE_BYTES = 4 * 1024 * 1024; +const DEFAULT_MAX_CHALLENGE_TRANSCRIPT_BYTES = 64 * 1024; +const DEFAULT_MAX_CALLER_SESSIONS = 256; +const UNIQUE_ID_RETRIES = 16; +const RANDOM_ID_BYTES = 24; +const MAX_ENDPOINT_LENGTH = 2_048; +const PRODUCER_OPERATIONS = new Set(["register", "reconnect"]); +const CALLER_OPERATIONS = new Set([ + "list", + "get", + "dispatch", + "diagnostics", + "shutdown", + "capability:issue", +]); + +export type SessionBrokerAuthenticationFailureCode = + | "authentication-required" + | "invalid-credential" + | "credential-expired" + | "credential-revoked" + | "challenge-expired" + | "challenge-used" + | "caller-session-expired" + | "invalid-signature" + | "replay-rejected" + | "authentication-capacity"; + +/** Report one stable, redacted authentication failure without credential or signature material. */ +export class SessionBrokerAuthenticationError extends Error { + constructor(readonly code: SessionBrokerAuthenticationFailureCode) { + super("Session broker authentication failed."); + this.name = "SessionBrokerAuthenticationError"; + } +} + +export interface SessionBrokerCredential { + readonly grant: Grant; + readonly publicKey: CryptoKey; +} + +export interface SessionBrokerDaemonIdentity { + readonly keyId: string; + readonly privateKey: CryptoKey; +} + +export interface SessionBrokerHelloChallengeRequest { + readonly role: "producer" | "caller"; + readonly appId: string; + readonly endpoint: string; + readonly keyId: string; + readonly grantId: string; + readonly initiatorNonce: string; + readonly proposal: BrokerHelloProposal; +} + +export interface SessionBrokerHelloChallenge { + readonly challengeId: string; + readonly responderNonce: string; + readonly expiresAt: number; + readonly daemonKeyId: string; + readonly daemonSignature: string; +} + +export interface SessionBrokerHelloProof { + readonly challengeId: string; + readonly signature: string; +} + +export interface AuthenticatedCallerSession { + readonly callerSessionId: string; + readonly principal: CallerPrincipal; + readonly expiresAt: number; + readonly initialSequence: "1"; + readonly brokerRevision: typeof SESSION_BROKER_PROTOCOL_REVISION; + readonly appRevision: number; + readonly features: readonly []; + readonly helloTranscriptHash: string; + readonly daemonKeyId: string; + readonly daemonSignature: string; +} + +export interface AuthenticatedProducerHello { + readonly principal: ProducerPrincipal; + readonly connectionId: string; + readonly brokerRevision: typeof SESSION_BROKER_PROTOCOL_REVISION; + readonly appRevision: number; + readonly features: readonly []; + readonly helloTranscriptHash: string; + readonly daemonKeyId: string; + readonly daemonSignature: string; +} + +export interface CallerRequestAuthenticationInput { + readonly request: Request; + readonly body: Uint8Array; +} + +export interface SessionBrokerResponseAuthentication { + readonly generation: string; + readonly brokerRevision: typeof SESSION_BROKER_PROTOCOL_REVISION; + readonly appContract?: BrokerAppContract; + readonly requestId: string; + readonly httpStatus: number; + readonly bodyDigest: string; + readonly daemonKeyId: string; + readonly daemonSignature: string; +} + +export interface CallerResponseSigningInput { + readonly httpStatus: number; + readonly body: CanonicalJsonValue; + readonly appContract?: BrokerAppContract; +} + +export interface AuthenticatedCallerRequest { + readonly principal: CallerPrincipal; + readonly requestId: string; + assertActive(): void; + signResponse(input: CallerResponseSigningInput): Promise; +} + +export interface CallerRequestAuthenticator { + authenticate(input: CallerRequestAuthenticationInput): Promise; +} + +interface PendingChallenge { + readonly request: SessionBrokerHelloChallengeRequest; + readonly transcript: Uint8Array; + readonly grant: BrokerGrant; + readonly publicKey: CryptoKey; + readonly expiresAt: number; + readonly retainedBytes: number; +} + +interface CallerSessionRecord { + readonly principal: CallerPrincipal; + readonly grant: CallerGrant; + readonly publicKey: CryptoKey; + readonly helloTranscriptHash: string; + readonly expiresAt: number; + readonly replay: CallerSequenceReplayWindow; +} + +export interface SessionBrokerAuthenticatorOptions { + readonly appId: string; + readonly appRevision: number; + readonly generation: string; + readonly daemonIdentity: SessionBrokerDaemonIdentity; + readonly credentials: readonly SessionBrokerCredential[]; + readonly crypto?: SessionBrokerCrypto; + readonly now?: () => number; + readonly isRevoked?: (revocationId: string) => boolean; + readonly challengeTtlMs?: number; + readonly callerSessionTtlMs?: number; + readonly maxChallenges?: number; + readonly maxChallengeBytes?: number; + readonly maxChallengeTranscriptBytes?: number; + readonly maxCallerSessions?: number; +} + +interface AuthenticatorSnapshot { + readonly appId: string; + readonly appRevision: number; + readonly generation: string; + readonly daemonIdentity: SessionBrokerDaemonIdentity; + readonly now: () => number; + readonly isRevoked?: (revocationId: string) => boolean; + readonly challengeTtlMs: number; + readonly callerSessionTtlMs: number; + readonly maxChallenges: number; + readonly maxChallengeBytes: number; + readonly maxChallengeTranscriptBytes: number; + readonly maxCallerSessions: number; +} + +function authenticationError(code: SessionBrokerAuthenticationFailureCode): never { + throw new SessionBrokerAuthenticationError(code); +} + +function invalidStartup(message: string): never { + throw new TypeError(`Invalid session broker authenticator configuration: ${message}`); +} + +/** Require a runtime key with the expected Ed25519 role and usage. */ +function assertCryptoKey( + value: unknown, + expected: { type: "public" | "private"; usage: "verify" | "sign" }, +): asserts value is CryptoKey { + if ( + !value || + typeof value !== "object" || + (value as CryptoKey).type !== expected.type || + (value as CryptoKey).algorithm?.name !== SESSION_BROKER_SIGNATURE_ALGORITHM || + !(value as CryptoKey).usages?.includes(expected.usage) || + (expected.type === "private" && (value as CryptoKey).extractable) + ) { + invalidStartup(`expected an Ed25519 ${expected.type} key with ${expected.usage} usage.`); + } +} + +/** Select and validate one immutable integer authenticator limit. */ +function configuredNumber( + value: number | undefined, + fallback: number, + name: string, + allowZero: boolean, +): number { + const selected = value ?? fallback; + if (!Number.isSafeInteger(selected) || selected < (allowZero ? 0 : 1)) { + invalidStartup(`${name} must be ${allowZero ? "a non-negative" : "a positive"} safe integer.`); + } + return selected; +} + +/** Bind the injected crypto methods so later mutations cannot alter authenticator behavior. */ +function copyCrypto(value: SessionBrokerCrypto | undefined): SessionBrokerCrypto { + const source = value ?? webSessionBrokerCrypto; + if ( + typeof source.randomBytes !== "function" || + typeof source.sha256 !== "function" || + typeof source.sign !== "function" || + typeof source.verify !== "function" + ) { + invalidStartup("crypto must implement randomBytes, sha256, sign, and verify."); + } + return Object.freeze({ + randomBytes: source.randomBytes.bind(source), + sha256: source.sha256.bind(source), + sign: source.sign.bind(source), + verify: source.verify.bind(source), + }); +} + +/** Validate and deeply snapshot one startup grant. */ +function copyGrant(input: BrokerGrant, appId: string): BrokerGrant { + if (!input || typeof input !== "object") invalidStartup("credential grant must be an object."); + const grant = input as BrokerGrant; + if (grant.kind !== "producer" && grant.kind !== "caller") { + invalidStartup("credential kind is not recognized."); + } + if (!isValidBrokerAppId(grant.appId) || grant.appId !== appId) { + invalidStartup("every credential must exactly match the configured appId."); + } + for (const [name, value] of [ + ["principalId", grant.principalId], + ["keyId", grant.keyId], + ["grantId", grant.grantId], + ["revocationId", grant.revocationId], + ] as const) { + if (!isValidBrokerIdentifier(value)) invalidStartup(`${name} has an invalid identifier.`); + } + if (grant.sessionId !== undefined && !isValidBrokerIdentifier(grant.sessionId)) { + invalidStartup("sessionId has an invalid identifier."); + } + if (grant.algorithm !== SESSION_BROKER_SIGNATURE_ALGORITHM) { + invalidStartup("credential algorithm must be Ed25519."); + } + if ( + !Number.isFinite(grant.issuedAt) || + !Number.isFinite(grant.expiresAt) || + grant.issuedAt >= grant.expiresAt + ) { + invalidStartup("credential timestamps must be finite and strictly ordered."); + } + if (typeof grant.mayDelegate !== "boolean" || !Array.isArray(grant.operations)) { + invalidStartup("credential delegation and operations are malformed."); + } + + const recognized = grant.kind === "producer" ? PRODUCER_OPERATIONS : CALLER_OPERATIONS; + const operations = [...grant.operations]; + if ( + new Set(operations).size !== operations.length || + operations.some((operation) => !recognized.has(operation as never)) + ) { + invalidStartup("credential operations must be recognized and unique."); + } + + if (grant.kind === "producer") { + return freezeBrokerGrant({ ...grant, operations } as ProducerGrant); + } + if (!Array.isArray(grant.commands) || grant.commands.length > MAX_BROKER_COMMAND_SCOPES) { + invalidStartup("caller command scopes exceed the configured bound."); + } + const commandKeys = new Set(); + const commands = grant.commands.map((scope) => { + if (!scope || typeof scope !== "object" || !isValidBrokerIdentifier(scope.name)) { + invalidStartup("caller command scope name is invalid."); + } + if (!isValidBrokerRevision(scope.version)) { + invalidStartup("caller command scope version is invalid."); + } + const key = `${scope.name}\u0000${scope.version}`; + if (commandKeys.has(key)) invalidStartup("caller command scopes must be unique."); + commandKeys.add(key); + return { name: scope.name, version: scope.version }; + }); + return freezeBrokerGrant({ ...grant, operations, commands } as CallerGrant); +} + +/** Validate and index immutable startup credentials by role and verifier identity. */ +function copyCredentials( + credentials: readonly SessionBrokerCredential[], + appId: string, +): Map { + if (!Array.isArray(credentials)) invalidStartup("credentials must be an array."); + const copied = new Map(); + for (const credential of credentials) { + if (!credential || typeof credential !== "object") { + invalidStartup("credential must be an object."); + } + assertCryptoKey(credential.publicKey, { type: "public", usage: "verify" }); + const grant = copyGrant(credential.grant, appId); + const credentialId = `${grant.kind}:${grant.keyId}:${grant.grantId}`; + if (copied.has(credentialId)) invalidStartup("credential identities must be unique."); + copied.set(credentialId, Object.freeze({ grant, publicKey: credential.publicKey })); + } + return copied; +} + +/** Validate and snapshot every non-credential authenticator startup option. */ +function copyOptions(options: SessionBrokerAuthenticatorOptions): AuthenticatorSnapshot { + if (!options || typeof options !== "object") invalidStartup("options must be an object."); + if (!isValidBrokerAppId(options.appId)) invalidStartup("appId has an invalid grammar."); + if (!isValidBrokerRevision(options.appRevision)) { + invalidStartup("appRevision must be a positive safe integer."); + } + if (!isValidBrokerIdentifier(options.generation)) { + invalidStartup("generation has an invalid identifier."); + } + if (!options.daemonIdentity || !isValidBrokerIdentifier(options.daemonIdentity.keyId)) { + invalidStartup("daemon keyId has an invalid identifier."); + } + assertCryptoKey(options.daemonIdentity.privateKey, { type: "private", usage: "sign" }); + if (options.now !== undefined && typeof options.now !== "function") { + invalidStartup("now must be a function."); + } + if (options.isRevoked !== undefined && typeof options.isRevoked !== "function") { + invalidStartup("isRevoked must be a function."); + } + return Object.freeze({ + appId: options.appId, + appRevision: options.appRevision, + generation: options.generation, + daemonIdentity: Object.freeze({ + keyId: options.daemonIdentity.keyId, + privateKey: options.daemonIdentity.privateKey, + }), + now: options.now ?? Date.now, + ...(options.isRevoked ? { isRevoked: options.isRevoked } : {}), + challengeTtlMs: configuredNumber( + options.challengeTtlMs, + DEFAULT_CHALLENGE_TTL_MS, + "challengeTtlMs", + false, + ), + callerSessionTtlMs: configuredNumber( + options.callerSessionTtlMs, + DEFAULT_CALLER_SESSION_TTL_MS, + "callerSessionTtlMs", + false, + ), + maxChallenges: configuredNumber( + options.maxChallenges, + DEFAULT_MAX_CHALLENGES, + "maxChallenges", + true, + ), + maxChallengeBytes: configuredNumber( + options.maxChallengeBytes, + DEFAULT_MAX_CHALLENGE_BYTES, + "maxChallengeBytes", + true, + ), + maxChallengeTranscriptBytes: configuredNumber( + options.maxChallengeTranscriptBytes, + DEFAULT_MAX_CHALLENGE_TRANSCRIPT_BYTES, + "maxChallengeTranscriptBytes", + true, + ), + maxCallerSessions: configuredNumber( + options.maxCallerSessions, + DEFAULT_MAX_CALLER_SESSIONS, + "maxCallerSessions", + true, + ), + }); +} + +/** Parse one bounded credential-free HTTP or websocket hello endpoint. */ +function parseEndpoint(value: string): URL | null { + if (value.length === 0 || value.length > MAX_ENDPOINT_LENGTH) return null; + try { + const url = new URL(value); + if ( + !["http:", "https:", "ws:", "wss:"].includes(url.protocol) || + !url.hostname || + url.username || + url.password || + url.hash + ) { + return null; + } + canonicalHttpTarget(url); + return url; + } catch { + return null; + } +} + +/** Build the only broker and application selection supported during Phase 1. */ +function fixedProposal(appRevision: number): BrokerHelloProposal { + return Object.freeze({ + brokerRevision: SESSION_BROKER_PROTOCOL_REVISION, + appRevision, + features: Object.freeze([]), + }); +} + +/** Produce the canonical path and RFC 3986 encoded sorted query covered by request signatures. */ +export function canonicalHttpTarget(url: URL): string { + const encode = (value: string) => + encodeURIComponent(value).replace( + /[!'()*]/g, + (character) => `%${character.charCodeAt(0).toString(16).toUpperCase()}`, + ); + const decode = (value: string, plusAsSpace: boolean) => + decodeURIComponent(plusAsSpace ? value.replaceAll("+", "%20") : value); + + const rawQuery = url.search.startsWith("?") ? url.search.slice(1) : url.search; + const pairs = rawQuery + ? rawQuery + .split("&") + .map((pair) => { + const separator = pair.indexOf("="); + const rawKey = separator < 0 ? pair : pair.slice(0, separator); + const rawValue = separator < 0 ? "" : pair.slice(separator + 1); + return [encode(decode(rawKey, true)), encode(decode(rawValue, true))] as const; + }) + .sort(([leftKey, leftValue], [rightKey, rightValue]) => + leftKey === rightKey + ? leftValue < rightValue + ? -1 + : leftValue > rightValue + ? 1 + : 0 + : leftKey < rightKey + ? -1 + : 1, + ) + : []; + const path = url.pathname + .split("/") + .map((segment) => encode(decode(segment, false))) + .join("/"); + const query = pairs.map(([key, value]) => `${key}=${value}`).join("&"); + return query ? `${path}?${query}` : path; +} + +/** Authenticate bounded producer hellos and generation-bound signed caller request sessions. */ +export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { + private readonly crypto: SessionBrokerCrypto; + private readonly config: AuthenticatorSnapshot; + private readonly credentials: Map; + private readonly challenges = new Map(); + private readonly callerSessions = new Map(); + private readonly reservedCallerSessionIds = new Set(); + private challengeBytes = 0; + private pendingCallerSessionAdmissions = 0; + + constructor(options: SessionBrokerAuthenticatorOptions) { + this.config = copyOptions(options); + this.crypto = copyCrypto(options.crypto); + this.credentials = copyCredentials(options.credentials, this.config.appId); + } + + /** Issue one bounded, expiring challenge signed by the daemon identity. */ + async issueChallenge( + request: SessionBrokerHelloChallengeRequest, + listenerEndpoint: string, + ): Promise { + this.pruneExpired(); + if (this.challenges.size >= this.config.maxChallenges) { + authenticationError("authentication-capacity"); + } + const normalized = this.validateHello(request, listenerEndpoint); + const credential = this.credentials.get( + `${normalized.role}:${normalized.keyId}:${normalized.grantId}`, + ); + if (!credential || credential.grant.kind !== normalized.role) { + authenticationError("invalid-credential"); + } + this.requireActiveGrant(credential.grant); + + const challengeId = this.uniqueId((id) => this.challenges.has(id)); + const responderNonce = this.randomId(); + const expiresAt = this.currentTime() + this.config.challengeTtlMs; + if (!Number.isFinite(expiresAt)) authenticationError("invalid-credential"); + const transcript = buildBrokerChallengeTranscript({ + ...normalized, + generation: this.config.generation, + responderNonce, + }); + if ( + transcript.byteLength > this.config.maxChallengeTranscriptBytes || + this.challengeBytes + transcript.byteLength > this.config.maxChallengeBytes + ) { + authenticationError("authentication-capacity"); + } + this.challengeBytes += transcript.byteLength; + this.challenges.set(challengeId, { + request: normalized, + transcript, + grant: credential.grant, + publicKey: credential.publicKey, + expiresAt, + retainedBytes: transcript.byteLength, + }); + try { + const daemonSignature = encodeBase64Url( + await this.crypto.sign(this.config.daemonIdentity.privateKey, transcript), + ); + return Object.freeze({ + challengeId, + responderNonce, + expiresAt, + daemonKeyId: this.config.daemonIdentity.keyId, + daemonSignature, + }); + } catch { + this.deleteChallenge(challengeId); + authenticationError("invalid-credential"); + } + } + + /** Consume one caller proof and issue a short-lived replay-protected caller session. */ + async completeCallerHello(proof: SessionBrokerHelloProof): Promise { + const pending = this.takeChallenge(proof.challengeId, "caller"); + await this.verifyProof(pending, proof.signature); + const grant = pending.grant as CallerGrant; + this.requireActiveGrant(grant); + this.pruneExpired(); + if ( + this.callerSessions.size + this.pendingCallerSessionAdmissions >= + this.config.maxCallerSessions + ) { + authenticationError("authentication-capacity"); + } + this.pendingCallerSessionAdmissions += 1; + + try { + return await this.createCallerSession(pending, grant); + } finally { + this.pendingCallerSessionAdmissions -= 1; + } + } + + private async createCallerSession( + pending: PendingChallenge, + grant: CallerGrant, + ): Promise { + const transcriptHash = encodeBase64Url(await this.crypto.sha256(pending.transcript)); + const callerSessionId = this.uniqueId( + (id) => this.callerSessions.has(id) || this.reservedCallerSessionIds.has(id), + ); + this.reservedCallerSessionIds.add(callerSessionId); + try { + const expiresAt = Math.min( + grant.expiresAt, + this.currentTime() + this.config.callerSessionTtlMs, + ); + const principal = principalFromGrant(grant); + const daemonSignature = encodeBase64Url( + await this.crypto.sign( + this.config.daemonIdentity.privateKey, + buildBrokerHelloAckTranscript({ + role: "caller", + appId: this.config.appId, + generation: this.config.generation, + keyId: grant.keyId, + grantId: grant.grantId, + helloTranscriptHash: transcriptHash, + selection: pending.request.proposal, + callerSessionId, + initialSequence: "1", + }), + ), + ); + this.callerSessions.set(callerSessionId, { + principal, + grant, + publicKey: pending.publicKey, + helloTranscriptHash: transcriptHash, + expiresAt, + replay: new CallerSequenceReplayWindow(), + }); + return Object.freeze({ + callerSessionId, + principal, + expiresAt, + initialSequence: "1", + brokerRevision: SESSION_BROKER_PROTOCOL_REVISION, + appRevision: this.config.appRevision, + features: Object.freeze([]) as readonly [], + helloTranscriptHash: transcriptHash, + daemonKeyId: this.config.daemonIdentity.keyId, + daemonSignature, + }); + } finally { + this.reservedCallerSessionIds.delete(callerSessionId); + } + } + + /** Consume one producer proof and sign the connection binding supplied by the adapter. */ + async completeProducerHello( + proof: SessionBrokerHelloProof, + connectionId: string, + ): Promise { + if (!isValidBrokerIdentifier(connectionId)) authenticationError("invalid-credential"); + const pending = this.takeChallenge(proof.challengeId, "producer"); + await this.verifyProof(pending, proof.signature); + const grant = pending.grant as ProducerGrant; + this.requireActiveGrant(grant); + const helloTranscriptHash = encodeBase64Url(await this.crypto.sha256(pending.transcript)); + const daemonSignature = encodeBase64Url( + await this.crypto.sign( + this.config.daemonIdentity.privateKey, + buildBrokerHelloAckTranscript({ + role: "producer", + appId: this.config.appId, + generation: this.config.generation, + keyId: grant.keyId, + grantId: grant.grantId, + helloTranscriptHash, + selection: pending.request.proposal, + connectionId, + }), + ), + ); + return Object.freeze({ + principal: principalFromGrant(grant), + connectionId, + brokerRevision: SESSION_BROKER_PROTOCOL_REVISION, + appRevision: this.config.appRevision, + features: Object.freeze([]) as readonly [], + helloTranscriptHash, + daemonKeyId: this.config.daemonIdentity.keyId, + daemonSignature, + }); + } + + /** Verify one signed HTTP request and atomically admit its sequence before returning authority. */ + async authenticate({ + request, + body, + }: CallerRequestAuthenticationInput): Promise { + this.pruneExpired(); + const callerSessionId = request.headers.get("x-session-broker-caller-session"); + const requestId = request.headers.get("x-session-broker-request-id"); + const sequence = request.headers.get("x-session-broker-sequence"); + const encodedSignature = request.headers.get("x-session-broker-signature"); + if (!callerSessionId || !requestId || !sequence || !encodedSignature) { + authenticationError("authentication-required"); + } + if (!isValidBrokerIdentifier(callerSessionId) || !isValidBrokerIdentifier(requestId)) { + authenticationError("invalid-signature"); + } + const session = this.callerSessions.get(callerSessionId); + if (!session) authenticationError("caller-session-expired"); + this.requireActiveGrant(session.grant); + if (this.currentTime() >= session.expiresAt) { + this.callerSessions.delete(callerSessionId); + authenticationError("caller-session-expired"); + } + + const signature = decodeBase64Url(encodedSignature); + if (!signature || signature.byteLength === 0) authenticationError("invalid-signature"); + const bodyDigest = encodeBase64Url(await this.crypto.sha256(body)); + let target: string; + try { + const url = new URL(request.url); + if ( + (url.protocol !== "http:" && url.protocol !== "https:") || + url.username || + url.password || + url.hash + ) { + authenticationError("invalid-signature"); + } + target = canonicalHttpTarget(url); + } catch { + authenticationError("invalid-signature"); + } + const transcript = buildCallerRequestTranscript({ + appId: this.config.appId, + generation: this.config.generation, + callerSessionId, + keyId: session.grant.keyId, + grantId: session.grant.grantId, + helloTranscriptHash: session.helloTranscriptHash, + method: request.method, + target, + bodyDigest, + requestId, + sequence, + }); + if (!(await this.crypto.verify(session.publicKey, signature, transcript))) { + authenticationError("invalid-signature"); + } + this.assertCallerSessionActive(callerSessionId, session); + if (session.replay.admit(sequence) !== "accepted") { + authenticationError("replay-rejected"); + } + + const principal = session.principal; + return Object.freeze({ + principal, + requestId, + assertActive: () => this.assertCallerSessionActive(callerSessionId, session), + signResponse: (input: CallerResponseSigningInput) => this.signResponse(requestId, input), + }); + } + + /** Revoke one in-memory caller session without exposing whether it previously existed. */ + revokeCallerSession(callerSessionId: string): void { + this.callerSessions.delete(callerSessionId); + } + + /** Recheck identity, revocation, and expiry after every asynchronous policy boundary. */ + private assertCallerSessionActive(callerSessionId: string, session: CallerSessionRecord): void { + if (this.callerSessions.get(callerSessionId) !== session) { + authenticationError("caller-session-expired"); + } + this.requireActiveGrant(session.grant); + if (this.currentTime() >= session.expiresAt) { + this.callerSessions.delete(callerSessionId); + authenticationError("caller-session-expired"); + } + } + + private async signResponse( + requestId: string, + input: CallerResponseSigningInput, + ): Promise { + if (!Number.isInteger(input.httpStatus) || input.httpStatus < 100 || input.httpStatus > 599) { + authenticationError("invalid-credential"); + } + if ( + input.appContract && + (input.appContract.appRevision !== this.config.appRevision || + !Array.isArray(input.appContract.features) || + input.appContract.features.length !== 0) + ) { + authenticationError("invalid-credential"); + } + const appContract = input.appContract + ? Object.freeze({ appRevision: this.config.appRevision, features: Object.freeze([]) }) + : undefined; + const bodyDigest = encodeBase64Url(await this.crypto.sha256(canonicalJsonBytes(input.body))); + const transcript = buildBrokerResponseTranscript({ + appId: this.config.appId, + generation: this.config.generation, + brokerRevision: SESSION_BROKER_PROTOCOL_REVISION, + requestId, + httpStatus: input.httpStatus, + bodyDigest, + ...(appContract ? { appContract } : {}), + }); + const daemonSignature = encodeBase64Url( + await this.crypto.sign(this.config.daemonIdentity.privateKey, transcript), + ); + return Object.freeze({ + generation: this.config.generation, + brokerRevision: SESSION_BROKER_PROTOCOL_REVISION, + ...(appContract ? { appContract } : {}), + requestId, + httpStatus: input.httpStatus, + bodyDigest, + daemonKeyId: this.config.daemonIdentity.keyId, + daemonSignature, + }); + } + + private validateHello( + request: SessionBrokerHelloChallengeRequest, + listenerEndpoint: string, + ): SessionBrokerHelloChallengeRequest { + if ( + !request || + typeof request !== "object" || + (request.role !== "producer" && request.role !== "caller") || + request.appId !== this.config.appId || + !isValidBrokerIdentifier(request.keyId) || + !isValidBrokerIdentifier(request.grantId) || + !isValidBrokerIdentifier(request.initiatorNonce) || + request.endpoint !== listenerEndpoint || + !parseEndpoint(request.endpoint) || + !parseEndpoint(listenerEndpoint) || + !request.proposal || + request.proposal.brokerRevision !== SESSION_BROKER_PROTOCOL_REVISION || + request.proposal.appRevision !== this.config.appRevision || + !Array.isArray(request.proposal.features) || + request.proposal.features.length !== 0 + ) { + authenticationError("invalid-credential"); + } + return Object.freeze({ + role: request.role, + appId: this.config.appId, + endpoint: request.endpoint, + keyId: request.keyId, + grantId: request.grantId, + initiatorNonce: request.initiatorNonce, + proposal: fixedProposal(this.config.appRevision), + }); + } + + private takeChallenge(challengeId: string, role: BrokerGrant["kind"]): PendingChallenge { + if (!isValidBrokerIdentifier(challengeId)) authenticationError("challenge-used"); + const pending = this.challenges.get(challengeId); + if (!pending) authenticationError("challenge-used"); + // Delete before any asynchronous verification so concurrent proofs cannot both consume it. + this.deleteChallenge(challengeId); + if (this.currentTime() >= pending.expiresAt) authenticationError("challenge-expired"); + if (pending.grant.kind !== role) authenticationError("invalid-credential"); + return pending; + } + + private async verifyProof(pending: PendingChallenge, encodedSignature: string): Promise { + const signature = decodeBase64Url(encodedSignature); + if ( + !signature || + signature.byteLength === 0 || + !(await this.crypto.verify(pending.publicKey, signature, pending.transcript)) + ) { + authenticationError("invalid-signature"); + } + } + + private requireActiveGrant(grant: BrokerGrant): void { + let revoked = false; + try { + revoked = this.config.isRevoked?.(grant.revocationId) ?? false; + } catch { + authenticationError("invalid-credential"); + } + if (revoked) authenticationError("credential-revoked"); + if (!isGrantActive(grant, { appId: this.config.appId, now: this.currentTime() })) { + authenticationError("credential-expired"); + } + } + + private currentTime(): number { + let now: number; + try { + now = this.config.now(); + } catch { + authenticationError("invalid-credential"); + } + if (!Number.isFinite(now)) authenticationError("invalid-credential"); + return now; + } + + private deleteChallenge(challengeId: string): void { + const challenge = this.challenges.get(challengeId); + if (!challenge) return; + this.challenges.delete(challengeId); + this.challengeBytes = Math.max(0, this.challengeBytes - challenge.retainedBytes); + } + + private pruneExpired(): void { + const now = this.currentTime(); + for (const [id, challenge] of this.challenges) { + if (now >= challenge.expiresAt) this.deleteChallenge(id); + } + for (const [id, session] of this.callerSessions) { + if (now >= session.expiresAt) this.callerSessions.delete(id); + } + } + + private uniqueId(isReserved: (id: string) => boolean): string { + for (let attempt = 0; attempt < UNIQUE_ID_RETRIES; attempt += 1) { + const id = this.randomId(); + if (!isReserved(id)) return id; + } + authenticationError("authentication-capacity"); + } + + private randomId(): string { + let bytes: Uint8Array; + try { + bytes = this.crypto.randomBytes(RANDOM_ID_BYTES); + } catch { + authenticationError("invalid-credential"); + } + if (!(bytes instanceof Uint8Array) || bytes.byteLength !== RANDOM_ID_BYTES) { + authenticationError("invalid-credential"); + } + // Fixed alphanumeric bookends keep generated values inside the public identifier grammar. + return `b_${encodeBase64Url(bytes)}_0`; + } +} + +/** Build the exact challenge transcript so clients can verify daemon identity before signing. */ +export function challengeTranscriptForClient( + request: SessionBrokerHelloChallengeRequest, + challenge: Pick, + generation: string, +): Uint8Array { + return buildBrokerChallengeTranscript({ + ...request, + generation, + responderNonce: challenge.responderNonce, + } satisfies BrokerChallengeTranscriptInput); +} diff --git a/packages/session-broker/src/broker.ts b/packages/session-broker/src/broker.ts index 37a4b6300..058c004d8 100644 --- a/packages/session-broker/src/broker.ts +++ b/packages/session-broker/src/broker.ts @@ -65,6 +65,7 @@ export interface SessionBrokerController< dispatchCommand(options: { selector: SessionTargetInput; command: ServerMessage["command"]; + commandVersion?: number; input: unknown; timeoutMessage: string; timeoutMs?: number; @@ -183,12 +184,14 @@ export class SessionBroker< dispatchCommand({ selector, command, + commandVersion, input, timeoutMessage, timeoutMs, }: { selector: SessionTargetInput; command: CommandName; + commandVersion?: number; input: Extract["input"]; timeoutMessage: string; timeoutMs?: number; @@ -196,12 +199,14 @@ export class SessionBroker< dispatchCommand({ selector, command, + commandVersion, input, timeoutMessage, timeoutMs, }: { selector: SessionTargetInput; command: ServerMessage["command"]; + commandVersion?: number; input: unknown; timeoutMessage: string; timeoutMs?: number; @@ -209,6 +214,7 @@ export class SessionBroker< return this.state.dispatchCommand({ selector, command, + commandVersion, input: input as Extract["input"], timeoutMessage, timeoutMs, diff --git a/packages/session-broker/src/crypto.test.ts b/packages/session-broker/src/crypto.test.ts new file mode 100644 index 000000000..df3f25246 --- /dev/null +++ b/packages/session-broker/src/crypto.test.ts @@ -0,0 +1,83 @@ +import { describe, expect, test } from "bun:test"; +import { + buildBrokerChallengeTranscript, + buildCallerRequestTranscript, +} from "@hunk/session-broker-core"; +import { + decodeBase64Url, + encodeBase64Url, + importEd25519PrivateKey, + importEd25519PublicKey, + webSessionBrokerCrypto, +} from "./crypto"; + +function hex(value: string): Uint8Array { + return Uint8Array.from(value.match(/../g) ?? [], (byte) => Number.parseInt(byte, 16)); +} + +describe("session broker Ed25519 crypto", () => { + test("round-trips canonical base64url empty and edge values", () => { + expect(encodeBase64Url(new Uint8Array())).toBe(""); + expect(decodeBase64Url("")).toEqual(new Uint8Array()); + expect(decodeBase64Url("A")).toBeNull(); + expect(decodeBase64Url("AA==")).toBeNull(); + expect(decodeBase64Url("AB")).toBeNull(); + expect(decodeBase64Url("AA")).toEqual(new Uint8Array([0])); + }); + + test("matches the golden transcript signature fixture", async () => { + // RFC 8032 test vector 1 key material wrapped in standard PKCS#8/SPKI containers. + const privateKey = await importEd25519PrivateKey( + new Uint8Array([ + ...hex("302e020100300506032b657004220420"), + ...hex("9d61b19deffd5a60ba844af492ec2cc44449c5697b326919703bac031cae7f60"), + ]), + ); + const publicKey = await importEd25519PublicKey( + new Uint8Array([ + ...hex("302a300506032b6570032100"), + ...hex("d75a980182b10ab7d54bfed3c964073a0ee172f3daa62325af021a68f707511a"), + ]), + ); + const transcript = buildBrokerChallengeTranscript({ + role: "caller", + appId: "dev.example", + generation: "generation-1", + endpoint: "http://127.0.0.1:47657/broker", + keyId: "caller-key-1", + grantId: "caller-grant-1", + initiatorNonce: "nonce-a", + responderNonce: "nonce-b", + proposal: { brokerRevision: 1, appRevision: 7, features: ["z", "a"] }, + }); + const signature = await webSessionBrokerCrypto.sign(privateKey, transcript); + + expect(encodeBase64Url(signature)).toBe( + "u1r3Q-Fji-3PYthUZ8TudnJm2Gw3b0jJqFYddzFMzpXmgshYVY9OL2iKD0zbEy2tWcm5L_evmmKAYClbc4_sAg", + ); + expect(await webSessionBrokerCrypto.verify(publicKey, signature, transcript)).toBe(true); + + expect( + encodeBase64Url( + await webSessionBrokerCrypto.sha256(new TextEncoder().encode('{"action":"list"}')), + ), + ).toBe("WE52AIFcTHUuQvjzwAqkvxWX8TuwjtXeRDszCX4aF1E"); + + const requestTranscript = buildCallerRequestTranscript({ + appId: "dev.example", + generation: "generation-1", + callerSessionId: "caller-session-1", + keyId: "caller-key-1", + grantId: "caller-grant-1", + helloTranscriptHash: "hello-hash", + method: "post", + target: "/broker?a=1&b=2", + bodyDigest: "body-hash", + requestId: "request-1", + sequence: "1", + }); + expect(encodeBase64Url(await webSessionBrokerCrypto.sign(privateKey, requestTranscript))).toBe( + "I9FPiLt4mNyEzZRIYNPFffwTRf-SutIK9ml9BdJFT-6JQs6i58G6sAfx-JJt8N3yI_VzDFJXFLv_4JbyaX5kDg", + ); + }); +}); diff --git a/packages/session-broker/src/crypto.ts b/packages/session-broker/src/crypto.ts new file mode 100644 index 000000000..5c774b86f --- /dev/null +++ b/packages/session-broker/src/crypto.ts @@ -0,0 +1,62 @@ +/** Encode bytes without padding for wire-safe identifiers and signatures. */ +export function encodeBase64Url(bytes: Uint8Array): string { + let binary = ""; + for (const byte of bytes) binary += String.fromCharCode(byte); + return btoa(binary).replaceAll("+", "-").replaceAll("/", "_").replace(/=+$/, ""); +} + +/** Decode a base64url value, rejecting non-canonical spellings. */ +export function decodeBase64Url(value: string): Uint8Array | null { + if (!/^[A-Za-z0-9_-]*$/.test(value) || value.length % 4 === 1) return null; + if (value.length === 0) return new Uint8Array(); + try { + const base64 = value.replaceAll("-", "+").replaceAll("_", "/"); + const binary = atob(base64.padEnd(Math.ceil(base64.length / 4) * 4, "=")); + const bytes = Uint8Array.from(binary, (character) => character.charCodeAt(0)); + return encodeBase64Url(bytes) === value ? bytes : null; + } catch { + return null; + } +} + +function asArrayBuffer(bytes: Uint8Array): ArrayBuffer { + return Uint8Array.from(bytes).buffer; +} + +export interface SessionBrokerCrypto { + randomBytes(length: number): Uint8Array; + sha256(value: Uint8Array): Promise; + sign(privateKey: CryptoKey, value: Uint8Array): Promise; + verify(publicKey: CryptoKey, signature: Uint8Array, value: Uint8Array): Promise; +} + +/** Use the runtime WebCrypto implementation for Ed25519 proof-of-possession and SHA-256. */ +export const webSessionBrokerCrypto: SessionBrokerCrypto = { + randomBytes(length) { + return crypto.getRandomValues(new Uint8Array(length)); + }, + async sha256(value) { + return new Uint8Array(await crypto.subtle.digest("SHA-256", asArrayBuffer(value))); + }, + async sign(privateKey, value) { + return new Uint8Array(await crypto.subtle.sign("Ed25519", privateKey, asArrayBuffer(value))); + }, + async verify(publicKey, signature, value) { + return crypto.subtle.verify( + "Ed25519", + publicKey, + asArrayBuffer(signature), + asArrayBuffer(value), + ); + }, +}; + +/** Import an Ed25519 SubjectPublicKeyInfo verifier without exposing private material. */ +export function importEd25519PublicKey(spki: Uint8Array): Promise { + return crypto.subtle.importKey("spki", asArrayBuffer(spki), "Ed25519", false, ["verify"]); +} + +/** Import an Ed25519 PKCS#8 signer into non-extractable runtime memory. */ +export function importEd25519PrivateKey(pkcs8: Uint8Array): Promise { + return crypto.subtle.importKey("pkcs8", asArrayBuffer(pkcs8), "Ed25519", false, ["sign"]); +} diff --git a/packages/session-broker/src/daemon.test.ts b/packages/session-broker/src/daemon.test.ts index cca6a166f..62e959897 100644 --- a/packages/session-broker/src/daemon.test.ts +++ b/packages/session-broker/src/daemon.test.ts @@ -4,12 +4,14 @@ import { brokerWireParsers, parseSessionRegistrationEnvelope, parseSessionSnapshotEnvelope, + type CallerPrincipal, type SessionRegistration, type SessionServerMessage, type SessionSnapshot, } from "@hunk/session-broker-core"; import { SessionBroker } from "./broker"; import { createSessionBrokerDaemon } from "./daemon"; +import type { AuthenticatedCallerRequest } from "./authentication"; interface TestSessionInfo { title: string; @@ -77,6 +79,52 @@ function createSnapshot( }; } +function authenticatedRequest(principal: CallerPrincipal): AuthenticatedCallerRequest { + return { + principal, + requestId: "request-1", + assertActive() {}, + async signResponse(input) { + return { + generation: "generation-1", + brokerRevision: 1, + ...(input.appContract ? { appContract: input.appContract } : {}), + requestId: "request-1", + httpStatus: input.httpStatus, + bodyDigest: "test-body-digest", + daemonKeyId: "daemon-key-1", + daemonSignature: "test-signature", + }; + }, + }; +} + +const authenticatedHttpApi = { + appId: "session-broker", + appRevision: 1, + callerAuthenticator: { + authenticate: async () => + authenticatedRequest({ + kind: "caller" as const, + appId: "session-broker", + principalId: "test-caller", + keyId: "test-key", + grantId: "test-grant", + operations: ["list", "get", "dispatch", "diagnostics"] as const, + commands: [ + { name: "annotate", version: 1 }, + { name: "annotate", version: 2 }, + ], + }), + }, + authorizer: async () => true, +}; + +async function authenticatedBody(response: Response | null) { + const envelope = (await response?.json()) as { body: unknown } | undefined; + return envelope?.body; +} + function createConnection() { const sent: string[] = []; let closed: { code?: number; reason?: string } | null = null; @@ -103,6 +151,7 @@ describe("session broker daemon", () => { broker: createBroker(), capabilities: { version: 1, name: "test-broker" }, exposeHttpApi: true, + ...authenticatedHttpApi, }); const { connection } = createConnection(); daemon.handleConnectionMessage( @@ -129,7 +178,7 @@ describe("session broker daemon", () => { }), ); expect(listResponse).toBeInstanceOf(Response); - await expect(listResponse?.json()).resolves.toMatchObject({ + await expect(authenticatedBody(listResponse)).resolves.toMatchObject({ sessions: [{ sessionId: "session-1", title: "repo working tree" }], }); @@ -140,7 +189,7 @@ describe("session broker daemon", () => { body: JSON.stringify({ action: "get", selector: { sessionId: "session-1" } }), }), ); - await expect(getResponse?.json()).resolves.toMatchObject({ + await expect(authenticatedBody(getResponse)).resolves.toMatchObject({ session: { registration: { sessionId: "session-1" }, snapshot: { state: { selectedIndex: 0 } }, @@ -150,6 +199,55 @@ describe("session broker daemon", () => { daemon.shutdown(); }); + test("refuses exposeHttpApi without both an explicit authenticator and authorizer", async () => { + const withoutAuthentication = createSessionBrokerDaemon({ + broker: createBroker(), + capabilities: { version: 1 }, + exposeHttpApi: true, + appId: "session-broker", + }); + const withoutAuthorization = createSessionBrokerDaemon({ + broker: createBroker(), + capabilities: { version: 1 }, + exposeHttpApi: true, + appId: "session-broker", + callerAuthenticator: authenticatedHttpApi.callerAuthenticator, + }); + + for (const daemon of [withoutAuthentication, withoutAuthorization]) { + expect(daemon.paths).toEqual({ health: "/health", socket: "/session" }); + await expect( + daemon.handleRequest( + new Request("http://broker.test/broker", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ action: "list" }), + }), + ), + ).resolves.toBeNull(); + daemon.shutdown(); + } + }); + + test("requires GET with an empty body for authenticated capabilities", async () => { + const daemon = createSessionBrokerDaemon({ + broker: createBroker(), + capabilities: { version: 1 }, + exposeHttpApi: true, + ...authenticatedHttpApi, + }); + + const response = await daemon.handleRequest( + new Request("http://broker.test/broker/capabilities", { + method: "POST", + body: "unsigned bytes", + }), + ); + + expect(response?.status).toBe(405); + daemon.shutdown(); + }); + test("does not expose the raw broker HTTP API by default", async () => { const daemon = createSessionBrokerDaemon({ broker: createBroker(), @@ -182,6 +280,7 @@ describe("session broker daemon", () => { broker: createBroker(), capabilities: { version: 1 }, exposeHttpApi: true, + ...authenticatedHttpApi, }); const response = await daemon.handleRequest( @@ -199,11 +298,53 @@ describe("session broker daemon", () => { daemon.shutdown(); }); + test("authenticates exact body bytes before strictly rejecting BOM and malformed UTF-8", async () => { + const authenticatedBodies: number[][] = []; + const daemon = createSessionBrokerDaemon({ + broker: createBroker(), + exposeHttpApi: true, + appId: "session-broker", + appRevision: 1, + callerAuthenticator: { + authenticate: async ({ body }) => { + authenticatedBodies.push([...body]); + return authenticatedRequest({ + kind: "caller", + appId: "session-broker", + principalId: "test-caller", + keyId: "test-key", + grantId: "test-grant", + operations: ["list"], + commands: [], + }); + }, + }, + authorizer: async () => true, + }); + const malformedBodies = [ + new Uint8Array([0xef, 0xbb, 0xbf, ...new TextEncoder().encode('{"action":"list"}')]), + new Uint8Array([0x7b, 0x22, 0x78, 0x22, 0x3a, 0xc0, 0xaf, 0x7d]), + ]; + for (const body of malformedBodies) { + const response = await daemon.handleRequest( + new Request("http://broker.test/broker", { + method: "POST", + headers: { "content-type": "application/json" }, + body, + }), + ); + expect(response?.status).toBe(400); + } + expect(authenticatedBodies).toEqual(malformedBodies.map((body) => [...body])); + daemon.shutdown(); + }); + test("rejects raw broker API bodies that exceed the size limit", async () => { const daemon = createSessionBrokerDaemon({ broker: createBroker(), capabilities: { version: 1 }, exposeHttpApi: true, + ...authenticatedHttpApi, }); const oversized = JSON.stringify({ action: "list", filler: "x".repeat(5 * 1024 * 1024) }); @@ -227,6 +368,7 @@ describe("session broker daemon", () => { broker: createBroker(), capabilities: { version: 1 }, exposeHttpApi: true, + ...authenticatedHttpApi, }); const session = createConnection(); const { connection, sent } = session; @@ -247,14 +389,19 @@ describe("session broker daemon", () => { action: "dispatch", selector: { sessionId: "session-1" }, command: "annotate", + commandVersion: 2, input: { summary: "Review note" }, }), }), ); await Bun.sleep(0); - const outgoing = JSON.parse(sent[sent.length - 1]!) as { requestId: string; command: string }; - expect(outgoing.command).toBe("annotate"); + const outgoing = JSON.parse(sent[sent.length - 1]!) as { + requestId: string; + command: string; + commandVersion: number; + }; + expect(outgoing).toMatchObject({ command: "annotate", commandVersion: 2 }); daemon.handleConnectionMessage( connection, @@ -267,7 +414,7 @@ describe("session broker daemon", () => { ); const response = await pendingResponse; - await expect(response?.json()).resolves.toEqual({ result: { applied: true } }); + await expect(authenticatedBody(response)).resolves.toEqual({ result: { applied: true } }); daemon.shutdown(); }); @@ -323,6 +470,7 @@ describe("session broker daemon", () => { broker: createBroker(), capabilities: { version: 1 }, exposeHttpApi: true, + ...authenticatedHttpApi, }); const owner = createConnection(); const snapshotPeer = createConnection(); @@ -403,7 +551,107 @@ describe("session broker daemon", () => { }), ); const response = await pendingResponse; - await expect(response?.json()).resolves.toEqual({ result: { applied: true } }); + await expect(authenticatedBody(response)).resolves.toEqual({ result: { applied: true } }); + daemon.shutdown(); + }); + + test("requires operation, command, and app authorization before broker control", async () => { + let appAuthorizerCalls = 0; + const daemon = createSessionBrokerDaemon({ + broker: createBroker(), + capabilities: { version: 1 }, + exposeHttpApi: true, + appId: "session-broker", + appRevision: 1, + callerAuthenticator: { + authenticate: async () => + authenticatedRequest({ + kind: "caller" as const, + appId: "session-broker", + principalId: "limited-caller", + keyId: "limited-key", + grantId: "limited-grant", + operations: ["list", "dispatch"] as const, + commands: [{ name: "allowed", version: 1 }], + }), + }, + authorizer: async () => { + appAuthorizerCalls += 1; + return true; + }, + }); + const post = (body: unknown) => + daemon.handleRequest( + new Request("http://broker.test/broker", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify(body), + }), + ); + + expect((await post({ action: "get", selector: { sessionId: "missing" } }))?.status).toBe(403); + expect(appAuthorizerCalls).toBe(0); + expect( + ( + await post({ + action: "dispatch", + selector: { sessionId: "missing" }, + command: "forbidden", + input: {}, + }) + )?.status, + ).toBe(403); + expect(appAuthorizerCalls).toBe(0); + expect( + ( + await post({ + action: "dispatch", + selector: { sessionId: "missing" }, + command: "allowed", + commandVersion: 0, + input: {}, + }) + )?.status, + ).toBe(400); + expect(appAuthorizerCalls).toBe(0); + expect((await post({ action: "list" }))?.status).toBe(200); + expect(appAuthorizerCalls).toBe(1); + daemon.shutdown(); + }); + + test("returns stable redacted authentication failures without invoking app authorization", async () => { + let authorized = false; + const daemon = createSessionBrokerDaemon({ + broker: createBroker(), + exposeHttpApi: true, + appId: "session-broker", + appRevision: 1, + callerAuthenticator: { + authenticate: async () => { + const { SessionBrokerAuthenticationError } = await import("./authentication"); + throw new SessionBrokerAuthenticationError("invalid-signature"); + }, + }, + authorizer: async () => { + authorized = true; + return true; + }, + }); + const response = await daemon.handleRequest( + new Request("http://broker.test/broker", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ action: "list" }), + }), + ); + expect(response?.status).toBe(401); + const responseText = await response?.text(); + expect(JSON.parse(responseText ?? "null")).toEqual({ + error: "authentication-failed", + code: "invalid-signature", + }); + expect(authorized).toBe(false); + expect(responseText).not.toContain("private"); daemon.shutdown(); }); diff --git a/packages/session-broker/src/daemon.ts b/packages/session-broker/src/daemon.ts index c03fa5e52..13df35d9c 100644 --- a/packages/session-broker/src/daemon.ts +++ b/packages/session-broker/src/daemon.ts @@ -1,11 +1,24 @@ import { MAX_HTTP_BODY_BYTES, PayloadTooLargeError, - readRequestTextWithLimit, + callerPrincipalAllows, + canonicalizeJson, + isValidBrokerAppId, + isValidBrokerIdentifier, + isValidBrokerRevision, + readRequestBytesWithLimit, + type CallerOperation, + type CallerPrincipal, + type CanonicalJsonValue, type SessionServerMessage, type SessionTargetSelector, } from "@hunk/session-broker-core"; import type { SessionBrokerController, SessionBrokerPeer } from "./broker"; +import { + SessionBrokerAuthenticationError, + type AuthenticatedCallerRequest, + type CallerRequestAuthenticator, +} from "./authentication"; import { DEFAULT_SESSION_BROKER_API_PATH, DEFAULT_SESSION_BROKER_CAPABILITIES_PATH, @@ -14,6 +27,10 @@ import { type SessionBrokerCapabilities, type SessionBrokerDaemonRequest, type SessionBrokerDaemonResponse, + type SessionBrokerAuthenticatedResponse, + type SessionBrokerAuditEvent, + type SessionBrokerAuditHook, + type SessionBrokerAuthorizer, type SessionBrokerHealth, type SessionBrokerHttpPaths, } from "./types"; @@ -32,6 +49,11 @@ export interface SessionBrokerDaemonOptions< capabilities?: SessionBrokerCapabilities; paths?: Partial; exposeHttpApi?: boolean; + callerAuthenticator?: CallerRequestAuthenticator; + authorizer?: SessionBrokerAuthorizer; + audit?: SessionBrokerAuditHook; + appId?: string; + appRevision?: number; idleTimeoutMs?: number; staleSessionTtlMs?: number; staleSessionSweepIntervalMs?: number; @@ -67,15 +89,68 @@ function hasJsonContentType(request: Request) { } /** Decode one raw broker API request body and surface a friendly transport-level error. */ -async function parseJsonRequest( - request: Request, -) { - const text = await readRequestTextWithLimit(request, MAX_HTTP_BODY_BYTES); +function parseJsonRequest( + body: Uint8Array, +): SessionBrokerDaemonRequest { + let parsed: unknown; try { - return JSON.parse(text) as SessionBrokerDaemonRequest; + if (body[0] === 0xef && body[1] === 0xbb && body[2] === 0xbf) { + throw new TypeError("UTF-8 BOM is not permitted."); + } + parsed = JSON.parse(new TextDecoder("utf-8", { fatal: true }).decode(body)); } catch { - throw new Error("Expected one JSON request body."); + throw new Error("Expected one strictly encoded JSON request body."); + } + if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) { + throw new Error("Expected one JSON request object."); + } + const record = parsed as Record; + if (record.action === "list") return { action: "list" }; + if (record.action !== "get" && record.action !== "dispatch") { + throw new Error("Unknown broker API action."); + } + if (!record.selector || typeof record.selector !== "object" || Array.isArray(record.selector)) { + throw new Error("Expected one broker session selector."); + } + const selector = record.selector as Record; + for (const key of ["sessionId", "sessionPath", "repoRoot", "repoBoundary"] as const) { + if (selector[key] !== undefined && typeof selector[key] !== "string") { + throw new Error("Expected one valid broker session selector."); + } } + if (selector.sessionId !== undefined && !isValidBrokerIdentifier(selector.sessionId)) { + throw new Error("Expected one valid broker session identifier."); + } + if (record.action === "get") { + return { action: "get", selector: selector as SessionTargetSelector }; + } + if (!isValidBrokerIdentifier(record.command)) { + throw new Error("Expected one valid broker command name."); + } + const commandVersion = record.commandVersion ?? 1; + if (!isValidBrokerRevision(commandVersion)) { + throw new Error("Expected one positive broker command version."); + } + if ( + record.timeoutMs !== undefined && + (!Number.isSafeInteger(record.timeoutMs) || (record.timeoutMs as number) <= 0) + ) { + throw new Error("Expected one positive command timeout."); + } + if (record.timeoutMessage !== undefined && typeof record.timeoutMessage !== "string") { + throw new Error("Expected one command timeout message."); + } + return { + action: "dispatch", + selector: selector as SessionTargetSelector, + command: record.command as CommandName, + commandVersion, + input: record.input as CommandInput, + ...(record.timeoutMs === undefined ? {} : { timeoutMs: record.timeoutMs as number }), + ...(record.timeoutMessage === undefined + ? {} + : { timeoutMessage: record.timeoutMessage as string }), + }; } /** Build the default dispatch timeout text so adapters can override only when they need to. */ @@ -100,6 +175,11 @@ export class SessionBrokerDaemon< private readonly idleTimeoutMs: number; private readonly staleSessionTtlMs: number; private readonly staleSessionSweepIntervalMs: number; + private readonly appId: string; + private readonly appRevision?: number; + private readonly callerAuthenticator?: CallerRequestAuthenticator; + private readonly authorizer?: SessionBrokerAuthorizer; + private readonly audit?: SessionBrokerAuditHook; private lastActivityAt = this.startedAt; private sweepTimer: ReturnType | null = null; private idleTimer: ReturnType | null = null; @@ -113,16 +193,28 @@ export class SessionBrokerDaemon< "broker" > = {}, ) { - const exposeHttpApi = options.exposeHttpApi ?? false; + const exposeAuthenticatedHttpApi = + (options.exposeHttpApi ?? false) && + isValidBrokerAppId(options.appId) && + isValidBrokerRevision(options.appRevision) && + !!options.callerAuthenticator && + !!options.authorizer; this.paths = { health: options.paths?.health ?? DEFAULT_SESSION_BROKER_HEALTH_PATH, socket: options.paths?.socket ?? DEFAULT_SESSION_BROKER_SOCKET_PATH, - api: exposeHttpApi ? (options.paths?.api ?? DEFAULT_SESSION_BROKER_API_PATH) : undefined, - capabilities: exposeHttpApi + api: exposeAuthenticatedHttpApi + ? (options.paths?.api ?? DEFAULT_SESSION_BROKER_API_PATH) + : undefined, + capabilities: exposeAuthenticatedHttpApi ? (options.paths?.capabilities ?? DEFAULT_SESSION_BROKER_CAPABILITIES_PATH) : undefined, }; this.capabilities = options.capabilities ?? { version: 1 }; + this.appId = options.appId ?? "session-broker"; + this.appRevision = options.appRevision; + this.callerAuthenticator = options.callerAuthenticator; + this.authorizer = options.authorizer; + this.audit = options.audit; this.idleTimeoutMs = options.idleTimeoutMs ?? DEFAULT_IDLE_TIMEOUT_MS; this.staleSessionTtlMs = options.staleSessionTtlMs ?? DEFAULT_STALE_SESSION_TTL_MS; this.staleSessionSweepIntervalMs = @@ -170,12 +262,35 @@ export class SessionBrokerDaemon< this.noteActivity(); } - return Response.json(this.getHealth()); + // Public health is deliberately liveness-only. Apps may expose authenticated diagnostics on + // a separate route, but broker identity, paths, counts, and process facts stay private. + return Response.json({ ok: true }); } if (this.paths.capabilities && url.pathname === this.paths.capabilities) { + if (request.method !== "GET") { + return jsonError("Broker capabilities requests must use GET.", 405); + } + let body: Uint8Array; + try { + body = await readRequestBytesWithLimit(request, MAX_HTTP_BODY_BYTES); + } catch (error) { + return error instanceof PayloadTooLargeError + ? jsonError(error.message, 413) + : jsonError("Could not read broker capabilities request body."); + } + if (body.byteLength !== 0) { + return jsonError("Broker capabilities requests must not include a body."); + } + const authenticated = await this.authenticateRequest(request, body, "diagnostics"); + if (authenticated instanceof Response) return authenticated; + if (!(await this.authorize(request, authenticated, { operation: "diagnostics" }))) { + return this.authenticatedResponse(authenticated, { error: "authorization-denied" }, 403); + } + const inactive = this.rejectInactiveRequest(authenticated); + if (inactive) return inactive; this.noteActivity(); - return Response.json(this.capabilities); + return this.authenticatedResponse(authenticated, this.capabilities, 200); } if (this.paths.api && url.pathname === this.paths.api) { @@ -357,19 +472,184 @@ export class SessionBrokerDaemon< }, remainingMs); } + private async authenticateRequest( + request: Request, + body: Uint8Array, + operation: CallerOperation, + ): Promise { + const requestId = request.headers.get("x-session-broker-request-id") ?? undefined; + try { + if (!this.callerAuthenticator || !this.authorizer) { + return jsonError("Broker control is unavailable.", 404); + } + const authenticated = await this.callerAuthenticator.authenticate({ request, body }); + if ( + !authenticated || + typeof authenticated !== "object" || + !authenticated.principal || + !isValidBrokerIdentifier(authenticated.requestId) || + typeof authenticated.assertActive !== "function" || + typeof authenticated.signResponse !== "function" + ) { + throw new SessionBrokerAuthenticationError("invalid-credential"); + } + return authenticated; + } catch (error) { + const code = + error instanceof SessionBrokerAuthenticationError ? error.code : "authentication-required"; + await this.emitAudit({ + appId: this.appId, + operation, + ...(requestId !== undefined ? { requestId } : {}), + decision: "deny", + outcome: "authentication-failed", + timestamp: Date.now(), + }); + return Response.json({ error: "authentication-failed", code }, { status: 401 }); + } + } + + /** Fail a request whose caller session changed while asynchronous app policy was running. */ + private rejectInactiveRequest(authenticated: AuthenticatedCallerRequest): Response | null { + try { + authenticated.assertActive(); + return null; + } catch (error) { + const code = + error instanceof SessionBrokerAuthenticationError ? error.code : "authentication-required"; + return Response.json({ error: "authentication-failed", code }, { status: 401 }); + } + } + + private async authorize( + request: Request, + authenticated: AuthenticatedCallerRequest, + facts: { + operation: CallerOperation; + sessionId?: string; + command?: string; + commandVersion?: number; + }, + ): Promise { + const principal: CallerPrincipal = authenticated.principal; + const allowedByGrant = callerPrincipalAllows(principal, { appId: this.appId, ...facts }); + let allowedByApp = false; + if (allowedByGrant && this.authorizer) { + try { + allowedByApp = await this.authorizer({ + principal, + ...facts, + requestId: authenticated.requestId, + signal: request.signal, + }); + } catch { + // App policy errors fail closed and never expose callback details to the caller. + allowedByApp = false; + } + } + await this.emitAudit({ + appId: this.appId, + principalId: principal.principalId, + keyId: principal.keyId, + ...(facts.sessionId !== undefined ? { sessionId: facts.sessionId } : {}), + operation: facts.operation, + ...(facts.command !== undefined ? { command: facts.command } : {}), + ...(facts.commandVersion === undefined ? {} : { commandVersion: facts.commandVersion }), + requestId: authenticated.requestId, + decision: allowedByApp ? "allow" : "deny", + outcome: allowedByApp ? "authenticated" : "authorization-failed", + timestamp: Date.now(), + }); + return allowedByApp; + } + + private async authenticatedResponse( + authenticated: AuthenticatedCallerRequest, + body: unknown, + status: number, + targetSpecific = false, + ): Promise { + // Normalize with JSON semantics first so optional undefined fields cannot create digest aliases. + const structuredBody = JSON.parse(JSON.stringify(body)) as CanonicalJsonValue; + canonicalizeJson(structuredBody); + const authentication = await authenticated.signResponse({ + httpStatus: status, + body: structuredBody, + ...(targetSpecific && this.appRevision !== undefined + ? { appContract: { appRevision: this.appRevision, features: [] } } + : {}), + }); + const envelope: SessionBrokerAuthenticatedResponse = { body: structuredBody, authentication }; + return new Response(canonicalizeJson(envelope as unknown as CanonicalJsonValue), { + status, + headers: { "content-type": "application/json" }, + }); + } + + private async emitAudit(event: SessionBrokerAuditEvent): Promise { + try { + await this.audit?.(event); + } catch { + // Audit sinks observe decisions but cannot weaken them or leak their failures to callers. + } + } + private async handleApiRequest(request: Request) { if (request.method !== "POST") { return jsonError("Broker API requests must use POST.", 405); } - if (!hasJsonContentType(request)) { return jsonError("Expected Content-Type application/json.", 415); } + let body: Uint8Array; try { - const input = await parseJsonRequest(request); - let response: SessionBrokerDaemonResponse; + body = await readRequestBytesWithLimit(request, MAX_HTTP_BODY_BYTES); + } catch (error) { + return error instanceof PayloadTooLargeError + ? jsonError(error.message, 413) + : jsonError("Could not read broker API request body."); + } + + // Authenticate the exact transport bytes before decoding or interpreting attacker-controlled JSON. + const authenticated = await this.authenticateRequest(request, body, "list"); + if (authenticated instanceof Response) return authenticated; + + let input: SessionBrokerDaemonRequest; + try { + input = parseJsonRequest(body); + } catch (error) { + return this.authenticatedResponse( + authenticated, + { error: error instanceof Error ? error.message : "Invalid broker API request." }, + 400, + ); + } + const operation = input.action as CallerOperation; + const selector = "selector" in input ? input.selector : undefined; + const sessionId = selector?.sessionId; + const command = input.action === "dispatch" ? input.command : undefined; + const commandVersion = input.action === "dispatch" ? (input.commandVersion ?? 1) : undefined; + const facts = { + operation, + ...(sessionId !== undefined ? { sessionId } : {}), + ...(command !== undefined ? { command, commandVersion } : {}), + }; + const targetSpecific = input.action !== "list"; + if (!(await this.authorize(request, authenticated, facts))) { + return this.authenticatedResponse( + authenticated, + { error: "authorization-denied" }, + 403, + targetSpecific, + ); + } + const inactive = this.rejectInactiveRequest(authenticated); + if (inactive) return inactive; + + try { + let response: SessionBrokerDaemonResponse; switch (input.action) { case "list": response = { sessions: this.broker.listSessions() }; @@ -379,11 +659,10 @@ export class SessionBrokerDaemon< break; case "dispatch": response = { - // The HTTP API stays generic JSON, while the broker keeps ownership of target - // resolution, timeout handling, and websocket command delivery. result: await this.broker.dispatchCommand({ selector: input.selector, command: input.command, + commandVersion: input.commandVersion ?? 1, input: input.input as Extract< ServerMessage, { command: ServerMessage["command"] } @@ -393,17 +672,15 @@ export class SessionBrokerDaemon< }), }; break; - default: - throw new Error("Unknown broker API action."); } - - return Response.json(response); + return this.authenticatedResponse(authenticated, response, 200, targetSpecific); } catch (error) { - if (error instanceof PayloadTooLargeError) { - return jsonError(error.message, 413); - } - - return jsonError(error instanceof Error ? error.message : "Unknown broker API error."); + return this.authenticatedResponse( + authenticated, + { error: error instanceof Error ? error.message : "Unknown broker API error." }, + 400, + targetSpecific, + ); } } } diff --git a/packages/session-broker/src/index.ts b/packages/session-broker/src/index.ts index 433013493..802081113 100644 --- a/packages/session-broker/src/index.ts +++ b/packages/session-broker/src/index.ts @@ -3,3 +3,5 @@ export * from "./types"; export * from "./broker"; export * from "./daemon"; export * from "./connection"; +export * from "./crypto"; +export * from "./authentication"; diff --git a/packages/session-broker/src/types.ts b/packages/session-broker/src/types.ts index cade3397e..ee5a35778 100644 --- a/packages/session-broker/src/types.ts +++ b/packages/session-broker/src/types.ts @@ -1,4 +1,10 @@ -import type { SessionTargetInput } from "@hunk/session-broker-core"; +import type { + BrokerAppContract, + CallerOperation, + CallerPrincipal, + SessionTargetInput, +} from "@hunk/session-broker-core"; +import type { SessionBrokerResponseAuthentication } from "./authentication"; export const DEFAULT_SESSION_BROKER_HEALTH_PATH = "/health"; export const DEFAULT_SESSION_BROKER_API_PATH = "/broker"; @@ -35,6 +41,7 @@ export type SessionBrokerDaemonRequest< action: "dispatch"; selector: SessionTargetInput; command: CommandName; + commandVersion?: number; input: CommandInput; timeoutMs?: number; timeoutMessage?: string; @@ -51,6 +58,17 @@ export type SessionBrokerDaemonResponse { + readonly body: Body; + readonly authentication: SessionBrokerResponseAuthentication; +} + +/** Select the fixed Phase-1 application contract for target-specific responses. */ +export interface SessionBrokerTargetContract extends BrokerAppContract { + readonly features: readonly []; +} + export interface SessionBrokerHealth { ok: boolean; pid: number; @@ -86,3 +104,35 @@ export interface SessionBrokerConnectionCloseDirective { reconnect?: boolean; warning?: string; } + +/** Facts supplied to the mandatory app authorization hook after signed authentication. */ +export interface SessionBrokerAuthorizationContext { + readonly principal: CallerPrincipal; + readonly operation: CallerOperation; + readonly sessionId?: string; + readonly command?: string; + readonly commandVersion?: number; + readonly requestId?: string; + readonly signal: AbortSignal; +} + +export type SessionBrokerAuthorizer = ( + context: SessionBrokerAuthorizationContext, +) => boolean | Promise; + +/** Redacted decision metadata suitable for an app-owned audit sink. */ +export interface SessionBrokerAuditEvent { + readonly appId: string; + readonly principalId?: string; + readonly keyId?: string; + readonly sessionId?: string; + readonly operation: CallerOperation; + readonly command?: string; + readonly commandVersion?: number; + readonly requestId?: string; + readonly decision: "allow" | "deny"; + readonly outcome: "authenticated" | "authentication-failed" | "authorization-failed"; + readonly timestamp: number; +} + +export type SessionBrokerAuditHook = (event: SessionBrokerAuditEvent) => void | Promise; From bba569a3d6d115838060b4a85108a1edaf1f5e82 Mon Sep 17 00:00:00 2001 From: Ben Vinegar Date: Sat, 29 Aug 2026 12:29:50 -0400 Subject: [PATCH 3/5] fix(session): validate broker runtime boundaries --- ...trict-session-broker-runtime-validation.md | 2 + packages/session-broker-bun/src/serve.test.ts | 41 +- .../src/brokerState.test.ts | 13 +- .../session-broker-core/src/brokerState.ts | 81 ++-- .../src/brokerWire.test.ts | 31 ++ .../session-broker-core/src/brokerWire.ts | 168 ++++---- packages/session-broker-core/src/index.ts | 1 + .../src/validation.test.ts | 67 ++++ .../session-broker-core/src/validation.ts | 166 ++++++++ .../session-broker-node/src/serve.test.ts | 28 +- packages/session-broker/README.md | 20 +- .../session-broker/src/authentication.test.ts | 6 +- packages/session-broker/src/authentication.ts | 95 +++-- packages/session-broker/src/broker.test.ts | 35 +- packages/session-broker/src/broker.ts | 45 ++- .../session-broker/src/connection.test.ts | 180 ++++++++- packages/session-broker/src/connection.ts | 38 +- packages/session-broker/src/daemon.test.ts | 288 +++++++++++++- packages/session-broker/src/daemon.ts | 208 ++++------ packages/session-broker/src/index.ts | 1 + .../src/protocolParsers.test.ts | 197 ++++++++++ .../session-broker/src/protocolParsers.ts | 357 +++++++++++++++++ packages/session-broker/src/types.ts | 4 +- src/session/agent/cliClient.test.ts | 28 ++ src/session/agent/cliClient.ts | 55 +-- src/session/broker/brokerClient.ts | 65 ++-- src/session/broker/brokerLauncher.test.ts | 19 + src/session/broker/brokerLauncher.ts | 61 ++- .../broker/brokerServer.helpers.test.ts | 31 ++ src/session/broker/brokerServer.ts | 2 + src/session/broker/protocolParsers.ts | 209 ++++++++++ src/session/broker/state.ts | 9 +- src/session/broker/wire.test.ts | 91 ++++- src/session/broker/wire.ts | 168 +++++--- src/session/client/capabilities.ts | 21 +- src/session/protocol.ts | 31 +- src/session/protocolSchemas.test.ts | 169 ++++++++- src/session/protocolSchemas.ts | 359 +++++++++++++++++- src/ui/runInteractiveApp.tsx | 13 +- 39 files changed, 2880 insertions(+), 523 deletions(-) create mode 100644 .changeset/strict-session-broker-runtime-validation.md create mode 100644 packages/session-broker-core/src/validation.test.ts create mode 100644 packages/session-broker-core/src/validation.ts create mode 100644 packages/session-broker/src/protocolParsers.test.ts create mode 100644 packages/session-broker/src/protocolParsers.ts create mode 100644 src/session/broker/protocolParsers.ts diff --git a/.changeset/strict-session-broker-runtime-validation.md b/.changeset/strict-session-broker-runtime-validation.md new file mode 100644 index 000000000..a845151cc --- /dev/null +++ b/.changeset/strict-session-broker-runtime-validation.md @@ -0,0 +1,2 @@ +--- +--- diff --git a/packages/session-broker-bun/src/serve.test.ts b/packages/session-broker-bun/src/serve.test.ts index 58086272d..1f0c61e50 100644 --- a/packages/session-broker-bun/src/serve.test.ts +++ b/packages/session-broker-bun/src/serve.test.ts @@ -8,7 +8,11 @@ import { type SessionRegistration, type SessionSnapshot, } from "@hunk/session-broker-core"; -import { SessionBroker, createSessionBrokerDaemon } from "@hunk/session-broker"; +import { + SessionBroker, + createSessionBrokerDaemon, + createSessionBrokerProtocolParsers, +} from "@hunk/session-broker"; import { serveSessionBrokerDaemon } from "./serve"; interface TestSessionInfo { @@ -53,7 +57,9 @@ function createRegistration(overrides: Partial["state"]> & { updatedAt?: string } = {}, + overrides: Partial["state"]> & { + updatedAt?: string; + } = {}, ) { const { updatedAt = "2026-04-15T00:00:00.000Z", ...stateOverrides } = overrides; return { @@ -65,6 +71,14 @@ function createSnapshot( } satisfies SessionSnapshot; } +const protocolParsers = createSessionBrokerProtocolParsers({ + appRevision: 1, + features: [], + parseRegistration: (value) => parseSessionRegistrationEnvelope(value, parseInfo), + parseSnapshot: (value) => parseSessionSnapshotEnvelope(value, parseState), + commands: [], +}); + async function reserveLoopbackPort() { const listener = createServer(() => undefined); await new Promise((resolve, reject) => { @@ -137,10 +151,7 @@ afterEach(() => { describe("session broker bun adapter", () => { test("serves the generic daemon API and websocket path through Bun", async () => { - const broker = new SessionBroker({ - parseRegistration: (value) => parseSessionRegistrationEnvelope(value, parseInfo), - parseSnapshot: (value) => parseSessionSnapshotEnvelope(value, parseState), - }); + const broker = new SessionBroker({ protocolParsers }); const daemon = createSessionBrokerDaemon({ broker, capabilities: { version: 1 }, @@ -221,7 +232,10 @@ describe("session broker bun adapter", () => { const response = await fetch(`http://127.0.0.1:${port}/broker`, { method: "POST", headers: { "content-type": "application/json" }, - body: JSON.stringify({ action: "get", selector: { sessionId: "session-1" } }), + body: JSON.stringify({ + action: "get", + selector: { sessionId: "session-1" }, + }), }); expect(response.status).toBe(200); await expect(response.json()).resolves.toMatchObject({ @@ -241,11 +255,11 @@ describe("session broker bun adapter", () => { }); test("lets custom request handlers override generic routes", async () => { - const broker = new SessionBroker({ - parseRegistration: (value) => parseSessionRegistrationEnvelope(value, parseInfo), - parseSnapshot: (value) => parseSessionSnapshotEnvelope(value, parseState), + const broker = new SessionBroker({ protocolParsers }); + const daemon = createSessionBrokerDaemon({ + broker, + capabilities: { version: 1 }, }); - const daemon = createSessionBrokerDaemon({ broker, capabilities: { version: 1 } }); const port = await reserveLoopbackPort(); const server = serveSessionBrokerDaemon({ daemon, @@ -263,7 +277,10 @@ describe("session broker bun adapter", () => { try { const response = await fetch(`http://127.0.0.1:${port}/health`); - await expect(response.json()).resolves.toEqual({ ok: true, overridden: true }); + await expect(response.json()).resolves.toEqual({ + ok: true, + overridden: true, + }); } finally { server.stop(true); await server.stopped; diff --git a/packages/session-broker-core/src/brokerState.test.ts b/packages/session-broker-core/src/brokerState.test.ts index 5d8a732ef..261ca24dd 100644 --- a/packages/session-broker-core/src/brokerState.test.ts +++ b/packages/session-broker-core/src/brokerState.test.ts @@ -102,10 +102,15 @@ const testBrokerView: SessionBrokerViewAdapter< TestListedSession, TestSelectedContext, TestSessionReview, - TestCommentSummary + TestCommentSummary, + TestServerMessage, + TestCommandResult > = { parseRegistration: (value) => parseSessionRegistrationEnvelope(value, parseTestInfo), parseSnapshot: (value) => parseSessionSnapshotEnvelope(value, parseTestState), + parseCommandInput: (_command, _version, value) => value, + parseCommandResult: (_command, _version, value) => + value && typeof value === "object" ? (value as TestCommandResult) : null, buildListedSession: (entry) => ({ sessionId: entry.registration.sessionId, pid: entry.registration.pid, @@ -315,6 +320,9 @@ describe("session broker state", () => { send() {}, }; + expect(state.registerSession(socket, createRegistration(), createSnapshot())).toBe( + "registered", + ); const accepted = state.registerSession( socket, { @@ -325,7 +333,8 @@ describe("session broker state", () => { ); expect(accepted).toBe("invalid"); - expect(state.listSessions()).toEqual([]); + expect(state.listSessions()).toHaveLength(1); + expect(state.getSession({ sessionId: "session-1" }).snapshot.state.selectedIndex).toBe(0); }); test("reports invalid snapshot updates without replacing the last valid selection", () => { diff --git a/packages/session-broker-core/src/brokerState.ts b/packages/session-broker-core/src/brokerState.ts index cb97dadb3..107d4ecf7 100644 --- a/packages/session-broker-core/src/brokerState.ts +++ b/packages/session-broker-core/src/brokerState.ts @@ -1,5 +1,6 @@ import { randomUUID } from "node:crypto"; import { isValidBrokerRevision } from "./auth"; +import { parseBrokerAppPayload } from "./validation"; import { matchesSessionSelector, repoSelectorDistance, type SelectableSession } from "./selectors"; import type { SessionRegistration, @@ -11,6 +12,8 @@ import type { interface PendingCommand { sessionId: string; socket: DaemonSessionSocket; + command: string; + commandVersion: number; resolve: (result: Result) => void; reject: (error: Error) => void; timeout: ReturnType; @@ -48,9 +51,21 @@ export interface SessionBrokerViewAdapter< SelectedContext, SessionReview, SessionCommentSummary, + ServerMessage extends SessionServerMessage = SessionServerMessage, + CommandResult = unknown, > { parseRegistration: (value: unknown) => SessionRegistration | null; parseSnapshot: (value: unknown) => SessionSnapshot | null; + parseCommandInput: ( + command: ServerMessage["command"], + version: number, + value: unknown, + ) => unknown; + parseCommandResult: ( + command: ServerMessage["command"], + version: number, + value: unknown, + ) => CommandResult | null; buildListedSession: (entry: SessionBrokerEntry) => ListedSession; buildSelectedContext: (session: ListedSession) => SelectedContext; buildSessionReview: ( @@ -63,7 +78,7 @@ export interface SessionBrokerViewAdapter< export type RegisterSessionResult = "registered" | "invalid" | "already-connected"; export type UpdateSnapshotResult = "updated" | "invalid" | "not-owner"; export type MarkSessionSeenResult = "seen" | "not-owner"; -export type HandleCommandResult = "handled" | "not-found" | "not-owner"; +export type HandleCommandResult = "handled" | "not-found" | "not-owner" | "invalid"; export type SessionTargetSelector = SessionTargetInput; @@ -170,7 +185,9 @@ export class SessionBrokerState< ListedSession, SelectedContext, SessionReview, - SessionCommentSummary + SessionCommentSummary, + ServerMessage, + CommandResult >, ) {} @@ -213,21 +230,15 @@ export class SessionBrokerState< registrationInput: unknown, snapshotInput: unknown, ): RegisterSessionResult { - const registration = this.view.parseRegistration(registrationInput); - const snapshot = this.view.parseSnapshot(snapshotInput); - if (!registration || !snapshot) { - const previousSessionId = this.sessionIdsBySocket.get(socket); - if (previousSessionId) { - // Drop any stale session already tied to this socket so an incompatible replacement - // payload cannot leave old review data behind after an upgrade or reload. - this.removeSession( - previousSessionId, - new Error("The session sent an incompatible registration payload."), - ); - } - + let registration: SessionRegistration | null; + let snapshot: SessionSnapshot | null; + try { + registration = this.view.parseRegistration(registrationInput); + snapshot = this.view.parseSnapshot(snapshotInput); + } catch { return "invalid"; } + if (!registration || !snapshot) return "invalid"; const existing = this.sessions.get(registration.sessionId); if (existing && existing.socket !== socket) { @@ -268,10 +279,13 @@ export class SessionBrokerState< return "not-owner"; } - const snapshot = this.view.parseSnapshot(snapshotInput); - if (!snapshot) { + let snapshot: SessionSnapshot | null; + try { + snapshot = this.view.parseSnapshot(snapshotInput); + } catch { return "invalid"; } + if (!snapshot) return "invalid"; this.sessions.set(ownedSessionId, { ...entry, @@ -360,6 +374,10 @@ export class SessionBrokerState< throw new TypeError("Command version must be a positive safe integer."); } const session = resolveSessionTarget(this.listSessions(), selector); + const parsedInput = parseBrokerAppPayload( + (value) => this.view.parseCommandInput(command, commandVersion, value), + input, + ) as Extract["input"]; const requestId = randomUUID(); return new Promise((resolve, reject) => { @@ -381,6 +399,8 @@ export class SessionBrokerState< this.pendingCommands.set(requestId, { sessionId: session.sessionId, socket: entry.socket, + command, + commandVersion, resolve: (result) => resolve(result as ResultType), reject, timeout, @@ -392,7 +412,7 @@ export class SessionBrokerState< requestId, command, commandVersion, - input, + input: parsedInput, } as Extract; entry.socket.send(JSON.stringify(message)); @@ -426,14 +446,31 @@ export class SessionBrokerState< return "not-owner"; } - clearTimeout(pending.timeout); - this.pendingCommands.delete(message.requestId); - if (message.ok) { - pending.resolve(message.result as CommandResult); + let result: CommandResult; + try { + result = parseBrokerAppPayload( + (value) => + this.view.parseCommandResult( + pending.command as ServerMessage["command"], + pending.commandVersion, + value, + ), + message.result, + ); + } catch { + // Keep the pending entry intact until the malformed producer is closed and normal + // connection cleanup rejects it. This avoids resolving work from an invalid contract. + return "invalid"; + } + clearTimeout(pending.timeout); + this.pendingCommands.delete(message.requestId); + pending.resolve(result); return "handled"; } + clearTimeout(pending.timeout); + this.pendingCommands.delete(message.requestId); pending.reject(new Error(message.error ?? "The session failed to handle the command.")); return "handled"; } diff --git a/packages/session-broker-core/src/brokerWire.test.ts b/packages/session-broker-core/src/brokerWire.test.ts index 8a1ef980f..6697c6365 100644 --- a/packages/session-broker-core/src/brokerWire.test.ts +++ b/packages/session-broker-core/src/brokerWire.test.ts @@ -22,6 +22,37 @@ describe("session broker wire parsing", () => { ).toBeNull(); }); + test("rejects arrays, unknown keys, malformed optionals, and throwing app parsers", () => { + const valid = { + registrationVersion: SESSION_BROKER_REGISTRATION_VERSION, + sessionId: "session-1", + pid: 123, + cwd: "/repo", + launchedAt: "2026-03-22T00:00:00.000Z", + info: { ok: true }, + }; + const parseInfo = (value: unknown) => (value && typeof value === "object" ? value : null); + for (const value of [ + null, + [], + { ...valid, extra: true }, + { ...valid, repoRoot: null }, + { ...valid, terminal: { locations: [], extra: true } }, + { ...valid, pid: Number.MAX_SAFE_INTEGER + 1 }, + { ...valid, sessionId: "bad id!" }, + ]) { + expect(parseSessionRegistrationEnvelope(value, parseInfo)).toBeNull(); + } + expect( + parseSessionRegistrationEnvelope(valid, () => { + throw new Error("parser internals"); + }), + ).toBeNull(); + expect( + parseSessionSnapshotEnvelope({ updatedAt: "now", state: {}, extra: true }, parseInfo), + ).toBeNull(); + }); + test("snapshot parsing delegates opaque app state validation", () => { const snapshot = parseSessionSnapshotEnvelope( { diff --git a/packages/session-broker-core/src/brokerWire.ts b/packages/session-broker-core/src/brokerWire.ts index 6667f95b2..9c695ed5c 100644 --- a/packages/session-broker-core/src/brokerWire.ts +++ b/packages/session-broker-core/src/brokerWire.ts @@ -4,74 +4,89 @@ import type { SessionTerminalLocation, SessionTerminalMetadata, } from "./types"; +import { + BrokerProtocolError, + parseBrokerAppPayload, + parseBrokerIdentifier, + parseBrokerSafeInteger, + parseBrokerString, + parseExactBrokerRecord, +} from "./validation"; /** Version the live broker registration payload separately from the public session CLI API. */ export const SESSION_BROKER_REGISTRATION_VERSION = 2; type JsonRecord = Record; -/** Return one JSON object record when the wire payload is object-shaped. */ +/** Return one plain JSON object record when the wire payload is object-shaped. */ function asRecord(value: unknown): JsonRecord | null { - return value && typeof value === "object" ? (value as JsonRecord) : null; + if (value === null || typeof value !== "object" || Array.isArray(value)) return null; + const prototype = Object.getPrototypeOf(value); + return prototype === Object.prototype || prototype === null ? (value as JsonRecord) : null; } -/** Parse one required non-empty string field from the websocket payload. */ +/** Parse one required bounded non-empty string field. */ function parseRequiredString(value: unknown) { - return typeof value === "string" && value.length > 0 ? value : null; + try { + return parseBrokerString(value); + } catch { + return null; + } } -/** Parse one optional string field, dropping malformed values instead of rejecting the payload. */ +/** Parse one optional bounded string, rejecting malformed present values. */ function parseOptionalString(value: unknown) { - return typeof value === "string" && value.length > 0 ? value : undefined; + if (value === undefined) return undefined; + return parseBrokerString(value); } -/** Parse one required non-negative integer field from the websocket payload. */ +/** Parse one required non-negative safe integer field. */ function parseNonNegativeInt(value: unknown) { - return typeof value === "number" && Number.isInteger(value) && value >= 0 ? value : null; -} - -/** Parse one required positive integer field from the websocket payload. */ -function parsePositiveInt(value: unknown) { - return typeof value === "number" && Number.isInteger(value) && value > 0 ? value : null; -} - -/** Parse one terminal location entry, skipping malformed optional metadata. */ -function parseSessionTerminalLocation(value: unknown): SessionTerminalLocation | null { - const record = asRecord(value); - if (!record) { + try { + return parseBrokerSafeInteger(value); + } catch { return null; } +} - const source = parseRequiredString(record.source); - if (source === null) { +/** Parse one required positive safe integer field. */ +function parsePositiveInt(value: unknown) { + try { + return parseBrokerSafeInteger(value, { minimum: 1 }); + } catch { return null; } +} +/** Parse one terminal location with exact keys and strict optional fields. */ +function parseSessionTerminalLocation(value: unknown): SessionTerminalLocation { + const record = parseExactBrokerRecord( + value, + ["source"] as const, + ["tty", "windowId", "tabId", "paneId", "terminalId", "sessionId"] as const, + ); return { - source, - tty: parseOptionalString(record.tty), - windowId: parseOptionalString(record.windowId), - tabId: parseOptionalString(record.tabId), - paneId: parseOptionalString(record.paneId), - terminalId: parseOptionalString(record.terminalId), - sessionId: parseOptionalString(record.sessionId), + source: parseBrokerString(record.source), + ...(record.tty === undefined ? {} : { tty: parseBrokerString(record.tty) }), + ...(record.windowId === undefined ? {} : { windowId: parseBrokerString(record.windowId) }), + ...(record.tabId === undefined ? {} : { tabId: parseBrokerString(record.tabId) }), + ...(record.paneId === undefined ? {} : { paneId: parseBrokerString(record.paneId) }), + ...(record.terminalId === undefined + ? {} + : { terminalId: parseBrokerString(record.terminalId) }), + ...(record.sessionId === undefined + ? {} + : { sessionId: parseBrokerIdentifier(record.sessionId) }), }; } -/** Parse terminal metadata while tolerating malformed optional location detail. */ -function parseSessionTerminalMetadata(value: unknown): SessionTerminalMetadata | undefined { - const record = asRecord(value); - if (!record || !Array.isArray(record.locations)) { - return undefined; - } - - const locations = record.locations - .map(parseSessionTerminalLocation) - .filter((location): location is SessionTerminalLocation => location !== null); - +/** Parse terminal metadata with exact keys and no partial location filtering. */ +function parseSessionTerminalMetadata(value: unknown): SessionTerminalMetadata { + const record = parseExactBrokerRecord(value, ["locations"] as const, ["program"] as const); + if (!Array.isArray(record.locations)) throw new BrokerProtocolError("invalid-field"); return { - program: parseOptionalString(record.program), - locations, + ...(record.program === undefined ? {} : { program: parseBrokerString(record.program) }), + locations: record.locations.map(parseSessionTerminalLocation), }; } @@ -80,38 +95,29 @@ export function parseSessionRegistrationEnvelope( value: unknown, parseInfo: (value: unknown) => Info | null, ): SessionRegistration | null { - const record = asRecord(value); - if (!record) { + try { + const record = parseExactBrokerRecord( + value, + ["registrationVersion", "sessionId", "pid", "cwd", "launchedAt", "info"] as const, + ["repoRoot", "terminal"] as const, + ); + const registrationVersion = parseBrokerSafeInteger(record.registrationVersion, { minimum: 1 }); + if (registrationVersion !== SESSION_BROKER_REGISTRATION_VERSION) return null; + return { + registrationVersion, + sessionId: parseBrokerIdentifier(record.sessionId), + pid: parseBrokerSafeInteger(record.pid, { minimum: 1 }), + cwd: parseBrokerString(record.cwd), + ...(record.repoRoot === undefined ? {} : { repoRoot: parseBrokerString(record.repoRoot) }), + launchedAt: parseBrokerString(record.launchedAt), + ...(record.terminal === undefined + ? {} + : { terminal: parseSessionTerminalMetadata(record.terminal) }), + info: parseBrokerAppPayload(parseInfo, record.info), + }; + } catch { return null; } - - const registrationVersion = parsePositiveInt(record.registrationVersion); - const sessionId = parseRequiredString(record.sessionId); - const pid = parsePositiveInt(record.pid); - const cwd = parseRequiredString(record.cwd); - const launchedAt = parseRequiredString(record.launchedAt); - const info = parseInfo(record.info); - if ( - registrationVersion !== SESSION_BROKER_REGISTRATION_VERSION || - sessionId === null || - pid === null || - cwd === null || - launchedAt === null || - info === null - ) { - return null; - } - - return { - registrationVersion, - sessionId, - pid, - cwd, - repoRoot: parseOptionalString(record.repoRoot), - launchedAt, - terminal: parseSessionTerminalMetadata(record.terminal), - info, - }; } /** Parse one broker snapshot envelope and delegate app-owned state parsing to the caller. */ @@ -119,21 +125,15 @@ export function parseSessionSnapshotEnvelope( value: unknown, parseState: (value: unknown) => State | null, ): SessionSnapshot | null { - const record = asRecord(value); - if (!record) { + try { + const record = parseExactBrokerRecord(value, ["updatedAt", "state"] as const); + return { + updatedAt: parseBrokerString(record.updatedAt), + state: parseBrokerAppPayload(parseState, record.state), + }; + } catch { return null; } - - const updatedAt = parseRequiredString(record.updatedAt); - const state = parseState(record.state); - if (updatedAt === null || state === null) { - return null; - } - - return { - updatedAt, - state, - }; } export const brokerWireParsers = { diff --git a/packages/session-broker-core/src/index.ts b/packages/session-broker-core/src/index.ts index c7f918b59..6ce0d1780 100644 --- a/packages/session-broker-core/src/index.ts +++ b/packages/session-broker-core/src/index.ts @@ -1,6 +1,7 @@ export * from "./types"; export * from "./canonicalJson"; export * from "./auth"; +export * from "./validation"; export * from "./brokerWire"; export * from "./limits"; export * from "./brokerState"; diff --git a/packages/session-broker-core/src/validation.test.ts b/packages/session-broker-core/src/validation.test.ts new file mode 100644 index 000000000..acbf855f6 --- /dev/null +++ b/packages/session-broker-core/src/validation.test.ts @@ -0,0 +1,67 @@ +import { describe, expect, test } from "bun:test"; +import { + BrokerProtocolError, + parseBrokerAppId, + parseBrokerAppPayload, + parseBrokerDeadline, + parseBrokerIdentifier, + parseBrokerRevision, + parseBrokerSelector, + parseBrokerTimeout, + parseBrokerUint64, + parseExactBrokerRecord, +} from "./validation"; + +function code(task: () => unknown) { + try { + task(); + return null; + } catch (error) { + return error instanceof BrokerProtocolError ? error.code : "unexpected"; + } +} + +describe("session broker runtime validation", () => { + test("requires exact plain records and strict optional fields", () => { + for (const value of [null, [], "record", { required: 1, extra: true }]) { + expect(code(() => parseExactBrokerRecord(value, ["required"] as const))).not.toBeNull(); + } + expect(parseExactBrokerRecord({ required: 1 }, ["required"] as const)).toEqual({ required: 1 }); + expect(code(() => parseBrokerSelector({ sessionId: "session-1", extra: true }))).toBe( + "invalid-keys", + ); + expect(code(() => parseBrokerSelector({ repoRoot: null }))).toBe("invalid-field"); + }); + + test("bounds identifiers, revisions, uint64 values, and deadlines", () => { + expect(parseBrokerAppId("dev.hunk")).toBe("dev.hunk"); + expect(parseBrokerIdentifier("session-1")).toBe("session-1"); + for (const value of [0, -1, 1.5, Number.MAX_SAFE_INTEGER + 1, NaN]) { + expect(code(() => parseBrokerRevision(value))).toBe("invalid-contract"); + } + expect(parseBrokerUint64("18446744073709551615")).toBe("18446744073709551615"); + expect(code(() => parseBrokerUint64("01"))).toBe("invalid-field"); + expect(parseBrokerTimeout(300_000)).toBe(300_000); + expect(code(() => parseBrokerTimeout(300_001))).toBe("invalid-deadline"); + expect(code(() => parseBrokerDeadline(Infinity))).toBe("invalid-deadline"); + }); + + test("redacts app parser rejection and exceptions", () => { + expect(code(() => parseBrokerAppPayload(() => null, {}))).toBe("invalid-app-payload"); + expect( + code(() => + parseBrokerAppPayload(() => { + throw new Error("secret parser internals"); + }, {}), + ), + ).toBe("app-parser-failed"); + try { + parseBrokerAppPayload(() => { + throw new Error("secret parser internals"); + }, {}); + } catch (error) { + expect(String(error)).not.toContain("secret parser internals"); + expect((error as Error).stack).not.toContain("secret parser internals"); + } + }); +}); diff --git a/packages/session-broker-core/src/validation.ts b/packages/session-broker-core/src/validation.ts new file mode 100644 index 000000000..9cf51b87c --- /dev/null +++ b/packages/session-broker-core/src/validation.ts @@ -0,0 +1,166 @@ +import { + isValidBrokerAppId, + isValidBrokerIdentifier, + isValidBrokerRevision, + parseCallerSequence, +} from "./auth"; +import type { SessionTargetInput } from "./types"; + +export const MAX_BROKER_STRING_BYTES = 4_096; +export const MAX_BROKER_ERROR_BYTES = 1_024; +export const MAX_BROKER_DEADLINE_MS = 5 * 60_000; + +export type BrokerProtocolFailureCode = + | "invalid-json" + | "invalid-record" + | "invalid-keys" + | "invalid-discriminant" + | "invalid-field" + | "invalid-selector" + | "invalid-deadline" + | "invalid-contract" + | "unknown-command" + | "invalid-app-payload" + | "app-parser-failed"; + +/** Reports one stable protocol failure without reflecting parser messages or attacker payloads. */ +export class BrokerProtocolError extends Error { + constructor(readonly code: BrokerProtocolFailureCode) { + super("Session broker protocol validation failed."); + this.name = "BrokerProtocolError"; + } +} + +/** Throw one stable protocol failure code. */ +export function failBrokerProtocol(code: BrokerProtocolFailureCode): never { + throw new BrokerProtocolError(code); +} + +/** Return an object record while rejecting null, arrays, and exotic prototypes. */ +export function parseBrokerRecord(value: unknown): Record { + if (value === null || typeof value !== "object" || Array.isArray(value)) { + return failBrokerProtocol("invalid-record"); + } + const prototype = Object.getPrototypeOf(value); + if (prototype !== Object.prototype && prototype !== null) { + return failBrokerProtocol("invalid-record"); + } + return value as Record; +} + +/** Require an exact key set with separately declared required and optional keys. */ +export function parseExactBrokerRecord( + value: unknown, + required: readonly Required[], + optional: readonly Optional[] = [], +): Record & Partial> { + const record = parseBrokerRecord(value); + const allowed = new Set([...required, ...optional]); + const keys = Object.keys(record); + if ( + required.some((key) => !Object.hasOwn(record, key)) || + keys.some((key) => !allowed.has(key)) + ) { + return failBrokerProtocol("invalid-keys"); + } + return record as Record & Partial>; +} + +/** Parse a bounded string by UTF-8 byte length. */ +export function parseBrokerString( + value: unknown, + options: { minBytes?: number; maxBytes?: number } = {}, +): string { + if (typeof value !== "string") return failBrokerProtocol("invalid-field"); + const bytes = new TextEncoder().encode(value).byteLength; + if (bytes < (options.minBytes ?? 1) || bytes > (options.maxBytes ?? MAX_BROKER_STRING_BYTES)) { + return failBrokerProtocol("invalid-field"); + } + return value; +} + +/** Parse one bounded broker identifier. */ +export function parseBrokerIdentifier(value: unknown): string { + if (!isValidBrokerIdentifier(value)) return failBrokerProtocol("invalid-field"); + return value; +} + +/** Parse one immutable application identifier. */ +export function parseBrokerAppId(value: unknown): string { + if (!isValidBrokerAppId(value)) return failBrokerProtocol("invalid-field"); + return value; +} + +/** Parse one non-negative or positive safe integer. */ +export function parseBrokerSafeInteger( + value: unknown, + options: { minimum?: number; maximum?: number } = {}, +): number { + const minimum = options.minimum ?? 0; + const maximum = options.maximum ?? Number.MAX_SAFE_INTEGER; + if (!Number.isSafeInteger(value) || (value as number) < minimum || (value as number) > maximum) { + return failBrokerProtocol("invalid-field"); + } + return value as number; +} + +/** Parse one positive protocol or command revision. */ +export function parseBrokerRevision(value: unknown): number { + if (!isValidBrokerRevision(value)) return failBrokerProtocol("invalid-contract"); + return value; +} + +/** Parse one canonical uint64 decimal value. */ +export function parseBrokerUint64(value: unknown, options: { allowZero?: boolean } = {}): string { + if (typeof value !== "string") return failBrokerProtocol("invalid-field"); + const parsed = parseCallerSequence(value); + if (parsed === null || (!options.allowZero && parsed === 0n)) { + return failBrokerProtocol("invalid-field"); + } + return value; +} + +/** Parse a strict generic session selector without app-owned selector semantics. */ +export function parseBrokerSelector(value: unknown): SessionTargetInput { + const record = parseExactBrokerRecord(value, [], [ + "sessionId", + "sessionPath", + "repoRoot", + "repoBoundary", + ] as const); + const selector: SessionTargetInput = {}; + if (record.sessionId !== undefined) selector.sessionId = parseBrokerIdentifier(record.sessionId); + for (const key of ["sessionPath", "repoRoot", "repoBoundary"] as const) { + if (record[key] !== undefined) selector[key] = parseBrokerString(record[key]); + } + return selector; +} + +/** Parse a positive relative timeout bounded by the public caller maximum. */ +export function parseBrokerTimeout(value: unknown): number { + try { + return parseBrokerSafeInteger(value, { minimum: 1, maximum: MAX_BROKER_DEADLINE_MS }); + } catch { + return failBrokerProtocol("invalid-deadline"); + } +} + +/** Parse a finite absolute deadline timestamp. */ +export function parseBrokerDeadline(value: unknown): number { + if (typeof value !== "number" || !Number.isSafeInteger(value) || value <= 0) { + return failBrokerProtocol("invalid-deadline"); + } + return value; +} + +/** Invoke an app parser and normalize null returns and throws into redacted failures. */ +export function parseBrokerAppPayload(parser: (value: unknown) => T | null, value: unknown): T { + let parsed: T | null; + try { + parsed = parser(value); + } catch { + return failBrokerProtocol("app-parser-failed"); + } + if (parsed === null || parsed === undefined) return failBrokerProtocol("invalid-app-payload"); + return parsed; +} diff --git a/packages/session-broker-node/src/serve.test.ts b/packages/session-broker-node/src/serve.test.ts index ec50e6ca1..bbcdce627 100644 --- a/packages/session-broker-node/src/serve.test.ts +++ b/packages/session-broker-node/src/serve.test.ts @@ -8,7 +8,11 @@ import { type SessionRegistration, type SessionSnapshot, } from "@hunk/session-broker-core"; -import { SessionBroker, createSessionBrokerDaemon } from "@hunk/session-broker"; +import { + SessionBroker, + createSessionBrokerDaemon, + createSessionBrokerProtocolParsers, +} from "@hunk/session-broker"; import { serveSessionBrokerDaemon } from "./serve"; interface TestSessionInfo { @@ -53,7 +57,9 @@ function createRegistration(overrides: Partial["state"]> & { updatedAt?: string } = {}, + overrides: Partial["state"]> & { + updatedAt?: string; + } = {}, ) { const { updatedAt = "2026-04-15T00:00:00.000Z", ...stateOverrides } = overrides; return { @@ -65,6 +71,14 @@ function createSnapshot( } satisfies SessionSnapshot; } +const protocolParsers = createSessionBrokerProtocolParsers({ + appRevision: 1, + features: [], + parseRegistration: (value) => parseSessionRegistrationEnvelope(value, parseInfo), + parseSnapshot: (value) => parseSessionSnapshotEnvelope(value, parseState), + commands: [], +}); + async function reserveLoopbackPort() { const listener = createServer(() => undefined); await new Promise((resolve, reject) => { @@ -102,10 +116,7 @@ async function waitUntil( describe("session broker node adapter", () => { test("serves the generic daemon API and websocket path through Node", async () => { - const broker = new SessionBroker({ - parseRegistration: (value) => parseSessionRegistrationEnvelope(value, parseInfo), - parseSnapshot: (value) => parseSessionSnapshotEnvelope(value, parseState), - }); + const broker = new SessionBroker({ protocolParsers }); const daemon = createSessionBrokerDaemon({ broker, capabilities: { version: 1 }, @@ -202,7 +213,10 @@ describe("session broker node adapter", () => { const response = await fetch(`http://127.0.0.1:${port}/broker`, { method: "POST", headers: { "content-type": "application/json" }, - body: JSON.stringify({ action: "get", selector: { sessionId: "session-1" } }), + body: JSON.stringify({ + action: "get", + selector: { sessionId: "session-1" }, + }), }); await expect(response.json()).resolves.toMatchObject({ body: { diff --git a/packages/session-broker/README.md b/packages/session-broker/README.md index 0c82d7680..9223a0dea 100644 --- a/packages/session-broker/README.md +++ b/packages/session-broker/README.md @@ -61,6 +61,7 @@ internal workspace usage only. import { SessionBroker, brokerWireParsers, + createSessionBrokerProtocolParsers, parseSessionRegistrationEnvelope, parseSessionSnapshotEnvelope, } from "@hunk/session-broker"; @@ -93,10 +94,22 @@ function parseState(value: unknown): SessionState | null { return selectedIndex === null ? null : { selectedIndex }; } -const broker = new SessionBroker({ +const protocolParsers = createSessionBrokerProtocolParsers({ + appRevision: 1, + features: [], parseRegistration: (value) => parseSessionRegistrationEnvelope(value, parseInfo), parseSnapshot: (value) => parseSessionSnapshotEnvelope(value, parseState), + commands: [ + { + command: "select", + version: 1, + parseInput: (value) => (brokerWireParsers.parseNonNegativeInt(value) === null ? null : value), + parseResult: (value) => (value === true ? true : null), + }, + ], }); + +const broker = new SessionBroker({ protocolParsers }); ``` ### 2. Create a daemon engine @@ -165,9 +178,12 @@ const connection = createSessionBrokerConnection({ createSocket: (url) => new WebSocket(url), registration, snapshot, + protocolParsers, bridge: { dispatchCommand: async (message) => { - return handleCommand(message); + if (message.command !== "select") throw new Error("Unsupported command."); + selectFile(message.input); + return true; }, }, }); diff --git a/packages/session-broker/src/authentication.test.ts b/packages/session-broker/src/authentication.test.ts index 0c42ab250..3ca2d550c 100644 --- a/packages/session-broker/src/authentication.test.ts +++ b/packages/session-broker/src/authentication.test.ts @@ -528,13 +528,17 @@ describe("session broker signed authentication", () => { const values = await setup(); const valid = challengeRequest(); for (const malformed of [ + null, + [], + { ...valid, extra: true }, { ...valid, initiatorNonce: "bad nonce" }, { ...valid, endpoint: "http://user@127.0.0.1/broker" }, { ...valid, proposal: { ...valid.proposal, appRevision: 2 } }, { ...valid, proposal: { ...valid.proposal, features: ["unexpected.feature"] } }, + { ...valid, proposal: { ...valid.proposal, extra: true } }, ]) { await expect( - values.authenticator.issueChallenge(malformed, malformed.endpoint), + values.authenticator.issueChallenge(malformed, valid.endpoint), ).rejects.toMatchObject({ code: "invalid-credential", }); diff --git a/packages/session-broker/src/authentication.ts b/packages/session-broker/src/authentication.ts index c9399fa26..dc0c3750e 100644 --- a/packages/session-broker/src/authentication.ts +++ b/packages/session-broker/src/authentication.ts @@ -14,6 +14,9 @@ import { isValidBrokerIdentifier, isValidBrokerRevision, principalFromGrant, + parseBrokerIdentifier, + parseBrokerString, + parseExactBrokerRecord, type BrokerAppContract, type BrokerChallengeTranscriptInput, type BrokerGrant, @@ -514,7 +517,7 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { /** Issue one bounded, expiring challenge signed by the daemon identity. */ async issueChallenge( - request: SessionBrokerHelloChallengeRequest, + request: unknown, listenerEndpoint: string, ): Promise { this.pruneExpired(); @@ -572,7 +575,8 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { } /** Consume one caller proof and issue a short-lived replay-protected caller session. */ - async completeCallerHello(proof: SessionBrokerHelloProof): Promise { + async completeCallerHello(proofInput: unknown): Promise { + const proof = this.parseHelloProof(proofInput); const pending = this.takeChallenge(proof.challengeId, "caller"); await this.verifyProof(pending, proof.signature); const grant = pending.grant as CallerGrant; @@ -651,9 +655,10 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { /** Consume one producer proof and sign the connection binding supplied by the adapter. */ async completeProducerHello( - proof: SessionBrokerHelloProof, - connectionId: string, + proofInput: unknown, + connectionId: unknown, ): Promise { + const proof = this.parseHelloProof(proofInput); if (!isValidBrokerIdentifier(connectionId)) authenticationError("invalid-credential"); const pending = this.takeChallenge(proof.challengeId, "producer"); await this.verifyProof(pending, proof.signature); @@ -820,37 +825,63 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { } private validateHello( - request: SessionBrokerHelloChallengeRequest, + value: unknown, listenerEndpoint: string, ): SessionBrokerHelloChallengeRequest { - if ( - !request || - typeof request !== "object" || - (request.role !== "producer" && request.role !== "caller") || - request.appId !== this.config.appId || - !isValidBrokerIdentifier(request.keyId) || - !isValidBrokerIdentifier(request.grantId) || - !isValidBrokerIdentifier(request.initiatorNonce) || - request.endpoint !== listenerEndpoint || - !parseEndpoint(request.endpoint) || - !parseEndpoint(listenerEndpoint) || - !request.proposal || - request.proposal.brokerRevision !== SESSION_BROKER_PROTOCOL_REVISION || - request.proposal.appRevision !== this.config.appRevision || - !Array.isArray(request.proposal.features) || - request.proposal.features.length !== 0 - ) { - authenticationError("invalid-credential"); + try { + const request = parseExactBrokerRecord(value, [ + "role", + "appId", + "endpoint", + "keyId", + "grantId", + "initiatorNonce", + "proposal", + ] as const); + const proposal = parseExactBrokerRecord(request.proposal, [ + "brokerRevision", + "appRevision", + "features", + ] as const); + if ( + (request.role !== "producer" && request.role !== "caller") || + request.appId !== this.config.appId || + request.endpoint !== listenerEndpoint || + typeof request.endpoint !== "string" || + !parseEndpoint(request.endpoint) || + !parseEndpoint(listenerEndpoint) || + proposal.brokerRevision !== SESSION_BROKER_PROTOCOL_REVISION || + proposal.appRevision !== this.config.appRevision || + !Array.isArray(proposal.features) || + proposal.features.length !== 0 + ) { + authenticationError("invalid-credential"); + } + return Object.freeze({ + role: request.role, + appId: this.config.appId, + endpoint: request.endpoint, + keyId: parseBrokerIdentifier(request.keyId), + grantId: parseBrokerIdentifier(request.grantId), + initiatorNonce: parseBrokerIdentifier(request.initiatorNonce), + proposal: fixedProposal(this.config.appRevision), + }); + } catch { + return authenticationError("invalid-credential"); + } + } + + /** Parse one exact proof envelope before consuming challenge state. */ + private parseHelloProof(value: unknown): SessionBrokerHelloProof { + try { + const proof = parseExactBrokerRecord(value, ["challengeId", "signature"] as const); + return { + challengeId: parseBrokerIdentifier(proof.challengeId), + signature: parseBrokerString(proof.signature, { maxBytes: 1_024 }), + }; + } catch { + return authenticationError("invalid-credential"); } - return Object.freeze({ - role: request.role, - appId: this.config.appId, - endpoint: request.endpoint, - keyId: request.keyId, - grantId: request.grantId, - initiatorNonce: request.initiatorNonce, - proposal: fixedProposal(this.config.appRevision), - }); } private takeChallenge(challengeId: string, role: BrokerGrant["kind"]): PendingChallenge { diff --git a/packages/session-broker/src/broker.test.ts b/packages/session-broker/src/broker.test.ts index 0e637e18d..d1612fa90 100644 --- a/packages/session-broker/src/broker.test.ts +++ b/packages/session-broker/src/broker.test.ts @@ -9,6 +9,7 @@ import { type SessionSnapshot, } from "@hunk/session-broker-core"; import { SessionBroker } from "./broker"; +import { createSessionBrokerProtocolParsers } from "./protocolParsers"; interface TestSessionInfo { title: string; @@ -57,15 +58,34 @@ function parseState(value: unknown): TestSessionState | null { return { selectedIndex, noteCount }; } +const protocolParsers = createSessionBrokerProtocolParsers< + TestSessionInfo, + TestSessionState, + TestServerMessage, + unknown +>({ + appRevision: 1, + features: [], + parseRegistration: (value) => parseSessionRegistrationEnvelope(value, parseInfo), + parseSnapshot: (value) => parseSessionSnapshotEnvelope(value, parseState), + commands: ["annotate", "reload_view"].map((command) => ({ + command: command as TestServerMessage["command"], + version: 1, + parseInput: (value: unknown) => (value && typeof value === "object" ? value : null), + parseResult: (value: unknown) => (value && typeof value === "object" ? value : null), + })), +}); + function createBroker() { return new SessionBroker({ - parseRegistration: (value) => parseSessionRegistrationEnvelope(value, parseInfo), - parseSnapshot: (value) => parseSessionSnapshotEnvelope(value, parseState), + protocolParsers, }); } function createRegistration( - overrides: Partial & { info?: Partial } = {}, + overrides: Partial & { + info?: Partial; + } = {}, ): TestRegistration { return { registrationVersion: SESSION_BROKER_REGISTRATION_VERSION, @@ -99,6 +119,10 @@ function createSnapshot( } describe("session broker wrapper", () => { + test("exposes the exact parser registry that owns its state contracts", () => { + expect(createBroker().protocolParsers).toBe(protocolParsers); + }); + test("stores raw registrations and snapshots without a custom projection adapter", () => { const broker = createBroker(); const connection = { send() {} }; @@ -156,7 +180,10 @@ describe("session broker wrapper", () => { timeoutMessage: "Timed out waiting for annotate.", }); - const outgoing = JSON.parse(sent[0]!) as { requestId: string; command: string }; + const outgoing = JSON.parse(sent[0]!) as { + requestId: string; + command: string; + }; expect(outgoing.command).toBe("annotate"); broker.handleCommandResult(connection, { diff --git a/packages/session-broker/src/broker.ts b/packages/session-broker/src/broker.ts index 058c004d8..c85798407 100644 --- a/packages/session-broker/src/broker.ts +++ b/packages/session-broker/src/broker.ts @@ -11,6 +11,7 @@ import { type SessionTargetSelector, type UpdateSnapshotResult, } from "@hunk/session-broker-core"; +import type { SessionBrokerProtocolParsers } from "./protocolParsers"; /** Minimal socket shape the broker needs in order to target one live session. */ export interface SessionBrokerPeer { @@ -30,9 +31,13 @@ export interface SessionBrokerRecord { snapshot: SessionSnapshot; } -export interface SessionBrokerOptions { - parseRegistration: (value: unknown) => SessionRegistration | null; - parseSnapshot: (value: unknown) => SessionSnapshot | null; +export interface SessionBrokerOptions< + Info, + State, + ServerMessage extends SessionServerMessage = SessionServerMessage, + CommandResult = unknown, +> { + protocolParsers: SessionBrokerProtocolParsers; describeSession?: ( registration: SessionRegistration, snapshot: SessionSnapshot, @@ -45,6 +50,13 @@ export interface SessionBrokerController< ServerMessage extends SessionServerMessage = SessionServerMessage, CommandResult = unknown, > { + /** The parser registry bound to this controller's state and command contracts. */ + readonly protocolParsers: SessionBrokerProtocolParsers< + unknown, + unknown, + ServerMessage, + CommandResult + >; listSessions(): SessionView[]; getSession(selector: SessionTargetSelector): SessionView; getSessionCount(): number; @@ -108,6 +120,8 @@ export class SessionBroker< ServerMessage, CommandResult > { + readonly protocolParsers: SessionBrokerProtocolParsers; + private readonly state: SessionBrokerState< Info, State, @@ -120,16 +134,33 @@ export class SessionBroker< >; private readonly describeSession: NonNullable< - SessionBrokerOptions["describeSession"] + SessionBrokerOptions["describeSession"] >; - constructor(options: SessionBrokerOptions) { + constructor(options: SessionBrokerOptions) { this.describeSession = options.describeSession ?? ((registration, _snapshot) => defaultSessionTitle(registration)); + this.protocolParsers = options.protocolParsers; this.state = new SessionBrokerState({ - parseRegistration: options.parseRegistration, - parseSnapshot: options.parseSnapshot, + parseRegistration: (value) => { + try { + return this.protocolParsers.parseRegistration(value); + } catch { + return null; + } + }, + parseSnapshot: (value) => { + try { + return this.protocolParsers.parseSnapshot(value); + } catch { + return null; + } + }, + parseCommandInput: (command, version, value) => + this.protocolParsers.parseCommandInput(command, version, value), + parseCommandResult: (command, version, value) => + this.protocolParsers.parseCommandResult(command, version, value), buildListedSession: (entry) => this.buildRecord(entry), buildSelectedContext: (session) => session, buildSessionReview: (entry) => this.buildRecord(entry), diff --git a/packages/session-broker/src/connection.test.ts b/packages/session-broker/src/connection.test.ts index 0cf19653d..69193e198 100644 --- a/packages/session-broker/src/connection.test.ts +++ b/packages/session-broker/src/connection.test.ts @@ -6,6 +6,7 @@ import type { } from "@hunk/session-broker-core"; import { SESSION_BROKER_REGISTRATION_VERSION } from "@hunk/session-broker-core"; import { createSessionBrokerConnection } from "./connection"; +import { createSessionBrokerProtocolParsers } from "./protocolParsers"; import type { SessionBrokerSocketLike } from "./types"; interface TestSessionInfo { @@ -22,6 +23,7 @@ class TestSocket implements SessionBrokerSocketLike { readyState = 0; sent: string[] = []; throwOnSend = false; + lastClose: { code?: number; reason?: string } | null = null; onopen: (() => void) | null = null; onmessage: ((event: { data: unknown }) => void) | null = null; onclose: ((event: { code: number; reason: string }) => void) | null = null; @@ -32,8 +34,9 @@ class TestSocket implements SessionBrokerSocketLike { this.sent.push(data); } - close() { - this.emitClose(); + close(code?: number, reason?: string) { + this.lastClose = { code, reason }; + this.emitClose(code, reason); } emitOpen() { @@ -69,6 +72,38 @@ function createSnapshot(): SessionSnapshot { }; } +const protocolParsers = createSessionBrokerProtocolParsers< + TestSessionInfo, + TestSessionState, + TestServerMessage, + { ok: true } +>({ + appRevision: 1, + features: [], + parseRegistration: (value) => + value && typeof value === "object" ? (value as SessionRegistration) : null, + parseSnapshot: (value) => + value && typeof value === "object" ? (value as SessionSnapshot) : null, + commands: [ + { + command: "annotate", + version: 1, + parseInput: (value) => { + if (!value || typeof value !== "object" || Array.isArray(value)) return null; + const record = value as Record; + return Object.keys(record).length === 1 && typeof record.summary === "string" + ? { summary: record.summary } + : null; + }, + parseResult: (value) => { + if (!value || typeof value !== "object" || Array.isArray(value)) return null; + const record = value as Record; + return Object.keys(record).length === 1 && record.ok === true ? { ok: true } : null; + }, + }, + ], +}); + describe("session broker connection", () => { test("registers on open and sends later snapshot updates", () => { const sockets: TestSocket[] = []; @@ -87,12 +122,15 @@ describe("session broker connection", () => { }, registration: createRegistration(), snapshot: createSnapshot(), + protocolParsers, }); connection.start(); sockets[0]?.emitOpen(); - const registerMessage = JSON.parse(sockets[0]!.sent[0]!) as { type: string }; + const registerMessage = JSON.parse(sockets[0]!.sent[0]!) as { + type: string; + }; expect(registerMessage.type).toBe("register"); connection.updateSnapshot({ @@ -100,7 +138,10 @@ describe("session broker connection", () => { state: { selectedIndex: 1 }, }); - const snapshotMessage = JSON.parse(sockets[0]!.sent[1]!) as { type: string; snapshot: unknown }; + const snapshotMessage = JSON.parse(sockets[0]!.sent[1]!) as { + type: string; + snapshot: unknown; + }; expect(snapshotMessage.type).toBe("snapshot"); expect(snapshotMessage.snapshot).toEqual({ updatedAt: "2026-04-15T00:00:01.000Z", @@ -122,6 +163,7 @@ describe("session broker connection", () => { createSocket: () => socket, registration, snapshot: createSnapshot(), + protocolParsers, }); connection.start(); socket.emitOpen(); @@ -150,6 +192,7 @@ describe("session broker connection", () => { createSocket: () => socket, registration: createRegistration(), snapshot: createSnapshot(), + protocolParsers, }); connection.start(); @@ -175,6 +218,131 @@ describe("session broker connection", () => { expect(resultMessage).toMatchObject({ type: "command-result", ok: true }); }); + test("parses and transforms each server command once before bridge dispatch", async () => { + let inputCalls = 0; + let resultCalls = 0; + const transformingParsers = createSessionBrokerProtocolParsers< + TestSessionInfo, + TestSessionState, + TestServerMessage, + { ok: true } + >({ + appRevision: 1, + features: [], + parseRegistration: protocolParsers.parseRegistration.bind(protocolParsers), + parseSnapshot: protocolParsers.parseSnapshot.bind(protocolParsers), + commands: [ + { + command: "annotate", + version: 1, + parseInput: (value) => { + inputCalls += 1; + const summary = (value as { summary?: unknown })?.summary; + return typeof summary === "string" ? { summary: summary.toUpperCase() } : null; + }, + parseResult: (value) => { + resultCalls += 1; + return (value as { ok?: unknown })?.ok === true ? { ok: true } : null; + }, + }, + ], + }); + const socket = new TestSocket(); + let bridgedInput: unknown; + const connection = createSessionBrokerConnection< + TestSessionInfo, + TestSessionState, + TestSocket, + TestServerMessage, + { ok: true } + >({ + url: "ws://broker.test/session", + createSocket: () => socket, + registration: createRegistration(), + snapshot: createSnapshot(), + protocolParsers: transformingParsers, + bridge: { + dispatchCommand: async (message) => { + bridgedInput = message.input; + return { ok: true }; + }, + }, + }); + + connection.start(); + socket.emitOpen(); + socket.emitMessage( + JSON.stringify({ + type: "command", + requestId: "request-1", + command: "annotate", + input: { summary: "review note" }, + }), + ); + await Bun.sleep(0); + + expect(bridgedInput).toEqual({ summary: "REVIEW NOTE" }); + expect({ inputCalls, resultCalls }).toEqual({ + inputCalls: 1, + resultCalls: 1, + }); + connection.stop(); + }); + + test("closes malformed and mismatched commands without invoking the bridge", async () => { + const socket = new TestSocket(); + let dispatched = 0; + const connection = createSessionBrokerConnection< + TestSessionInfo, + TestSessionState, + TestSocket, + TestServerMessage, + { ok: true } + >({ + url: "ws://broker.test/session", + createSocket: () => socket, + registration: createRegistration(), + snapshot: createSnapshot(), + protocolParsers, + bridge: { + dispatchCommand: async () => { + dispatched += 1; + return { ok: true }; + }, + }, + reconnectDelayMs: 1_000, + }); + connection.start(); + socket.emitOpen(); + + for (const message of [ + null, + [], + { + type: "command", + requestId: "request-1", + command: "unknown", + input: {}, + }, + { + type: "command", + requestId: "request-1", + command: "annotate", + input: { summary: "note", extra: true }, + }, + ]) { + socket.readyState = 1; + socket.emitMessage(JSON.stringify(message)); + expect(socket.lastClose).toEqual({ + code: 1008, + reason: "Malformed session broker command.", + }); + } + await Bun.sleep(0); + expect(dispatched).toBe(0); + connection.stop(); + }); + test("does not migrate a late command result onto a replacement socket", async () => { const sockets: TestSocket[] = []; let resolveCommand!: (result: { ok: true }) => void; @@ -196,6 +364,7 @@ describe("session broker connection", () => { }, registration: createRegistration(), snapshot: createSnapshot(), + protocolParsers, bridge: { dispatchCommand: () => commandResult }, reconnectDelayMs: 1, }); @@ -240,6 +409,7 @@ describe("session broker connection", () => { }, registration: createRegistration(), snapshot: createSnapshot(), + protocolParsers, reconnectDelayMs: 1, }); @@ -288,6 +458,7 @@ describe("session broker connection", () => { createSocket: () => socket, registration: createRegistration(), snapshot: createSnapshot(), + protocolParsers, reconnectDelayMs: 1_000, }); @@ -339,6 +510,7 @@ describe("session broker connection", () => { }, registration: createRegistration(), snapshot: createSnapshot(), + protocolParsers, reconnectDelayMs: 5, resolveClose: (event) => event.reason === "stop" diff --git a/packages/session-broker/src/connection.ts b/packages/session-broker/src/connection.ts index 4eab9db2e..e0772b4ab 100644 --- a/packages/session-broker/src/connection.ts +++ b/packages/session-broker/src/connection.ts @@ -1,9 +1,12 @@ -import type { - SessionClientMessage, - SessionRegistration, - SessionServerMessage, - SessionSnapshot, +import { + BrokerProtocolError, + type SessionClientMessage, + type SessionRegistration, + type SessionServerMessage, + type SessionSnapshot, } from "@hunk/session-broker-core"; +import type { SessionBrokerProtocolParsers } from "./protocolParsers"; +import { parseSessionBrokerJsonText } from "./protocolParsers"; import type { SessionBrokerConnectionCloseDirective, SessionBrokerSocketCloseEvent, @@ -33,6 +36,7 @@ export interface SessionBrokerConnectionOptions< registration: SessionRegistration; snapshot: SessionSnapshot; bridge?: SessionBrokerConnectionBridge | null; + protocolParsers: SessionBrokerProtocolParsers; heartbeatIntervalMs?: number; reconnectDelayMs?: number; openState?: number; @@ -145,14 +149,15 @@ export class SessionBrokerConnection< }; socket.onmessage = (event) => { - if (typeof event.data !== "string") { - return; - } - let parsed: ServerMessage; try { - parsed = JSON.parse(event.data) as ServerMessage; + parsed = this.options.protocolParsers.parseServerMessage( + parseSessionBrokerJsonText(event.data), + ); } catch { + // Never invoke the app bridge after a malformed or mismatched command contract. Closing + // prevents this producer from retaining daemon assumptions that were not actually parsed. + socket.close(1008, "Malformed session broker command."); return; } @@ -254,13 +259,24 @@ export class SessionBrokerConnection< try { const result = await this.bridge.dispatchCommand(message); + const parsedResult = this.options.protocolParsers.parseCommandResult( + message.command, + message.commandVersion ?? 1, + result, + ); this.sendToSocket(socket, { type: "command-result", requestId: message.requestId, ok: true, - result, + result: parsedResult, }); } catch (error) { + // Parser failures invalidate the selected command contract and cannot be represented as an + // app command rejection. Close without reflecting callback details. + if (error instanceof BrokerProtocolError) { + socket.close(1008, "Malformed session broker command result."); + return; + } this.sendToSocket(socket, { type: "command-result", requestId: message.requestId, diff --git a/packages/session-broker/src/daemon.test.ts b/packages/session-broker/src/daemon.test.ts index 62e959897..5877e3735 100644 --- a/packages/session-broker/src/daemon.test.ts +++ b/packages/session-broker/src/daemon.test.ts @@ -11,6 +11,7 @@ import { } from "@hunk/session-broker-core"; import { SessionBroker } from "./broker"; import { createSessionBrokerDaemon } from "./daemon"; +import { createSessionBrokerProtocolParsers } from "./protocolParsers"; import type { AuthenticatedCallerRequest } from "./authentication"; interface TestSessionInfo { @@ -45,10 +46,51 @@ function parseState(value: unknown): TestSessionState | null { return selectedIndex === null ? null : { selectedIndex }; } +const protocolParsers = createSessionBrokerProtocolParsers< + TestSessionInfo, + TestSessionState, + TestServerMessage, + unknown +>({ + appRevision: 1, + features: [], + parseRegistration: (value) => parseSessionRegistrationEnvelope(value, parseInfo), + parseSnapshot: (value) => parseSessionSnapshotEnvelope(value, parseState), + commands: [ + { + command: "annotate", + version: 1, + parseInput: (value) => { + const record = brokerWireParsers.asRecord(value); + return record && typeof record.summary === "string" ? { summary: record.summary } : null; + }, + parseResult: (value) => { + const record = brokerWireParsers.asRecord(value); + return record && Object.keys(record).length === 1 && record.applied === true + ? { applied: true } + : null; + }, + }, + { + command: "annotate", + version: 2, + parseInput: (value) => { + const record = brokerWireParsers.asRecord(value); + return record && typeof record.summary === "string" ? { summary: record.summary } : null; + }, + parseResult: (value) => { + const record = brokerWireParsers.asRecord(value); + return record && Object.keys(record).length === 1 && record.applied === true + ? { applied: true } + : null; + }, + }, + ], +}); + function createBroker() { return new SessionBroker({ - parseRegistration: (value) => parseSessionRegistrationEnvelope(value, parseInfo), - parseSnapshot: (value) => parseSessionSnapshotEnvelope(value, parseState), + protocolParsers, }); } @@ -186,7 +228,10 @@ describe("session broker daemon", () => { new Request("http://broker.test/broker", { method: "POST", headers: { "content-type": "application/json" }, - body: JSON.stringify({ action: "get", selector: { sessionId: "session-1" } }), + body: JSON.stringify({ + action: "get", + selector: { sessionId: "session-1" }, + }), }), ); await expect(authenticatedBody(getResponse)).resolves.toMatchObject({ @@ -347,7 +392,10 @@ describe("session broker daemon", () => { ...authenticatedHttpApi, }); - const oversized = JSON.stringify({ action: "list", filler: "x".repeat(5 * 1024 * 1024) }); + const oversized = JSON.stringify({ + action: "list", + filler: "x".repeat(5 * 1024 * 1024), + }); const response = await daemon.handleRequest( new Request("http://broker.test/broker", { method: "POST", @@ -414,7 +462,121 @@ describe("session broker daemon", () => { ); const response = await pendingResponse; - await expect(authenticatedBody(response)).resolves.toEqual({ result: { applied: true } }); + await expect(authenticatedBody(response)).resolves.toEqual({ + result: { applied: true }, + }); + daemon.shutdown(); + }); + + test("executes each app parser once per daemon boundary and forwards transformed input", async () => { + const calls = { registration: 0, snapshot: 0, input: 0, result: 0 }; + const countingParsers = createSessionBrokerProtocolParsers< + TestSessionInfo, + TestSessionState, + TestServerMessage, + unknown + >({ + appRevision: 1, + features: [], + parseRegistration: (value) => { + calls.registration += 1; + return parseSessionRegistrationEnvelope(value, parseInfo); + }, + parseSnapshot: (value) => { + calls.snapshot += 1; + return parseSessionSnapshotEnvelope(value, parseState); + }, + commands: [ + { + command: "annotate", + version: 1, + parseInput: (value) => { + calls.input += 1; + const record = brokerWireParsers.asRecord(value); + return typeof record?.summary === "string" + ? { summary: record.summary.toUpperCase() } + : null; + }, + parseResult: (value) => { + calls.result += 1; + const record = brokerWireParsers.asRecord(value); + return record?.applied === true ? { applied: true } : null; + }, + }, + ], + }); + const broker = new SessionBroker({ + protocolParsers: countingParsers, + }); + const daemon = createSessionBrokerDaemon({ + broker, + exposeHttpApi: true, + ...authenticatedHttpApi, + }); + const owner = createConnection(); + daemon.handleConnectionMessage( + owner.connection, + JSON.stringify({ + type: "register", + registration: createRegistration(), + snapshot: createSnapshot(), + }), + ); + expect(calls).toEqual({ + registration: 1, + snapshot: 1, + input: 0, + result: 0, + }); + + daemon.handleConnectionMessage( + owner.connection, + JSON.stringify({ + type: "snapshot", + sessionId: "session-1", + snapshot: createSnapshot({ selectedIndex: 1 }), + }), + ); + expect(calls.snapshot).toBe(2); + + const pending = daemon.handleRequest( + new Request("http://broker.test/broker", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + action: "dispatch", + selector: { sessionId: "session-1" }, + command: "annotate", + input: { summary: "review note" }, + }), + }), + ); + await Bun.sleep(0); + expect(calls.input).toBe(1); + const command = JSON.parse(owner.sent.at(-1)!) as { + requestId: string; + input: { summary: string }; + }; + expect(command.input).toEqual({ summary: "REVIEW NOTE" }); + + daemon.handleConnectionMessage( + owner.connection, + JSON.stringify({ + type: "command-result", + requestId: command.requestId, + ok: true, + result: { applied: true }, + }), + ); + await expect(authenticatedBody(await pending)).resolves.toEqual({ + result: { applied: true }, + }); + expect(calls).toEqual({ + registration: 1, + snapshot: 2, + input: 1, + result: 1, + }); daemon.shutdown(); }); @@ -458,7 +620,10 @@ describe("session broker daemon", () => { daemon.handleConnectionMessage(owner.connection, register); daemon.handleConnectionMessage(duplicate.connection, register); - expect(duplicate.closed).toEqual({ code: 1008, reason: "Session registration rejected." }); + expect(duplicate.closed).toEqual({ + code: 1008, + reason: "Session registration rejected.", + }); daemon.handleConnectionClose(duplicate.connection); expect(daemon.listSessions()).toHaveLength(1); expect(daemon.listSessions()[0]).toMatchObject({ sessionId: "session-1" }); @@ -498,20 +663,31 @@ describe("session broker daemon", () => { JSON.stringify({ type: "snapshot", sessionId: "session-1", - snapshot: createSnapshot({ updatedAt: "2026-04-15T00:00:01.000Z", selectedIndex: 1 }), + snapshot: createSnapshot({ + updatedAt: "2026-04-15T00:00:01.000Z", + selectedIndex: 1, + }), }), ); - expect(snapshotPeer.closed).toEqual({ code: 1008, reason: "Session ownership rejected." }); + expect(snapshotPeer.closed).toEqual({ + code: 1008, + reason: "Session ownership rejected.", + }); expect(daemon.getSession({ sessionId: "session-1" })).toMatchObject({ snapshot: { state: { selectedIndex: 0 } }, }); - const ownerSeenAt = daemon.getSession({ sessionId: "session-1" }).lastSeenAt; + const ownerSeenAt = daemon.getSession({ + sessionId: "session-1", + }).lastSeenAt; daemon.handleConnectionMessage( heartbeatPeer.connection, JSON.stringify({ type: "heartbeat", sessionId: "session-1" }), ); - expect(heartbeatPeer.closed).toEqual({ code: 1008, reason: "Session ownership rejected." }); + expect(heartbeatPeer.closed).toEqual({ + code: 1008, + reason: "Session ownership rejected.", + }); expect(daemon.getSession({ sessionId: "session-1" }).lastSeenAt).toBe(ownerSeenAt); const pendingResponse = daemon.handleRequest( @@ -538,7 +714,10 @@ describe("session broker daemon", () => { result: { applied: "forged" }, }), ); - expect(resultPeer.closed).toEqual({ code: 1008, reason: "Command ownership rejected." }); + expect(resultPeer.closed).toEqual({ + code: 1008, + reason: "Command ownership rejected.", + }); expect(daemon.getHealth().pendingCommands).toBe(1); daemon.handleConnectionMessage( @@ -551,7 +730,92 @@ describe("session broker daemon", () => { }), ); const response = await pendingResponse; - await expect(authenticatedBody(response)).resolves.toEqual({ result: { applied: true } }); + await expect(authenticatedBody(response)).resolves.toEqual({ + result: { applied: true }, + }); + daemon.shutdown(); + }); + + test("closes malformed results without resolving pending work or leaking parser details", async () => { + const daemon = createSessionBrokerDaemon({ + broker: createBroker(), + exposeHttpApi: true, + ...authenticatedHttpApi, + }); + const owner = createConnection(); + daemon.handleConnectionMessage( + owner.connection, + JSON.stringify({ + type: "register", + registration: createRegistration(), + snapshot: createSnapshot(), + }), + ); + const pendingResponse = daemon.handleRequest( + new Request("http://broker.test/broker", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + action: "dispatch", + selector: { sessionId: "session-1" }, + command: "annotate", + input: { summary: "note" }, + }), + }), + ); + await Bun.sleep(0); + const outgoing = JSON.parse(owner.sent.at(-1)!) as { requestId: string }; + + daemon.handleConnectionMessage( + owner.connection, + JSON.stringify({ + type: "command-result", + requestId: outgoing.requestId, + ok: true, + result: { applied: false, parserStack: "secret" }, + }), + ); + + expect(owner.closed).toEqual({ + code: 1008, + reason: "Malformed command result.", + }); + expect(daemon.getHealth().pendingCommands).toBe(1); + daemon.handleConnectionClose(owner.connection); + const response = await pendingResponse; + const text = await response?.text(); + expect(text).not.toContain("parserStack"); + expect(daemon.getHealth().pendingCommands).toBe(0); + daemon.shutdown(); + }); + + test("preserves the prior registration when a replacement parser rejects", () => { + const daemon = createSessionBrokerDaemon({ + broker: createBroker(), + }); + const owner = createConnection(); + daemon.handleConnectionMessage( + owner.connection, + JSON.stringify({ + type: "register", + registration: createRegistration(), + snapshot: createSnapshot(), + }), + ); + daemon.handleConnectionMessage( + owner.connection, + JSON.stringify({ + type: "register", + registration: { ...createRegistration(), unexpected: true }, + snapshot: createSnapshot(), + }), + ); + expect(owner.closed).toEqual({ + code: 1008, + reason: "Incompatible session registration.", + }); + expect(daemon.listSessions()).toHaveLength(1); + expect(daemon.listSessions()[0]).toMatchObject({ sessionId: "session-1" }); daemon.shutdown(); }); diff --git a/packages/session-broker/src/daemon.ts b/packages/session-broker/src/daemon.ts index 13df35d9c..5ad60be22 100644 --- a/packages/session-broker/src/daemon.ts +++ b/packages/session-broker/src/daemon.ts @@ -1,4 +1,5 @@ import { + BrokerProtocolError, MAX_HTTP_BODY_BYTES, PayloadTooLargeError, callerPrincipalAllows, @@ -19,13 +20,17 @@ import { type AuthenticatedCallerRequest, type CallerRequestAuthenticator, } from "./authentication"; +import { + parseSessionBrokerJsonBytes, + parseSessionBrokerJsonText, + type SessionBrokerProtocolParsers, +} from "./protocolParsers"; import { DEFAULT_SESSION_BROKER_API_PATH, DEFAULT_SESSION_BROKER_CAPABILITIES_PATH, DEFAULT_SESSION_BROKER_HEALTH_PATH, DEFAULT_SESSION_BROKER_SOCKET_PATH, type SessionBrokerCapabilities, - type SessionBrokerDaemonRequest, type SessionBrokerDaemonResponse, type SessionBrokerAuthenticatedResponse, type SessionBrokerAuditEvent, @@ -63,23 +68,12 @@ function jsonError(message: string, status = 400) { return Response.json({ error: message }, { status }); } -/** Parse one websocket envelope without committing the daemon to any runtime socket type. */ -function parseSocketEnvelope(message: string) { - let parsed: unknown; - try { - parsed = JSON.parse(message); - } catch { - return null; - } - - if (!parsed || typeof parsed !== "object") { - return null; - } - - const type = (parsed as { type?: unknown }).type; - return typeof type === "string" - ? (parsed as object as { type: string } & Record) - : null; +/** Build one redacted protocol failure body without reflecting parser details. */ +function protocolError(error: unknown) { + return { + error: "protocol-validation-failed", + code: error instanceof BrokerProtocolError ? error.code : "invalid-app-payload", + } as const; } /** Return whether one raw broker API request body was explicitly sent as JSON. */ @@ -88,71 +82,6 @@ function hasJsonContentType(request: Request) { return contentType?.split(";", 1)[0]?.trim().toLowerCase() === "application/json"; } -/** Decode one raw broker API request body and surface a friendly transport-level error. */ -function parseJsonRequest( - body: Uint8Array, -): SessionBrokerDaemonRequest { - let parsed: unknown; - try { - if (body[0] === 0xef && body[1] === 0xbb && body[2] === 0xbf) { - throw new TypeError("UTF-8 BOM is not permitted."); - } - parsed = JSON.parse(new TextDecoder("utf-8", { fatal: true }).decode(body)); - } catch { - throw new Error("Expected one strictly encoded JSON request body."); - } - if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) { - throw new Error("Expected one JSON request object."); - } - const record = parsed as Record; - if (record.action === "list") return { action: "list" }; - if (record.action !== "get" && record.action !== "dispatch") { - throw new Error("Unknown broker API action."); - } - if (!record.selector || typeof record.selector !== "object" || Array.isArray(record.selector)) { - throw new Error("Expected one broker session selector."); - } - const selector = record.selector as Record; - for (const key of ["sessionId", "sessionPath", "repoRoot", "repoBoundary"] as const) { - if (selector[key] !== undefined && typeof selector[key] !== "string") { - throw new Error("Expected one valid broker session selector."); - } - } - if (selector.sessionId !== undefined && !isValidBrokerIdentifier(selector.sessionId)) { - throw new Error("Expected one valid broker session identifier."); - } - if (record.action === "get") { - return { action: "get", selector: selector as SessionTargetSelector }; - } - if (!isValidBrokerIdentifier(record.command)) { - throw new Error("Expected one valid broker command name."); - } - const commandVersion = record.commandVersion ?? 1; - if (!isValidBrokerRevision(commandVersion)) { - throw new Error("Expected one positive broker command version."); - } - if ( - record.timeoutMs !== undefined && - (!Number.isSafeInteger(record.timeoutMs) || (record.timeoutMs as number) <= 0) - ) { - throw new Error("Expected one positive command timeout."); - } - if (record.timeoutMessage !== undefined && typeof record.timeoutMessage !== "string") { - throw new Error("Expected one command timeout message."); - } - return { - action: "dispatch", - selector: selector as SessionTargetSelector, - command: record.command as CommandName, - commandVersion, - input: record.input as CommandInput, - ...(record.timeoutMs === undefined ? {} : { timeoutMs: record.timeoutMs as number }), - ...(record.timeoutMessage === undefined - ? {} - : { timeoutMessage: record.timeoutMessage as string }), - }; -} - /** Build the default dispatch timeout text so adapters can override only when they need to. */ function defaultTimeoutMessage(command: string) { return `Timed out waiting for the session to handle ${command}.`; @@ -172,6 +101,12 @@ export class SessionBrokerDaemon< private readonly startedAt = Date.now(); private readonly capabilities: SessionBrokerCapabilities; + private readonly protocolParsers: SessionBrokerProtocolParsers< + unknown, + unknown, + ServerMessage, + CommandResult + >; private readonly idleTimeoutMs: number; private readonly staleSessionTtlMs: number; private readonly staleSessionSweepIntervalMs: number; @@ -188,10 +123,7 @@ export class SessionBrokerDaemon< constructor( private readonly broker: SessionBrokerController, - options: Omit< - SessionBrokerDaemonOptions, - "broker" - > = {}, + options: Omit, "broker">, ) { const exposeAuthenticatedHttpApi = (options.exposeHttpApi ?? false) && @@ -210,8 +142,15 @@ export class SessionBrokerDaemon< : undefined, }; this.capabilities = options.capabilities ?? { version: 1 }; + this.protocolParsers = broker.protocolParsers; + if ( + options.appRevision !== undefined && + options.appRevision !== this.protocolParsers.appRevision + ) { + throw new TypeError("Session broker app revision does not match its parser registry."); + } this.appId = options.appId ?? "session-broker"; - this.appRevision = options.appRevision; + this.appRevision = this.protocolParsers.appRevision; this.callerAuthenticator = options.callerAuthenticator; this.authorizer = options.authorizer; this.audit = options.audit; @@ -257,7 +196,9 @@ export class SessionBrokerDaemon< if (url.pathname === this.paths.health) { // Treat health checks as a cheap maintenance pulse so stale sessions disappear even when the // daemon is mostly idle and no websocket traffic is flowing. - const removed = this.broker.pruneStaleSessions({ ttlMs: this.staleSessionTtlMs }); + const removed = this.broker.pruneStaleSessions({ + ttlMs: this.staleSessionTtlMs, + }); if (removed > 0) { this.noteActivity(); } @@ -284,7 +225,11 @@ export class SessionBrokerDaemon< } const authenticated = await this.authenticateRequest(request, body, "diagnostics"); if (authenticated instanceof Response) return authenticated; - if (!(await this.authorize(request, authenticated, { operation: "diagnostics" }))) { + if ( + !(await this.authorize(request, authenticated, { + operation: "diagnostics", + })) + ) { return this.authenticatedResponse(authenticated, { error: "authorization-denied" }, 403); } const inactive = this.rejectInactiveRequest(authenticated); @@ -301,9 +246,12 @@ export class SessionBrokerDaemon< return null; } - handleConnectionMessage(connection: SessionBrokerPeer, message: string) { - const parsed = parseSocketEnvelope(message); - if (!parsed) { + handleConnectionMessage(connection: SessionBrokerPeer, message: unknown) { + let parsed; + try { + parsed = this.protocolParsers.parseClientMessage(parseSessionBrokerJsonText(message)); + } catch { + connection.close?.(INCOMPATIBLE_PAYLOAD_CLOSE_CODE, "Malformed session broker protocol."); return; } @@ -330,10 +278,6 @@ export class SessionBrokerDaemon< break; } case "snapshot": { - if (typeof parsed.sessionId !== "string") { - return; - } - // Snapshot updates are only valid after registration. Closing missing or invalid sessions // keeps the broker state single-sourced instead of guessing how to recover. const updateResult = this.broker.updateSnapshot( @@ -355,10 +299,6 @@ export class SessionBrokerDaemon< break; } case "heartbeat": { - if (typeof parsed.sessionId !== "string") { - return; - } - const seenResult = this.broker.markSessionSeen(connection, parsed.sessionId); if (seenResult === "not-owner") { connection.close?.(INCOMPATIBLE_PAYLOAD_CLOSE_CODE, "Session ownership rejected."); @@ -369,24 +309,20 @@ export class SessionBrokerDaemon< break; } case "command-result": { - if (typeof parsed.requestId !== "string" || typeof parsed.ok !== "boolean") { - return; - } - - const result = this.broker.handleCommandResult(connection, { - requestId: parsed.requestId, - ok: parsed.ok, - result: parsed.result as CommandResult | undefined, - error: typeof parsed.error === "string" ? parsed.error : undefined, - }); + const result = this.broker.handleCommandResult(connection, parsed); if (result === "not-owner") { connection.close?.(INCOMPATIBLE_PAYLOAD_CLOSE_CODE, "Command ownership rejected."); return; } - if (result === "handled") { - this.noteActivity(); + if (result === "invalid") { + // A result that violates the pending command contract invalidates the producer's current + // assumptions. The broker keeps pending state coherent until disconnect cleanup runs. + connection.close?.(INCOMPATIBLE_PAYLOAD_CLOSE_CODE, "Malformed command result."); + return; } + + if (result === "handled") this.noteActivity(); break; } } @@ -420,7 +356,9 @@ export class SessionBrokerDaemon< private startLifecycle() { this.sweepTimer = setInterval(() => { - const removed = this.broker.pruneStaleSessions({ ttlMs: this.staleSessionTtlMs }); + const removed = this.broker.pruneStaleSessions({ + ttlMs: this.staleSessionTtlMs, + }); if (removed > 0) { this.noteActivity(); } @@ -482,7 +420,10 @@ export class SessionBrokerDaemon< if (!this.callerAuthenticator || !this.authorizer) { return jsonError("Broker control is unavailable.", 404); } - const authenticated = await this.callerAuthenticator.authenticate({ request, body }); + const authenticated = await this.callerAuthenticator.authenticate({ + request, + body, + }); if ( !authenticated || typeof authenticated !== "object" || @@ -532,7 +473,10 @@ export class SessionBrokerDaemon< }, ): Promise { const principal: CallerPrincipal = authenticated.principal; - const allowedByGrant = callerPrincipalAllows(principal, { appId: this.appId, ...facts }); + const allowedByGrant = callerPrincipalAllows(principal, { + appId: this.appId, + ...facts, + }); let allowedByApp = false; if (allowedByGrant && this.authorizer) { try { @@ -579,7 +523,10 @@ export class SessionBrokerDaemon< ? { appContract: { appRevision: this.appRevision, features: [] } } : {}), }); - const envelope: SessionBrokerAuthenticatedResponse = { body: structuredBody, authentication }; + const envelope: SessionBrokerAuthenticatedResponse = { + body: structuredBody, + authentication, + }; return new Response(canonicalizeJson(envelope as unknown as CanonicalJsonValue), { status, headers: { "content-type": "application/json" }, @@ -615,15 +562,11 @@ export class SessionBrokerDaemon< const authenticated = await this.authenticateRequest(request, body, "list"); if (authenticated instanceof Response) return authenticated; - let input: SessionBrokerDaemonRequest; + let input; try { - input = parseJsonRequest(body); + input = this.protocolParsers.parseDaemonRequest(parseSessionBrokerJsonBytes(body)); } catch (error) { - return this.authenticatedResponse( - authenticated, - { error: error instanceof Error ? error.message : "Invalid broker API request." }, - 400, - ); + return this.authenticatedResponse(authenticated, protocolError(error), 400); } const operation = input.action as CallerOperation; @@ -657,27 +600,32 @@ export class SessionBrokerDaemon< case "get": response = { session: this.broker.getSession(input.selector) }; break; - case "dispatch": + case "dispatch": { + // Resolve the target before invoking app-owned parsing so the exact target contract is + // selected first. This read-only lookup happens only after authentication/authorization. + this.broker.getSession(input.selector); response = { result: await this.broker.dispatchCommand({ selector: input.selector, command: input.command, commandVersion: input.commandVersion ?? 1, - input: input.input as Extract< - ServerMessage, - { command: ServerMessage["command"] } - >["input"], + input: input.input, timeoutMessage: input.timeoutMessage ?? defaultTimeoutMessage(input.command), timeoutMs: input.timeoutMs, }), }; break; + } } return this.authenticatedResponse(authenticated, response, 200, targetSpecific); } catch (error) { return this.authenticatedResponse( authenticated, - { error: error instanceof Error ? error.message : "Unknown broker API error." }, + error instanceof BrokerProtocolError + ? protocolError(error) + : { + error: error instanceof Error ? error.message : "Unknown broker API error.", + }, 400, targetSpecific, ); diff --git a/packages/session-broker/src/index.ts b/packages/session-broker/src/index.ts index 802081113..7e0be3777 100644 --- a/packages/session-broker/src/index.ts +++ b/packages/session-broker/src/index.ts @@ -5,3 +5,4 @@ export * from "./daemon"; export * from "./connection"; export * from "./crypto"; export * from "./authentication"; +export * from "./protocolParsers"; diff --git a/packages/session-broker/src/protocolParsers.test.ts b/packages/session-broker/src/protocolParsers.test.ts new file mode 100644 index 000000000..a50fc5d62 --- /dev/null +++ b/packages/session-broker/src/protocolParsers.test.ts @@ -0,0 +1,197 @@ +import { describe, expect, test } from "bun:test"; +import { + BrokerProtocolError, + type SessionRegistration, + type SessionServerMessage, + type SessionSnapshot, +} from "@hunk/session-broker-core"; +import type { SessionBrokerDaemonRequest } from "./types"; +import { + createSessionBrokerProtocolParsers, + type StructuralSessionBrokerDaemonRequest, +} from "./protocolParsers"; + +type TestMessage = SessionServerMessage<"annotate", { summary: string }>; +type TestResult = { applied: true }; + +function exactObject(value: unknown, keys: readonly string[]) { + return ( + !!value && + typeof value === "object" && + !Array.isArray(value) && + Object.keys(value).length === keys.length && + keys.every((key) => Object.hasOwn(value, key)) + ); +} + +const parsers = createSessionBrokerProtocolParsers< + { title: string }, + { selected: number }, + TestMessage, + TestResult +>({ + appRevision: 7, + features: [], + parseRegistration: (value) => + exactObject(value, ["registrationVersion", "sessionId", "pid", "cwd", "launchedAt", "info"]) + ? (value as SessionRegistration<{ title: string }>) + : null, + parseSnapshot: (value) => + exactObject(value, ["updatedAt", "state"]) + ? (value as SessionSnapshot<{ selected: number }>) + : null, + commands: [ + { + command: "annotate", + version: 2, + parseInput: (value) => + exactObject(value, ["summary"]) && + typeof (value as Record).summary === "string" + ? { summary: (value as Record).summary! } + : null, + parseResult: (value) => + exactObject(value, ["applied"]) && (value as Record).applied === true + ? { applied: true } + : null, + }, + ], +}); + +const generatedMalformedValues = Array.from({ length: 12 }, (_, index) => { + const values: unknown[] = [null, false, index + 0.5, `bad value ${index}`, [], { extra: index }]; + return values[index % values.length]; +}); + +const malformedCorpus: readonly unknown[] = [ + ...generatedMalformedValues, + null, + [], + "message", + {}, + { type: "command", requestId: "request-1", command: "annotate", input: null }, + { + type: "command", + requestId: "request-1", + command: "annotate", + commandVersion: 0, + input: { summary: "note" }, + }, + { + type: "command", + requestId: "request-1", + command: "annotate", + commandVersion: 2, + input: { summary: "note", extra: true }, + }, + { + type: "command", + requestId: "bad id!", + command: "annotate", + commandVersion: 2, + input: { summary: "note" }, + }, + { + type: "command", + requestId: "request-1", + command: "mismatched", + commandVersion: 2, + input: { summary: "note" }, + }, + { type: "register", registration: null, snapshot: {}, extra: true }, + { type: "snapshot", sessionId: "bad id!", snapshot: {} }, + { + action: "dispatch", + selector: { nested: true }, + command: "annotate", + input: { summary: 1 }, + }, + { action: "get", selector: { sessionId: "session-1", nested: true } }, + { action: "list", selector: {} }, +]; + +function failureCode(task: () => unknown) { + try { + task(); + return null; + } catch (error) { + expect(error).toBeInstanceOf(BrokerProtocolError); + return (error as BrokerProtocolError).code; + } +} + +describe("session broker authoritative protocol parsers", () => { + test("type-locks complete parser outputs to the public unions", () => { + const client = parsers.parseClientMessage({ + type: "heartbeat", + sessionId: "session-1", + }); + const server: TestMessage = parsers.parseServerMessage({ + type: "command", + requestId: "request-1", + command: "annotate", + commandVersion: 2, + input: { summary: "note" }, + }); + const request: StructuralSessionBrokerDaemonRequest<"annotate"> = parsers.parseDaemonRequest({ + action: "list", + }); + const publicRequest: SessionBrokerDaemonRequest<"annotate"> = { + action: "list", + }; + expect([client.type, server.command, request.action, publicRequest.action]).toEqual([ + "heartbeat", + "annotate", + "list", + "list", + ]); + }); + + test("runs one deterministic malformed corpus through websocket and HTTP parsers", () => { + for (const value of malformedCorpus) { + expect(failureCode(() => parsers.parseServerMessage(value))).not.toBeNull(); + expect(failureCode(() => parsers.parseClientMessage(value))).not.toBeNull(); + expect(failureCode(() => parsers.parseDaemonRequest(value))).not.toBeNull(); + } + }); + + test("enforces complete client envelopes and exact app result contracts", () => { + expect( + failureCode(() => + parsers.parseClientMessage({ + type: "heartbeat", + sessionId: "session-1", + extra: true, + }), + ), + ).toBe("invalid-keys"); + expect(failureCode(() => parsers.parseCommandResult("annotate", 2, { applied: false }))).toBe( + "invalid-app-payload", + ); + expect(failureCode(() => parsers.parseCommandResult("annotate", 1, { applied: true }))).toBe( + "unknown-command", + ); + }); + + test("leaves registration and snapshot app payloads untouched for controller parsing", () => { + const registration = { malformedForApp: true }; + const snapshot = { transformedLater: true }; + expect(parsers.parseClientMessage({ type: "register", registration, snapshot })).toEqual({ + type: "register", + registration, + snapshot, + }); + }); + + test("normalizes parser throws without exposing their messages", () => { + const throwing = createSessionBrokerProtocolParsers({ + appRevision: 1, + features: [], + parseRegistration: () => { + throw new Error("registration secret"); + }, + parseSnapshot: () => null, + commands: [], + }); + expect(failureCode(() => throwing.parseRegistration({}))).toBe("app-parser-failed"); + }); +}); diff --git a/packages/session-broker/src/protocolParsers.ts b/packages/session-broker/src/protocolParsers.ts new file mode 100644 index 000000000..e8aeb9a9b --- /dev/null +++ b/packages/session-broker/src/protocolParsers.ts @@ -0,0 +1,357 @@ +import { + BrokerProtocolError, + SESSION_BROKER_PROTOCOL_REVISION, + failBrokerProtocol, + parseBrokerAppPayload, + parseBrokerDeadline, + parseBrokerIdentifier, + parseBrokerRevision, + parseBrokerSelector, + parseBrokerString, + parseBrokerTimeout, + parseExactBrokerRecord, + type SessionRegistration, + type SessionServerMessage, + type SessionSnapshot, +} from "@hunk/session-broker-core"; +import type { SessionBrokerDaemonRequest } from "./types"; + +export type SessionBrokerRuntimeParser = (value: unknown) => T | null; + +export interface SessionBrokerCommandParsers< + CommandName extends string = string, + Input = unknown, + Result = unknown, +> { + readonly command: CommandName; + readonly version: number; + readonly parseInput: SessionBrokerRuntimeParser; + readonly parseResult: SessionBrokerRuntimeParser; +} + +export interface SessionBrokerAppParserRegistry< + Info = unknown, + State = unknown, + ServerMessage extends SessionServerMessage = SessionServerMessage, + Result = unknown, +> { + readonly brokerRevision?: typeof SESSION_BROKER_PROTOCOL_REVISION; + readonly appRevision: number; + readonly features: readonly []; + readonly parseRegistration: SessionBrokerRuntimeParser>; + readonly parseSnapshot: SessionBrokerRuntimeParser>; + readonly commands: readonly SessionBrokerCommandParsers< + ServerMessage["command"], + unknown, + Result + >[]; +} + +interface StructuralDispatchRequest extends Omit< + Extract, { action: "dispatch" }>, + "input" +> { + input: unknown; +} + +export type StructuralSessionBrokerDaemonRequest = + | Exclude, { action: "dispatch" }> + | StructuralDispatchRequest; + +/** A producer envelope whose app-owned payloads remain unknown until controller state parses them. */ +export type StructuralSessionClientMessage = + | { type: "register"; registration: unknown; snapshot: unknown } + | { type: "snapshot"; sessionId: string; snapshot: unknown } + | { type: "heartbeat"; sessionId: string } + | { type: "command-result"; requestId: string; ok: true; result: Result } + | { type: "command-result"; requestId: string; ok: false; error: string }; + +/** Decode strict UTF-8 JSON without allowing a BOM or surfacing decoder details. */ +export function parseSessionBrokerJsonBytes(body: Uint8Array): unknown { + try { + if (body[0] === 0xef && body[1] === 0xbb && body[2] === 0xbf) { + return failBrokerProtocol("invalid-json"); + } + return JSON.parse(new TextDecoder("utf-8", { fatal: true }).decode(body)) as unknown; + } catch (error) { + if (error instanceof BrokerProtocolError) throw error; + return failBrokerProtocol("invalid-json"); + } +} + +/** Parse text JSON at a websocket boundary. */ +export function parseSessionBrokerJsonText(message: unknown): unknown { + if (typeof message !== "string") return failBrokerProtocol("invalid-json"); + try { + return JSON.parse(message) as unknown; + } catch { + return failBrokerProtocol("invalid-json"); + } +} + +/** Owns the fixed Phase-1 app contract and every app parser selected by exact command identity. */ +export class SessionBrokerProtocolParsers< + Info = unknown, + State = unknown, + ServerMessage extends SessionServerMessage = SessionServerMessage, + Result = unknown, +> { + readonly brokerRevision = SESSION_BROKER_PROTOCOL_REVISION; + readonly appRevision: number; + readonly features = Object.freeze([]) as readonly []; + private readonly parseRegistrationValue: SessionBrokerRuntimeParser>; + private readonly parseSnapshotValue: SessionBrokerRuntimeParser>; + private readonly commandParsers = new Map< + string, + SessionBrokerCommandParsers + >(); + + constructor(registry: SessionBrokerAppParserRegistry) { + if ( + registry.brokerRevision !== undefined && + registry.brokerRevision !== SESSION_BROKER_PROTOCOL_REVISION + ) { + throw new TypeError("Invalid session broker parser registry."); + } + this.appRevision = parseBrokerRevision(registry.appRevision); + if (!Array.isArray(registry.features) || registry.features.length !== 0) { + throw new TypeError("Invalid session broker parser registry."); + } + if ( + typeof registry.parseRegistration !== "function" || + typeof registry.parseSnapshot !== "function" || + !Array.isArray(registry.commands) + ) { + throw new TypeError("Invalid session broker parser registry."); + } + this.parseRegistrationValue = registry.parseRegistration; + this.parseSnapshotValue = registry.parseSnapshot; + for (const descriptor of registry.commands) { + if ( + !descriptor || + typeof descriptor !== "object" || + typeof descriptor.parseInput !== "function" || + typeof descriptor.parseResult !== "function" + ) { + throw new TypeError("Invalid session broker parser registry."); + } + const command = parseBrokerIdentifier(descriptor.command); + const version = parseBrokerRevision(descriptor.version); + const key = this.commandKey(command, version); + if (this.commandParsers.has(key)) + throw new TypeError("Invalid session broker parser registry."); + this.commandParsers.set(key, Object.freeze({ ...descriptor, command, version })); + } + } + + /** Parse app registration through the one registered callback. */ + parseRegistration(value: unknown): SessionRegistration { + return parseBrokerAppPayload(this.parseRegistrationValue, value); + } + + /** Parse app snapshot through the one registered callback. */ + parseSnapshot(value: unknown): SessionSnapshot { + return parseBrokerAppPayload(this.parseSnapshotValue, value); + } + + /** Parse exact command input after the caller has selected a target contract. */ + parseCommandInput( + command: CommandName, + version: number, + value: unknown, + ): Extract["input"] { + const descriptor = this.lookupCommand(command, version); + return parseBrokerAppPayload(descriptor.parseInput, value) as Extract< + ServerMessage, + { command: CommandName } + >["input"]; + } + + /** Parse exact command result before pending broker work is resolved. */ + parseCommandResult(command: string, version: number, value: unknown): Result { + return parseBrokerAppPayload(this.lookupCommand(command, version).parseResult, value); + } + + /** Parse one producer envelope structurally, leaving app payloads for controller state. */ + parseClientMessage(value: unknown): StructuralSessionClientMessage { + const base = parseExactBrokerRecord( + value, + ["type"] as const, + ["registration", "snapshot", "sessionId", "requestId", "ok", "result", "error"] as const, + ); + if (typeof base.type !== "string") return failBrokerProtocol("invalid-discriminant"); + switch (base.type) { + case "register": { + const record = parseExactBrokerRecord(value, ["type", "registration", "snapshot"] as const); + return { + type: "register", + registration: record.registration, + snapshot: record.snapshot, + }; + } + case "snapshot": { + const record = parseExactBrokerRecord(value, ["type", "sessionId", "snapshot"] as const); + return { + type: "snapshot", + sessionId: parseBrokerIdentifier(record.sessionId), + snapshot: record.snapshot, + }; + } + case "heartbeat": { + const record = parseExactBrokerRecord(value, ["type", "sessionId"] as const); + return { + type: "heartbeat", + sessionId: parseBrokerIdentifier(record.sessionId), + }; + } + case "command-result": { + const common = parseExactBrokerRecord( + value, + ["type", "requestId", "ok"] as const, + ["result", "error"] as const, + ); + const requestId = parseBrokerIdentifier(common.requestId); + if (common.ok === true) { + const record = parseExactBrokerRecord(value, [ + "type", + "requestId", + "ok", + "result", + ] as const); + return { + type: "command-result", + requestId, + ok: true, + result: record.result as Result, + }; + } + if (common.ok === false) { + const record = parseExactBrokerRecord(value, [ + "type", + "requestId", + "ok", + "error", + ] as const); + return { + type: "command-result", + requestId, + ok: false, + error: parseBrokerString(record.error, { maxBytes: 1_024 }), + }; + } + return failBrokerProtocol("invalid-field"); + } + default: + return failBrokerProtocol("invalid-discriminant"); + } + } + + /** Parse one complete daemon-to-producer command envelope before bridge dispatch. */ + parseServerMessage(value: unknown): ServerMessage { + const record = parseExactBrokerRecord( + value, + ["type", "requestId", "command", "input"] as const, + ["commandVersion"] as const, + ); + if (record.type !== "command") return failBrokerProtocol("invalid-discriminant"); + const requestId = parseBrokerIdentifier(record.requestId); + const command = parseBrokerIdentifier(record.command) as ServerMessage["command"]; + const commandVersion = + record.commandVersion === undefined ? 1 : parseBrokerRevision(record.commandVersion); + const input = this.parseCommandInput(command, commandVersion, record.input); + return { + type: "command", + requestId, + command, + commandVersion, + input, + } as ServerMessage; + } + + /** Parse generic HTTP structure while leaving app input unknown until target selection. */ + parseDaemonRequest( + value: unknown, + ): StructuralSessionBrokerDaemonRequest { + const base = parseExactBrokerRecord( + value, + ["action"] as const, + [ + "selector", + "command", + "commandVersion", + "input", + "timeoutMs", + "timeoutMessage", + "deadline", + "idempotencyKey", + ] as const, + ); + switch (base.action) { + case "list": + parseExactBrokerRecord(value, ["action"] as const); + return { action: "list" }; + case "get": { + const record = parseExactBrokerRecord(value, ["action", "selector"] as const); + return { + action: "get", + selector: parseBrokerSelector(record.selector), + }; + } + case "dispatch": { + const record = parseExactBrokerRecord( + value, + ["action", "selector", "command", "input"] as const, + ["commandVersion", "timeoutMs", "timeoutMessage", "deadline", "idempotencyKey"] as const, + ); + const command = parseBrokerIdentifier(record.command) as ServerMessage["command"]; + const commandVersion = + record.commandVersion === undefined ? 1 : parseBrokerRevision(record.commandVersion); + return { + action: "dispatch", + selector: parseBrokerSelector(record.selector), + command, + commandVersion, + input: record.input, + ...(record.timeoutMs === undefined + ? {} + : { timeoutMs: parseBrokerTimeout(record.timeoutMs) }), + ...(record.timeoutMessage === undefined + ? {} + : { + timeoutMessage: parseBrokerString(record.timeoutMessage, { + maxBytes: 1_024, + }), + }), + ...(record.deadline === undefined + ? {} + : { deadline: parseBrokerDeadline(record.deadline) }), + ...(record.idempotencyKey === undefined + ? {} + : { idempotencyKey: parseBrokerIdentifier(record.idempotencyKey) }), + }; + } + default: + return failBrokerProtocol("invalid-discriminant"); + } + } + + private lookupCommand(command: string, version: number) { + const descriptor = this.commandParsers.get(this.commandKey(command, version)); + if (!descriptor) return failBrokerProtocol("unknown-command"); + return descriptor; + } + + private commandKey(command: string, version: number) { + return `${command}\u0000${version}`; + } +} + +/** Snapshot and validate one authoritative fixed-contract parser registry. */ +export function createSessionBrokerProtocolParsers< + Info = unknown, + State = unknown, + ServerMessage extends SessionServerMessage = SessionServerMessage, + Result = unknown, +>(registry: SessionBrokerAppParserRegistry) { + return new SessionBrokerProtocolParsers(registry); +} diff --git a/packages/session-broker/src/types.ts b/packages/session-broker/src/types.ts index ee5a35778..54e3d3aa3 100644 --- a/packages/session-broker/src/types.ts +++ b/packages/session-broker/src/types.ts @@ -45,6 +45,8 @@ export type SessionBrokerDaemonRequest< input: CommandInput; timeoutMs?: number; timeoutMessage?: string; + deadline?: number; + idempotencyKey?: string; }; export type SessionBrokerDaemonResponse = @@ -93,7 +95,7 @@ export interface SessionBrokerSocketMessageEvent { export interface SessionBrokerSocketLike { readonly readyState: number; send(data: string): void; - close(): void; + close(code?: number, reason?: string): void; onopen: (() => void) | null; onmessage: ((event: SessionBrokerSocketMessageEvent) => void) | null; onclose: ((event: SessionBrokerSocketCloseEvent) => void) | null; diff --git a/src/session/agent/cliClient.test.ts b/src/session/agent/cliClient.test.ts index 79431964a..ab4738856 100644 --- a/src/session/agent/cliClient.test.ts +++ b/src/session/agent/cliClient.test.ts @@ -334,6 +334,34 @@ describe("HTTP Hunk session CLI client", () => { ); }); + test("rejects malformed successful responses instead of returning asserted types", async () => { + globalThis.fetch = (async () => + Response.json({ + sessions: [{ sessionId: "partial", unknown: true }], + })) as unknown as typeof fetch; + + const client = createHttpHunkSessionCliClient(); + await expect(client.listSessions()).rejects.toThrow( + "Invalid Hunk session daemon response for list.", + ); + + globalThis.fetch = (async () => new Response("not json")) as unknown as typeof fetch; + await expect(client.listSessions()).rejects.toThrow( + "Invalid Hunk session daemon response for list.", + ); + }); + + test("returns schema-transformed response objects", async () => { + const session = createTestListedSession(); + globalThis.fetch = (async () => + Response.json({ sessions: [session] })) as unknown as typeof fetch; + + const client = createHttpHunkSessionCliClient(); + const result = await client.listSessions(); + expect(result).toEqual([session]); + expect(result[0]).not.toBe(session); + }); + test("throws daemon response errors with JSON messages or status text fallbacks", async () => { globalThis.fetch = (async () => Response.json( diff --git a/src/session/agent/cliClient.ts b/src/session/agent/cliClient.ts index 511bf0eed..ca22c0f56 100644 --- a/src/session/agent/cliClient.ts +++ b/src/session/agent/cliClient.ts @@ -8,9 +8,12 @@ import { } from "../client/daemonHttp"; import { HUNK_SESSION_API_PATH, + type SessionDaemonAction, type SessionDaemonCapabilities, type SessionDaemonRequest, + type SessionDaemonResponses, } from "../protocol"; +import { parseSessionDaemonResponse } from "../protocolSchemas"; import type { AppliedCommentBatchResult, AppliedCommentResult, @@ -78,7 +81,9 @@ class HttpHunkSessionCliClient implements HunkSessionCliClient { constructor(private readonly timeoutMs = HUNK_SESSION_DAEMON_HTTP_TIMEOUT_MS) {} - private async request(input: SessionDaemonRequest) { + private async request( + input: Extract, + ): Promise { return requestSessionDaemonHttp({ config: this.config, path: HUNK_SESSION_API_PATH, @@ -96,7 +101,13 @@ class HttpHunkSessionCliClient implements HunkSessionCliClient { throw new Error(await extractResponseError(response)); } - return (await response.json()) as ResultType; + let value: unknown; + try { + value = await response.json(); + } catch { + throw new Error(`Invalid Hunk session daemon response for ${input.action}.`); + } + return parseSessionDaemonResponse(input.action, value); }, }); } @@ -106,22 +117,20 @@ class HttpHunkSessionCliClient implements HunkSessionCliClient { } async listSessions() { - return (await this.request<{ sessions: ListedSession[] }>({ action: "list" })).sessions; + return (await this.request({ action: "list" })).sessions; } async getSession(selector: SessionSelectorInput) { - return (await this.request<{ session: ListedSession }>({ action: "get", selector })).session; + return (await this.request({ action: "get", selector })).session; } async getSelectedContext(selector: SessionSelectorInput) { - return ( - await this.request<{ context: SelectedSessionContext }>({ action: "context", selector }) - ).context; + return (await this.request({ action: "context", selector })).context; } async getSessionReview(input: SessionReviewCommandInput) { return ( - await this.request<{ review: SessionReview }>({ + await this.request({ action: "review", selector: input.selector, includePatch: input.includePatch, @@ -132,7 +141,7 @@ class HttpHunkSessionCliClient implements HunkSessionCliClient { async navigateToHunk(input: SessionNavigateCommandInput) { return ( - await this.request<{ result: NavigatedSelectionResult }>({ + await this.request({ action: "navigate", selector: input.selector, filePath: input.filePath, @@ -147,7 +156,7 @@ class HttpHunkSessionCliClient implements HunkSessionCliClient { async reloadSession(input: SessionReloadCommandInput) { return ( - await this.request<{ result: ReloadedSessionResult }>({ + await this.request({ action: "reload", selector: input.selector, nextInput: input.nextInput, @@ -158,7 +167,7 @@ class HttpHunkSessionCliClient implements HunkSessionCliClient { async addComment(input: SessionCommentAddCommandInput) { return ( - await this.request<{ result: AppliedCommentResult }>({ + await this.request({ action: "comment-add", selector: input.selector, filePath: input.filePath, @@ -175,7 +184,7 @@ class HttpHunkSessionCliClient implements HunkSessionCliClient { async applyComments(input: SessionCommentApplyCommandInput) { return ( - await this.request<{ result: AppliedCommentBatchResult }>({ + await this.request({ action: "comment-apply", selector: input.selector, comments: input.comments, @@ -186,20 +195,18 @@ class HttpHunkSessionCliClient implements HunkSessionCliClient { async listComments(input: SessionCommentListCommandInput) { return ( - await this.request<{ comments: Array }>( - { - action: "comment-list", - selector: input.selector, - filePath: input.filePath, - type: input.type, - }, - ) + await this.request({ + action: "comment-list", + selector: input.selector, + filePath: input.filePath, + type: input.type, + }) ).comments; } async removeComment(input: SessionCommentRemoveCommandInput) { return ( - await this.request<{ result: RemovedCommentResult }>({ + await this.request({ action: "comment-rm", selector: input.selector, commentId: input.commentId, @@ -209,7 +216,7 @@ class HttpHunkSessionCliClient implements HunkSessionCliClient { async clearComments(input: SessionCommentClearCommandInput) { return ( - await this.request<{ result: ClearedCommentsResult }>({ + await this.request({ action: "comment-clear", selector: input.selector, filePath: input.filePath, @@ -220,7 +227,7 @@ class HttpHunkSessionCliClient implements HunkSessionCliClient { async addHighlight(input: SessionHighlightAddCommandInput) { return ( - await this.request<{ result: AppliedHighlightResult }>({ + await this.request({ action: "highlight-add", selector: input.selector, filePath: input.filePath, @@ -236,7 +243,7 @@ class HttpHunkSessionCliClient implements HunkSessionCliClient { async clearHighlights(input: SessionHighlightClearCommandInput) { return ( - await this.request<{ result: ClearedHighlightsResult }>({ + await this.request({ action: "highlight-clear", selector: input.selector, filePath: input.filePath, diff --git a/src/session/broker/brokerClient.ts b/src/session/broker/brokerClient.ts index 256a38b9e..53ab54330 100644 --- a/src/session/broker/brokerClient.ts +++ b/src/session/broker/brokerClient.ts @@ -4,11 +4,7 @@ import { type SessionBrokerConnectionBridge, type SessionBrokerSocketLike, } from "@hunk/session-broker"; -import type { - SessionRegistration, - SessionServerMessage, - SessionSnapshot, -} from "@hunk/session-broker-core"; +import type { SessionRegistration, SessionSnapshot } from "@hunk/session-broker-core"; import { SESSION_BROKER_SOCKET_PATH, resolveSessionBrokerConfig, @@ -19,6 +15,7 @@ import { readSessionBrokerHealth, waitForSessionBrokerShutdown, } from "./brokerLauncher"; +import { hunkSessionProtocolParsers } from "./protocolParsers"; import { readHunkSessionDaemonCapabilities, reportHunkDaemonUpgradeRestart, @@ -38,47 +35,37 @@ const INCOMPATIBLE_SESSION_CLOSE_REASON_PREFIX = "Incompatible session "; const INCOMPATIBLE_SESSION_CLOSE_MESSAGE = "This window is too old for the refreshed session broker daemon. Restart the window to reconnect."; -type SessionAppBridge< - ServerMessage extends SessionServerMessage = SessionServerMessage, - Result = unknown, -> = SessionBrokerConnectionBridge; +type SessionAppBridge = SessionBrokerConnectionBridge< + HunkSessionServerMessage, + HunkSessionCommandResult +>; interface SessionBrokerClientTiming { daemonStartupTimeoutMs?: number; reconnectDelayMs?: number; } -/** The broker client bound to Hunk's session info, state, message, and result types. */ -export type HunkSessionBrokerClient = SessionBrokerClient< - HunkSessionInfo, - HunkSessionState, - HunkSessionServerMessage, - HunkSessionCommandResult ->; +/** The concrete broker client bound to Hunk's session contracts. */ +export type HunkSessionBrokerClient = SessionBrokerClient; -/** Keep one running app session registered with the local session broker daemon. */ -export class SessionBrokerClient< - Info = unknown, - State = unknown, - ServerMessage extends SessionServerMessage = SessionServerMessage, - Result = unknown, -> { +/** Keep one running Hunk session registered with the local session broker daemon. */ +export class SessionBrokerClient { private connection: GenericSessionBrokerConnection< - Info, - State, + HunkSessionInfo, + HunkSessionState, SessionBrokerSocketLike, - ServerMessage, - Result + HunkSessionServerMessage, + HunkSessionCommandResult > | null = null; - private bridge: SessionAppBridge | null = null; + private bridge: SessionAppBridge | null = null; private reconnectTimer: ReturnType | null = null; private stopped = false; private startupPromise: Promise | null = null; private lastConnectionWarning: string | null = null; constructor( - private registration: SessionRegistration, - private snapshot: SessionSnapshot, + private registration: SessionRegistration, + private snapshot: SessionSnapshot, private timing: SessionBrokerClientTiming = {}, ) {} @@ -122,7 +109,10 @@ export class SessionBrokerClient< return this.registration; } - replaceSession(registration: SessionRegistration, snapshot: SessionSnapshot) { + replaceSession( + registration: SessionRegistration, + snapshot: SessionSnapshot, + ) { // Let the connection validate/send first. If it throws, the client keeps // serving the previous registration and snapshot as one coherent pair. this.connection?.replaceSession(registration, snapshot); @@ -201,12 +191,12 @@ export class SessionBrokerClient< } } - setBridge(bridge: SessionAppBridge | null) { + setBridge(bridge: SessionAppBridge | null) { this.bridge = bridge; this.connection?.setBridge(bridge); } - updateSnapshot(snapshot: SessionSnapshot) { + updateSnapshot(snapshot: SessionSnapshot) { this.snapshot = snapshot; this.connection?.updateSnapshot(snapshot); } @@ -217,17 +207,18 @@ export class SessionBrokerClient< } this.connection = createSessionBrokerConnection< - Info, - State, + HunkSessionInfo, + HunkSessionState, SessionBrokerSocketLike, - ServerMessage, - Result + HunkSessionServerMessage, + HunkSessionCommandResult >({ url: `${config.wsOrigin}${SESSION_BROKER_SOCKET_PATH}`, createSocket: (url) => new WebSocket(url) as unknown as SessionBrokerSocketLike, registration: this.registration, snapshot: this.snapshot, bridge: this.bridge, + protocolParsers: hunkSessionProtocolParsers, heartbeatIntervalMs: HEARTBEAT_INTERVAL_MS, reconnectDelayMs: this.timing.reconnectDelayMs ?? RECONNECT_DELAY_MS, resolveClose: (event) => diff --git a/src/session/broker/brokerLauncher.test.ts b/src/session/broker/brokerLauncher.test.ts index 937b7babc..e3b9d6097 100644 --- a/src/session/broker/brokerLauncher.test.ts +++ b/src/session/broker/brokerLauncher.test.ts @@ -6,6 +6,7 @@ import { join } from "node:path"; import { ensureSessionBrokerAvailable, isLoopbackPortReachable, + parseSessionBrokerHealth, resolveDaemonLaunchCommand, resolveSessionBrokerRuntimePaths, } from "./brokerLauncher"; @@ -34,6 +35,24 @@ afterEach(() => { }); describe("session daemon launcher", () => { + test("strictly parses minimal and legacy health responses", () => { + expect(parseSessionBrokerHealth({ ok: true })).toEqual({ ok: true }); + expect( + parseSessionBrokerHealth({ + ok: true, + pid: 123, + sessions: 1, + pendingCommands: 0, + startedAt: "2026-04-15T00:00:00.000Z", + uptimeMs: 10, + staleSessionTtlMs: 45_000, + paths: { health: "/health", socket: "/session" }, + }), + ).toMatchObject({ ok: true, pid: 123, sessions: 1 }); + for (const value of [null, [], { ok: "yes" }, { ok: true, pid: 1.5 }, { ok: true, extra: 1 }]) { + expect(parseSessionBrokerHealth(value)).toBeNull(); + } + }); test("reuses the current script entrypoint when Hunk is running from source or a JS wrapper", () => { expect(resolveDaemonLaunchCommand(["bun", "src/main.tsx", "diff"], "/usr/bin/bun")).toEqual({ command: "/usr/bin/bun", diff --git a/src/session/broker/brokerLauncher.ts b/src/session/broker/brokerLauncher.ts index 1c8691875..3976484df 100644 --- a/src/session/broker/brokerLauncher.ts +++ b/src/session/broker/brokerLauncher.ts @@ -4,6 +4,11 @@ import { existsSync, mkdirSync, readFileSync, rmSync, statSync, writeFileSync } import { connect } from "node:net"; import { tmpdir } from "node:os"; import { join } from "node:path"; +import { + parseBrokerSafeInteger, + parseBrokerString, + parseExactBrokerRecord, +} from "@hunk/session-broker-core"; import { resolveSessionBrokerConfig, type ResolvedSessionBrokerConfig } from "./brokerConfig"; const SCRIPT_ENTRYPOINT_PATTERN = /[\\/]|\.(?:[cm]?js|tsx?)$/; @@ -317,6 +322,60 @@ export interface SessionBrokerHealth { staleSessionTtlMs?: number; } +/** Parse the minimal or legacy-rich health response without trusting cross-process JSON. */ +export function parseSessionBrokerHealth(value: unknown): SessionBrokerHealth | null { + try { + const record = parseExactBrokerRecord( + value, + ["ok"] as const, + [ + "pid", + "sessions", + "pendingCommands", + "startedAt", + "uptimeMs", + "sessionApi", + "sessionCapabilities", + "sessionSocket", + "staleSessionTtlMs", + "paths", + ] as const, + ); + if (record.ok !== true) return null; + const parsed: SessionBrokerHealth = { ok: true }; + for (const key of [ + "pid", + "sessions", + "pendingCommands", + "uptimeMs", + "staleSessionTtlMs", + ] as const) { + if (record[key] !== undefined) parsed[key] = parseBrokerSafeInteger(record[key]); + } + for (const key of [ + "startedAt", + "sessionApi", + "sessionCapabilities", + "sessionSocket", + ] as const) { + if (record[key] !== undefined) parsed[key] = parseBrokerString(record[key]); + } + // Generic rich health used to carry a paths object. It is accepted only as one exact bounded + // compatibility shape and intentionally not projected into caller authority. + if (record.paths !== undefined) { + const paths = parseExactBrokerRecord( + record.paths, + ["health", "socket"] as const, + ["api", "capabilities"] as const, + ); + for (const path of Object.values(paths)) parseBrokerString(path); + } + return parsed; + } catch { + return null; + } +} + /** Read the daemon's health payload when one is reachable on the configured loopback port. */ export async function readSessionBrokerHealth( config: ResolvedSessionBrokerConfig = resolveSessionBrokerConfig(), @@ -334,7 +393,7 @@ export async function readSessionBrokerHealth( return null; } - return (await response.json()) as SessionBrokerHealth; + return parseSessionBrokerHealth(await response.json()); } catch { return null; } finally { diff --git a/src/session/broker/brokerServer.helpers.test.ts b/src/session/broker/brokerServer.helpers.test.ts index 77b2520c5..c7040431a 100644 --- a/src/session/broker/brokerServer.helpers.test.ts +++ b/src/session/broker/brokerServer.helpers.test.ts @@ -212,6 +212,37 @@ describe("handleSessionApiRequest", () => { expect(response.status).toBe(400); }); + test("rejects malformed nested HTTP daemon request bodies before state dispatch", async () => { + const { state, calls } = createFakeState(); + const malformed = [ + { action: "get", selector: { sessionId: "s-1", extra: true } }, + { + action: "reload", + selector: { sessionId: "s-1" }, + nextInput: { kind: "vcs", staged: false, options: { tabWidth: 0 } }, + }, + { + action: "comment-apply", + selector: { sessionId: "s-1" }, + comments: [{ filePath: "a.ts", summary: "note", hunkNumber: 0 }], + revealMode: "first", + }, + ]; + + for (const body of malformed) { + const response = await handleSessionApiRequest( + state, + new Request(`http://127.0.0.1:${PORT}/session-api`, { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify(body), + }), + ); + expect(response.status).toBe(400); + } + expect(calls).toEqual([]); + }); + test("routes list/get/context/review to the matching state methods", async () => { const { state, calls } = createFakeState(); for (const action of ["list", "get", "context", "review"] as const) { diff --git a/src/session/broker/brokerServer.ts b/src/session/broker/brokerServer.ts index dd44cbcb7..4bd580aa0 100644 --- a/src/session/broker/brokerServer.ts +++ b/src/session/broker/brokerServer.ts @@ -40,6 +40,7 @@ import { type SessionDaemonResponse, } from "../protocol"; import { parseSessionDaemonRequest } from "../protocolSchemas"; +import { hunkSessionProtocolParsers } from "./protocolParsers"; const DEFAULT_STALE_SESSION_TTL_MS = 45_000; const DEFAULT_STALE_SESSION_SWEEP_INTERVAL_MS = 15_000; @@ -460,6 +461,7 @@ function createHunkBrokerController( state: HunkSessionBrokerState, ): SessionBrokerController { return { + protocolParsers: hunkSessionProtocolParsers, listSessions: () => state.listSessions(), getSession: (selector) => state.getSession(selector), getSessionCount: () => state.getSessionCount(), diff --git a/src/session/broker/protocolParsers.ts b/src/session/broker/protocolParsers.ts new file mode 100644 index 000000000..783990d67 --- /dev/null +++ b/src/session/broker/protocolParsers.ts @@ -0,0 +1,209 @@ +import { createSessionBrokerProtocolParsers } from "@hunk/session-broker"; +import { z } from "zod"; +import { parseSessionRegistration, parseSessionSnapshot } from "./wire"; +import { + HUNK_REVIEW_PROTOCOL_VERSION, + parseHunkReviewActionEnvelope, + parseHunkReviewResourceReadEnvelope, +} from "../reviewProtocol"; +import type { + HunkSessionCommandResult, + HunkSessionInfo, + HunkSessionServerMessage, + HunkSessionState, +} from "../types"; +import { HUNK_SESSION_DAEMON_VERSION } from "../protocol"; +import { cliInputSchema, hunkCommandResultSchemas } from "../protocolSchemas"; + +const selectorFields = { + sessionId: z.string().min(1).max(128).optional(), + sessionPath: z.string().min(1).max(4096).optional(), + repoRoot: z.string().min(1).max(4096).optional(), + repoBoundary: z.string().min(1).max(4096).optional(), +}; +const side = z.enum(["old", "new"]); +const positive = z.int().positive(); +const nonnegative = z.int().nonnegative(); +const optionalString = z.string().min(1).max(4096).optional(); +const commentItem = z.strictObject({ + filePath: z.string().min(1).max(4096), + hunkIndex: nonnegative.optional(), + side: side.optional(), + line: positive.optional(), + summary: z.string().min(1).max(4096), + rationale: optionalString, + markup: optionalString, + author: optionalString, +}); + +const commandInputs = { + comment: z.strictObject({ + ...selectorFields, + filePath: z.string().min(1).max(4096), + hunkIndex: nonnegative.optional(), + side: side.optional(), + line: positive.optional(), + summary: z.string().min(1).max(4096), + rationale: optionalString, + markup: optionalString, + author: optionalString, + reveal: z.boolean().optional(), + }), + comment_batch: z.strictObject({ + ...selectorFields, + comments: z.array(commentItem), + revealMode: z.enum(["none", "first"]).optional(), + }), + navigate_to_hunk: z.strictObject({ + ...selectorFields, + filePath: optionalString, + hunkIndex: nonnegative.optional(), + side: side.optional(), + line: positive.optional(), + commentDirection: z.enum(["next", "prev"]).optional(), + }), + reload_session: z.strictObject({ + ...selectorFields, + nextInput: cliInputSchema, + sourcePath: optionalString, + }), + remove_comment: z.strictObject({ + ...selectorFields, + commentId: z.string().min(1).max(128), + }), + clear_comments: z.strictObject({ + ...selectorFields, + filePath: optionalString, + includeUser: z.boolean().optional(), + }), + read_review_resource: z + .strictObject({ + ...selectorFields, + protocolVersion: z.literal(HUNK_REVIEW_PROTOCOL_VERSION), + actor: z.unknown(), + request: z.unknown(), + }) + .refine( + ({ protocolVersion, actor, request }) => + parseHunkReviewResourceReadEnvelope({ protocolVersion, actor, request }).ok, + ), + apply_review_action: z + .strictObject({ + ...selectorFields, + protocolVersion: z.literal(HUNK_REVIEW_PROTOCOL_VERSION), + generation: z.string(), + expectedStateRevision: nonnegative.optional(), + actor: z.unknown(), + action: z.unknown(), + }) + .refine( + ({ protocolVersion, generation, expectedStateRevision, actor, action }) => + parseHunkReviewActionEnvelope({ + protocolVersion, + generation, + ...(expectedStateRevision === undefined ? {} : { expectedStateRevision }), + actor, + action, + }).ok, + ), + highlight: z.strictObject({ + ...selectorFields, + filePath: z.string().min(1).max(4096), + side, + line: positive, + start: nonnegative, + end: positive, + tone: z.enum(["match", "current", "info", "warning", "error"]).optional(), + reveal: z.boolean().optional(), + }), + clear_highlights: z.strictObject({ + ...selectorFields, + filePath: optionalString, + }), +} as const; + +const failure = z.strictObject({ + ok: z.literal(false), + code: z.enum([ + "unknown-resource", + "resource-unavailable", + "resource-too-large", + "resource-integrity", + "invalid-range", + "stale-generation", + "invalid-request", + "file-not-found", + "hunk-not-found", + "gap-not-found", + "draft-missing", + "note-not-found", + "missing-fact", + ]), + message: z.string(), + currentGeneration: z.string(), +}); +const results = { + comment: hunkCommandResultSchemas.comment, + comment_batch: hunkCommandResultSchemas.comment_batch, + navigate_to_hunk: hunkCommandResultSchemas.navigate_to_hunk, + reload_session: hunkCommandResultSchemas.reload_session, + remove_comment: hunkCommandResultSchemas.remove_comment, + clear_comments: hunkCommandResultSchemas.clear_comments, + read_review_resource: z.union([ + failure, + z.strictObject({ + ok: z.literal(true), + chunk: z.strictObject({ + generation: z.string(), + resourceId: z.string(), + offset: nonnegative, + byteLength: nonnegative, + encoding: z.literal("base64"), + data: z.string(), + contentDigest: z.string(), + contentSize: nonnegative, + eof: z.boolean(), + }), + }), + ]), + apply_review_action: z.union([ + failure, + z.strictObject({ + ok: z.literal(true), + generation: z.string(), + stateRevision: nonnegative, + }), + ]), + highlight: hunkCommandResultSchemas.highlight, + clear_highlights: hunkCommandResultSchemas.clear_highlights, +} as const; + +/** Parse with one strict Hunk schema while keeping Zod diagnostics inside the app boundary. */ +function schemaParser(schema: z.ZodType) { + return (value: unknown): unknown | null => { + const parsed = schema.safeParse(value); + return parsed.success ? parsed.data : null; + }; +} + +/** Fixed Hunk app parser registry shared by the daemon and every producer connection. */ +export const hunkSessionProtocolParsers = createSessionBrokerProtocolParsers< + HunkSessionInfo, + HunkSessionState, + HunkSessionServerMessage, + HunkSessionCommandResult +>({ + appRevision: HUNK_SESSION_DAEMON_VERSION, + features: [], + parseRegistration: parseSessionRegistration, + parseSnapshot: parseSessionSnapshot, + commands: (Object.keys(commandInputs) as Array).map((command) => ({ + command, + version: HUNK_REVIEW_PROTOCOL_VERSION, + parseInput: schemaParser(commandInputs[command]), + parseResult: (value: unknown): HunkSessionCommandResult | null => { + const parsed = results[command].safeParse(value); + return parsed.success ? (parsed.data as HunkSessionCommandResult) : null; + }, + })), +}); diff --git a/src/session/broker/state.ts b/src/session/broker/state.ts index a4478acdf..48c8321c6 100644 --- a/src/session/broker/state.ts +++ b/src/session/broker/state.ts @@ -35,6 +35,7 @@ import type { SessionReview, SessionReviewFile, } from "../types"; +import { hunkSessionProtocolParsers } from "./protocolParsers"; import { parseSessionRegistration, parseSessionSnapshot } from "./wire"; import { ReviewMirror, @@ -75,10 +76,16 @@ const hunkSessionBrokerView: SessionBrokerViewAdapter< ListedSession, SelectedSessionContext, SessionReview, - SessionLiveCommentSummary + SessionLiveCommentSummary, + HunkSessionServerMessage, + HunkSessionCommandResult > = { parseRegistration: parseSessionRegistration, parseSnapshot: parseSessionSnapshot, + parseCommandInput: (command, version, value) => + hunkSessionProtocolParsers.parseCommandInput(command, version, value), + parseCommandResult: (command, version, value) => + hunkSessionProtocolParsers.parseCommandResult(command, version, value), buildListedSession: buildListedHunkSession, buildSelectedContext: buildSelectedHunkSessionContext, buildSessionReview: buildHunkSessionReview, diff --git a/src/session/broker/wire.test.ts b/src/session/broker/wire.test.ts index f97332308..5e933c6f1 100644 --- a/src/session/broker/wire.test.ts +++ b/src/session/broker/wire.test.ts @@ -16,7 +16,12 @@ function createRegistration(files: unknown[]) { pid: 123, cwd: "/repo", launchedAt: "2026-03-22T00:00:00.000Z", - info: { inputKind: "vcs", title: "repo working tree", sourceLabel: "/repo", files }, + info: { + inputKind: "vcs", + title: "repo working tree", + sourceLabel: "/repo", + files, + }, }; } @@ -45,7 +50,7 @@ function createValidComment(overrides: Record = {}) { } describe("hunk session wire parsing", () => { - test("snapshot comment counts only include validated comment summaries", () => { + test("snapshot rejects malformed comment summaries instead of partially filtering them", () => { const snapshot = parseSessionSnapshot({ updatedAt: "2026-03-22T00:00:00.000Z", state: { @@ -64,12 +69,10 @@ describe("hunk session wire parsing", () => { }, }); - expect(snapshot).not.toBeNull(); - expect(snapshot?.state.liveComments).toHaveLength(1); - expect(snapshot?.state.liveCommentCount).toBe(1); + expect(snapshot).toBeNull(); }); - test("snapshot carries the live note markup width and drops invalid values", () => { + test("snapshot carries the live note markup width and rejects invalid values", () => { const parse = (noteMarkupWidth: unknown) => parseSessionSnapshot({ updatedAt: "2026-03-22T00:00:00.000Z", @@ -82,7 +85,7 @@ describe("hunk session wire parsing", () => { }); expect(parse(112)?.state.noteMarkupWidth).toBe(112); - expect(parse("wide")?.state.noteMarkupWidth).toBeUndefined(); + expect(parse("wide")).toBeNull(); expect(parse(undefined)?.state.noteMarkupWidth).toBeUndefined(); }); @@ -110,7 +113,7 @@ describe("hunk session wire parsing", () => { }); }); - test("registration preserves only recognized experimental feature ids", () => { + test("registration rejects malformed or unknown experimental feature ids", () => { const registration = parseSessionRegistration({ registrationVersion: SESSION_BROKER_REGISTRATION_VERSION, sessionId: "session-1", @@ -126,7 +129,7 @@ describe("hunk session wire parsing", () => { }, }); - expect(registration?.info.experimentalFeatures).toEqual(["stml"]); + expect(registration).toBeNull(); }); test("rejects registrations with more files than the cap", () => { @@ -146,12 +149,73 @@ describe("hunk session wire parsing", () => { expect(parseSessionRegistration(createRegistration([createFile({ hunks })]))).toBeNull(); }); - test("rejects files whose patch exceeds the byte cap", () => { + test("accepts legacy embedded patches beyond the generic string ceiling through the exact cap", () => { + for (const size of [4_097, MAX_REGISTRATION_PATCH_BYTES]) { + const patch = "x".repeat(size); + expect( + parseSessionRegistration(createRegistration([createFile({ patch })]))?.info.files[0]?.patch, + ).toHaveLength(size); + } + }); + + test("rejects a legacy embedded patch one byte beyond its cap", () => { const patch = "x".repeat(MAX_REGISTRATION_PATCH_BYTES + 1); expect(parseSessionRegistration(createRegistration([createFile({ patch })]))).toBeNull(); }); + test("accepts zero-based pure-add, pure-delete, and new-file hunk ranges", () => { + const ranges: Array<{ + oldRange: [number, number]; + newRange: [number, number]; + }> = [ + { oldRange: [0, 0], newRange: [1, 3] }, + { oldRange: [4, 2], newRange: [0, 0] }, + { oldRange: [0, 0], newRange: [0, 4] }, + ]; + const files = ranges.map((range, index) => + createFile({ + id: `file-${index}`, + path: `src/file-${index}.ts`, + hunks: [{ index: 0, header: "@@", ...range }], + }), + ); + + expect( + parseSessionRegistration(createRegistration(files))?.info.files.map((file) => file.hunks[0]), + ).toEqual(ranges.map((range) => ({ index: 0, header: "@@", ...range }))); + }); + + test("accepts zero-based selected ranges and review-note ranges in snapshots", () => { + const snapshot = parseSessionSnapshot({ + updatedAt: "2026-03-22T00:00:00.000Z", + state: { + selectedHunkIndex: 0, + selectedHunkOldRange: [0, 0], + selectedHunkNewRange: [0, 3], + showAgentNotes: true, + liveComments: [], + reviewNotes: [ + { + noteId: "note-1", + source: "user", + filePath: "new-file.ts", + oldRange: [0, 0], + newRange: [0, 3], + body: "New file", + createdAt: "2026-03-22T00:00:00.000Z", + }, + ], + }, + }); + + expect(snapshot?.state).toMatchObject({ + selectedHunkOldRange: [0, 0], + selectedHunkNewRange: [0, 3], + reviewNotes: [{ oldRange: [0, 0], newRange: [0, 3] }], + }); + }); + test("rejects snapshots with more live comments than the cap", () => { const liveComments = Array.from({ length: MAX_SNAPSHOT_LIVE_COMMENTS + 1 }, (_, index) => createValidComment({ commentId: `comment-${index}` }), @@ -176,7 +240,12 @@ describe("hunk session wire parsing", () => { const snapshot = parseSessionSnapshot({ updatedAt: "2026-03-22T00:00:00.000Z", - state: { selectedHunkIndex: 0, showAgentNotes: true, liveComments: [], reviewNotes }, + state: { + selectedHunkIndex: 0, + showAgentNotes: true, + liveComments: [], + reviewNotes, + }, }); expect(snapshot).toBeNull(); diff --git a/src/session/broker/wire.ts b/src/session/broker/wire.ts index 4ff45ce24..c92405b04 100644 --- a/src/session/broker/wire.ts +++ b/src/session/broker/wire.ts @@ -5,11 +5,13 @@ import { MAX_REGISTRATION_HUNKS_PER_FILE, MAX_REGISTRATION_PATCH_BYTES, MAX_SNAPSHOT_LIVE_COMMENTS, + BrokerProtocolError, MAX_SNAPSHOT_REVIEW_NOTES, brokerWireParsers, + parseBrokerString, + parseExactBrokerRecord, parseSessionRegistrationEnvelope, parseSessionSnapshotEnvelope, - utf8ByteLength, } from "@hunk/session-broker-core"; import { parseHunkReviewPublicationAddress, @@ -36,35 +38,42 @@ const REVIEW_INPUT_KINDS = new Set([ ]); const EXPERIMENTAL_FEATURE_SET = new Set(EXPERIMENTAL_FEATURES); -/** Preserve only recognized experimental feature ids from a session registration. */ +/** Parse unique recognized experimental feature ids without silently dropping malformed entries. */ function parseExperimentalFeatures(value: unknown): ExperimentalFeature[] { - if (!Array.isArray(value)) { - return []; + if (value === undefined) return []; + if (!Array.isArray(value)) throw new BrokerProtocolError("invalid-app-payload"); + if ( + value.some((feature) => typeof feature !== "string" || !EXPERIMENTAL_FEATURE_SET.has(feature)) + ) { + throw new BrokerProtocolError("invalid-app-payload"); } + return [...new Set(value)] as ExperimentalFeature[]; +} - return [...new Set(value)].filter( - (feature): feature is ExperimentalFeature => - typeof feature === "string" && EXPERIMENTAL_FEATURE_SET.has(feature), - ); +/** Read one app-owned object with an exact field set. */ +function exactRecord( + value: unknown, + required: readonly string[], + optional: readonly string[] = [], +) { + return parseExactBrokerRecord(value, required, optional); } /** Parse one optional diff-side line range tuple when the payload shape matches. */ function parseOptionalRange(value: unknown): [number, number] | undefined { + if (value === undefined) return undefined; if (!Array.isArray(value) || value.length !== 2) { - return undefined; + throw new BrokerProtocolError("invalid-app-payload"); } - - const start = brokerWireParsers.parsePositiveInt(value[0]); - const end = brokerWireParsers.parsePositiveInt(value[1]); - return start !== null && end !== null ? [start, end] : undefined; + const start = brokerWireParsers.parseNonNegativeInt(value[0]); + const end = brokerWireParsers.parseNonNegativeInt(value[1]); + if (start === null || end === null) throw new BrokerProtocolError("invalid-app-payload"); + return [start, end]; } /** Parse one registered review hunk from the app-owned session payload. */ function parseSessionReviewHunk(value: unknown): SessionReviewHunk | null { - const record = brokerWireParsers.asRecord(value); - if (!record) { - return null; - } + const record = exactRecord(value, ["index", "header"], ["oldRange", "newRange"]); const index = brokerWireParsers.parseNonNegativeInt(record.index); const header = brokerWireParsers.parseRequiredString(record.header); @@ -82,10 +91,11 @@ function parseSessionReviewHunk(value: unknown): SessionReviewHunk | null { /** Parse one registered review file from the app-owned session payload. */ function parseSessionReviewFile(value: unknown): SessionReviewFile | null { - const record = brokerWireParsers.asRecord(value); - if (!record) { - return null; - } + const record = exactRecord( + value, + ["id", "path", "additions", "deletions", "hunks"], + ["previousPath", "patch", "hunkCount"], + ); const id = brokerWireParsers.parseRequiredString(record.id); const path = brokerWireParsers.parseRequiredString(record.path); @@ -98,6 +108,16 @@ function parseSessionReviewFile(value: unknown): SessionReviewFile | null { if (!Array.isArray(record.hunks) || record.hunks.length > MAX_REGISTRATION_HUNKS_PER_FILE) { return null; } + const assertedHunkCount = + record.hunkCount === undefined + ? undefined + : brokerWireParsers.parseNonNegativeInt(record.hunkCount); + if ( + record.hunkCount !== undefined && + (assertedHunkCount === null || assertedHunkCount !== record.hunks.length) + ) { + return null; + } const hunks = record.hunks.map(parseSessionReviewHunk); if (hunks.some((hunk) => hunk === null)) { @@ -106,10 +126,12 @@ function parseSessionReviewFile(value: unknown): SessionReviewFile | null { // Reject files whose patch text alone would blow the per-file memory budget instead of // silently dropping it, so an oversized registration fails loudly rather than half-loading. - const patch = brokerWireParsers.parseOptionalString(record.patch); - if (patch !== undefined && utf8ByteLength(patch) > MAX_REGISTRATION_PATCH_BYTES) { - return null; - } + const patch = + record.patch === undefined + ? undefined + : parseBrokerString(record.patch, { + maxBytes: MAX_REGISTRATION_PATCH_BYTES, + }); return { id, @@ -134,10 +156,11 @@ function parseReviewInputKind(value: unknown): CliInput["kind"] | null { /** Parse one live comment summary from the app-owned snapshot payload. */ function parseSessionLiveCommentSummary(value: unknown): SessionLiveCommentSummary | null { - const record = brokerWireParsers.asRecord(value); - if (!record) { - return null; - } + const record = exactRecord( + value, + ["commentId", "filePath", "hunkIndex", "summary", "createdAt", "line", "side"], + ["rationale", "author"], + ); const commentId = brokerWireParsers.parseRequiredString(record.commentId); const filePath = brokerWireParsers.parseRequiredString(record.filePath); @@ -173,10 +196,11 @@ function parseSessionLiveCommentSummary(value: unknown): SessionLiveCommentSumma /** Parse one review note summary from the app-owned snapshot payload. */ function parseSessionReviewNoteSummary(value: unknown): SessionReviewNoteSummary | null { - const record = brokerWireParsers.asRecord(value); - if (!record) { - return null; - } + const record = exactRecord( + value, + ["noteId", "source", "filePath", "body", "createdAt"], + ["hunkIndex", "oldRange", "newRange", "title", "author", "updatedAt", "editable"], + ); const noteId = brokerWireParsers.parseRequiredString(record.noteId); const filePath = brokerWireParsers.parseRequiredString(record.filePath); @@ -196,11 +220,18 @@ function parseSessionReviewNoteSummary(value: unknown): SessionReviewNoteSummary return null; } + const hunkIndex = + record.hunkIndex === undefined + ? undefined + : brokerWireParsers.parseNonNegativeInt(record.hunkIndex); + if (record.hunkIndex !== undefined && hunkIndex === null) return null; + if (record.editable !== undefined && typeof record.editable !== "boolean") return null; + return { noteId, source, filePath, - hunkIndex: brokerWireParsers.parseNonNegativeInt(record.hunkIndex) ?? undefined, + hunkIndex: hunkIndex ?? undefined, oldRange: parseOptionalRange(record.oldRange), newRange: parseOptionalRange(record.newRange), body, @@ -214,10 +245,12 @@ function parseSessionReviewNoteSummary(value: unknown): SessionReviewNoteSummary /** Parse the app-owned registration info embedded inside one broker registration envelope. */ function parseHunkSessionInfo(value: unknown): HunkSessionInfo | null { - const record = brokerWireParsers.asRecord(value); - if (!record || !Array.isArray(record.files) || record.files.length > MAX_REGISTRATION_FILES) { - return null; - } + const record = exactRecord( + value, + ["inputKind", "title", "sourceLabel", "files"], + ["experimentalFeatures", "reviewCatalog", "reviewCapabilityDigest"], + ); + if (!Array.isArray(record.files) || record.files.length > MAX_REGISTRATION_FILES) return null; const inputKind = parseReviewInputKind(record.inputKind); const title = brokerWireParsers.parseRequiredString(record.title); @@ -265,9 +298,22 @@ function parseHunkSessionInfo(value: unknown): HunkSessionInfo | null { /** Parse the app-owned snapshot state embedded inside one broker snapshot envelope. */ function parseHunkSessionState(value: unknown): HunkSessionState | null { - const record = brokerWireParsers.asRecord(value); + const record = exactRecord( + value, + ["liveComments", "selectedHunkIndex", "showAgentNotes"], + [ + "selectedFileId", + "selectedFilePath", + "selectedHunkOldRange", + "selectedHunkNewRange", + "noteMarkupWidth", + "liveCommentCount", + "reviewNoteCount", + "reviewNotes", + "reviewPublication", + ], + ); if ( - !record || !Array.isArray(record.liveComments) || record.liveComments.length > MAX_SNAPSHOT_LIVE_COMMENTS || (Array.isArray(record.reviewNotes) && record.reviewNotes.length > MAX_SNAPSHOT_REVIEW_NOTES) @@ -291,12 +337,36 @@ function parseHunkSessionState(value: unknown): HunkSessionState | null { return null; } - const liveComments = record.liveComments - .map(parseSessionLiveCommentSummary) - .filter((comment): comment is SessionLiveCommentSummary => comment !== null); - const reviewNotes = (Array.isArray(record.reviewNotes) ? record.reviewNotes : []) - .map(parseSessionReviewNoteSummary) - .filter((note): note is SessionReviewNoteSummary => note !== null); + if (record.reviewNotes !== undefined && !Array.isArray(record.reviewNotes)) return null; + const assertedLiveCommentCount = + record.liveCommentCount === undefined + ? undefined + : brokerWireParsers.parseNonNegativeInt(record.liveCommentCount); + const assertedReviewNoteCount = + record.reviewNoteCount === undefined + ? undefined + : brokerWireParsers.parseNonNegativeInt(record.reviewNoteCount); + if ( + (record.liveCommentCount !== undefined && assertedLiveCommentCount === null) || + (record.reviewNoteCount !== undefined && assertedReviewNoteCount === null) + ) { + return null; + } + const liveComments = record.liveComments.map(parseSessionLiveCommentSummary); + const reviewNotes = (record.reviewNotes ?? []).map(parseSessionReviewNoteSummary); + if ( + liveComments.some((comment) => comment === null) || + reviewNotes.some((note) => note === null) || + (assertedLiveCommentCount !== undefined && assertedLiveCommentCount !== liveComments.length) || + (assertedReviewNoteCount !== undefined && assertedReviewNoteCount !== reviewNotes.length) + ) { + return null; + } + const noteMarkupWidth = + record.noteMarkupWidth === undefined + ? undefined + : brokerWireParsers.parseNonNegativeInt(record.noteMarkupWidth); + if (record.noteMarkupWidth !== undefined && noteMarkupWidth === null) return null; return { selectedFileId: brokerWireParsers.parseOptionalString(record.selectedFileId), @@ -305,11 +375,11 @@ function parseHunkSessionState(value: unknown): HunkSessionState | null { selectedHunkOldRange: parseOptionalRange(record.selectedHunkOldRange), selectedHunkNewRange: parseOptionalRange(record.selectedHunkNewRange), showAgentNotes, - noteMarkupWidth: brokerWireParsers.parseNonNegativeInt(record.noteMarkupWidth) ?? undefined, + noteMarkupWidth: noteMarkupWidth ?? undefined, liveCommentCount: liveComments.length, - liveComments, + liveComments: liveComments as SessionLiveCommentSummary[], reviewNoteCount: reviewNotes.length, - reviewNotes, + reviewNotes: reviewNotes as SessionReviewNoteSummary[], ...(reviewPublication ? { reviewPublication } : {}), }; } diff --git a/src/session/client/capabilities.ts b/src/session/client/capabilities.ts index 0f7cc0118..2892b5e98 100644 --- a/src/session/client/capabilities.ts +++ b/src/session/client/capabilities.ts @@ -2,12 +2,8 @@ import { resolveSessionBrokerConfig, type ResolvedSessionBrokerConfig, } from "../broker/brokerConfig"; -import { - HUNK_SESSION_API_VERSION, - HUNK_SESSION_CAPABILITIES_PATH, - HUNK_SESSION_DAEMON_VERSION, - type SessionDaemonCapabilities, -} from "../protocol"; +import { HUNK_SESSION_CAPABILITIES_PATH, type SessionDaemonCapabilities } from "../protocol"; +import { parseSessionDaemonCapabilities } from "../protocolSchemas"; import { HUNK_SESSION_DAEMON_HTTP_TIMEOUT_MS, requestSessionDaemonHttp } from "./daemonHttp"; export const HUNK_DAEMON_UPGRADE_RESTART_NOTICE = @@ -47,18 +43,7 @@ export async function readHunkSessionDaemonCapabilities( return null; } - if ( - !capabilities || - typeof capabilities !== "object" || - (capabilities as { version?: unknown }).version !== HUNK_SESSION_API_VERSION || - (capabilities as { daemonVersion?: unknown }).daemonVersion !== - HUNK_SESSION_DAEMON_VERSION || - !Array.isArray((capabilities as { actions?: unknown }).actions) - ) { - return null; - } - - return capabilities as SessionDaemonCapabilities; + return parseSessionDaemonCapabilities(capabilities); }, }); } diff --git a/src/session/protocol.ts b/src/session/protocol.ts index ae56001da..073142f9a 100644 --- a/src/session/protocol.ts +++ b/src/session/protocol.ts @@ -145,17 +145,20 @@ export type SessionDaemonRequest = filePath?: string; }; -export type SessionDaemonResponse = - | { sessions: ListedSession[] } - | { session: ListedSession } - | { context: SelectedSessionContext } - | { review: SessionReview } - | { result: NavigatedSelectionResult } - | { result: ReloadedSessionResult } - | { result: AppliedCommentResult } - | { result: AppliedCommentBatchResult } - | { comments: Array } - | { result: RemovedCommentResult } - | { result: ClearedCommentsResult } - | { result: AppliedHighlightResult } - | { result: ClearedHighlightsResult }; +export interface SessionDaemonResponses { + list: { sessions: ListedSession[] }; + get: { session: ListedSession }; + context: { context: SelectedSessionContext }; + review: { review: SessionReview }; + navigate: { result: NavigatedSelectionResult }; + reload: { result: ReloadedSessionResult }; + "comment-add": { result: AppliedCommentResult }; + "comment-apply": { result: AppliedCommentBatchResult }; + "comment-list": { comments: Array }; + "comment-rm": { result: RemovedCommentResult }; + "comment-clear": { result: ClearedCommentsResult }; + "highlight-add": { result: AppliedHighlightResult }; + "highlight-clear": { result: ClearedHighlightsResult }; +} + +export type SessionDaemonResponse = SessionDaemonResponses[SessionDaemonAction]; diff --git a/src/session/protocolSchemas.test.ts b/src/session/protocolSchemas.test.ts index 6504965f4..fff4b59f2 100644 --- a/src/session/protocolSchemas.test.ts +++ b/src/session/protocolSchemas.test.ts @@ -1,7 +1,18 @@ import { describe, expect, test } from "bun:test"; import type { z } from "zod"; -import type { SessionDaemonRequest } from "./protocol"; -import { parseSessionDaemonRequest, sessionDaemonRequestSchema } from "./protocolSchemas"; +import type { CliInput } from "../core/run/commandInputs"; +import { + HUNK_SESSION_API_VERSION, + HUNK_SESSION_DAEMON_VERSION, + type SessionDaemonRequest, +} from "./protocol"; +import { + cliInputSchema, + parseSessionDaemonCapabilities, + parseSessionDaemonRequest, + parseSessionDaemonResponse, + sessionDaemonRequestSchema, +} from "./protocolSchemas"; /** Strict structural equality; `true` only when A and B are the same type. */ type Equal = @@ -14,8 +25,40 @@ const _schemaMatchesProtocol: Equal< SessionDaemonRequest > = true; void _schemaMatchesProtocol; +const _cliInputSchemaMatchesProtocol: Equal, CliInput> = true; +void _cliInputSchemaMatchesProtocol; describe("session daemon request validation", () => { + test("strictly parses cross-process capabilities", () => { + expect( + parseSessionDaemonCapabilities({ + version: HUNK_SESSION_API_VERSION, + daemonVersion: HUNK_SESSION_DAEMON_VERSION, + actions: ["list", "get"], + }), + ).toEqual({ + version: HUNK_SESSION_API_VERSION, + daemonVersion: HUNK_SESSION_DAEMON_VERSION, + actions: ["list", "get"], + }); + for (const value of [ + null, + [], + { + version: HUNK_SESSION_API_VERSION, + daemonVersion: HUNK_SESSION_DAEMON_VERSION, + actions: ["unknown"], + }, + { + version: HUNK_SESSION_API_VERSION, + daemonVersion: HUNK_SESSION_DAEMON_VERSION, + actions: ["list"], + extra: true, + }, + ]) { + expect(parseSessionDaemonCapabilities(value)).toBeNull(); + } + }); test("accepts every wire-shaped action payload", () => { const requests: unknown[] = [ { action: "list" }, @@ -25,7 +68,12 @@ describe("session daemon request validation", () => { selector: { repoRoot: "/repo/nested", repoBoundary: "/repo" }, }, { action: "review", selector: { sessionId: "s-1" } }, - { action: "review", selector: { sessionId: "s-1" }, includePatch: true, includeNotes: true }, + { + action: "review", + selector: { sessionId: "s-1" }, + includePatch: true, + includeNotes: true, + }, { action: "navigate", selector: { sessionId: "s-1" }, hunkNumber: 2 }, { action: "navigate", @@ -34,8 +82,16 @@ describe("session daemon request validation", () => { side: "new", line: 12, }, - { action: "navigate", selector: { sessionId: "s-1" }, commentDirection: "next" }, - { action: "navigate", selector: { sessionId: "s-1" }, commentId: "comment-1" }, + { + action: "navigate", + selector: { sessionId: "s-1" }, + commentDirection: "next", + }, + { + action: "navigate", + selector: { sessionId: "s-1" }, + commentId: "comment-1", + }, { action: "reload", selector: { sessionId: "s-1" }, @@ -57,8 +113,16 @@ describe("session daemon request validation", () => { revealMode: "first", }, { action: "comment-list", selector: { sessionId: "s-1" }, type: "user" }, - { action: "comment-rm", selector: { sessionId: "s-1" }, commentId: "c-1" }, - { action: "comment-clear", selector: { sessionId: "s-1" }, includeUser: true }, + { + action: "comment-rm", + selector: { sessionId: "s-1" }, + commentId: "c-1", + }, + { + action: "comment-clear", + selector: { sessionId: "s-1" }, + includeUser: true, + }, { action: "highlight-add", selector: { sessionId: "s-1" }, @@ -80,7 +144,11 @@ describe("session daemon request validation", () => { end: 9, reveal: false, }, - { action: "highlight-clear", selector: { sessionId: "s-1" }, filePath: "a.ts" }, + { + action: "highlight-clear", + selector: { sessionId: "s-1" }, + filePath: "a.ts", + }, { action: "highlight-clear", selector: { sessionId: "s-1" } }, ]; @@ -89,6 +157,52 @@ describe("session daemon request validation", () => { } }); + test("accepts zero-based ranges in navigation responses", () => { + expect( + parseSessionDaemonResponse("navigate", { + result: { + fileId: "file-1", + filePath: "new-file.ts", + hunkIndex: 0, + selectedHunk: { index: 0, oldRange: [0, 0], newRange: [0, 4] }, + }, + }), + ).toEqual({ + result: { + fileId: "file-1", + filePath: "new-file.ts", + hunkIndex: 0, + selectedHunk: { index: 0, oldRange: [0, 0], newRange: [0, 4] }, + }, + }); + }); + + test("rejects malformed action-specific responses with stable errors", () => { + for (const [action, body] of [ + ["list", { sessions: "not-an-array" }], + ["get", { session: { sessionId: "partial" } }], + ["context", { context: { sessionId: "partial" } }], + ["review", { review: { files: [] } }], + ["navigate", { result: { fileId: "file-1", filePath: "a.ts", hunkIndex: -1 } }], + ["comment-list", { comments: [{ commentId: "partial" }] }], + ["highlight-clear", { result: { removedCount: "two", remainingCount: 0 } }], + ] as const) { + expect(() => parseSessionDaemonResponse(action, body)).toThrow( + `Invalid Hunk session daemon response for ${action}.`, + ); + } + expect(() => + parseSessionDaemonResponse("navigate", { + result: { + fileId: "file-1", + filePath: "a.ts", + hunkIndex: 0, + unknown: true, + }, + }), + ).toThrow("Invalid Hunk session daemon response for navigate."); + }); + test("rejects malformed highlight payloads", () => { expect(() => parseSessionDaemonRequest({ @@ -152,14 +266,49 @@ describe("session daemon request validation", () => { ).toThrow(/side/); }); + test("rejects deterministic nested Hunk command mutations", () => { + const malformed = [ + { + action: "reload", + selector: { sessionId: "s-1" }, + nextInput: { kind: "vcs", staged: false, options: { tabWidth: 0 } }, + }, + { + action: "reload", + selector: { sessionId: "s-1" }, + nextInput: { kind: "patch", options: {}, unknown: true }, + }, + { + action: "comment-apply", + selector: { sessionId: "s-1" }, + comments: [{ filePath: "a.ts", summary: "note", hunkNumber: 0 }], + revealMode: "first", + }, + ...Array.from({ length: 8 }, (_, index) => ({ + action: "navigate", + selector: index % 2 === 0 ? { sessionId: index } : { sessionId: "s-1", extra: index }, + hunkNumber: index + 1, + })), + ]; + for (const value of malformed) { + expect(() => parseSessionDaemonRequest(value)).toThrow("Invalid session API request:"); + } + }); + test("rejects non-object payloads and missing required fields", () => { expect(() => parseSessionDaemonRequest("list")).toThrow(/Invalid session API request/); expect(() => parseSessionDaemonRequest(null)).toThrow(/Invalid session API request/); expect(() => - parseSessionDaemonRequest({ action: "comment-rm", selector: { sessionId: "s-1" } }), + parseSessionDaemonRequest({ + action: "comment-rm", + selector: { sessionId: "s-1" }, + }), ).toThrow(/commentId/); expect(() => - parseSessionDaemonRequest({ action: "reload", selector: { sessionId: "s-1" } }), + parseSessionDaemonRequest({ + action: "reload", + selector: { sessionId: "s-1" }, + }), ).toThrow(/nextInput/); }); }); diff --git a/src/session/protocolSchemas.ts b/src/session/protocolSchemas.ts index faa4836e1..9e09fd87d 100644 --- a/src/session/protocolSchemas.ts +++ b/src/session/protocolSchemas.ts @@ -1,6 +1,14 @@ import { z } from "zod"; import type { CliInput } from "../core/run/commandInputs"; -import type { SessionDaemonRequest } from "./protocol"; +import { EXPERIMENTAL_FEATURES } from "../core/run/experimental"; +import { + HUNK_SESSION_API_VERSION, + HUNK_SESSION_DAEMON_VERSION, + type SessionDaemonAction, + type SessionDaemonCapabilities, + type SessionDaemonRequest, + type SessionDaemonResponses, +} from "./protocol"; /** * Runtime validation for the session daemon's HTTP action surface. @@ -20,20 +28,96 @@ const selectorSchema = z.strictObject({ }); const sideSchema = z.enum(["old", "new"]); +const sessionDaemonActionSchema = z.enum([ + "list", + "get", + "context", + "review", + "navigate", + "reload", + "comment-add", + "comment-apply", + "comment-list", + "comment-rm", + "comment-clear", + "highlight-add", + "highlight-clear", +]); -/** - * Reload payloads embed a full parsed CLI input tree whose deep shape is owned by the CLI - * parser, so this envelope check is intentionally shallow: an object with a string `kind` - * discriminant. A well-formed-but-wrong tree still reaches the reload path, where thrown errors - * surface through the daemon's JSON error response rather than a schema rejection. - */ -const nextInputSchema = z.custom( - (value) => - typeof value === "object" && - value !== null && - !Array.isArray(value) && - typeof (value as { kind?: unknown }).kind === "string", -); +const sessionDaemonCapabilitiesSchema = z.strictObject({ + version: z.literal(HUNK_SESSION_API_VERSION), + daemonVersion: z.literal(HUNK_SESSION_DAEMON_VERSION), + actions: z.array(sessionDaemonActionSchema), +}); + +const commonOptionsSchema = z.strictObject({ + mode: z.enum(["auto", "split", "stack"]).optional(), + cursorLine: z.enum(["row", "number", "off"]).optional(), + vcs: z.string().optional(), + theme: z.string().optional(), + agentContext: z.string().optional(), + pager: z.boolean().optional(), + watch: z.boolean().optional(), + experimental: z.boolean().optional(), + fast: z.boolean().optional(), + excludeUntracked: z.boolean().optional(), + lineNumbers: z.boolean().optional(), + tabWidth: z.int().positive().optional(), + fileGap: z.int().nonnegative().optional(), + hunkGap: z.int().nonnegative().optional(), + wrapLines: z.boolean().optional(), + hunkHeaders: z.boolean().optional(), + menuBar: z.boolean().optional(), + sidebar: z.union([z.boolean(), z.literal("auto")]).optional(), + agentNotes: z.boolean().optional(), + copyDecorations: z.boolean().optional(), + promptSaveViewPreferences: z.boolean().optional(), + transparentBackground: z.boolean().optional(), + colorMoved: z.boolean().optional(), + extensions: z.boolean().optional(), + extensionPaths: z.array(z.string()).optional(), +}); + +/** Parses the complete reloadable CLI input tree carried inside a command. */ +export const cliInputSchema = z.discriminatedUnion("kind", [ + z.strictObject({ + kind: z.literal("vcs"), + range: z.string().optional(), + staged: z.boolean(), + pathspecs: z.array(z.string()).optional(), + options: commonOptionsSchema, + }), + z.strictObject({ + kind: z.literal("show"), + ref: z.string().optional(), + pathspecs: z.array(z.string()).optional(), + options: commonOptionsSchema, + }), + z.strictObject({ + kind: z.literal("stash-show"), + ref: z.string().optional(), + options: commonOptionsSchema, + }), + z.strictObject({ + kind: z.literal("diff"), + left: z.string(), + right: z.string(), + options: commonOptionsSchema, + }), + z.strictObject({ + kind: z.literal("patch"), + file: z.string().optional(), + text: z.string().optional(), + options: commonOptionsSchema, + }), + z.strictObject({ + kind: z.literal("difftool"), + left: z.string(), + right: z.string(), + path: z.string().optional(), + options: commonOptionsSchema, + }), +]) satisfies z.ZodType; const commentApplyItemSchema = z.strictObject({ filePath: z.string(), @@ -69,7 +153,7 @@ export const sessionDaemonRequestSchema = z.discriminatedUnion("action", [ z.strictObject({ action: z.literal("reload"), selector: selectorSchema, - nextInput: nextInputSchema, + nextInput: cliInputSchema, sourcePath: z.string().optional(), }), z.strictObject({ @@ -125,6 +209,251 @@ export const sessionDaemonRequestSchema = z.discriminatedUnion("action", [ }), ]); +const nonnegative = z.int().nonnegative(); +const positive = z.int().positive(); +const lineRangeSchema = z.tuple([nonnegative, nonnegative]); +const inputKindSchema = z.enum(["vcs", "show", "stash-show", "diff", "patch", "difftool"]); +const experimentalFeaturesSchema = z.array(z.enum(EXPERIMENTAL_FEATURES)); +const terminalLocationSchema = z.strictObject({ + source: z.string(), + tty: z.string().optional(), + windowId: z.string().optional(), + tabId: z.string().optional(), + paneId: z.string().optional(), + terminalId: z.string().optional(), + sessionId: z.string().optional(), +}); +const terminalSchema = z.strictObject({ + program: z.string().optional(), + locations: z.array(terminalLocationSchema), +}); +const fileSummarySchema = z.strictObject({ + id: z.string(), + path: z.string(), + previousPath: z.string().optional(), + additions: nonnegative, + deletions: nonnegative, + hunkCount: nonnegative, +}); +const reviewHunkSchema = z.strictObject({ + index: nonnegative, + header: z.string(), + oldRange: lineRangeSchema.optional(), + newRange: lineRangeSchema.optional(), +}); +const selectedHunkSchema = reviewHunkSchema.omit({ header: true }); +const reviewFileSchema = fileSummarySchema.extend({ + patch: z.string().optional(), + hunks: z.array(reviewHunkSchema), +}); +const liveCommentSchema = z.strictObject({ + commentId: z.string(), + filePath: z.string(), + hunkIndex: nonnegative, + side: sideSchema, + line: positive, + summary: z.string(), + rationale: z.string().optional(), + author: z.string().optional(), + createdAt: z.string(), +}); +const reviewNoteSchema = z.strictObject({ + noteId: z.string(), + source: z.enum(["ai", "agent", "user"]), + filePath: z.string(), + hunkIndex: nonnegative.optional(), + oldRange: lineRangeSchema.optional(), + newRange: lineRangeSchema.optional(), + body: z.string(), + title: z.string().optional(), + author: z.string().optional(), + createdAt: z.string(), + updatedAt: z.string().optional(), + editable: z.boolean(), +}); +const snapshotSchema = z.strictObject({ + updatedAt: z.string(), + state: z.strictObject({ + selectedFileId: z.string().optional(), + selectedFilePath: z.string().optional(), + selectedHunkIndex: nonnegative, + selectedHunkOldRange: lineRangeSchema.optional(), + selectedHunkNewRange: lineRangeSchema.optional(), + showAgentNotes: z.boolean(), + noteMarkupWidth: nonnegative.optional(), + liveCommentCount: nonnegative, + liveComments: z.array(liveCommentSchema), + reviewNoteCount: nonnegative.optional(), + reviewNotes: z.array(reviewNoteSchema).optional(), + reviewPublication: z + .strictObject({ generation: z.string(), stateRevision: nonnegative }) + .optional(), + }), +}); +const listedSessionSchema = z.strictObject({ + sessionId: z.string(), + pid: positive, + cwd: z.string(), + repoRoot: z.string().optional(), + launchedAt: z.string(), + terminal: terminalSchema.optional(), + inputKind: inputKindSchema, + title: z.string(), + sourceLabel: z.string(), + experimentalFeatures: experimentalFeaturesSchema.optional(), + fileCount: nonnegative, + files: z.array(fileSummarySchema), + snapshot: snapshotSchema, +}); +const selectedContextSchema = z.strictObject({ + sessionId: z.string(), + title: z.string(), + sourceLabel: z.string(), + cwd: z.string().optional(), + repoRoot: z.string().optional(), + inputKind: inputKindSchema, + experimentalFeatures: experimentalFeaturesSchema.optional(), + selectedFile: fileSummarySchema.nullable(), + selectedHunk: selectedHunkSchema.nullable(), + showAgentNotes: z.boolean(), + noteMarkupWidth: nonnegative.optional(), + liveCommentCount: nonnegative, +}); +const reviewSchema = z.strictObject({ + sessionId: z.string(), + title: z.string(), + sourceLabel: z.string(), + cwd: z.string().optional(), + repoRoot: z.string().optional(), + inputKind: inputKindSchema, + experimentalFeatures: experimentalFeaturesSchema.optional(), + selectedFile: reviewFileSchema.nullable(), + selectedHunk: reviewHunkSchema.nullable(), + showAgentNotes: z.boolean(), + liveCommentCount: nonnegative, + reviewNoteCount: nonnegative.optional(), + reviewNotes: z.array(reviewNoteSchema).optional(), + files: z.array(reviewFileSchema), +}); + +const appliedCommentSchema = z.strictObject({ + commentId: z.string(), + fileId: z.string(), + filePath: z.string(), + hunkIndex: nonnegative, + side: sideSchema, + line: positive, + markupWidth: nonnegative.optional(), + markupNotes: z.array(z.string()).optional(), +}); + +/** Strict Hunk command result schemas shared by broker and HTTP response validation. */ +export const hunkCommandResultSchemas = { + comment: appliedCommentSchema, + comment_batch: z.strictObject({ applied: z.array(appliedCommentSchema) }), + navigate_to_hunk: z.strictObject({ + fileId: z.string(), + filePath: z.string(), + hunkIndex: nonnegative, + selectedHunk: selectedHunkSchema.optional(), + revealed: z.enum(["line", "hunk"]).optional(), + side: sideSchema.optional(), + line: positive.optional(), + }), + reload_session: z.strictObject({ + sessionId: z.string(), + inputKind: inputKindSchema, + title: z.string(), + sourceLabel: z.string(), + fileCount: nonnegative, + selectedFilePath: z.string().optional(), + selectedHunkIndex: nonnegative, + }), + remove_comment: z.strictObject({ + commentId: z.string(), + removed: z.boolean(), + remainingCommentCount: nonnegative, + source: z.enum(["ai", "agent", "user"]).optional(), + }), + clear_comments: z.strictObject({ + removedCount: nonnegative, + remainingCommentCount: nonnegative, + filePath: z.string().optional(), + includeUser: z.boolean().optional(), + removedLiveCommentCount: nonnegative.optional(), + removedUserNoteCount: nonnegative.optional(), + remainingLiveCommentCount: nonnegative.optional(), + remainingUserNoteCount: nonnegative.optional(), + }), + highlight: z.strictObject({ + fileId: z.string(), + filePath: z.string(), + hunkIndex: nonnegative, + side: sideSchema, + line: positive, + start: nonnegative, + end: positive, + tone: z.enum(["match", "current", "info", "warning", "error"]), + fileMarkCount: nonnegative, + revealed: z.enum(["line", "hunk"]).optional(), + }), + clear_highlights: z.strictObject({ + removedCount: nonnegative, + remainingCount: nonnegative, + filePath: z.string().optional(), + }), +} as const; + +const daemonResponseSchemas = { + list: z.strictObject({ sessions: z.array(listedSessionSchema) }), + get: z.strictObject({ session: listedSessionSchema }), + context: z.strictObject({ context: selectedContextSchema }), + review: z.strictObject({ review: reviewSchema }), + navigate: z.strictObject({ + result: hunkCommandResultSchemas.navigate_to_hunk, + }), + reload: z.strictObject({ result: hunkCommandResultSchemas.reload_session }), + "comment-add": z.strictObject({ result: hunkCommandResultSchemas.comment }), + "comment-apply": z.strictObject({ + result: hunkCommandResultSchemas.comment_batch, + }), + "comment-list": z.strictObject({ + comments: z.array(z.union([liveCommentSchema, reviewNoteSchema])), + }), + "comment-rm": z.strictObject({ + result: hunkCommandResultSchemas.remove_comment, + }), + "comment-clear": z.strictObject({ + result: hunkCommandResultSchemas.clear_comments, + }), + "highlight-add": z.strictObject({ + result: hunkCommandResultSchemas.highlight, + }), + "highlight-clear": z.strictObject({ + result: hunkCommandResultSchemas.clear_highlights, + }), +} satisfies { + [Action in SessionDaemonAction]: z.ZodType; +}; + +/** Strictly parse the response associated with one Hunk session API action. */ +export function parseSessionDaemonResponse( + action: Action, + value: unknown, +): SessionDaemonResponses[Action] { + const result = daemonResponseSchemas[action].safeParse(value); + if (!result.success) { + throw new Error(`Invalid Hunk session daemon response for ${action}.`); + } + return result.data as SessionDaemonResponses[Action]; +} + +/** Parse one exact cross-process Hunk daemon capability response. */ +export function parseSessionDaemonCapabilities(value: unknown): SessionDaemonCapabilities | null { + const result = sessionDaemonCapabilitiesSchema.safeParse(value); + return result.success ? result.data : null; +} + /** Compose one readable rejection reason from the first schema issue. */ function describeFirstIssue(error: z.ZodError) { const issue = error.issues[0]; diff --git a/src/ui/runInteractiveApp.tsx b/src/ui/runInteractiveApp.tsx index 6a2d40aed..bb6eae7ae 100644 --- a/src/ui/runInteractiveApp.tsx +++ b/src/ui/runInteractiveApp.tsx @@ -20,12 +20,6 @@ import { createInitialSessionSnapshot, createSessionRegistration, } from "../app/session/registration"; -import type { - HunkSessionCommandResult, - HunkSessionInfo, - HunkSessionServerMessage, - HunkSessionState, -} from "../session/types"; import { SessionBrokerClient } from "../session/broker/brokerClient"; import { AppHost } from "./AppHost"; import { disposeHighlightWorker } from "./diff/worker"; @@ -54,12 +48,7 @@ export async function runInteractiveApp({ sourceLabel: bootstrap.changeset.sourceLabel, }); const publication = reviewProducer.getPublication(); - const hostClient = new SessionBrokerClient< - HunkSessionInfo, - HunkSessionState, - HunkSessionServerMessage, - HunkSessionCommandResult - >( + const hostClient = new SessionBrokerClient( createSessionRegistration(bootstrap, publication), createInitialSessionSnapshot(bootstrap, publication), ); From 9cd85c89776f61f348e87a8537b271cc21cb07b5 Mon Sep 17 00:00:00 2001 From: Ben Vinegar Date: Sat, 29 Aug 2026 13:41:27 -0400 Subject: [PATCH 4/5] fix(session): bound broker scheduling and transport --- .changeset/calm-brokers-bound.md | 2 + package.json | 3 +- packages/session-broker-bun/package.json | 2 +- packages/session-broker-bun/src/serve.test.ts | 128 ++++++ packages/session-broker-bun/src/serve.ts | 178 +++++++- packages/session-broker-core/package.json | 2 +- .../src/brokerState.test.ts | 240 ++++++++++- .../session-broker-core/src/brokerState.ts | 389 +++++++++++++---- .../session-broker-core/src/budgets.test.ts | 102 +++++ packages/session-broker-core/src/budgets.ts | 258 +++++++++++ packages/session-broker-core/src/index.ts | 1 + .../session-broker-core/src/limits.test.ts | 26 ++ packages/session-broker-core/src/limits.ts | 134 +++++- packages/session-broker-node/package.json | 2 +- .../session-broker-node/src/serve.test.ts | 8 + packages/session-broker-node/src/serve.ts | 270 ++++++++++-- packages/session-broker/package.json | 2 +- .../session-broker/src/authentication.test.ts | 61 +++ packages/session-broker/src/authentication.ts | 255 +++++++---- packages/session-broker/src/broker.ts | 59 ++- .../session-broker/src/connection.test.ts | 180 ++++++++ packages/session-broker/src/connection.ts | 172 ++++++-- packages/session-broker/src/daemon.test.ts | 50 ++- packages/session-broker/src/daemon.ts | 401 +++++++++++++----- scripts/test-session-broker-node.ts | 38 ++ src/session/broker/brokerServer.ts | 28 +- .../sessionBrokerAdapterConformance.json | 13 + test/session-broker-node/adapter.test.mjs | 197 +++++++++ 28 files changed, 2767 insertions(+), 434 deletions(-) create mode 100644 .changeset/calm-brokers-bound.md create mode 100644 packages/session-broker-core/src/budgets.test.ts create mode 100644 packages/session-broker-core/src/budgets.ts create mode 100644 scripts/test-session-broker-node.ts create mode 100644 test/fixtures/sessionBrokerAdapterConformance.json create mode 100644 test/session-broker-node/adapter.test.mjs diff --git a/.changeset/calm-brokers-bound.md b/.changeset/calm-brokers-bound.md new file mode 100644 index 000000000..a845151cc --- /dev/null +++ b/.changeset/calm-brokers-bound.md @@ -0,0 +1,2 @@ +--- +--- diff --git a/package.json b/package.json index 6089d8ddf..ff949b4e2 100644 --- a/package.json +++ b/package.json @@ -86,6 +86,7 @@ "test": "bun run ./scripts/run-test-suite.ts", "test:theme-contrast": "bun test src/ui/themes.test.ts --test-name-pattern contrast", "test:integration": "\"${npm_execpath:-bun}\" test ./test/pty", + "test:session-broker-node": "bun run ./scripts/test-session-broker-node.ts", "test:tty-smoke": "HUNK_RUN_TTY_SMOKE=1 \"${npm_execpath:-bun}\" test ./test/smoke", "check:pack": "bun run ./scripts/check-pack.ts", "check:prebuilt-pack": "bun run ./scripts/check-prebuilt-pack.ts", @@ -170,7 +171,7 @@ "pre-commit": "bunx lint-staged" }, "engines": { - "node": ">=18" + "node": ">=22" }, "packageManager": "bun@1.3.14", "pi": { diff --git a/packages/session-broker-bun/package.json b/packages/session-broker-bun/package.json index 597fc3cec..1873c83c4 100644 --- a/packages/session-broker-bun/package.json +++ b/packages/session-broker-bun/package.json @@ -20,6 +20,6 @@ }, "engines": { "bun": ">=1.0.0", - "node": ">=18" + "node": ">=22" } } diff --git a/packages/session-broker-bun/src/serve.test.ts b/packages/session-broker-bun/src/serve.test.ts index 1f0c61e50..16d5e6af1 100644 --- a/packages/session-broker-bun/src/serve.test.ts +++ b/packages/session-broker-bun/src/serve.test.ts @@ -13,6 +13,7 @@ import { createSessionBrokerDaemon, createSessionBrokerProtocolParsers, } from "@hunk/session-broker"; +import SESSION_BROKER_ADAPTER_CONFORMANCE from "../../../test/fixtures/sessionBrokerAdapterConformance.json" with { type: "json" }; import { serveSessionBrokerDaemon } from "./serve"; interface TestSessionInfo { @@ -114,6 +115,51 @@ async function waitUntil( } } +async function openTestSocket(url: string) { + const socket = new WebSocket(url); + await new Promise((resolve, reject) => { + const timer = setTimeout( + () => reject(new Error("Timed out waiting for websocket open.")), + 1_000, + ); + socket.addEventListener( + "open", + () => { + clearTimeout(timer); + resolve(); + }, + { once: true }, + ); + socket.addEventListener( + "error", + () => { + clearTimeout(timer); + reject(new Error("Websocket failed to open.")); + }, + { once: true }, + ); + }); + return socket; +} + +function testSocketCloseCode(socket: WebSocket) { + return new Promise((resolve, reject) => { + const timer = setTimeout( + () => reject(new Error("Timed out waiting for websocket close.")), + 1_000, + ); + socket.addEventListener( + "close", + (event) => { + clearTimeout(timer); + resolve(event.code); + }, + { once: true }, + ); + socket.addEventListener("error", () => {}, { once: true }); + }); +} + async function readHealth(port: number) { try { const response = await fetch(`http://127.0.0.1:${port}/health`); @@ -150,6 +196,68 @@ afterEach(() => { }); describe("session broker bun adapter", () => { + test("uses the shared binary, oversize, and pressure close corpus", () => { + expect(SESSION_BROKER_ADAPTER_CONFORMANCE).toMatchObject({ + textOnly: { binaryCloseCode: 1003 }, + inbound: { oversizedCloseCode: 1009, pressureCloseCode: 1013 }, + }); + }); + + test("closes binary, oversized, and aggregate-pressure messages per the shared corpus", async () => { + const broker = new SessionBroker({ protocolParsers }); + const daemon = createSessionBrokerDaemon({ + broker, + limits: { maxWsMessageBytes: 8, maxInFlightWsBytes: 0 }, + }); + const port = await reserveLoopbackPort(); + const server = serveSessionBrokerDaemon({ daemon, hostname: "127.0.0.1", port }); + try { + const binary = await openTestSocket(`ws://127.0.0.1:${port}/session`); + const binaryClosed = testSocketCloseCode(binary); + binary.send(new Uint8Array([1])); + expect(await binaryClosed).toBe(SESSION_BROKER_ADAPTER_CONFORMANCE.textOnly.binaryCloseCode); + + const oversized = await openTestSocket(`ws://127.0.0.1:${port}/session`); + const oversizedClosed = testSocketCloseCode(oversized); + oversized.send("123456789"); + expect(await oversizedClosed).toBe( + SESSION_BROKER_ADAPTER_CONFORMANCE.inbound.oversizedCloseCode, + ); + + const pressure = await openTestSocket(`ws://127.0.0.1:${port}/session`); + const pressureClosed = testSocketCloseCode(pressure); + pressure.send("{}"); + expect(await pressureClosed).toBe( + SESSION_BROKER_ADAPTER_CONFORMANCE.inbound.pressureCloseCode, + ); + } finally { + server.stop(true); + await server.stopped; + } + }); + + test("admits exactly the configured number of active websocket peers", async () => { + const broker = new SessionBroker({ protocolParsers }); + const daemon = createSessionBrokerDaemon({ + broker, + limits: { maxUnauthenticatedSockets: 1 }, + }); + const port = await reserveLoopbackPort(); + const server = serveSessionBrokerDaemon({ daemon, hostname: "127.0.0.1", port }); + try { + const first = await openTestSocket(`ws://127.0.0.1:${port}/session`); + await expect(openTestSocket(`ws://127.0.0.1:${port}/session`)).rejects.toThrow(); + const closed = testSocketCloseCode(first); + first.close(); + await closed; + const afterRelease = await openTestSocket(`ws://127.0.0.1:${port}/session`); + afterRelease.close(); + } finally { + server.stop(true); + await server.stopped; + } + }); + test("serves the generic daemon API and websocket path through Bun", async () => { const broker = new SessionBroker({ protocolParsers }); const daemon = createSessionBrokerDaemon({ @@ -254,6 +362,26 @@ describe("session broker bun adapter", () => { } }); + test("falls back to an empty 503 when even the capacity envelope exceeds the response cap", async () => { + const broker = new SessionBroker({ protocolParsers }); + const daemon = createSessionBrokerDaemon({ broker, limits: { maxHttpResponseBytes: 1 } }); + const port = await reserveLoopbackPort(); + const server = serveSessionBrokerDaemon({ + daemon, + hostname: "127.0.0.1", + port, + handleRequest: () => new Response("too large"), + }); + try { + const response = await fetch(`http://127.0.0.1:${port}/large`); + expect(response.status).toBe(503); + expect(await response.text()).toBe(""); + } finally { + server.stop(true); + await server.stopped; + } + }); + test("lets custom request handlers override generic routes", async () => { const broker = new SessionBroker({ protocolParsers }); const daemon = createSessionBrokerDaemon({ diff --git a/packages/session-broker-bun/src/serve.ts b/packages/session-broker-bun/src/serve.ts index 67cdd920b..9bbfd883c 100644 --- a/packages/session-broker-bun/src/serve.ts +++ b/packages/session-broker-bun/src/serve.ts @@ -1,9 +1,16 @@ import { - MAX_WS_MESSAGE_BYTES, + BrokerCapacityError, + ResourceBudget, + boundHttpResponse, utf8ByteLength, + type BudgetReservation, type SessionServerMessage, } from "@hunk/session-broker-core"; -import type { SessionBrokerDaemon } from "@hunk/session-broker"; +import type { SessionBrokerDaemon, SessionBrokerPeer } from "@hunk/session-broker"; + +interface BrokerWebSocketData { + admission: BudgetReservation; +} export interface ServeSessionBrokerDaemonOptions< SessionView = unknown, @@ -15,13 +22,13 @@ export interface ServeSessionBrokerDaemonOptions< port: number; handleRequest?: ( request: Request, - server: ReturnType>, + server: ReturnType>, ) => Response | Promise | undefined; notFound?: (request: Request) => Response | Promise; formatServeError?: (error: unknown, address: { hostname: string; port: number }) => Error; } -export type RunningSessionBrokerDaemon = ReturnType> & { +export type RunningSessionBrokerDaemon = ReturnType> & { stopped: Promise; }; @@ -44,6 +51,82 @@ export function serveSessionBrokerDaemon< >( options: ServeSessionBrokerDaemonOptions, ): RunningSessionBrokerDaemon { + const inboundBudget = new ResourceBudget( + options.daemon.limits.maxInFlightWsBytes, + "maxInFlightWsBytes", + ); + const outboundBudget = new ResourceBudget( + options.daemon.limits.maxOutboundBytesTotal, + "maxOutboundBytesTotal", + "busy", + ); + const unauthenticatedSocketBudget = new ResourceBudget( + options.daemon.limits.maxUnauthenticatedSockets, + "maxUnauthenticatedSockets", + "busy", + ); + const responseBudget = new ResourceBudget( + options.daemon.limits.maxHttpResponseBytes, + "maxHttpResponseBytes", + "busy", + ); + const bufferedReservations = new Map(); + const activeAdmissions = new Set(); + const peers = new WeakMap(); + const peerFor = (socket: { + send(data: string): number; + close(code?: number, reason?: string): void; + getBufferedAmount?(): number; + }): SessionBrokerPeer => { + const key = socket as object; + const existing = peers.get(key); + if (existing) return existing; + const peer: SessionBrokerPeer = { + send(data) { + const bytes = utf8ByteLength(data); + const buffered = socket.getBufferedAmount?.() ?? 0; + if (bytes > options.daemon.limits.maxOutboundBytesPerPeer - buffered) { + socket.close(1013, "Session broker outbound pressure exceeded."); + throw new BrokerCapacityError("busy", "maxOutboundBytesPerPeer"); + } + const previous = bufferedReservations.get(key); + let provisional: BudgetReservation; + try { + provisional = previous + ? outboundBudget.resize(previous, buffered + bytes) + : outboundBudget.reserve(buffered + bytes); + } catch { + socket.close(1013, "Session broker outbound pressure exceeded."); + throw new BrokerCapacityError("busy", "maxOutboundBytesTotal"); + } + bufferedReservations.set(key, provisional); + try { + const sent = socket.send(data); + if (sent === 0) { + socket.close(1013, "Session broker outbound pressure exceeded."); + throw new BrokerCapacityError("busy", "maxOutboundBytesPerPeer"); + } + const remaining = socket.getBufferedAmount?.() ?? 0; + if (remaining === 0) { + bufferedReservations.delete(key); + provisional.release(); + } else { + bufferedReservations.set(key, outboundBudget.resize(provisional, remaining)); + } + } catch (error) { + // A dropped send retains the provisional bound until close; other failures release now. + if (socket.getBufferedAmount?.() === 0) { + bufferedReservations.delete(key); + provisional.release(); + } + throw error; + } + }, + close: (code, reason) => socket.close(code, reason), + }; + peers.set(key, peer); + return peer; + }; let resolved = false; let resolveStopped: (() => void) | null = null; const stopped = new Promise((resolve) => { @@ -59,9 +142,9 @@ export function serveSessionBrokerDaemon< resolveStopped = null; }; - let server: ReturnType>; + let server: ReturnType>; try { - server = Bun.serve<{}>({ + server = Bun.serve({ hostname: options.hostname, port: options.port, fetch: async (request, bunServer) => { @@ -69,19 +152,32 @@ export function serveSessionBrokerDaemon< // Let host apps extend or override routes first; the generic daemon only handles the // broker's shared HTTP surface plus the websocket upgrade path. if (customResponse !== undefined) { - return customResponse; + return boundHttpResponse( + customResponse, + options.daemon.limits.maxHttpResponseBytes, + responseBudget, + ); } const daemonResponse = await options.daemon.handleRequest(request); if (daemonResponse) { - return daemonResponse; + return boundHttpResponse( + daemonResponse, + options.daemon.limits.maxHttpResponseBytes, + responseBudget, + ); } const url = new URL(request.url); if (options.daemon.matchesSocketPath(url.pathname)) { - if (bunServer.upgrade(request, { data: {} })) { + const admission = unauthenticatedSocketBudget.tryReserve(); + if (!admission) return new Response(null, { status: 503 }); + activeAdmissions.add(admission); + if (bunServer.upgrade(request, { data: { admission } })) { return undefined; } + activeAdmissions.delete(admission); + admission.release(); // Bun signals failed upgrades by returning false from upgrade rather than by throwing, // so surface that as one explicit HTTP response here. @@ -89,27 +185,65 @@ export function serveSessionBrokerDaemon< return new Response("Expected websocket upgrade.", { status: 426 }); } - return (await options.notFound?.(request)) ?? defaultNotFound(); + return boundHttpResponse( + (await options.notFound?.(request)) ?? defaultNotFound(), + options.daemon.limits.maxHttpResponseBytes, + responseBudget, + ); }, websocket: { - // Let Bun reject oversized frames at the protocol layer before they are ever buffered. - maxPayloadLength: MAX_WS_MESSAGE_BYTES, + // Bun cannot customize the close code of its native payload rejection. Keep the native cap + // at the fixed aggregate ceiling so decoded messages above the per-message limit reach the + // portable 1009 path while runtime buffering remains bounded. + maxPayloadLength: Math.min( + Number.MAX_SAFE_INTEGER, + Math.max( + options.daemon.limits.maxWsMessageBytes + 1, + options.daemon.limits.maxInFlightWsBytes, + ), + ), message: (socket, message) => { + const peer = peerFor(socket); if (typeof message !== "string") { + socket.close(1003, "Session broker accepts text messages only."); return; } - // Defense in depth: Bun's maxPayloadLength already bounds raw frames, but guard the - // decoded string too so a registration payload cannot be parsed unbounded here. - if (utf8ByteLength(message) > MAX_WS_MESSAGE_BYTES) { + const bytes = utf8ByteLength(message); + if (bytes > options.daemon.limits.maxWsMessageBytes) { socket.close(1009, "Message exceeds the session broker size limit."); return; } - - options.daemon.handleConnectionMessage(socket, message); + const reservation = inboundBudget.tryReserve(bytes); + if (!reservation) { + socket.close(1013, "Session broker inbound pressure exceeded."); + return; + } + try { + options.daemon.handleConnectionMessage(peer, message); + } finally { + reservation.release(); + } + }, + drain: (socket) => { + const key = socket as object; + const previous = bufferedReservations.get(key); + if (!previous) return; + const remaining = socket.getBufferedAmount(); + if (remaining === 0) { + bufferedReservations.delete(key); + previous.release(); + } else { + bufferedReservations.set(key, outboundBudget.resize(previous, remaining)); + } }, close: (socket) => { - options.daemon.handleConnectionClose(socket); + const key = socket as object; + bufferedReservations.get(key)?.release(); + bufferedReservations.delete(key); + activeAdmissions.delete(socket.data.admission); + socket.data.admission.release(); + options.daemon.handleConnectionClose(peerFor(socket)); }, }, }); @@ -125,6 +259,10 @@ export function serveSessionBrokerDaemon< // Wrap Bun's stop so callers do not need to remember that the daemon and transport have to be // torn down together. options.daemon.shutdown(); + for (const reservation of bufferedReservations.values()) reservation.release(); + bufferedReservations.clear(); + for (const admission of activeAdmissions) admission.release(); + activeAdmissions.clear(); const result = originalStop(closeActiveConnections); finish(); return result; @@ -140,6 +278,10 @@ export function serveSessionBrokerDaemon< void options.daemon.stopped.then(() => { // Idle shutdown and manual stop share one completion promise, but the Bun server only needs // the original transport stop here because the daemon has already transitioned to stopped. + for (const reservation of bufferedReservations.values()) reservation.release(); + bufferedReservations.clear(); + for (const admission of activeAdmissions) admission.release(); + activeAdmissions.clear(); originalStop(true); finish(); }); diff --git a/packages/session-broker-core/package.json b/packages/session-broker-core/package.json index b86888c3c..36e39e37e 100644 --- a/packages/session-broker-core/package.json +++ b/packages/session-broker-core/package.json @@ -17,6 +17,6 @@ }, "engines": { "bun": ">=1.0.0", - "node": ">=18" + "node": ">=22" } } diff --git a/packages/session-broker-core/src/brokerState.test.ts b/packages/session-broker-core/src/brokerState.test.ts index 261ca24dd..6c2855d06 100644 --- a/packages/session-broker-core/src/brokerState.test.ts +++ b/packages/session-broker-core/src/brokerState.test.ts @@ -5,6 +5,7 @@ import { type SessionBrokerListedSession, type SessionBrokerViewAdapter, } from "./brokerState"; +import type { SessionBrokerLimitOptions } from "./budgets"; import { SESSION_BROKER_REGISTRATION_VERSION, brokerWireParsers, @@ -134,7 +135,7 @@ const testBrokerView: SessionBrokerViewAdapter< listComments: (_session, filter) => [{ id: "note-1", filePath: filter.filePath }], }; -function createState() { +function createState(limitOptions: SessionBrokerLimitOptions = {}) { return new SessionBrokerState< TestSessionInfo, TestSessionState, @@ -144,7 +145,7 @@ function createState() { TestSelectedContext, TestSessionReview, TestCommentSummary - >(testBrokerView); + >(testBrokerView, limitOptions); } function createRegistration( @@ -651,4 +652,239 @@ describe("session broker state", () => { expect(state.pruneStaleSessions({ ttlMs, now: lastSeenAt + wallClockJumpMs + 15_000 })).toBe(1); expect(state.listSessions()).toHaveLength(0); }); + + test("schedules commands FIFO with one active command per session", async () => { + const state = createState(); + const sent: string[] = []; + const socket = { send: (data: string) => sent.push(data) }; + state.registerSession(socket, createRegistration(), createSnapshot()); + + const first = state.dispatchCommand({ + selector: { sessionId: "session-1" }, + command: "annotate", + input: { filePath: "a", summary: "first" }, + timeoutMessage: "first timeout", + }); + const second = state.dispatchCommand({ + selector: { sessionId: "session-1" }, + command: "annotate", + input: { filePath: "b", summary: "second" }, + timeoutMessage: "second timeout", + }); + expect(sent).toHaveLength(1); + const firstId = JSON.parse(sent[0]!).requestId as string; + state.handleCommandResult(socket, { + requestId: firstId, + ok: true, + result: { kind: "annotated", annotationId: "one" }, + }); + expect(sent).toHaveLength(2); + const secondId = JSON.parse(sent[1]!).requestId as string; + state.handleCommandResult(socket, { + requestId: secondId, + ok: true, + result: { kind: "annotated", annotationId: "two" }, + }); + await expect(first).resolves.toMatchObject({ annotationId: "one" }); + await expect(second).resolves.toMatchObject({ annotationId: "two" }); + }); + + test("lets different sessions progress independently", async () => { + const state = createState(); + const firstSent: string[] = []; + const secondSent: string[] = []; + const firstSocket = { send: (data: string) => firstSent.push(data) }; + const secondSocket = { send: (data: string) => secondSent.push(data) }; + state.registerSession(firstSocket, createRegistration(), createSnapshot()); + state.registerSession( + secondSocket, + createRegistration({ sessionId: "session-2", cwd: "/two", repoRoot: "/two" }), + createSnapshot(), + ); + const first = state.dispatchCommand({ + selector: { sessionId: "session-1" }, + command: "annotate", + input: { filePath: "a", summary: "one" }, + timeoutMessage: "timeout", + }); + const second = state.dispatchCommand({ + selector: { sessionId: "session-2" }, + command: "annotate", + input: { filePath: "b", summary: "two" }, + timeoutMessage: "timeout", + }); + expect([firstSent.length, secondSent.length]).toEqual([1, 1]); + for (const [socket, raw, pending, id] of [ + [firstSocket, firstSent[0]!, first, "one"], + [secondSocket, secondSent[0]!, second, "two"], + ] as const) { + state.handleCommandResult(socket, { + requestId: JSON.parse(raw).requestId, + ok: true, + result: { kind: "annotated", annotationId: id }, + }); + await expect(pending).resolves.toMatchObject({ annotationId: id }); + } + }); + + test("rejects exact count boundaries plus one without dropping admitted work", async () => { + const state = createState({ + limits: { maxCommandsPerSession: 2, maxCommandsTotal: 2 }, + }); + const sent: string[] = []; + const socket = { send: (data: string) => sent.push(data) }; + state.registerSession(socket, createRegistration(), createSnapshot()); + const commands = ["one", "two"].map((summary) => + state.dispatchCommand({ + selector: { sessionId: "session-1" }, + command: "annotate", + input: { filePath: "a", summary }, + timeoutMessage: "timeout", + }), + ); + expect(() => + state.dispatchCommand({ + selector: { sessionId: "session-1" }, + command: "annotate", + input: { filePath: "a", summary: "three" }, + timeoutMessage: "timeout", + }), + ).toThrow("queue-full"); + state.shutdown(); + for (const command of commands) await expect(command).rejects.toThrow("shut down"); + }); + + test("accounts queued command UTF-8 bytes at the exact daemon boundary", async () => { + const input = { filePath: "é", summary: "😀" }; + const bytes = + new TextEncoder().encode( + JSON.stringify({ + type: "command", + requestId: "0".repeat(36), + command: "annotate", + commandVersion: 1, + input, + }), + ).byteLength + 128; + const state = createState({ limits: { maxQueuedCommandBytes: bytes } }); + const socket = { send() {} }; + state.registerSession(socket, createRegistration(), createSnapshot()); + const admitted = state.dispatchCommand({ + selector: { sessionId: "session-1" }, + command: "annotate", + input, + timeoutMessage: "timeout", + }); + expect(() => + state.dispatchCommand({ + selector: { sessionId: "session-1" }, + command: "annotate", + input, + timeoutMessage: "timeout", + }), + ).toThrow("queue-full"); + state.shutdown(); + await expect(admitted).rejects.toThrow("shut down"); + }); + + test("releases a timed-out active command and advances its session FIFO", async () => { + const state = createState({ limits: { defaultCommandTimeoutMs: 5, maxCommandTimeoutMs: 100 } }); + const sent: string[] = []; + const socket = { send: (data: string) => sent.push(data) }; + state.registerSession(socket, createRegistration(), createSnapshot()); + const first = state.dispatchCommand({ + selector: { sessionId: "session-1" }, + command: "annotate", + input: { filePath: "a", summary: "one" }, + timeoutMessage: "timed out", + }); + const second = state.dispatchCommand({ + selector: { sessionId: "session-1" }, + command: "annotate", + input: { filePath: "b", summary: "two" }, + timeoutMessage: "second timeout", + timeoutMs: 100, + }); + await expect(first).rejects.toThrow("timed out"); + expect(sent).toHaveLength(2); + state.handleCommandResult(socket, { + requestId: JSON.parse(sent[1]!).requestId, + ok: true, + result: { kind: "annotated", annotationId: "two" }, + }); + await expect(second).resolves.toMatchObject({ annotationId: "two" }); + expect(() => + state.dispatchCommand({ + selector: { sessionId: "session-1" }, + command: "annotate", + input: { filePath: "c", summary: "three" }, + timeoutMessage: "timeout", + timeoutMs: 101, + }), + ).toThrow("capacity-exceeded"); + }); + + test("rejects a new session at capacity without evicting the existing owner", () => { + const state = createState({ limits: { maxSessions: 1 } }); + const first = { send() {} }; + const second = { send() {} }; + expect(state.registerSession(first, createRegistration(), createSnapshot())).toBe("registered"); + expect( + state.registerSession( + second, + createRegistration({ sessionId: "session-2" }), + createSnapshot(), + ), + ).toBe("capacity-exceeded"); + expect(state.listSessions().map((session) => session.sessionId)).toEqual(["session-1"]); + }); + + test("transfers a same-socket count reservation when the session id changes at capacity", () => { + const state = createState({ limits: { maxSessions: 1 } }); + const socket = { send() {} }; + expect(state.registerSession(socket, createRegistration(), createSnapshot())).toBe( + "registered", + ); + expect( + state.registerSession( + socket, + createRegistration({ sessionId: "session-2", cwd: "/two", repoRoot: "/two" }), + createSnapshot(), + ), + ).toBe("registered"); + expect(state.listSessions().map((session) => session.sessionId)).toEqual(["session-2"]); + }); + + test("accepts identical registration and snapshot replacement at the exact retained ceiling", () => { + const registration = createRegistration(); + const snapshot = createSnapshot(); + const retainedBytes = + new TextEncoder().encode(JSON.stringify({ registration, snapshot })).byteLength + 256; + const state = createState({ + limits: { maxRetainedSessionBytes: retainedBytes, maxRetainedBytes: retainedBytes }, + }); + const socket = { send() {} }; + expect(state.registerSession(socket, registration, snapshot)).toBe("registered"); + expect(state.registerSession(socket, registration, snapshot)).toBe("registered"); + expect(state.updateSnapshot(socket, "session-1", snapshot)).toBe("updated"); + }); + + test("preserves retained state when a replacement cannot reserve capacity", () => { + const registration = createRegistration(); + const snapshot = createSnapshot(); + const retainedBytes = + new TextEncoder().encode(JSON.stringify({ registration, snapshot })).byteLength + 256; + const state = createState({ + limits: { maxRetainedSessionBytes: retainedBytes, maxRetainedBytes: retainedBytes }, + }); + const socket = { send() {} }; + expect(state.registerSession(socket, registration, snapshot)).toBe("registered"); + expect( + state.updateSnapshot(socket, "session-1", { + ...snapshot, + state: { selectedIndex: 123_456, noteCount: 0 }, + }), + ).toBe("capacity-exceeded"); + expect(state.getSession({ sessionId: "session-1" }).snapshot.state.selectedIndex).toBe(0); + }); }); diff --git a/packages/session-broker-core/src/brokerState.ts b/packages/session-broker-core/src/brokerState.ts index 107d4ecf7..2f2f52bff 100644 --- a/packages/session-broker-core/src/brokerState.ts +++ b/packages/session-broker-core/src/brokerState.ts @@ -1,6 +1,16 @@ import { randomUUID } from "node:crypto"; import { isValidBrokerRevision } from "./auth"; -import { parseBrokerAppPayload } from "./validation"; +import { + BrokerCapacityError, + ReservationGroup, + ResourceBudget, + resolveSessionBrokerLimits, + type BudgetReservation, + type SessionBrokerLimitOptions, + type SessionBrokerLimits, +} from "./budgets"; +import { utf8ByteLength } from "./limits"; +import { BrokerProtocolError, parseBrokerAppPayload } from "./validation"; import { matchesSessionSelector, repoSelectorDistance, type SelectableSession } from "./selectors"; import type { SessionRegistration, @@ -10,13 +20,17 @@ import type { } from "./types"; interface PendingCommand { + requestId: string; sessionId: string; socket: DaemonSessionSocket; command: string; commandVersion: number; + serializedMessage: string; + reservation: BudgetReservation; resolve: (result: Result) => void; reject: (error: Error) => void; timeout: ReturnType; + active: boolean; } interface DaemonSessionSocket { @@ -75,13 +89,33 @@ export interface SessionBrokerViewAdapter< listComments: (session: ListedSession, filter: { filePath?: string }) => SessionCommentSummary[]; } -export type RegisterSessionResult = "registered" | "invalid" | "already-connected"; -export type UpdateSnapshotResult = "updated" | "invalid" | "not-owner"; +export type RegisterSessionResult = + | "registered" + | "invalid" + | "already-connected" + | "capacity-exceeded"; +export type UpdateSnapshotResult = "updated" | "invalid" | "not-owner" | "capacity-exceeded"; export type MarkSessionSeenResult = "seen" | "not-owner"; export type HandleCommandResult = "handled" | "not-found" | "not-owner" | "invalid"; export type SessionTargetSelector = SessionTargetInput; +const RETAINED_SESSION_OVERHEAD_BYTES = 256; +const QUEUED_COMMAND_OVERHEAD_BYTES = 128; + +/** Measure one JSON-safe retained value in UTF-8 plus its fixed broker bookkeeping overhead. */ +function retainedJsonBytes(value: unknown, overhead: number): number { + let serialized: string | undefined; + try { + serialized = JSON.stringify(value); + } catch { + throw new TypeError("Session broker data is not JSON serializable."); + } + if (serialized === undefined) + throw new TypeError("Session broker data is not JSON serializable."); + return utf8ByteLength(serialized) + overhead; +} + function describeSessionChoices( sessions: ListedSession[], ) { @@ -173,9 +207,18 @@ export class SessionBrokerState< SessionReview = unknown, SessionCommentSummary = unknown, > { + readonly limits: Readonly; + private sessions = new Map>(); private sessionIdsBySocket = new Map(); private pendingCommands = new Map>(); + private commandQueues = new Map(); + private retainedReservations = new Map(); + private sessionReservations = new Map(); + private readonly sessionBudget: ResourceBudget; + private readonly commandBudget: ResourceBudget; + private readonly queuedCommandByteBudget: ResourceBudget; + private readonly retainedByteBudget: ResourceBudget; private lastPruneAt: number | null = null; constructor( @@ -189,7 +232,22 @@ export class SessionBrokerState< ServerMessage, CommandResult >, - ) {} + limitOptions: SessionBrokerLimitOptions = {}, + ) { + this.limits = resolveSessionBrokerLimits(limitOptions); + this.sessionBudget = new ResourceBudget(this.limits.maxSessions, "maxSessions"); + this.commandBudget = new ResourceBudget( + this.limits.maxCommandsTotal, + "maxCommandsTotal", + "queue-full", + ); + this.queuedCommandByteBudget = new ResourceBudget( + this.limits.maxQueuedCommandBytes, + "maxQueuedCommandBytes", + "queue-full", + ); + this.retainedByteBudget = new ResourceBudget(this.limits.maxRetainedBytes, "maxRetainedBytes"); + } listSessions(): ListedSession[] { return [...this.sessions.values()] @@ -240,28 +298,70 @@ export class SessionBrokerState< } if (!registration || !snapshot) return "invalid"; - const existing = this.sessions.get(registration.sessionId); - if (existing && existing.socket !== socket) { - // Reconnect proof is not available yet, so an unauthenticated peer cannot supersede a live - // owner. Once the owner closes, normal unregister cleanup makes this ID available again. - return "already-connected"; + let retainedBytes: number; + try { + // Measure the values the parser actually retains so transforming parsers cannot expand past + // either the per-session or aggregate ceiling. + retainedBytes = retainedJsonBytes( + { registration, snapshot }, + RETAINED_SESSION_OVERHEAD_BYTES, + ); + } catch { + return "invalid"; } + if (retainedBytes > this.limits.maxRetainedSessionBytes) return "capacity-exceeded"; + const existing = this.sessions.get(registration.sessionId); + if (existing && existing.socket !== socket) return "already-connected"; const previousSessionId = this.sessionIdsBySocket.get(socket); - if (previousSessionId && previousSessionId !== registration.sessionId) { - this.unregisterSocket(socket); - } + const transferSessionId = existing ? registration.sessionId : previousSessionId; + const previousRetained = transferSessionId + ? this.retainedReservations.get(transferSessionId) + : undefined; + const previousCount = transferSessionId + ? this.sessionReservations.get(transferSessionId) + : undefined; + + let retainedReservation: BudgetReservation | null = null; + let sessionReservation: BudgetReservation | null = null; + try { + try { + retainedReservation = previousRetained + ? this.retainedByteBudget.resize(previousRetained, retainedBytes) + : this.retainedByteBudget.reserve(retainedBytes); + sessionReservation = previousCount ?? this.sessionBudget.reserve(); + } catch { + return "capacity-exceeded"; + } - const now = new Date().toISOString(); - this.sessions.set(registration.sessionId, { - registration, - snapshot, - socket, - connectedAt: now, - lastSeenAt: now, - }); - this.sessionIdsBySocket.set(socket, registration.sessionId); - return "registered"; + const now = new Date().toISOString(); + if (previousSessionId && previousSessionId !== registration.sessionId) { + // Detach the old identity without releasing the reservations transferred to its replacement. + this.sessions.delete(previousSessionId); + this.retainedReservations.delete(previousSessionId); + this.sessionReservations.delete(previousSessionId); + this.rejectPendingCommandsForSession( + previousSessionId, + new Error("The session registration was replaced."), + ); + } + this.sessions.set(registration.sessionId, { + registration, + snapshot, + socket, + connectedAt: existing?.connectedAt ?? now, + lastSeenAt: now, + }); + this.sessionIdsBySocket.set(socket, registration.sessionId); + this.retainedReservations.set(registration.sessionId, retainedReservation); + this.sessionReservations.set(registration.sessionId, sessionReservation); + retainedReservation = null; + sessionReservation = null; + return "registered"; + } finally { + retainedReservation?.release(); + if (sessionReservation && sessionReservation !== previousCount) sessionReservation.release(); + } } updateSnapshot( @@ -287,12 +387,37 @@ export class SessionBrokerState< } if (!snapshot) return "invalid"; - this.sessions.set(ownedSessionId, { - ...entry, - snapshot, - lastSeenAt: new Date().toISOString(), - }); - return "updated"; + let retainedBytes: number; + try { + retainedBytes = retainedJsonBytes( + { registration: entry.registration, snapshot }, + RETAINED_SESSION_OVERHEAD_BYTES, + ); + } catch { + return "invalid"; + } + if (retainedBytes > this.limits.maxRetainedSessionBytes) return "capacity-exceeded"; + + const previous = this.retainedReservations.get(ownedSessionId); + if (!previous) return "capacity-exceeded"; + let reservation: BudgetReservation | null; + try { + reservation = this.retainedByteBudget.resize(previous, retainedBytes); + } catch { + return "capacity-exceeded"; + } + try { + this.sessions.set(ownedSessionId, { + ...entry, + snapshot, + lastSeenAt: new Date().toISOString(), + }); + this.retainedReservations.set(ownedSessionId, reservation); + reservation = null; + return "updated"; + } finally { + reservation?.release(); + } } markSessionSeen(socket: DaemonSessionSocket, sessionIdAssertion: string): MarkSessionSeenResult { @@ -354,14 +479,14 @@ export class SessionBrokerState< return removed; } - /** Dispatch one app-owned command through the generic broker transport. */ + /** Admit one command and schedule it through the target session's capacity-one FIFO. */ dispatchCommand({ selector, command, commandVersion = 1, input, timeoutMessage, - timeoutMs = 15_000, + timeoutMs = this.limits.defaultCommandTimeoutMs, }: { selector: SessionTargetInput; command: CommandName; @@ -373,59 +498,81 @@ export class SessionBrokerState< if (!isValidBrokerRevision(commandVersion)) { throw new TypeError("Command version must be a positive safe integer."); } + if ( + !Number.isSafeInteger(timeoutMs) || + timeoutMs < 1 || + timeoutMs > this.limits.maxCommandTimeoutMs + ) { + throw new BrokerCapacityError("capacity-exceeded", "maxCommandTimeoutMs"); + } const session = resolveSessionTarget(this.listSessions(), selector); - const parsedInput = parseBrokerAppPayload( - (value) => this.view.parseCommandInput(command, commandVersion, value), - input, - ) as Extract["input"]; - const requestId = randomUUID(); - - return new Promise((resolve, reject) => { - const timeout = setTimeout(() => { - this.pendingCommands.delete(requestId); - reject(new Error(timeoutMessage)); - }, timeoutMs); - - // Record the pending request before sending so synchronous transport failures and later close - // events can both resolve the same command bookkeeping path. - - const entry = this.sessions.get(session.sessionId); - if (!entry) { - clearTimeout(timeout); - reject(new Error("The targeted session is no longer connected.")); - return; - } + const entry = this.sessions.get(session.sessionId); + if (!entry) return Promise.reject(new Error("The targeted session is no longer connected.")); + const sessionCount = this.commandQueues.get(session.sessionId)?.length ?? 0; + if (sessionCount >= this.limits.maxCommandsPerSession) { + throw new BrokerCapacityError("queue-full", "maxCommandsPerSession"); + } - this.pendingCommands.set(requestId, { - sessionId: session.sessionId, - socket: entry.socket, + // Measure the untrusted app input before its parser and hold aggregate capacity until terminal. + let inputBytes: number; + try { + inputBytes = retainedJsonBytes(input, 0); + } catch { + throw new BrokerProtocolError("invalid-app-payload"); + } + if (inputBytes > this.limits.maxCommandInputBytes) { + throw new BrokerCapacityError("capacity-exceeded", "maxCommandInputBytes"); + } + const reservations = new ReservationGroup(); + try { + reservations.add(this.commandBudget.reserve()); + const parsedInput = parseBrokerAppPayload( + (value) => this.view.parseCommandInput(command, commandVersion, value), + input, + ) as Extract["input"]; + if (retainedJsonBytes(parsedInput, 0) > this.limits.maxCommandInputBytes) { + throw new BrokerCapacityError("capacity-exceeded", "maxCommandInputBytes"); + } + const requestId = randomUUID(); + const serializedMessage = JSON.stringify({ + type: "command", + requestId, command, commandVersion, - resolve: (result) => resolve(result as ResultType), - reject, - timeout, + input: parsedInput, }); - - try { - const message = { - type: "command", + const queuedBytes = utf8ByteLength(serializedMessage) + QUEUED_COMMAND_OVERHEAD_BYTES; + reservations.add(this.queuedCommandByteBudget.reserve(queuedBytes)); + + return new Promise((resolve, reject) => { + const timeout = setTimeout(() => { + const pending = this.pendingCommands.get(requestId); + if (!pending) return; + this.finishPending(pending, () => reject(new Error(timeoutMessage))); + }, timeoutMs); + const pending: PendingCommand = { requestId, + sessionId: session.sessionId, + socket: entry.socket, command, commandVersion, - input: parsedInput, - } as Extract; - - entry.socket.send(JSON.stringify(message)); - } catch (error) { - clearTimeout(timeout); - this.pendingCommands.delete(requestId); - reject( - error instanceof Error - ? error - : new Error("The targeted session could not receive the command."), - ); - } - }); + serializedMessage, + reservation: reservations, + resolve: (result) => resolve(result as ResultType), + reject, + timeout, + active: false, + }; + this.pendingCommands.set(requestId, pending); + const queue = this.commandQueues.get(session.sessionId) ?? []; + queue.push(requestId); + this.commandQueues.set(session.sessionId, queue); + this.advanceSessionQueue(session.sessionId); + }); + } catch (error) { + reservations.release(); + throw error; + } } handleCommandResult( @@ -449,6 +596,9 @@ export class SessionBrokerState< if (message.ok) { let result: CommandResult; try { + if (retainedJsonBytes(message.result, 0) > this.limits.maxCommandResultBytes) { + return "invalid"; + } result = parseBrokerAppPayload( (value) => this.view.parseCommandResult( @@ -458,32 +608,85 @@ export class SessionBrokerState< ), message.result, ); + if (retainedJsonBytes(result, 0) > this.limits.maxCommandResultBytes) return "invalid"; } catch { // Keep the pending entry intact until the malformed producer is closed and normal // connection cleanup rejects it. This avoids resolving work from an invalid contract. return "invalid"; } - clearTimeout(pending.timeout); - this.pendingCommands.delete(message.requestId); - pending.resolve(result); + this.finishPending(pending, () => pending.resolve(result)); return "handled"; } - clearTimeout(pending.timeout); - this.pendingCommands.delete(message.requestId); - pending.reject(new Error(message.error ?? "The session failed to handle the command.")); + this.finishPending(pending, () => + pending.reject(new Error(message.error ?? "The session failed to handle the command.")), + ); return "handled"; } shutdown(error = new Error("The session broker daemon shut down.")) { - for (const [requestId, pending] of this.pendingCommands.entries()) { - clearTimeout(pending.timeout); - this.pendingCommands.delete(requestId); - pending.reject(error); + for (const pending of this.pendingCommands.values()) { + this.finishPending(pending, () => pending.reject(error), false); } + this.commandQueues.clear(); this.sessionIdsBySocket.clear(); this.sessions.clear(); + for (const reservation of this.retainedReservations.values()) reservation.release(); + for (const reservation of this.sessionReservations.values()) reservation.release(); + this.retainedReservations.clear(); + this.sessionReservations.clear(); + } + + /** Write the oldest queued command only when the session has no active command. */ + private advanceSessionQueue(sessionId: string): void { + const queue = this.commandQueues.get(sessionId); + if (!queue?.length) { + this.commandQueues.delete(sessionId); + return; + } + const first = this.pendingCommands.get(queue[0]!); + if (!first || first.active) return; + const entry = this.sessions.get(sessionId); + if (!entry || entry.socket !== first.socket) { + this.finishPending(first, () => + first.reject(new Error("The targeted session is no longer connected.")), + ); + return; + } + first.active = true; + try { + const accepted = entry.socket.send(first.serializedMessage); + if (accepted === false) throw new BrokerCapacityError("busy", "outbound"); + } catch (error) { + this.finishPending(first, () => + first.reject( + error instanceof Error + ? error + : new Error("The targeted session could not receive the command."), + ), + ); + } + } + + /** Complete one command exactly once, release reservations, and advance its session FIFO. */ + private finishPending( + pending: PendingCommand, + settle: () => void, + advance = true, + ): void { + if (this.pendingCommands.get(pending.requestId) !== pending) return; + clearTimeout(pending.timeout); + this.pendingCommands.delete(pending.requestId); + const queue = this.commandQueues.get(pending.sessionId); + if (queue) { + const index = queue.indexOf(pending.requestId); + if (index >= 0) queue.splice(index, 1); + if (queue.length === 0) this.commandQueues.delete(pending.sessionId); + } + pending.reservation.release(); + settle(); + if (advance) this.advanceSessionQueue(pending.sessionId); } /** Resolve one live session selector into the full in-memory registration entry. */ @@ -506,6 +709,10 @@ export class SessionBrokerState< } this.sessions.delete(sessionId); + this.retainedReservations.get(sessionId)?.release(); + this.retainedReservations.delete(sessionId); + this.sessionReservations.get(sessionId)?.release(); + this.sessionReservations.delete(sessionId); if (this.sessionIdsBySocket.get(entry.socket) === sessionId) { this.sessionIdsBySocket.delete(entry.socket); } @@ -514,14 +721,10 @@ export class SessionBrokerState< } private rejectPendingCommandsForSession(sessionId: string, error: Error) { - for (const [requestId, pending] of this.pendingCommands.entries()) { - if (pending.sessionId !== sessionId) { - continue; - } - - clearTimeout(pending.timeout); - this.pendingCommands.delete(requestId); - pending.reject(error); + for (const pending of this.pendingCommands.values()) { + if (pending.sessionId !== sessionId) continue; + this.finishPending(pending, () => pending.reject(error), false); } + this.commandQueues.delete(sessionId); } } diff --git a/packages/session-broker-core/src/budgets.test.ts b/packages/session-broker-core/src/budgets.test.ts new file mode 100644 index 000000000..86f189d68 --- /dev/null +++ b/packages/session-broker-core/src/budgets.test.ts @@ -0,0 +1,102 @@ +import { describe, expect, test } from "bun:test"; +import { + DEFAULT_SESSION_BROKER_LIMITS, + BrokerCapacityError, + ReservationGroup, + mergeSessionBrokerLimits, + ResourceBudget, + resolveSessionBrokerLimits, +} from "./budgets"; + +describe("session broker limits", () => { + test("publishes immutable contract defaults", () => { + expect(DEFAULT_SESSION_BROKER_LIMITS).toMatchObject({ + maxSessions: 256, + maxCommandsPerSession: 64, + maxCommandsTotal: 1_024, + maxPreBridgeCommands: 32, + maxCommandInputBytes: 1024 * 1024, + maxQueuedCommandBytes: 64 * 1024 * 1024, + maxRetainedSessionBytes: 4 * 1024 * 1024, + maxRetainedBytes: 256 * 1024 * 1024, + defaultCommandTimeoutMs: 15_000, + maxCommandTimeoutMs: 300_000, + maxConcurrentHttpControls: 32, + maxInFlightHttpBodyBytes: 64 * 1024 * 1024, + maxHttpBodyBytes: 4 * 1024 * 1024, + maxHttpResponseBytes: 8 * 1024 * 1024, + maxWsMessageBytes: 8 * 1024 * 1024, + maxInFlightWsBytes: 64 * 1024 * 1024, + }); + expect(Object.isFrozen(DEFAULT_SESSION_BROKER_LIMITS)).toBe(true); + }); + + test("merges partial daemon lowerings onto a complete broker snapshot", () => { + const broker = resolveSessionBrokerLimits({ + limits: { maxSessions: 4, maxWsMessageBytes: 64 }, + }); + const merged = mergeSessionBrokerLimits(broker, { + limits: { maxSessions: 2 }, + unsafeLimits: { maxHttpBodyBytes: broker.maxHttpBodyBytes + 1 }, + }); + expect(merged.maxSessions).toBe(2); + expect(merged.maxWsMessageBytes).toBe(64); + expect(merged.maxHttpBodyBytes).toBe(broker.maxHttpBodyBytes + 1); + expect(() => mergeSessionBrokerLimits(broker, { limits: { maxSessions: 5 } })).toThrow( + "unsafeLimits", + ); + }); + + test("allows supported lowering and requires unsafeLimits for raising", () => { + expect(resolveSessionBrokerLimits({ limits: { maxSessions: 1 } }).maxSessions).toBe(1); + expect(() => resolveSessionBrokerLimits({ limits: { maxSessions: 257 } })).toThrow( + "unsafeLimits", + ); + expect(resolveSessionBrokerLimits({ unsafeLimits: { maxSessions: 257 } }).maxSessions).toBe( + 257, + ); + expect(() => resolveSessionBrokerLimits({ limits: { maxSessions: -1 } })).toThrow( + "non-negative safe integer", + ); + expect(() => resolveSessionBrokerLimits({ limits: { unknown: 1 } as never })).toThrow( + "Unknown session broker limit", + ); + }); +}); + +describe("resource reservations", () => { + test("accepts the exact boundary, rejects boundary plus one, and releases idempotently", () => { + const budget = new ResourceBudget(4, "test"); + const reservation = budget.reserve(4); + expect(budget.used).toBe(4); + expect(() => budget.reserve(1)).toThrow(BrokerCapacityError); + reservation.release(); + reservation.release(); + expect(budget.used).toBe(0); + }); + + test("resizes retained records by their delta and transfers release ownership", () => { + const budget = new ResourceBudget(4, "bytes"); + const original = budget.reserve(4); + const identical = budget.resize(original, 4); + expect(budget.used).toBe(4); + original.release(); + expect(budget.used).toBe(4); + const smaller = budget.resize(identical, 2); + expect(budget.used).toBe(2); + smaller.release(); + expect(budget.used).toBe(0); + }); + + test("rolls grouped parser/send reservations back exactly once", () => { + const count = new ResourceBudget(1, "count"); + const bytes = new ResourceBudget(4, "bytes"); + const group = new ReservationGroup(); + group.add(count.reserve()); + group.add(bytes.reserve(4)); + expect(group.amount).toBe(5); + group.release(); + group.release(); + expect({ count: count.used, bytes: bytes.used }).toEqual({ count: 0, bytes: 0 }); + }); +}); diff --git a/packages/session-broker-core/src/budgets.ts b/packages/session-broker-core/src/budgets.ts new file mode 100644 index 000000000..51372b8bc --- /dev/null +++ b/packages/session-broker-core/src/budgets.ts @@ -0,0 +1,258 @@ +/** Defines the supported Phase-1 broker resource ceilings and reusable reservations. */ + +/** Enumerate every broker-owned resource ceiling used by current Phase-1 surfaces. */ +export interface SessionBrokerLimits { + readonly maxSessions: number; + readonly maxCommandsPerSession: number; + readonly maxCommandsTotal: number; + readonly maxPreBridgeCommands: number; + readonly maxCommandInputBytes: number; + readonly maxCommandResultBytes: number; + readonly maxQueuedCommandBytes: number; + readonly maxRetainedSessionBytes: number; + readonly maxRetainedBytes: number; + readonly defaultCommandTimeoutMs: number; + readonly maxCommandTimeoutMs: number; + readonly maxConcurrentHttpControls: number; + readonly maxInFlightHttpBodyBytes: number; + readonly maxHttpBodyBytes: number; + readonly maxHttpResponseBytes: number; + readonly maxWsMessageBytes: number; + readonly maxInFlightWsBytes: number; + readonly maxOutboundBytesPerPeer: number; + readonly maxOutboundBytesTotal: number; + readonly maxUnauthenticatedSockets: number; + readonly maxIncompleteHandshakes: number; + readonly maxIncompleteHandshakeBytes: number; + readonly maxHandshakeProposalBytes: number; + readonly maxCallerSessions: number; + readonly maxCallerSessionBytes: number; + readonly maxCallerSessionsBytes: number; +} + +/** Hold the immutable supported broker ceilings. Hosts may lower these without opting out. */ +export const DEFAULT_SESSION_BROKER_LIMITS: Readonly = Object.freeze({ + maxSessions: 256, + maxCommandsPerSession: 64, + maxCommandsTotal: 1_024, + maxPreBridgeCommands: 32, + maxCommandInputBytes: 1024 * 1024, + maxCommandResultBytes: 1024 * 1024, + maxQueuedCommandBytes: 64 * 1024 * 1024, + maxRetainedSessionBytes: 4 * 1024 * 1024, + maxRetainedBytes: 256 * 1024 * 1024, + defaultCommandTimeoutMs: 15_000, + maxCommandTimeoutMs: 5 * 60_000, + maxConcurrentHttpControls: 32, + maxInFlightHttpBodyBytes: 64 * 1024 * 1024, + maxHttpBodyBytes: 4 * 1024 * 1024, + maxHttpResponseBytes: 8 * 1024 * 1024, + maxWsMessageBytes: 8 * 1024 * 1024, + maxInFlightWsBytes: 64 * 1024 * 1024, + maxOutboundBytesPerPeer: 8 * 1024 * 1024, + maxOutboundBytesTotal: 64 * 1024 * 1024, + maxUnauthenticatedSockets: 64, + maxIncompleteHandshakes: 128, + maxIncompleteHandshakeBytes: 4 * 1024 * 1024, + maxHandshakeProposalBytes: 64 * 1024, + maxCallerSessions: 256, + maxCallerSessionBytes: 8 * 1024, + maxCallerSessionsBytes: 2 * 1024 * 1024, +}); + +export interface SessionBrokerLimitOptions { + /** Supported configuration may only lower the immutable defaults. */ + readonly limits?: Partial; + /** Explicitly opts out of supported ceilings. Hosts assume the resulting resource risk. */ + readonly unsafeLimits?: Partial; +} + +export type BrokerCapacityCode = "busy" | "queue-full" | "capacity-exceeded"; + +/** Report a stable resource-admission failure without exposing internal accounting. */ +export class BrokerCapacityError extends Error { + constructor( + readonly code: BrokerCapacityCode, + readonly resource: keyof SessionBrokerLimits | string, + ) { + super(code); + this.name = "BrokerCapacityError"; + } +} + +function assertLimit(value: unknown, name: string): asserts value is number { + if (!Number.isSafeInteger(value) || (value as number) < 0) { + throw new TypeError(`Session broker limit ${name} must be a non-negative safe integer.`); + } +} + +/** Merge named lowerings and explicit unsafe overrides onto one complete limit snapshot. */ +export function mergeSessionBrokerLimits( + base: Readonly, + options: SessionBrokerLimitOptions = {}, +): Readonly { + const supported = options.limits ?? {}; + const unsafe = options.unsafeLimits ?? {}; + if (!supported || typeof supported !== "object" || !unsafe || typeof unsafe !== "object") { + throw new TypeError("Session broker limits must be objects."); + } + + const resolved = { ...base } as SessionBrokerLimits; + for (const name of Object.keys(DEFAULT_SESSION_BROKER_LIMITS) as (keyof SessionBrokerLimits)[]) { + assertLimit(base[name], name); + const supportedValue = supported[name]; + if (supportedValue !== undefined) { + assertLimit(supportedValue, name); + if (supportedValue > base[name]) { + throw new TypeError( + `Session broker limit ${name} may only be raised through unsafeLimits.`, + ); + } + (resolved as Record)[name] = supportedValue; + } + const unsafeValue = unsafe[name]; + if (unsafeValue !== undefined) { + assertLimit(unsafeValue, name); + (resolved as Record)[name] = unsafeValue; + } + } + + for (const key of [...Object.keys(supported), ...Object.keys(unsafe)]) { + if (!(key in DEFAULT_SESSION_BROKER_LIMITS)) { + throw new TypeError(`Unknown session broker limit ${key}.`); + } + } + if (resolved.defaultCommandTimeoutMs > resolved.maxCommandTimeoutMs) { + throw new TypeError( + "Session broker defaultCommandTimeoutMs must not exceed maxCommandTimeoutMs.", + ); + } + if (resolved.maxRetainedSessionBytes > resolved.maxRetainedBytes) { + throw new TypeError("Session broker per-session retained bytes must not exceed daemon bytes."); + } + if (resolved.maxCallerSessionBytes > resolved.maxCallerSessionsBytes) { + throw new TypeError("Session broker per-caller retained bytes must not exceed daemon bytes."); + } + return Object.freeze(resolved); +} + +/** Validate and snapshot limits relative to the immutable supported defaults. */ +export function resolveSessionBrokerLimits( + options: SessionBrokerLimitOptions = {}, +): Readonly { + return mergeSessionBrokerLimits(DEFAULT_SESSION_BROKER_LIMITS, options); +} + +/** Release one successful budget reservation at most once. */ +export interface BudgetReservation { + readonly amount: number; + readonly released: boolean; + release(): void; +} + +/** Track one count or byte budget with reserve-before-work admission. */ +export class ResourceBudget { + private reserved = 0; + private readonly reservationStates = new WeakMap< + BudgetReservation, + { amount: number; released: boolean } + >(); + + constructor( + readonly capacity: number, + private readonly resource: string, + private readonly code: BrokerCapacityCode = "capacity-exceeded", + ) { + assertLimit(capacity, resource); + } + + get used(): number { + return this.reserved; + } + + tryReserve(amount = 1): BudgetReservation | null { + assertLimit(amount, this.resource); + if (amount > this.capacity - this.reserved) return null; + this.reserved += amount; + const state = { amount, released: false }; + const reservation: BudgetReservation = { + amount, + get released() { + return state.released; + }, + release: () => { + if (state.released) return; + state.released = true; + this.reservationStates.delete(reservation); + this.reserved -= state.amount; + }, + }; + this.reservationStates.set(reservation, state); + return reservation; + } + + reserve(amount = 1): BudgetReservation { + const reservation = this.tryReserve(amount); + if (!reservation) throw new BrokerCapacityError(this.code, this.resource); + return reservation; + } + + /** Atomically replace one live reservation, charging only its positive size delta. */ + resize(previous: BudgetReservation, amount: number): BudgetReservation { + assertLimit(amount, this.resource); + const previousState = this.reservationStates.get(previous); + if (!previousState || previousState.released) { + throw new TypeError(`Cannot resize an inactive ${this.resource} reservation.`); + } + const delta = amount - previousState.amount; + if (delta > this.capacity - this.reserved) { + throw new BrokerCapacityError(this.code, this.resource); + } + this.reserved += delta; + const replacementState = { amount, released: false }; + const replacement: BudgetReservation = { + amount, + get released() { + return replacementState.released; + }, + release: () => { + if (replacementState.released) return; + replacementState.released = true; + this.reservationStates.delete(replacement); + this.reserved -= replacementState.amount; + }, + }; + previousState.released = true; + this.reservationStates.delete(previous); + this.reservationStates.set(replacement, replacementState); + return replacement; + } +} + +/** Own several incremental reservations and release all of them idempotently. */ +export class ReservationGroup implements BudgetReservation { + private reservations: BudgetReservation[] = []; + private done = false; + + get amount(): number { + return this.reservations.reduce((sum, reservation) => sum + reservation.amount, 0); + } + + get released(): boolean { + return this.done; + } + + add(reservation: BudgetReservation): void { + if (this.done) { + reservation.release(); + throw new Error("Cannot add to a released reservation group."); + } + this.reservations.push(reservation); + } + + release(): void { + if (this.done) return; + this.done = true; + for (const reservation of this.reservations.splice(0)) reservation.release(); + } +} diff --git a/packages/session-broker-core/src/index.ts b/packages/session-broker-core/src/index.ts index 6ce0d1780..9e2f37555 100644 --- a/packages/session-broker-core/src/index.ts +++ b/packages/session-broker-core/src/index.ts @@ -4,6 +4,7 @@ export * from "./auth"; export * from "./validation"; export * from "./brokerWire"; export * from "./limits"; +export * from "./budgets"; export * from "./brokerState"; export * from "./selectors"; export * from "./sessionTerminalMetadata"; diff --git a/packages/session-broker-core/src/limits.test.ts b/packages/session-broker-core/src/limits.test.ts index fef455aa3..84838e94b 100644 --- a/packages/session-broker-core/src/limits.test.ts +++ b/packages/session-broker-core/src/limits.test.ts @@ -1,10 +1,13 @@ import { describe, expect, test } from "bun:test"; import { + InvalidContentLengthError, PayloadTooLargeError, readRequestBytesWithLimit, + readRequestBytesWithReservation, readRequestTextWithLimit, utf8ByteLength, } from "./limits"; +import { ResourceBudget } from "./budgets"; /** Build a streaming request body so the read path runs without a Content-Length header. */ function streamingRequest(byteLength: number, chunkSize = 64 * 1024) { @@ -52,6 +55,29 @@ describe("readRequestTextWithLimit", () => { ); }); + test("rejects malformed Content-Length instead of treating it as undeclared", async () => { + const request = new Request("http://broker.test/api", { + method: "POST", + headers: { "content-length": "01" }, + body: "x", + }); + await expect(readRequestBytesWithLimit(request, 1024)).rejects.toBeInstanceOf( + InvalidContentLengthError, + ); + }); + + test("retains aggregate byte capacity through parse and releases it idempotently", async () => { + const budget = new ResourceBudget(4, "http"); + const request = new Request("http://broker.test/api", { method: "POST", body: "éé" }); + const read = await readRequestBytesWithReservation(request, 4, budget); + expect(read.bytes.byteLength).toBe(4); + expect(budget.used).toBe(4); + expect(budget.tryReserve(1)).toBeNull(); + read.reservation.release(); + read.reservation.release(); + expect(budget.used).toBe(0); + }); + test("returns the decoded body when it stays under the limit", async () => { const request = new Request("http://broker.test/api", { method: "POST", diff --git a/packages/session-broker-core/src/limits.ts b/packages/session-broker-core/src/limits.ts index 275f9f04d..90be29a56 100644 --- a/packages/session-broker-core/src/limits.ts +++ b/packages/session-broker-core/src/limits.ts @@ -7,11 +7,19 @@ * or patch bytes. These caps keep memory bounded while staying far above any realistic review. */ +import { + DEFAULT_SESSION_BROKER_LIMITS, + BrokerCapacityError, + ReservationGroup, + type BudgetReservation, + type ResourceBudget, +} from "./budgets"; + /** Maximum decoded byte length accepted for one HTTP API request body. */ -export const MAX_HTTP_BODY_BYTES = 4 * 1024 * 1024; +export const MAX_HTTP_BODY_BYTES = DEFAULT_SESSION_BROKER_LIMITS.maxHttpBodyBytes; /** Maximum byte length accepted for one inbound websocket message. */ -export const MAX_WS_MESSAGE_BYTES = 8 * 1024 * 1024; +export const MAX_WS_MESSAGE_BYTES = DEFAULT_SESSION_BROKER_LIMITS.maxWsMessageBytes; /** Maximum number of files accepted in one session registration payload. */ export const MAX_REGISTRATION_FILES = 5_000; @@ -36,6 +44,14 @@ export class PayloadTooLargeError extends Error { } } +/** Raised when Content-Length is ambiguous instead of a canonical non-negative integer. */ +export class InvalidContentLengthError extends Error { + constructor() { + super("Content-Length must be a canonical non-negative integer."); + this.name = "InvalidContentLengthError"; + } +} + // Reused across every websocket message, HTTP body, and patch check to avoid a per-call alloc. const sharedTextEncoder = new TextEncoder(); const fatalTextDecoder = new TextDecoder("utf-8", { fatal: true }); @@ -52,19 +68,80 @@ export function utf8ByteLength(value: string): number { * stream is aborted mid-read so a missing or lying Content-Length cannot force the daemon to * buffer an unbounded body before the cap is noticed. */ -export async function readRequestBytesWithLimit( +export async function readRequestBytesWithReservation( request: Request, maxBytes: number, -): Promise { - const declared = request.headers.get("content-length"); - if (declared && /^(?:0|[1-9][0-9]*)$/.test(declared) && Number(declared) > maxBytes) { + aggregateBudget?: ResourceBudget, +): Promise<{ bytes: Uint8Array; reservation: BudgetReservation }> { + const declaredHeader = request.headers.get("content-length"); + if (declaredHeader !== null && !/^(?:0|[1-9][0-9]*)$/.test(declaredHeader)) { + throw new InvalidContentLengthError(); + } + const declared = declaredHeader === null ? null : Number(declaredHeader); + if (declared !== null && (!Number.isSafeInteger(declared) || declared > maxBytes)) { throw new PayloadTooLargeError(maxBytes); } - const body = request.body; - if (!body) return new Uint8Array(); + const reservations = new ReservationGroup(); + try { + if (aggregateBudget && declared !== null) reservations.add(aggregateBudget.reserve(declared)); + const body = request.body; + if (!body) return { bytes: new Uint8Array(), reservation: reservations }; + + const reader = body.getReader(); + const chunks: Uint8Array[] = []; + let total = 0; + try { + for (;;) { + const { done, value } = await reader.read(); + if (done) break; + if (!value) continue; + + const nextTotal = total + value.byteLength; + if (nextTotal > maxBytes) { + await reader.cancel().catch(() => {}); + throw new PayloadTooLargeError(maxBytes); + } + if (aggregateBudget && nextTotal > (declared ?? 0)) { + reservations.add(aggregateBudget.reserve(nextTotal - Math.max(total, declared ?? 0))); + } + total = nextTotal; + chunks.push(value); + } + } finally { + reader.releaseLock(); + } + + const merged = new Uint8Array(total); + let offset = 0; + for (const chunk of chunks) { + merged.set(chunk, offset); + offset += chunk.byteLength; + } + return { bytes: merged, reservation: reservations }; + } catch (error) { + reservations.release(); + throw error; + } +} - const reader = body.getReader(); +/** Buffer one finite non-SSE response under the shared hard response and in-flight budget. */ +export async function boundHttpResponse( + response: Response, + maxBytes: number, + aggregateBudget?: ResourceBudget, +): Promise { + if (response.headers.get("content-type")?.toLowerCase().startsWith("text/event-stream")) { + return response; + } + const declared = response.headers.get("content-length"); + if (declared && /^(?:0|[1-9][0-9]*)$/.test(declared) && Number(declared) > maxBytes) { + return new Response(null, { status: 503 }); + } + if (!response.body) return response; + + const reader = response.body.getReader(); + const reservations = new ReservationGroup(); const chunks: Uint8Array[] = []; let total = 0; try { @@ -72,25 +149,46 @@ export async function readRequestBytesWithLimit( const { done, value } = await reader.read(); if (done) break; if (!value) continue; - total += value.byteLength; if (total > maxBytes) { await reader.cancel().catch(() => {}); - throw new PayloadTooLargeError(maxBytes); + return new Response(null, { status: 503 }); } + if (aggregateBudget) reservations.add(aggregateBudget.reserve(value.byteLength)); chunks.push(value); } + const body = new Uint8Array(total); + let offset = 0; + for (const chunk of chunks) { + body.set(chunk, offset); + offset += chunk.byteLength; + } + const headers = new Headers(response.headers); + headers.set("content-length", String(total)); + return new Response(body, { + status: response.status, + statusText: response.statusText, + headers, + }); + } catch (error) { + if (error instanceof BrokerCapacityError) { + await reader.cancel().catch(() => {}); + return new Response(null, { status: 503 }); + } + throw error; } finally { + reservations.release(); reader.releaseLock(); } +} - const merged = new Uint8Array(total); - let offset = 0; - for (const chunk of chunks) { - merged.set(chunk, offset); - offset += chunk.byteLength; - } - return merged; +export async function readRequestBytesWithLimit( + request: Request, + maxBytes: number, +): Promise { + const { bytes, reservation } = await readRequestBytesWithReservation(request, maxBytes); + reservation.release(); + return bytes; } /** Read and strictly decode one bounded request body as UTF-8 text. */ diff --git a/packages/session-broker-node/package.json b/packages/session-broker-node/package.json index e8bd9b1f9..c417693e6 100644 --- a/packages/session-broker-node/package.json +++ b/packages/session-broker-node/package.json @@ -21,6 +21,6 @@ }, "engines": { "bun": ">=1.0.0", - "node": ">=18" + "node": ">=22" } } diff --git a/packages/session-broker-node/src/serve.test.ts b/packages/session-broker-node/src/serve.test.ts index bbcdce627..876c568dd 100644 --- a/packages/session-broker-node/src/serve.test.ts +++ b/packages/session-broker-node/src/serve.test.ts @@ -13,6 +13,7 @@ import { createSessionBrokerDaemon, createSessionBrokerProtocolParsers, } from "@hunk/session-broker"; +import SESSION_BROKER_ADAPTER_CONFORMANCE from "../../../test/fixtures/sessionBrokerAdapterConformance.json" with { type: "json" }; import { serveSessionBrokerDaemon } from "./serve"; interface TestSessionInfo { @@ -115,6 +116,13 @@ async function waitUntil( } describe("session broker node adapter", () => { + test("uses the shared binary, oversize, and pressure close corpus", () => { + expect(SESSION_BROKER_ADAPTER_CONFORMANCE).toMatchObject({ + textOnly: { binaryCloseCode: 1003 }, + inbound: { oversizedCloseCode: 1009, pressureCloseCode: 1013 }, + }); + }); + test("serves the generic daemon API and websocket path through Node", async () => { const broker = new SessionBroker({ protocolParsers }); const daemon = createSessionBrokerDaemon({ diff --git a/packages/session-broker-node/src/serve.ts b/packages/session-broker-node/src/serve.ts index 82f63b725..fc5b6a52a 100644 --- a/packages/session-broker-node/src/serve.ts +++ b/packages/session-broker-node/src/serve.ts @@ -1,7 +1,13 @@ import { createServer, type IncomingMessage, type ServerResponse } from "node:http"; import { Readable } from "node:stream"; -import type { AddressInfo } from "node:net"; -import type { SessionServerMessage } from "@hunk/session-broker-core"; +import type { AddressInfo, Socket } from "node:net"; +import { + BrokerCapacityError, + ResourceBudget, + boundHttpResponse, + type BudgetReservation, + type SessionServerMessage, +} from "@hunk/session-broker-core"; import type { SessionBrokerDaemon, SessionBrokerPeer } from "@hunk/session-broker"; import { WebSocketServer, type WebSocket } from "ws"; @@ -39,10 +45,34 @@ function defaultServeError(error: unknown, address: { hostname: string; port: nu ); } -function toNodeConnection(socket: WebSocket): SessionBrokerPeer { +function toNodeConnection( + socket: WebSocket, + outboundBudget: ResourceBudget, + maxPeerBytes: number, +): SessionBrokerPeer { return { send(data: string) { - socket.send(data); + const bytes = Buffer.byteLength(data); + if (bytes > maxPeerBytes - socket.bufferedAmount) { + socket.close(1013, "Session broker outbound pressure exceeded."); + throw new BrokerCapacityError("busy", "maxOutboundBytesPerPeer"); + } + const reservation = outboundBudget.tryReserve(bytes); + if (!reservation) { + socket.close(1013, "Session broker outbound pressure exceeded."); + throw new BrokerCapacityError("busy", "maxOutboundBytesTotal"); + } + try { + socket.send(data, (error) => { + reservation.release(); + if (error && socket.readyState < 2) { + socket.close(1013, "Session broker outbound delivery failed."); + } + }); + } catch (error) { + reservation.release(); + throw error; + } }, close(code?: number, reason?: string) { socket.close(code, reason); @@ -67,21 +97,56 @@ async function toRequest(request: IncomingMessage, hostname: string, port: numbe } as RequestInit & { duplex?: "half" }); } -async function writeResponse(nodeResponse: ServerResponse, response: Response) { +async function writeResponse( + nodeResponse: ServerResponse, + sourceResponse: Response, + responseBudget: ResourceBudget, + maxResponseBytes: number, +) { + const response = await boundHttpResponse(sourceResponse, maxResponseBytes, responseBudget); + const streaming = response.headers + .get("content-type") + ?.toLowerCase() + .startsWith("text/event-stream"); nodeResponse.statusCode = response.status; nodeResponse.statusMessage = response.statusText; - - response.headers.forEach((value, key) => { - nodeResponse.setHeader(key, value); - }); - + response.headers.forEach((value, key) => nodeResponse.setHeader(key, value)); if (!response.body) { nodeResponse.end(); return; } - const body = Buffer.from(await response.arrayBuffer()); - nodeResponse.end(body); + const reader = response.body.getReader(); + let total = 0; + try { + for (;;) { + const { done, value } = await reader.read(); + if (done) break; + if (!value) continue; + total += value.byteLength; + if (!streaming && total > maxResponseBytes) { + await reader.cancel().catch(() => {}); + nodeResponse.destroy(new BrokerCapacityError("capacity-exceeded", "maxHttpResponseBytes")); + return; + } + const reservation = streaming ? null : responseBudget.tryReserve(value.byteLength); + if (!streaming && !reservation) { + await reader.cancel().catch(() => {}); + nodeResponse.destroy(new BrokerCapacityError("busy", "maxHttpResponseBytes")); + return; + } + await new Promise((resolve, reject) => { + nodeResponse.write(value, (error) => { + reservation?.release(); + if (error) reject(error); + else resolve(); + }); + }); + } + nodeResponse.end(); + } finally { + reader.releaseLock(); + } } /** Serve one runtime-neutral broker daemon through Node HTTP and ws. */ @@ -92,26 +157,70 @@ export async function serveSessionBrokerDaemon< >( options: ServeSessionBrokerDaemonOptions, ): Promise { + const inboundBudget = new ResourceBudget( + options.daemon.limits.maxInFlightWsBytes, + "maxInFlightWsBytes", + ); + const outboundBudget = new ResourceBudget( + options.daemon.limits.maxOutboundBytesTotal, + "maxOutboundBytesTotal", + "busy", + ); + const responseBudget = new ResourceBudget( + options.daemon.limits.maxHttpResponseBytes, + "maxHttpResponseBytes", + "busy", + ); const server = createServer(async (incoming, outgoing) => { const request = await toRequest(incoming, options.hostname, options.port); const customResponse = await options.handleRequest?.(request, server); if (customResponse !== undefined) { - await writeResponse(outgoing, customResponse); + await writeResponse( + outgoing, + customResponse, + responseBudget, + options.daemon.limits.maxHttpResponseBytes, + ); return; } const daemonResponse = await options.daemon.handleRequest(request); if (daemonResponse) { - await writeResponse(outgoing, daemonResponse); + await writeResponse( + outgoing, + daemonResponse, + responseBudget, + options.daemon.limits.maxHttpResponseBytes, + ); return; } - await writeResponse(outgoing, (await options.notFound?.(request)) ?? defaultNotFound()); + await writeResponse( + outgoing, + (await options.notFound?.(request)) ?? defaultNotFound(), + responseBudget, + options.daemon.limits.maxHttpResponseBytes, + ); + }); + const unauthenticatedSocketBudget = new ResourceBudget( + options.daemon.limits.maxUnauthenticatedSockets, + "maxUnauthenticatedSockets", + "busy", + ); + const webSocketServer = new WebSocketServer({ + noServer: true, + maxPayload: Math.max(1, options.daemon.limits.maxWsMessageBytes), }); - const webSocketServer = new WebSocketServer({ noServer: true }); // Reuse one stable peer wrapper per websocket so close events unregister the same logical // connection object that registration and message handling used earlier. const peerBySocket = new WeakMap(); + const admissionBySocket = new WeakMap(); + const activeWebSockets = new Set(); + const activeSockets = new Set(); + server.on("connection", (socket) => { + activeSockets.add(socket); + socket.once("close", () => activeSockets.delete(socket)); + }); let resolved = false; let resolveStopped: (() => void) | null = null; const stopped = new Promise((resolve) => { @@ -128,20 +237,58 @@ export async function serveSessionBrokerDaemon< }; webSocketServer.on("connection", (socket: WebSocket) => { - const peer = toNodeConnection(socket); + activeWebSockets.add(socket); + const peer = toNodeConnection( + socket, + outboundBudget, + options.daemon.limits.maxOutboundBytesPerPeer, + ); peerBySocket.set(socket, peer); - socket.on("message", (message: string | Buffer | ArrayBuffer | Buffer[]) => { - const text = - typeof message === "string" - ? message - : Array.isArray(message) - ? Buffer.concat(message).toString() - : message instanceof ArrayBuffer - ? Buffer.from(new Uint8Array(message)).toString() - : Buffer.from(message).toString(); - options.daemon.handleConnectionMessage(peer, text); + socket.on("message", (message: Buffer | ArrayBuffer | Buffer[], isBinary: boolean) => { + if (isBinary) { + socket.close(1003, "Session broker accepts text messages only."); + return; + } + const bytes = Array.isArray(message) + ? Buffer.concat(message) + : message instanceof ArrayBuffer + ? Buffer.from(new Uint8Array(message)) + : Buffer.from(message); + if (bytes.byteLength > options.daemon.limits.maxWsMessageBytes) { + socket.close(1009, "Message exceeds the session broker size limit."); + return; + } + const reservation = inboundBudget.tryReserve(bytes.byteLength); + if (!reservation) { + socket.close(1013, "Session broker inbound pressure exceeded."); + return; + } + try { + let text: string; + try { + text = new TextDecoder("utf-8", { fatal: true }).decode(bytes); + } catch { + socket.close(1007, "Malformed UTF-8 session broker message."); + return; + } + try { + options.daemon.handleConnectionMessage(peer, text); + } catch (error) { + socket.close( + error instanceof BrokerCapacityError ? 1013 : 1011, + "Session broker message handling failed.", + ); + } + } finally { + reservation.release(); + } }); + // ws reports maxPayload violations through an error event before its protocol close. Keep the + // process alive while ws completes the required 1009 close handshake. + socket.on("error", () => {}); socket.on("close", (code: number, reason: Buffer) => { + activeWebSockets.delete(socket); + admissionBySocket.get(socket)?.release(); options.daemon.handleConnectionClose(peerBySocket.get(socket) ?? peer); // The runtime-neutral daemon only cares that the transport closed; Node-specific close data // stays ignored here instead of leaking into the shared broker API. @@ -158,9 +305,22 @@ export async function serveSessionBrokerDaemon< return; } - webSocketServer.handleUpgrade(request, socket, head, (webSocket: WebSocket) => { - webSocketServer.emit("connection", webSocket, request); - }); + const admission = unauthenticatedSocketBudget.tryReserve(); + if (!admission) { + socket.write("HTTP/1.1 503 Service Unavailable\r\nContent-Length: 0\r\n\r\n"); + socket.destroy(); + return; + } + socket.once("close", () => admission.release()); + try { + webSocketServer.handleUpgrade(request, socket, head, (webSocket: WebSocket) => { + admissionBySocket.set(webSocket, admission); + webSocketServer.emit("connection", webSocket, request); + }); + } catch { + admission.release(); + socket.destroy(); + } }); await new Promise((resolve, reject) => { @@ -183,21 +343,41 @@ export async function serveSessionBrokerDaemon< server.listen(options.port, options.hostname); }); - const stop = async () => { - // Shut down the daemon first so pending broker commands reject before the transport disappears. - options.daemon.shutdown(); - await new Promise((resolve) => webSocketServer.close(() => resolve())); - await new Promise((resolve, reject) => { - server.close((error) => { - if (error) { - reject(error); - return; - } - - resolve(); - }); - }); - finish(); + let stopPromise: Promise | null = null; + const stop = () => { + if (stopPromise) return stopPromise; + stopPromise = (async () => { + // Reject broker work, then terminate active peers before waiting for ws/server close callbacks. + options.daemon.shutdown(); + const peerClosures = [...activeWebSockets].map( + (socket) => + new Promise((resolve) => { + if (socket.readyState === socket.CLOSED) return resolve(); + socket.once("close", () => resolve()); + socket.terminate(); + }), + ); + for (const socket of activeSockets) socket.destroy(); + await Promise.all(peerClosures); + await new Promise((resolve) => webSocketServer.close(() => resolve())); + const closeServer = () => { + server.close(); + server.closeAllConnections(); + }; + if ((globalThis as { Bun?: unknown }).Bun) { + // Bun's Node compatibility layer does not consistently emit Server's close callback after + // upgraded sockets are terminated; the real Node path below still awaits native closure. + closeServer(); + } else { + await new Promise((resolve, reject) => { + server.once("close", resolve); + server.once("error", reject); + closeServer(); + }); + } + finish(); + })(); + return stopPromise; }; void options.daemon.stopped.then(async () => { diff --git a/packages/session-broker/package.json b/packages/session-broker/package.json index 0a9573322..27dcc8bbc 100644 --- a/packages/session-broker/package.json +++ b/packages/session-broker/package.json @@ -20,6 +20,6 @@ }, "engines": { "bun": ">=1.0.0", - "node": ">=18" + "node": ">=22" } } diff --git a/packages/session-broker/src/authentication.test.ts b/packages/session-broker/src/authentication.test.ts index 3ca2d550c..190853f92 100644 --- a/packages/session-broker/src/authentication.test.ts +++ b/packages/session-broker/src/authentication.test.ts @@ -57,8 +57,10 @@ async function setup( options: { revoked?: () => boolean; maxChallenges?: number; + maxChallengeBytes?: number; maxChallengeTranscriptBytes?: number; maxCallerSessions?: number; + limits?: { maxCallerSessionBytes?: number; maxCallerSessionsBytes?: number }; callerSessionTtlMs?: number; crypto?: SessionBrokerCrypto; } = {}, @@ -79,9 +81,11 @@ async function setup( now: () => now, isRevoked: options.revoked, maxChallenges: options.maxChallenges, + maxChallengeBytes: options.maxChallengeBytes, maxChallengeTranscriptBytes: options.maxChallengeTranscriptBytes, maxCallerSessions: options.maxCallerSessions, callerSessionTtlMs: options.callerSessionTtlMs, + limits: options.limits, crypto: options.crypto, }); return { @@ -551,6 +555,40 @@ describe("session broker signed authentication", () => { } }); + test("retains incomplete-handshake capacity through asynchronous proof verification", async () => { + let blockVerify = false; + let releaseVerify!: () => void; + const verifyGate = new Promise((resolve) => { + releaseVerify = resolve; + }); + const cryptoWithGate: SessionBrokerCrypto = { + ...webSessionBrokerCrypto, + async verify(publicKey, signature, value) { + if (blockVerify) await verifyGate; + return webSessionBrokerCrypto.verify(publicKey, signature, value); + }, + }; + const values = await setup({ maxChallenges: 1, crypto: cryptoWithGate }); + const request = challengeRequest(); + const challenge = await values.authenticator.issueChallenge(request, request.endpoint); + const transcript = challengeTranscriptForClient(request, challenge, "generation-1"); + const signature = encodeBase64Url( + await webSessionBrokerCrypto.sign(values.caller.privateKey, transcript), + ); + + blockVerify = true; + const completing = values.authenticator.completeCallerHello({ + challengeId: challenge.challengeId, + signature, + }); + await Bun.sleep(0); + await expect( + values.authenticator.issueChallenge(request, request.endpoint), + ).rejects.toMatchObject({ code: "authentication-capacity" }); + releaseVerify(); + await expect(completing).resolves.toMatchObject({ brokerRevision: 1 }); + }); + test("bounds pending challenge counts and retained transcript bytes", async () => { const values = await setup({ maxChallenges: 1 }); await values.authenticator.issueChallenge(challengeRequest(), challengeRequest().endpoint); @@ -565,9 +603,32 @@ describe("session broker signed authentication", () => { byteBound.authenticator.issueChallenge(challengeRequest(), challengeRequest().endpoint), ).rejects.toMatchObject({ code: "authentication-capacity" }); + const completeRecordBound = await setup({ maxChallengeBytes: 512 }); + await expect( + completeRecordBound.authenticator.issueChallenge( + challengeRequest(), + challengeRequest().endpoint, + ), + ).rejects.toMatchObject({ code: "authentication-capacity" }); + const noCallerCapacity = await setup({ maxCallerSessions: 0 }); await expect(openCallerSession(noCallerCapacity)).rejects.toMatchObject({ code: "authentication-capacity", }); + + const noCallerByteCapacity = await setup({ + limits: { maxCallerSessionBytes: 1, maxCallerSessionsBytes: 1 }, + }); + await expect(openCallerSession(noCallerByteCapacity)).rejects.toMatchObject({ + code: "authentication-capacity", + }); + + const reusableCallerCapacity = await setup({ maxCallerSessions: 1 }); + const first = await openCallerSession(reusableCallerCapacity); + reusableCallerCapacity.authenticator.revokeCallerSession(first.session.callerSessionId); + await expect(openCallerSession(reusableCallerCapacity)).resolves.toMatchObject({ + session: { initialSequence: "1" }, + }); + reusableCallerCapacity.authenticator.clear(); }); }); diff --git a/packages/session-broker/src/authentication.ts b/packages/session-broker/src/authentication.ts index dc0c3750e..08e36cb01 100644 --- a/packages/session-broker/src/authentication.ts +++ b/packages/session-broker/src/authentication.ts @@ -17,7 +17,12 @@ import { parseBrokerIdentifier, parseBrokerString, parseExactBrokerRecord, + ReservationGroup, + ResourceBudget, + resolveSessionBrokerLimits, + utf8ByteLength, type BrokerAppContract, + type BudgetReservation, type BrokerChallengeTranscriptInput, type BrokerGrant, type BrokerHelloProposal, @@ -28,6 +33,7 @@ import { type ProducerGrant, type ProducerOperation, type ProducerPrincipal, + type SessionBrokerLimitOptions, } from "@hunk/session-broker-core"; import { decodeBase64Url, @@ -38,13 +44,27 @@ import { const DEFAULT_CHALLENGE_TTL_MS = 15_000; const DEFAULT_CALLER_SESSION_TTL_MS = 5 * 60_000; -const DEFAULT_MAX_CHALLENGES = 128; -const DEFAULT_MAX_CHALLENGE_BYTES = 4 * 1024 * 1024; -const DEFAULT_MAX_CHALLENGE_TRANSCRIPT_BYTES = 64 * 1024; -const DEFAULT_MAX_CALLER_SESSIONS = 256; const UNIQUE_ID_RETRIES = 16; const RANDOM_ID_BYTES = 24; const MAX_ENDPOINT_LENGTH = 2_048; +const CHALLENGE_RECORD_OVERHEAD_BYTES = 320; +const CALLER_SESSION_RECORD_OVERHEAD_BYTES = 384; +const CRYPTO_KEY_REFERENCE_BYTES = 64; + +/** Measure retained JSON fields plus opaque runtime references and map bookkeeping. */ +function retainedRecordBytes(values: readonly unknown[], overhead: number): number { + let total = overhead + CRYPTO_KEY_REFERENCE_BYTES; + for (const value of values) { + if (value instanceof Uint8Array) { + total += value.byteLength; + continue; + } + const serialized = JSON.stringify(value); + if (serialized === undefined) authenticationError("invalid-credential"); + total += utf8ByteLength(serialized); + } + return total; +} const PRODUCER_OPERATIONS = new Set(["register", "reconnect"]); const CALLER_OPERATIONS = new Set([ "list", @@ -163,6 +183,7 @@ export interface AuthenticatedCallerRequest { export interface CallerRequestAuthenticator { authenticate(input: CallerRequestAuthenticationInput): Promise; + clear?(): void; } interface PendingChallenge { @@ -171,7 +192,7 @@ interface PendingChallenge { readonly grant: BrokerGrant; readonly publicKey: CryptoKey; readonly expiresAt: number; - readonly retainedBytes: number; + readonly reservation: BudgetReservation; } interface CallerSessionRecord { @@ -181,6 +202,7 @@ interface CallerSessionRecord { readonly helloTranscriptHash: string; readonly expiresAt: number; readonly replay: CallerSequenceReplayWindow; + readonly reservation: BudgetReservation; } export interface SessionBrokerAuthenticatorOptions { @@ -198,6 +220,8 @@ export interface SessionBrokerAuthenticatorOptions { readonly maxChallengeBytes?: number; readonly maxChallengeTranscriptBytes?: number; readonly maxCallerSessions?: number; + readonly limits?: SessionBrokerLimitOptions["limits"]; + readonly unsafeLimits?: SessionBrokerLimitOptions["unsafeLimits"]; } interface AuthenticatorSnapshot { @@ -213,6 +237,8 @@ interface AuthenticatorSnapshot { readonly maxChallengeBytes: number; readonly maxChallengeTranscriptBytes: number; readonly maxCallerSessions: number; + readonly maxCallerSessionBytes: number; + readonly maxCallerSessionsBytes: number; } function authenticationError(code: SessionBrokerAuthenticationFailureCode): never { @@ -246,11 +272,13 @@ function configuredNumber( fallback: number, name: string, allowZero: boolean, + maximum = fallback, ): number { const selected = value ?? fallback; if (!Number.isSafeInteger(selected) || selected < (allowZero ? 0 : 1)) { invalidStartup(`${name} must be ${allowZero ? "a non-negative" : "a positive"} safe integer.`); } + if (selected > maximum) invalidStartup(`${name} may only be raised through unsafeLimits.`); return selected; } @@ -379,6 +407,15 @@ function copyOptions(options: SessionBrokerAuthenticatorOptions): AuthenticatorS if (options.isRevoked !== undefined && typeof options.isRevoked !== "function") { invalidStartup("isRevoked must be a function."); } + const limits = resolveSessionBrokerLimits({ + ...(options.limits ? { limits: options.limits } : {}), + ...(options.unsafeLimits ? { unsafeLimits: options.unsafeLimits } : {}), + }); + const retainedCallerCapacity = + limits.maxCallerSessionBytes === 0 + ? 0 + : Math.floor(limits.maxCallerSessionsBytes / limits.maxCallerSessionBytes); + const maxCallerSessions = Math.min(limits.maxCallerSessions, retainedCallerCapacity); return Object.freeze({ appId: options.appId, appRevision: options.appRevision, @@ -403,28 +440,34 @@ function copyOptions(options: SessionBrokerAuthenticatorOptions): AuthenticatorS ), maxChallenges: configuredNumber( options.maxChallenges, - DEFAULT_MAX_CHALLENGES, + limits.maxIncompleteHandshakes, "maxChallenges", true, + limits.maxIncompleteHandshakes, ), maxChallengeBytes: configuredNumber( options.maxChallengeBytes, - DEFAULT_MAX_CHALLENGE_BYTES, + limits.maxIncompleteHandshakeBytes, "maxChallengeBytes", true, + limits.maxIncompleteHandshakeBytes, ), maxChallengeTranscriptBytes: configuredNumber( options.maxChallengeTranscriptBytes, - DEFAULT_MAX_CHALLENGE_TRANSCRIPT_BYTES, + limits.maxHandshakeProposalBytes, "maxChallengeTranscriptBytes", true, + limits.maxHandshakeProposalBytes, ), maxCallerSessions: configuredNumber( options.maxCallerSessions, - DEFAULT_MAX_CALLER_SESSIONS, + maxCallerSessions, "maxCallerSessions", true, + maxCallerSessions, ), + maxCallerSessionBytes: limits.maxCallerSessionBytes, + maxCallerSessionsBytes: limits.maxCallerSessionsBytes, }); } @@ -506,13 +549,28 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { private readonly challenges = new Map(); private readonly callerSessions = new Map(); private readonly reservedCallerSessionIds = new Set(); - private challengeBytes = 0; - private pendingCallerSessionAdmissions = 0; + private readonly challengeCountBudget: ResourceBudget; + private readonly challengeBudget: ResourceBudget; + private readonly callerSessionCountBudget: ResourceBudget; + private readonly callerSessionByteBudget: ResourceBudget; constructor(options: SessionBrokerAuthenticatorOptions) { this.config = copyOptions(options); this.crypto = copyCrypto(options.crypto); this.credentials = copyCredentials(options.credentials, this.config.appId); + this.challengeCountBudget = new ResourceBudget( + this.config.maxChallenges, + "maxIncompleteHandshakes", + ); + this.challengeBudget = new ResourceBudget(this.config.maxChallengeBytes, "challengeBytes"); + this.callerSessionCountBudget = new ResourceBudget( + this.config.maxCallerSessions, + "maxCallerSessions", + ); + this.callerSessionByteBudget = new ResourceBudget( + this.config.maxCallerSessionsBytes, + "maxCallerSessionsBytes", + ); } /** Issue one bounded, expiring challenge signed by the daemon identity. */ @@ -521,9 +579,6 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { listenerEndpoint: string, ): Promise { this.pruneExpired(); - if (this.challenges.size >= this.config.maxChallenges) { - authenticationError("authentication-capacity"); - } const normalized = this.validateHello(request, listenerEndpoint); const credential = this.credentials.get( `${normalized.role}:${normalized.keyId}:${normalized.grantId}`, @@ -542,20 +597,28 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { generation: this.config.generation, responderNonce, }); - if ( - transcript.byteLength > this.config.maxChallengeTranscriptBytes || - this.challengeBytes + transcript.byteLength > this.config.maxChallengeBytes - ) { + const retainedBytes = retainedRecordBytes( + [challengeId, normalized, credential.grant, transcript, expiresAt], + CHALLENGE_RECORD_OVERHEAD_BYTES, + ); + if (retainedBytes > this.config.maxChallengeTranscriptBytes) { + authenticationError("authentication-capacity"); + } + const reservation = new ReservationGroup(); + try { + reservation.add(this.challengeCountBudget.reserve()); + reservation.add(this.challengeBudget.reserve(retainedBytes)); + } catch { + reservation.release(); authenticationError("authentication-capacity"); } - this.challengeBytes += transcript.byteLength; this.challenges.set(challengeId, { request: normalized, transcript, grant: credential.grant, publicKey: credential.publicKey, expiresAt, - retainedBytes: transcript.byteLength, + reservation, }); try { const daemonSignature = encodeBase64Url( @@ -578,22 +641,14 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { async completeCallerHello(proofInput: unknown): Promise { const proof = this.parseHelloProof(proofInput); const pending = this.takeChallenge(proof.challengeId, "caller"); - await this.verifyProof(pending, proof.signature); - const grant = pending.grant as CallerGrant; - this.requireActiveGrant(grant); - this.pruneExpired(); - if ( - this.callerSessions.size + this.pendingCallerSessionAdmissions >= - this.config.maxCallerSessions - ) { - authenticationError("authentication-capacity"); - } - this.pendingCallerSessionAdmissions += 1; - try { + await this.verifyProof(pending, proof.signature); + const grant = pending.grant as CallerGrant; + this.requireActiveGrant(grant); + this.pruneExpired(); return await this.createCallerSession(pending, grant); } finally { - this.pendingCallerSessionAdmissions -= 1; + pending.reservation.release(); } } @@ -606,12 +661,27 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { (id) => this.callerSessions.has(id) || this.reservedCallerSessionIds.has(id), ); this.reservedCallerSessionIds.add(callerSessionId); + const reservations = new ReservationGroup(); + let committed = false; try { const expiresAt = Math.min( grant.expiresAt, this.currentTime() + this.config.callerSessionTtlMs, ); const principal = principalFromGrant(grant); + const retainedBytes = retainedRecordBytes( + [callerSessionId, principal, grant, transcriptHash, expiresAt], + CALLER_SESSION_RECORD_OVERHEAD_BYTES, + ); + if (retainedBytes > this.config.maxCallerSessionBytes) { + authenticationError("authentication-capacity"); + } + try { + reservations.add(this.callerSessionCountBudget.reserve()); + reservations.add(this.callerSessionByteBudget.reserve(retainedBytes)); + } catch { + authenticationError("authentication-capacity"); + } const daemonSignature = encodeBase64Url( await this.crypto.sign( this.config.daemonIdentity.privateKey, @@ -635,7 +705,9 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { helloTranscriptHash: transcriptHash, expiresAt, replay: new CallerSequenceReplayWindow(), + reservation: reservations, }); + committed = true; return Object.freeze({ callerSessionId, principal, @@ -650,6 +722,7 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { }); } finally { this.reservedCallerSessionIds.delete(callerSessionId); + if (!committed) reservations.release(); } } @@ -661,35 +734,39 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { const proof = this.parseHelloProof(proofInput); if (!isValidBrokerIdentifier(connectionId)) authenticationError("invalid-credential"); const pending = this.takeChallenge(proof.challengeId, "producer"); - await this.verifyProof(pending, proof.signature); - const grant = pending.grant as ProducerGrant; - this.requireActiveGrant(grant); - const helloTranscriptHash = encodeBase64Url(await this.crypto.sha256(pending.transcript)); - const daemonSignature = encodeBase64Url( - await this.crypto.sign( - this.config.daemonIdentity.privateKey, - buildBrokerHelloAckTranscript({ - role: "producer", - appId: this.config.appId, - generation: this.config.generation, - keyId: grant.keyId, - grantId: grant.grantId, - helloTranscriptHash, - selection: pending.request.proposal, - connectionId, - }), - ), - ); - return Object.freeze({ - principal: principalFromGrant(grant), - connectionId, - brokerRevision: SESSION_BROKER_PROTOCOL_REVISION, - appRevision: this.config.appRevision, - features: Object.freeze([]) as readonly [], - helloTranscriptHash, - daemonKeyId: this.config.daemonIdentity.keyId, - daemonSignature, - }); + try { + await this.verifyProof(pending, proof.signature); + const grant = pending.grant as ProducerGrant; + this.requireActiveGrant(grant); + const helloTranscriptHash = encodeBase64Url(await this.crypto.sha256(pending.transcript)); + const daemonSignature = encodeBase64Url( + await this.crypto.sign( + this.config.daemonIdentity.privateKey, + buildBrokerHelloAckTranscript({ + role: "producer", + appId: this.config.appId, + generation: this.config.generation, + keyId: grant.keyId, + grantId: grant.grantId, + helloTranscriptHash, + selection: pending.request.proposal, + connectionId, + }), + ), + ); + return Object.freeze({ + principal: principalFromGrant(grant), + connectionId, + brokerRevision: SESSION_BROKER_PROTOCOL_REVISION, + appRevision: this.config.appRevision, + features: Object.freeze([]) as readonly [], + helloTranscriptHash, + daemonKeyId: this.config.daemonIdentity.keyId, + daemonSignature, + }); + } finally { + pending.reservation.release(); + } } /** Verify one signed HTTP request and atomically admit its sequence before returning authority. */ @@ -710,9 +787,14 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { } const session = this.callerSessions.get(callerSessionId); if (!session) authenticationError("caller-session-expired"); - this.requireActiveGrant(session.grant); + try { + this.requireActiveGrant(session.grant); + } catch (error) { + this.deleteCallerSession(callerSessionId); + throw error; + } if (this.currentTime() >= session.expiresAt) { - this.callerSessions.delete(callerSessionId); + this.deleteCallerSession(callerSessionId); authenticationError("caller-session-expired"); } @@ -766,7 +848,14 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { /** Revoke one in-memory caller session without exposing whether it previously existed. */ revokeCallerSession(callerSessionId: string): void { - this.callerSessions.delete(callerSessionId); + this.deleteCallerSession(callerSessionId); + } + + /** Release every retained authentication record during shutdown or credential reload. */ + clear(): void { + for (const id of this.challenges.keys()) this.deleteChallenge(id); + for (const id of this.callerSessions.keys()) this.deleteCallerSession(id); + this.reservedCallerSessionIds.clear(); } /** Recheck identity, revocation, and expiry after every asynchronous policy boundary. */ @@ -774,9 +863,14 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { if (this.callerSessions.get(callerSessionId) !== session) { authenticationError("caller-session-expired"); } - this.requireActiveGrant(session.grant); + try { + this.requireActiveGrant(session.grant); + } catch (error) { + this.deleteCallerSession(callerSessionId); + throw error; + } if (this.currentTime() >= session.expiresAt) { - this.callerSessions.delete(callerSessionId); + this.deleteCallerSession(callerSessionId); authenticationError("caller-session-expired"); } } @@ -888,11 +982,17 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { if (!isValidBrokerIdentifier(challengeId)) authenticationError("challenge-used"); const pending = this.challenges.get(challengeId); if (!pending) authenticationError("challenge-used"); - // Delete before any asynchronous verification so concurrent proofs cannot both consume it. - this.deleteChallenge(challengeId); - if (this.currentTime() >= pending.expiresAt) authenticationError("challenge-expired"); - if (pending.grant.kind !== role) authenticationError("invalid-credential"); - return pending; + // Remove lookup authority before asynchronous verification so a second proof cannot consume it. + // The count/byte reservation remains live until the consuming proof completes. + this.challenges.delete(challengeId); + try { + if (this.currentTime() >= pending.expiresAt) authenticationError("challenge-expired"); + if (pending.grant.kind !== role) authenticationError("invalid-credential"); + return pending; + } catch (error) { + pending.reservation.release(); + throw error; + } } private async verifyProof(pending: PendingChallenge, encodedSignature: string): Promise { @@ -934,7 +1034,14 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { const challenge = this.challenges.get(challengeId); if (!challenge) return; this.challenges.delete(challengeId); - this.challengeBytes = Math.max(0, this.challengeBytes - challenge.retainedBytes); + challenge.reservation.release(); + } + + private deleteCallerSession(callerSessionId: string): void { + const session = this.callerSessions.get(callerSessionId); + if (!session) return; + this.callerSessions.delete(callerSessionId); + session.reservation.release(); } private pruneExpired(): void { @@ -943,7 +1050,7 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { if (now >= challenge.expiresAt) this.deleteChallenge(id); } for (const [id, session] of this.callerSessions) { - if (now >= session.expiresAt) this.callerSessions.delete(id); + if (now >= session.expiresAt) this.deleteCallerSession(id); } } diff --git a/packages/session-broker/src/broker.ts b/packages/session-broker/src/broker.ts index c85798407..44ff0752a 100644 --- a/packages/session-broker/src/broker.ts +++ b/packages/session-broker/src/broker.ts @@ -4,6 +4,8 @@ import { type MarkSessionSeenResult, type RegisterSessionResult, type SessionBrokerEntry, + type SessionBrokerLimitOptions, + type SessionBrokerLimits, type SessionRegistration, type SessionServerMessage, type SessionSnapshot, @@ -38,6 +40,8 @@ export interface SessionBrokerOptions< CommandResult = unknown, > { protocolParsers: SessionBrokerProtocolParsers; + limits?: SessionBrokerLimitOptions["limits"]; + unsafeLimits?: SessionBrokerLimitOptions["unsafeLimits"]; describeSession?: ( registration: SessionRegistration, snapshot: SessionSnapshot, @@ -57,6 +61,7 @@ export interface SessionBrokerController< ServerMessage, CommandResult >; + readonly limits?: Readonly; listSessions(): SessionView[]; getSession(selector: SessionTargetSelector): SessionView; getSessionCount(): number; @@ -142,30 +147,40 @@ export class SessionBroker< options.describeSession ?? ((registration, _snapshot) => defaultSessionTitle(registration)); this.protocolParsers = options.protocolParsers; - this.state = new SessionBrokerState({ - parseRegistration: (value) => { - try { - return this.protocolParsers.parseRegistration(value); - } catch { - return null; - } + this.state = new SessionBrokerState( + { + parseRegistration: (value) => { + try { + return this.protocolParsers.parseRegistration(value); + } catch { + return null; + } + }, + parseSnapshot: (value) => { + try { + return this.protocolParsers.parseSnapshot(value); + } catch { + return null; + } + }, + parseCommandInput: (command, version, value) => + this.protocolParsers.parseCommandInput(command, version, value), + parseCommandResult: (command, version, value) => + this.protocolParsers.parseCommandResult(command, version, value), + buildListedSession: (entry) => this.buildRecord(entry), + buildSelectedContext: (session) => session, + buildSessionReview: (entry) => this.buildRecord(entry), + listComments: () => [], }, - parseSnapshot: (value) => { - try { - return this.protocolParsers.parseSnapshot(value); - } catch { - return null; - } + { + ...(options.limits ? { limits: options.limits } : {}), + ...(options.unsafeLimits ? { unsafeLimits: options.unsafeLimits } : {}), }, - parseCommandInput: (command, version, value) => - this.protocolParsers.parseCommandInput(command, version, value), - parseCommandResult: (command, version, value) => - this.protocolParsers.parseCommandResult(command, version, value), - buildListedSession: (entry) => this.buildRecord(entry), - buildSelectedContext: (session) => session, - buildSessionReview: (entry) => this.buildRecord(entry), - listComments: () => [], - }); + ); + } + + get limits() { + return this.state.limits; } listSessions() { diff --git a/packages/session-broker/src/connection.test.ts b/packages/session-broker/src/connection.test.ts index 69193e198..fe3dc41fb 100644 --- a/packages/session-broker/src/connection.test.ts +++ b/packages/session-broker/src/connection.test.ts @@ -492,6 +492,186 @@ describe("session broker connection", () => { connection.stop(); }); + test("keeps 32 missing-bridge commands FIFO and explicitly rejects the 33rd", async () => { + const socket = new TestSocket(); + const dispatched: string[] = []; + const connection = createSessionBrokerConnection< + TestSessionInfo, + TestSessionState, + TestSocket, + TestServerMessage, + { ok: true } + >({ + url: "ws://broker.test/session", + createSocket: () => socket, + registration: createRegistration(), + snapshot: createSnapshot(), + protocolParsers, + }); + connection.start(); + socket.emitOpen(); + for (let index = 1; index <= 33; index += 1) { + socket.emitMessage( + JSON.stringify({ + type: "command", + requestId: `request-${index}`, + command: "annotate", + input: { summary: `note-${index}` }, + }), + ); + } + const overflow = JSON.parse(socket.sent.at(-1)!) as { requestId: string; error: string }; + expect(overflow).toMatchObject({ requestId: "request-33", error: "queue-full" }); + + connection.setBridge({ + dispatchCommand: async (message) => { + dispatched.push(message.requestId); + return { ok: true }; + }, + }); + await Bun.sleep(10); + expect(dispatched).toEqual(Array.from({ length: 32 }, (_, index) => `request-${index + 1}`)); + connection.stop(); + }); + + test("keeps a hung bridge plus queued commands within the same 32-command budget", async () => { + const socket = new TestSocket(); + const never = new Promise<{ ok: true }>(() => {}); + const connection = createSessionBrokerConnection< + TestSessionInfo, + TestSessionState, + TestSocket, + TestServerMessage, + { ok: true } + >({ + url: "ws://broker.test/session", + createSocket: () => socket, + registration: createRegistration(), + snapshot: createSnapshot(), + protocolParsers, + bridge: { dispatchCommand: () => never }, + }); + connection.start(); + socket.emitOpen(); + for (let index = 1; index <= 33; index += 1) { + socket.emitMessage( + JSON.stringify({ + type: "command", + requestId: `request-${index}`, + command: "annotate", + input: { summary: `note-${index}` }, + }), + ); + } + await Bun.sleep(0); + expect(JSON.parse(socket.sent.at(-1)!)).toMatchObject({ + requestId: "request-33", + error: "queue-full", + }); + connection.stop(); + }); + + test("retains a hung command reservation across disconnect and reconnect", async () => { + const sockets: TestSocket[] = []; + const never = new Promise<{ ok: true }>(() => {}); + const connection = createSessionBrokerConnection< + TestSessionInfo, + TestSessionState, + TestSocket, + TestServerMessage, + { ok: true } + >({ + url: "ws://broker.test/session", + createSocket: () => { + const socket = new TestSocket(); + sockets.push(socket); + return socket; + }, + registration: createRegistration(), + snapshot: createSnapshot(), + protocolParsers, + bridge: { dispatchCommand: () => never }, + reconnectDelayMs: 1, + limits: { maxPreBridgeCommands: 1 }, + }); + + connection.start(); + sockets[0]!.emitOpen(); + sockets[0]!.emitMessage( + JSON.stringify({ + type: "command", + requestId: "request-hung", + command: "annotate", + input: { summary: "hung" }, + }), + ); + await Bun.sleep(0); + sockets[0]!.emitClose(); + await Bun.sleep(5); + sockets[1]!.emitOpen(); + sockets[1]!.emitMessage( + JSON.stringify({ + type: "command", + requestId: "request-new", + command: "annotate", + input: { summary: "new" }, + }), + ); + + expect(JSON.parse(sockets[1]!.sent.at(-1)!)).toMatchObject({ + requestId: "request-new", + error: "queue-full", + }); + connection.stop(); + }); + + test("serializes bridge execution while preserving arrival order", async () => { + const socket = new TestSocket(); + const started: string[] = []; + let releaseFirst!: () => void; + const firstGate = new Promise((resolve) => { + releaseFirst = resolve; + }); + const connection = createSessionBrokerConnection< + TestSessionInfo, + TestSessionState, + TestSocket, + TestServerMessage, + { ok: true } + >({ + url: "ws://broker.test/session", + createSocket: () => socket, + registration: createRegistration(), + snapshot: createSnapshot(), + protocolParsers, + bridge: { + dispatchCommand: async (message) => { + started.push(message.requestId); + if (message.requestId === "request-1") await firstGate; + return { ok: true }; + }, + }, + }); + connection.start(); + socket.emitOpen(); + for (const requestId of ["request-1", "request-2"]) { + socket.emitMessage( + JSON.stringify({ + type: "command", + requestId, + command: "annotate", + input: { summary: requestId }, + }), + ); + } + await Bun.sleep(0); + expect(started).toEqual(["request-1"]); + releaseFirst(); + await Bun.sleep(0); + expect(started).toEqual(["request-1", "request-2"]); + connection.stop(); + }); + test("reconnects after socket close unless a close directive disables it", async () => { const sockets: TestSocket[] = []; const warnings: string[] = []; diff --git a/packages/session-broker/src/connection.ts b/packages/session-broker/src/connection.ts index e0772b4ab..dbc72ac89 100644 --- a/packages/session-broker/src/connection.ts +++ b/packages/session-broker/src/connection.ts @@ -1,5 +1,13 @@ import { + BrokerCapacityError, BrokerProtocolError, + ReservationGroup, + ResourceBudget, + resolveSessionBrokerLimits, + utf8ByteLength, + type BudgetReservation, + type SessionBrokerLimitOptions, + type SessionBrokerLimits, type SessionClientMessage, type SessionRegistration, type SessionServerMessage, @@ -16,6 +24,20 @@ import type { const DEFAULT_RECONNECT_DELAY_MS = 3_000; const DEFAULT_HEARTBEAT_INTERVAL_MS = 10_000; const DEFAULT_SOCKET_OPEN_STATE = 1; +const PRODUCER_COMMAND_OVERHEAD_BYTES = 128; + +interface QueuedProducerCommand { + socket: Socket; + message: Message; + reservation: BudgetReservation; +} + +/** Measure one JSON-safe command value without relying on UTF-16 string length. */ +function commandValueBytes(value: unknown): number { + const serialized = JSON.stringify(value); + if (serialized === undefined) throw new BrokerProtocolError("invalid-app-payload"); + return utf8ByteLength(serialized); +} export interface SessionBrokerConnectionBridge< ServerMessage extends SessionServerMessage = SessionServerMessage, @@ -42,6 +64,8 @@ export interface SessionBrokerConnectionOptions< openState?: number; resolveClose?: (event: SessionBrokerSocketCloseEvent) => SessionBrokerConnectionCloseDirective; onWarning?: (message: string) => void; + limits?: SessionBrokerLimitOptions["limits"]; + unsafeLimits?: SessionBrokerLimitOptions["unsafeLimits"]; } /** @@ -57,7 +81,13 @@ export class SessionBrokerConnection< > { private socket: Socket | null = null; private bridge: SessionBrokerConnectionBridge | null; - private queuedMessages: Array<{ socket: Socket; message: ServerMessage }> = []; + readonly limits: Readonly; + + private queuedMessages: Array> = []; + private executingMessages = new Set>(); + private readonly queuedCommandCountBudget: ResourceBudget; + private readonly queuedCommandByteBudget: ResourceBudget; + private draining = false; private reconnectTimer: ReturnType | null = null; private heartbeatTimer: ReturnType | null = null; private stopped = false; @@ -73,9 +103,23 @@ export class SessionBrokerConnection< Result >, ) { + this.limits = resolveSessionBrokerLimits({ + ...(options.limits ? { limits: options.limits } : {}), + ...(options.unsafeLimits ? { unsafeLimits: options.unsafeLimits } : {}), + }); this.bridge = options.bridge ?? null; this.registration = options.registration; this.snapshot = options.snapshot; + this.queuedCommandCountBudget = new ResourceBudget( + this.limits.maxPreBridgeCommands, + "maxPreBridgeCommands", + "queue-full", + ); + this.queuedCommandByteBudget = new ResourceBudget( + this.limits.maxQueuedCommandBytes, + "maxQueuedCommandBytes", + "queue-full", + ); } start() { @@ -88,6 +132,9 @@ export class SessionBrokerConnection< stop() { this.stopped = true; + for (const entry of this.queuedMessages.splice(0)) entry.reservation.release(); + for (const entry of this.executingMessages) entry.reservation.release(); + this.executingMessages.clear(); if (this.reconnectTimer) { clearTimeout(this.reconnectTimer); this.reconnectTimer = null; @@ -151,9 +198,14 @@ export class SessionBrokerConnection< socket.onmessage = (event) => { let parsed: ServerMessage; try { - parsed = this.options.protocolParsers.parseServerMessage( - parseSessionBrokerJsonText(event.data), - ); + const raw = parseSessionBrokerJsonText(event.data) as { input?: unknown }; + if (commandValueBytes(raw?.input) > this.limits.maxCommandInputBytes) { + throw new BrokerCapacityError("capacity-exceeded", "maxCommandInputBytes"); + } + parsed = this.options.protocolParsers.parseServerMessage(raw); + if (commandValueBytes(parsed.input) > this.limits.maxCommandInputBytes) { + throw new BrokerCapacityError("capacity-exceeded", "maxCommandInputBytes"); + } } catch { // Never invoke the app bridge after a malformed or mismatched command contract. Closing // prevents this producer from retaining daemon assumptions that were not actually parsed. @@ -170,7 +222,13 @@ export class SessionBrokerConnection< this.stopHeartbeat(); } - this.queuedMessages = this.queuedMessages.filter((queued) => queued.socket !== socket); + this.queuedMessages = this.queuedMessages.filter((queued) => { + if (queued.socket !== socket) return true; + queued.reservation.release(); + return false; + }); + // Executing bridge work retains its reservation until its promise settles. A disconnect + // prevents its response from migrating but does not make the retained input or work vanish. if (this.stopped) { return; } @@ -250,20 +308,90 @@ export class SessionBrokerConnection< } private async handleServerMessage(socket: Socket, message: ServerMessage) { - if (!this.bridge) { - // Sessions may connect before the host app has finished wiring its command bridge. Bind each - // queued command to its source so a reconnect cannot inherit work from a disconnected socket. - this.queuedMessages.push({ socket, message }); + const reservations = new ReservationGroup(); + try { + reservations.add(this.queuedCommandCountBudget.reserve()); + reservations.add( + this.queuedCommandByteBudget.reserve( + commandValueBytes(message) + PRODUCER_COMMAND_OVERHEAD_BYTES, + ), + ); + } catch { + reservations.release(); + try { + this.sendToSocket(socket, { + type: "command-result", + requestId: message.requestId, + ok: false, + error: "queue-full", + }); + } catch { + socket.close(1013, "Session broker queue pressure exceeded."); + } return; } + // Every admitted command, including one executing in a hung bridge, retains its reservation. + this.queuedMessages.push({ socket, message, reservation: reservations }); + if (this.bridge) await this.flushQueuedMessages(socket); + } + + private async flushQueuedMessages(socket = this.socket) { + if (!this.bridge || !socket || this.draining) return; + this.draining = true; try { - const result = await this.bridge.dispatchCommand(message); + // Take one command at a time so newly received work joins the same FIFO and a disconnect can + // prevent every command that has not started from executing on a replacement transport. + for (;;) { + if (!this.bridge) return; + const index = this.queuedMessages.findIndex((entry) => entry.socket === socket); + if (index < 0) return; + const [entry] = this.queuedMessages.splice(index, 1); + if (!entry) return; + if ( + this.socket !== entry.socket || + entry.socket.readyState !== (this.options.openState ?? DEFAULT_SOCKET_OPEN_STATE) + ) { + entry.reservation.release(); + return; + } + this.executingMessages.add(entry); + try { + await this.executeServerMessage(entry.socket, entry.message); + } finally { + this.executingMessages.delete(entry); + entry.reservation.release(); + } + } + } finally { + this.draining = false; + if ( + this.bridge && + this.socket && + this.queuedMessages.some((entry) => entry.socket === this.socket) + ) { + void this.flushQueuedMessages(this.socket); + } + } + } + + /** Execute one already-admitted command without re-entering the producer FIFO. */ + private async executeServerMessage(socket: Socket, message: ServerMessage) { + const bridge = this.bridge; + if (!bridge) return; + try { + const result = await bridge.dispatchCommand(message); + if (commandValueBytes(result) > this.limits.maxCommandResultBytes) { + throw new BrokerProtocolError("invalid-app-payload"); + } const parsedResult = this.options.protocolParsers.parseCommandResult( message.command, message.commandVersion ?? 1, result, ); + if (commandValueBytes(parsedResult) > this.limits.maxCommandResultBytes) { + throw new BrokerProtocolError("invalid-app-payload"); + } this.sendToSocket(socket, { type: "command-result", requestId: message.requestId, @@ -271,8 +399,6 @@ export class SessionBrokerConnection< result: parsedResult, }); } catch (error) { - // Parser failures invalidate the selected command contract and cannot be represented as an - // app command rejection. Close without reflecting callback details. if (error instanceof BrokerProtocolError) { socket.close(1008, "Malformed session broker command result."); return; @@ -285,28 +411,6 @@ export class SessionBrokerConnection< }); } } - - private async flushQueuedMessages(socket = this.socket) { - if (!this.bridge || !socket || this.queuedMessages.length === 0) { - return; - } - - // Snapshot only this transport's queue so commands cannot cross a disconnect. Commands received - // while replay runs stay in the queue for a later pass and preserve their original ordering. - const queued = this.queuedMessages.filter((entry) => entry.socket === socket); - this.queuedMessages = this.queuedMessages.filter((entry) => entry.socket !== socket); - - for (const entry of queued) { - if ( - this.socket !== entry.socket || - entry.socket.readyState !== (this.options.openState ?? DEFAULT_SOCKET_OPEN_STATE) - ) { - break; - } - - await this.handleServerMessage(entry.socket, entry.message); - } - } } /** Create one runtime-neutral session connection around a browser-like websocket factory. */ diff --git a/packages/session-broker/src/daemon.test.ts b/packages/session-broker/src/daemon.test.ts index 5877e3735..11cfb03e3 100644 --- a/packages/session-broker/src/daemon.test.ts +++ b/packages/session-broker/src/daemon.test.ts @@ -405,8 +405,9 @@ describe("session broker daemon", () => { ); expect(response?.status).toBe(413); - await expect(response?.json()).resolves.toMatchObject({ - error: expect.stringContaining("session broker limit"), + await expect(response?.json()).resolves.toEqual({ + error: "capacity-exceeded", + resource: "maxHttpBodyBytes", }); daemon.shutdown(); }); @@ -919,6 +920,51 @@ describe("session broker daemon", () => { daemon.shutdown(); }); + test("rejects daemon state limits that were not configured on its broker controller", () => { + expect(() => + createSessionBrokerDaemon({ + broker: createBroker(), + limits: { maxSessions: 0 }, + }), + ).toThrow("state limits must be configured on the broker controller"); + }); + + test("returns busy before admitting more than the configured concurrent controls", async () => { + let release!: () => void; + const gate = new Promise((resolve) => { + release = resolve; + }); + const daemon = createSessionBrokerDaemon({ + broker: createBroker(), + exposeHttpApi: true, + ...authenticatedHttpApi, + limits: { maxConcurrentHttpControls: 1 }, + callerAuthenticator: { + authenticate: async () => { + await gate; + return authenticatedHttpApi.callerAuthenticator.authenticate(); + }, + }, + }); + const request = () => + new Request("http://broker.test/broker", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ action: "list" }), + }); + const first = daemon.handleRequest(request()); + await Bun.sleep(0); + const overflow = await daemon.handleRequest(request()); + expect(overflow?.status).toBe(503); + await expect(overflow?.json()).resolves.toEqual({ + error: "busy", + resource: "maxConcurrentHttpControls", + }); + release(); + expect((await first)?.status).toBe(200); + daemon.shutdown(); + }); + test("requests shutdown after the idle timeout when no sessions remain", async () => { const daemon = createSessionBrokerDaemon({ broker: createBroker(), diff --git a/packages/session-broker/src/daemon.ts b/packages/session-broker/src/daemon.ts index 5ad60be22..a555b3a6e 100644 --- a/packages/session-broker/src/daemon.ts +++ b/packages/session-broker/src/daemon.ts @@ -1,16 +1,24 @@ import { + BrokerCapacityError, BrokerProtocolError, - MAX_HTTP_BODY_BYTES, + InvalidContentLengthError, PayloadTooLargeError, + ResourceBudget, + readRequestBytesWithReservation, + mergeSessionBrokerLimits, + DEFAULT_SESSION_BROKER_LIMITS, callerPrincipalAllows, canonicalizeJson, isValidBrokerAppId, isValidBrokerIdentifier, isValidBrokerRevision, - readRequestBytesWithLimit, + utf8ByteLength, + type BudgetReservation, type CallerOperation, type CallerPrincipal, type CanonicalJsonValue, + type SessionBrokerLimitOptions, + type SessionBrokerLimits, type SessionServerMessage, type SessionTargetSelector, } from "@hunk/session-broker-core"; @@ -44,6 +52,18 @@ const DEFAULT_STALE_SESSION_TTL_MS = 45_000; const DEFAULT_STALE_SESSION_SWEEP_INTERVAL_MS = 15_000; const DEFAULT_IDLE_TIMEOUT_MS = 60_000; const INCOMPATIBLE_PAYLOAD_CLOSE_CODE = 1008; +const BROKER_STATE_LIMITS = [ + "maxSessions", + "maxCommandsPerSession", + "maxCommandsTotal", + "maxCommandInputBytes", + "maxCommandResultBytes", + "maxQueuedCommandBytes", + "maxRetainedSessionBytes", + "maxRetainedBytes", + "defaultCommandTimeoutMs", + "maxCommandTimeoutMs", +] as const satisfies readonly (keyof SessionBrokerLimits)[]; export interface SessionBrokerDaemonOptions< SessionView = unknown, @@ -62,12 +82,19 @@ export interface SessionBrokerDaemonOptions< idleTimeoutMs?: number; staleSessionTtlMs?: number; staleSessionSweepIntervalMs?: number; + limits?: SessionBrokerLimitOptions["limits"]; + unsafeLimits?: SessionBrokerLimitOptions["unsafeLimits"]; } function jsonError(message: string, status = 400) { return Response.json({ error: message }, { status }); } +/** Map resource admission failures to stable retryable HTTP errors. */ +function capacityResponse(error: BrokerCapacityError) { + return Response.json({ error: error.code, resource: error.resource }, { status: 503 }); +} + /** Build one redacted protocol failure body without reflecting parser details. */ function protocolError(error: unknown) { return { @@ -99,6 +126,8 @@ export class SessionBrokerDaemon< readonly paths: SessionBrokerHttpPaths; readonly stopped: Promise; + readonly limits: Readonly; + private readonly startedAt = Date.now(); private readonly capabilities: SessionBrokerCapabilities; private readonly protocolParsers: SessionBrokerProtocolParsers< @@ -115,6 +144,8 @@ export class SessionBrokerDaemon< private readonly callerAuthenticator?: CallerRequestAuthenticator; private readonly authorizer?: SessionBrokerAuthorizer; private readonly audit?: SessionBrokerAuditHook; + private readonly httpControlBudget: ResourceBudget; + private readonly httpBodyBudget: ResourceBudget; private lastActivityAt = this.startedAt; private sweepTimer: ReturnType | null = null; private idleTimer: ReturnType | null = null; @@ -125,6 +156,25 @@ export class SessionBrokerDaemon< private readonly broker: SessionBrokerController, options: Omit, "broker">, ) { + const brokerLimits = broker.limits ?? DEFAULT_SESSION_BROKER_LIMITS; + this.limits = mergeSessionBrokerLimits(brokerLimits, { + ...(options.limits ? { limits: options.limits } : {}), + ...(options.unsafeLimits ? { unsafeLimits: options.unsafeLimits } : {}), + }); + if (BROKER_STATE_LIMITS.some((name) => this.limits[name] !== brokerLimits[name])) { + throw new TypeError( + "Session broker state limits must be configured on the broker controller before daemon composition.", + ); + } + this.httpControlBudget = new ResourceBudget( + this.limits.maxConcurrentHttpControls, + "maxConcurrentHttpControls", + "busy", + ); + this.httpBodyBudget = new ResourceBudget( + this.limits.maxInFlightHttpBodyBytes, + "maxInFlightHttpBodyBytes", + ); const exposeAuthenticatedHttpApi = (options.exposeHttpApi ?? false) && isValidBrokerAppId(options.appId) && @@ -190,6 +240,46 @@ export class SessionBrokerDaemon< return pathname === this.paths.socket; } + /** Run one app-specific finite HTTP control through the daemon's shared count/body budgets. */ + async handleBoundedControl( + request: Request, + handler: (body: Uint8Array) => Response | Promise, + responses: { payloadTooLarge?: (error: PayloadTooLargeError) => Response } = {}, + ): Promise { + let control: BudgetReservation; + try { + control = this.httpControlBudget.reserve(); + } catch (error) { + return capacityResponse(error as BrokerCapacityError); + } + let bodyReservation: BudgetReservation | null = null; + try { + try { + const read = await readRequestBytesWithReservation( + request, + this.limits.maxHttpBodyBytes, + this.httpBodyBudget, + ); + bodyReservation = read.reservation; + return await handler(read.bytes); + } catch (error) { + if (error instanceof BrokerCapacityError) return capacityResponse(error); + if (error instanceof PayloadTooLargeError) { + if (responses.payloadTooLarge) return responses.payloadTooLarge(error); + return Response.json( + { error: "capacity-exceeded", resource: "maxHttpBodyBytes" }, + { status: 413 }, + ); + } + if (error instanceof InvalidContentLengthError) return jsonError("Invalid Content-Length."); + return jsonError("Could not read broker control request body."); + } + } finally { + bodyReservation?.release(); + control.release(); + } + } + async handleRequest(request: Request) { const url = new URL(request.url); @@ -209,33 +299,7 @@ export class SessionBrokerDaemon< } if (this.paths.capabilities && url.pathname === this.paths.capabilities) { - if (request.method !== "GET") { - return jsonError("Broker capabilities requests must use GET.", 405); - } - let body: Uint8Array; - try { - body = await readRequestBytesWithLimit(request, MAX_HTTP_BODY_BYTES); - } catch (error) { - return error instanceof PayloadTooLargeError - ? jsonError(error.message, 413) - : jsonError("Could not read broker capabilities request body."); - } - if (body.byteLength !== 0) { - return jsonError("Broker capabilities requests must not include a body."); - } - const authenticated = await this.authenticateRequest(request, body, "diagnostics"); - if (authenticated instanceof Response) return authenticated; - if ( - !(await this.authorize(request, authenticated, { - operation: "diagnostics", - })) - ) { - return this.authenticatedResponse(authenticated, { error: "authorization-denied" }, 403); - } - const inactive = this.rejectInactiveRequest(authenticated); - if (inactive) return inactive; - this.noteActivity(); - return this.authenticatedResponse(authenticated, this.capabilities, 200); + return this.handleCapabilitiesRequest(request); } if (this.paths.api && url.pathname === this.paths.api) { @@ -273,6 +337,10 @@ export class SessionBrokerDaemon< connection.close?.(INCOMPATIBLE_PAYLOAD_CLOSE_CODE, "Session registration rejected."); return; } + if (registrationResult === "capacity-exceeded") { + connection.close?.(1013, "Session broker capacity exceeded."); + return; + } this.noteActivity(); break; @@ -294,6 +362,10 @@ export class SessionBrokerDaemon< connection.close?.(INCOMPATIBLE_PAYLOAD_CLOSE_CODE, "Incompatible session snapshot."); return; } + if (updateResult === "capacity-exceeded") { + connection.close?.(1013, "Session broker capacity exceeded."); + return; + } this.noteActivity(); break; @@ -350,6 +422,7 @@ export class SessionBrokerDaemon< } this.broker.shutdown(error); + this.callerAuthenticator?.clear?.(); this.resolveStopped?.(); this.resolveStopped = null; } @@ -514,21 +587,44 @@ export class SessionBrokerDaemon< targetSpecific = false, ): Promise { // Normalize with JSON semantics first so optional undefined fields cannot create digest aliases. - const structuredBody = JSON.parse(JSON.stringify(body)) as CanonicalJsonValue; - canonicalizeJson(structuredBody); + let structuredBody = JSON.parse(JSON.stringify(body)) as CanonicalJsonValue; + let responseStatus = status; + const targetContract = + targetSpecific && this.appRevision !== undefined + ? { appContract: { appRevision: this.appRevision, features: [] as const } } + : {}; + if (utf8ByteLength(canonicalizeJson(structuredBody)) > this.limits.maxHttpResponseBytes) { + structuredBody = { error: "capacity-exceeded", resource: "maxHttpResponseBytes" }; + responseStatus = 503; + } const authentication = await authenticated.signResponse({ - httpStatus: status, + httpStatus: responseStatus, body: structuredBody, - ...(targetSpecific && this.appRevision !== undefined - ? { appContract: { appRevision: this.appRevision, features: [] } } - : {}), + ...targetContract, }); - const envelope: SessionBrokerAuthenticatedResponse = { + let envelope: SessionBrokerAuthenticatedResponse = { body: structuredBody, authentication, }; - return new Response(canonicalizeJson(envelope as unknown as CanonicalJsonValue), { - status, + let serializedEnvelope = canonicalizeJson(envelope as unknown as CanonicalJsonValue); + if (utf8ByteLength(serializedEnvelope) > this.limits.maxHttpResponseBytes) { + structuredBody = { error: "capacity-exceeded", resource: "maxHttpResponseBytes" }; + responseStatus = 503; + envelope = { + body: structuredBody, + authentication: await authenticated.signResponse({ + httpStatus: responseStatus, + body: structuredBody, + ...targetContract, + }), + }; + serializedEnvelope = canonicalizeJson(envelope as unknown as CanonicalJsonValue); + } + if (utf8ByteLength(serializedEnvelope) > this.limits.maxHttpResponseBytes) { + return new Response(null, { status: 503 }); + } + return new Response(serializedEnvelope, { + status: responseStatus, headers: { "content-type": "application/json" }, }); } @@ -541,94 +637,171 @@ export class SessionBrokerDaemon< } } - private async handleApiRequest(request: Request) { - if (request.method !== "POST") { - return jsonError("Broker API requests must use POST.", 405); + /** Run one authenticated capabilities control under the shared HTTP budgets. */ + private async handleCapabilitiesRequest(request: Request): Promise { + let control: BudgetReservation; + try { + control = this.httpControlBudget.reserve(); + } catch (error) { + return capacityResponse(error as BrokerCapacityError); } - if (!hasJsonContentType(request)) { - return jsonError("Expected Content-Type application/json.", 415); + let bodyReservation: BudgetReservation | null = null; + try { + if (request.method !== "GET") { + return jsonError("Broker capabilities requests must use GET.", 405); + } + let body: Uint8Array; + try { + const read = await readRequestBytesWithReservation( + request, + this.limits.maxHttpBodyBytes, + this.httpBodyBudget, + ); + body = read.bytes; + bodyReservation = read.reservation; + } catch (error) { + if (error instanceof BrokerCapacityError) return capacityResponse(error); + return error instanceof PayloadTooLargeError + ? Response.json( + { error: "capacity-exceeded", resource: "maxHttpBodyBytes" }, + { status: 413 }, + ) + : error instanceof InvalidContentLengthError + ? jsonError("Invalid Content-Length.") + : jsonError("Could not read broker capabilities request body."); + } + if (body.byteLength !== 0) { + return jsonError("Broker capabilities requests must not include a body."); + } + const authenticated = await this.authenticateRequest(request, body, "diagnostics"); + if (authenticated instanceof Response) return authenticated; + if (!(await this.authorize(request, authenticated, { operation: "diagnostics" }))) { + return this.authenticatedResponse(authenticated, { error: "authorization-denied" }, 403); + } + const inactive = this.rejectInactiveRequest(authenticated); + if (inactive) return inactive; + this.noteActivity(); + return this.authenticatedResponse(authenticated, this.capabilities, 200); + } finally { + bodyReservation?.release(); + control.release(); } + } - let body: Uint8Array; + private async handleApiRequest(request: Request) { + let control: BudgetReservation; try { - body = await readRequestBytesWithLimit(request, MAX_HTTP_BODY_BYTES); + control = this.httpControlBudget.reserve(); } catch (error) { - return error instanceof PayloadTooLargeError - ? jsonError(error.message, 413) - : jsonError("Could not read broker API request body."); + return capacityResponse(error as BrokerCapacityError); } + let bodyReservation: BudgetReservation | null = null; + try { + if (request.method !== "POST") { + return jsonError("Broker API requests must use POST.", 405); + } + if (!hasJsonContentType(request)) { + return jsonError("Expected Content-Type application/json.", 415); + } + + let body: Uint8Array; + try { + const read = await readRequestBytesWithReservation( + request, + this.limits.maxHttpBodyBytes, + this.httpBodyBudget, + ); + body = read.bytes; + bodyReservation = read.reservation; + } catch (error) { + if (error instanceof BrokerCapacityError) return capacityResponse(error); + return error instanceof PayloadTooLargeError + ? Response.json( + { error: "capacity-exceeded", resource: "maxHttpBodyBytes" }, + { status: 413 }, + ) + : error instanceof InvalidContentLengthError + ? jsonError("Invalid Content-Length.") + : jsonError("Could not read broker API request body."); + } - // Authenticate the exact transport bytes before decoding or interpreting attacker-controlled JSON. - const authenticated = await this.authenticateRequest(request, body, "list"); - if (authenticated instanceof Response) return authenticated; + // Authenticate the exact transport bytes before decoding or interpreting attacker-controlled JSON. + const authenticated = await this.authenticateRequest(request, body, "list"); + if (authenticated instanceof Response) return authenticated; - let input; - try { - input = this.protocolParsers.parseDaemonRequest(parseSessionBrokerJsonBytes(body)); - } catch (error) { - return this.authenticatedResponse(authenticated, protocolError(error), 400); - } + let input; + try { + input = this.protocolParsers.parseDaemonRequest(parseSessionBrokerJsonBytes(body)); + } catch (error) { + return this.authenticatedResponse(authenticated, protocolError(error), 400); + } - const operation = input.action as CallerOperation; - const selector = "selector" in input ? input.selector : undefined; - const sessionId = selector?.sessionId; - const command = input.action === "dispatch" ? input.command : undefined; - const commandVersion = input.action === "dispatch" ? (input.commandVersion ?? 1) : undefined; - const facts = { - operation, - ...(sessionId !== undefined ? { sessionId } : {}), - ...(command !== undefined ? { command, commandVersion } : {}), - }; - const targetSpecific = input.action !== "list"; - if (!(await this.authorize(request, authenticated, facts))) { - return this.authenticatedResponse( - authenticated, - { error: "authorization-denied" }, - 403, - targetSpecific, - ); - } - const inactive = this.rejectInactiveRequest(authenticated); - if (inactive) return inactive; + const operation = input.action as CallerOperation; + const selector = "selector" in input ? input.selector : undefined; + const sessionId = selector?.sessionId; + const command = input.action === "dispatch" ? input.command : undefined; + const commandVersion = input.action === "dispatch" ? (input.commandVersion ?? 1) : undefined; + const facts = { + operation, + ...(sessionId !== undefined ? { sessionId } : {}), + ...(command !== undefined ? { command, commandVersion } : {}), + }; + const targetSpecific = input.action !== "list"; + if (!(await this.authorize(request, authenticated, facts))) { + return this.authenticatedResponse( + authenticated, + { error: "authorization-denied" }, + 403, + targetSpecific, + ); + } + const inactive = this.rejectInactiveRequest(authenticated); + if (inactive) return inactive; - try { - let response: SessionBrokerDaemonResponse; - switch (input.action) { - case "list": - response = { sessions: this.broker.listSessions() }; - break; - case "get": - response = { session: this.broker.getSession(input.selector) }; - break; - case "dispatch": { - // Resolve the target before invoking app-owned parsing so the exact target contract is - // selected first. This read-only lookup happens only after authentication/authorization. - this.broker.getSession(input.selector); - response = { - result: await this.broker.dispatchCommand({ - selector: input.selector, - command: input.command, - commandVersion: input.commandVersion ?? 1, - input: input.input, - timeoutMessage: input.timeoutMessage ?? defaultTimeoutMessage(input.command), - timeoutMs: input.timeoutMs, - }), - }; - break; + try { + let response: SessionBrokerDaemonResponse; + switch (input.action) { + case "list": + response = { sessions: this.broker.listSessions() }; + break; + case "get": + response = { session: this.broker.getSession(input.selector) }; + break; + case "dispatch": { + // Resolve the target before invoking app-owned parsing so the exact target contract is + // selected first. This read-only lookup happens only after authentication/authorization. + this.broker.getSession(input.selector); + response = { + result: await this.broker.dispatchCommand({ + selector: input.selector, + command: input.command, + commandVersion: input.commandVersion ?? 1, + input: input.input, + timeoutMessage: input.timeoutMessage ?? defaultTimeoutMessage(input.command), + timeoutMs: input.timeoutMs, + }), + }; + break; + } } + return this.authenticatedResponse(authenticated, response, 200, targetSpecific); + } catch (error) { + return this.authenticatedResponse( + authenticated, + error instanceof BrokerCapacityError + ? { error: error.code, resource: error.resource } + : error instanceof BrokerProtocolError + ? protocolError(error) + : { + error: error instanceof Error ? error.message : "Unknown broker API error.", + }, + error instanceof BrokerCapacityError ? 503 : 400, + targetSpecific, + ); } - return this.authenticatedResponse(authenticated, response, 200, targetSpecific); - } catch (error) { - return this.authenticatedResponse( - authenticated, - error instanceof BrokerProtocolError - ? protocolError(error) - : { - error: error instanceof Error ? error.message : "Unknown broker API error.", - }, - 400, - targetSpecific, - ); + } finally { + bodyReservation?.release(); + control.release(); } } } diff --git a/scripts/test-session-broker-node.ts b/scripts/test-session-broker-node.ts new file mode 100644 index 000000000..b772d3e56 --- /dev/null +++ b/scripts/test-session-broker-node.ts @@ -0,0 +1,38 @@ +import { mkdtemp, rm } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; + +/** Build the TypeScript workspace adapter, then execute its conformance suite in real Node. */ +async function main() { + const directory = await mkdtemp(join(tmpdir(), "hunk-session-broker-node-")); + const outfile = join(directory, "adapter.mjs"); + try { + const build = await Bun.build({ + entrypoints: [join(process.cwd(), "packages/session-broker-node/src/index.ts")], + outdir: directory, + naming: "adapter.mjs", + target: "node", + format: "esm", + minify: false, + sourcemap: "none", + }); + if (!build.success) { + for (const log of build.logs) console.error(log); + process.exitCode = 1; + return; + } + const child = Bun.spawn({ + cmd: ["node", "--test", "test/session-broker-node/adapter.test.mjs"], + cwd: process.cwd(), + env: { ...process.env, HUNK_NODE_ADAPTER_BUNDLE: outfile }, + stdout: "inherit", + stderr: "inherit", + }); + const exitCode = await child.exited; + if (exitCode !== 0) process.exitCode = exitCode; + } finally { + await rm(directory, { force: true, recursive: true }); + } +} + +await main(); diff --git a/src/session/broker/brokerServer.ts b/src/session/broker/brokerServer.ts index 4bd580aa0..c9ee84f2c 100644 --- a/src/session/broker/brokerServer.ts +++ b/src/session/broker/brokerServer.ts @@ -25,9 +25,10 @@ import type { RemovedCommentResult, } from "../types"; import { + BrokerCapacityError, MAX_HTTP_BODY_BYTES, PayloadTooLargeError, - readRequestTextWithLimit, + readRequestBytesWithLimit, } from "@hunk/session-broker-core"; import { listHunkSessionNotes } from "./projections"; import { @@ -209,11 +210,10 @@ export function validateOriginHeader(request: Request, expectedPort: number, all return null; } -async function parseJsonRequest(request: Request) { - const text = await readRequestTextWithLimit(request, MAX_HTTP_BODY_BYTES); +function parseJsonRequestBytes(bytes: Uint8Array) { let raw: unknown; try { - raw = JSON.parse(text); + raw = JSON.parse(new TextDecoder("utf-8", { fatal: true }).decode(bytes)); } catch { throw new Error("Expected one JSON request body."); } @@ -221,7 +221,11 @@ async function parseJsonRequest(request: Request) { return parseSessionDaemonRequest(raw); } -export async function handleSessionApiRequest(state: HunkSessionBrokerState, request: Request) { +export async function handleSessionApiRequest( + state: HunkSessionBrokerState, + request: Request, + bodyBytes?: Uint8Array, +) { if (request.method !== "POST") { return jsonError("Session API requests must use POST.", 405); } @@ -231,7 +235,9 @@ export async function handleSessionApiRequest(state: HunkSessionBrokerState, req } try { - const input = await parseJsonRequest(request); + const input = parseJsonRequestBytes( + bodyBytes ?? (await readRequestBytesWithLimit(request, MAX_HTTP_BODY_BYTES)), + ); let response: SessionDaemonResponse; switch (input.action) { @@ -446,6 +452,9 @@ export async function handleSessionApiRequest(state: HunkSessionBrokerState, req if (error instanceof PayloadTooLargeError) { return jsonError(error.message, 413); } + if (error instanceof BrokerCapacityError) { + return Response.json({ error: error.code, resource: error.resource }, { status: 503 }); + } return jsonError(error instanceof Error ? error.message : "Unknown session API error."); } @@ -462,6 +471,7 @@ function createHunkBrokerController( ): SessionBrokerController { return { protocolParsers: hunkSessionProtocolParsers, + limits: state.limits, listSessions: () => state.listSessions(), getSession: (selector) => state.getSession(selector), getSessionCount: () => state.getSessionCount(), @@ -546,7 +556,11 @@ export function serveSessionBrokerDaemon( // Keep the richer Hunk session API here rather than in the shared package so commands like // review, reload, and comment flows stay app-specific. if (url.pathname === HUNK_SESSION_API_PATH) { - return handleSessionApiRequest(state, request); + return daemon.handleBoundedControl( + request, + (body) => handleSessionApiRequest(state, request, body), + { payloadTooLarge: (error) => jsonError(error.message, 413) }, + ); } // The review surface authorizes every one of its own routes with a per-session diff --git a/test/fixtures/sessionBrokerAdapterConformance.json b/test/fixtures/sessionBrokerAdapterConformance.json new file mode 100644 index 000000000..381186040 --- /dev/null +++ b/test/fixtures/sessionBrokerAdapterConformance.json @@ -0,0 +1,13 @@ +{ + "textOnly": { + "binaryCloseCode": 1003 + }, + "inbound": { + "oversizedCloseCode": 1009, + "pressureCloseCode": 1013, + "maxMessageBytes": 8388608 + }, + "outbound": { + "pressureCloseCode": 1013 + } +} diff --git a/test/session-broker-node/adapter.test.mjs b/test/session-broker-node/adapter.test.mjs new file mode 100644 index 000000000..9d9a27480 --- /dev/null +++ b/test/session-broker-node/adapter.test.mjs @@ -0,0 +1,197 @@ +import assert from "node:assert/strict"; +import { readFile } from "node:fs/promises"; +import { createServer } from "node:net"; +import { createRequire } from "node:module"; +import { test } from "node:test"; +import { fileURLToPath, pathToFileURL } from "node:url"; + +const require = createRequire( + new URL("../../packages/session-broker-node/package.json", import.meta.url), +); +const WebSocket = require("ws"); +const corpus = JSON.parse( + await readFile( + fileURLToPath(new URL("../fixtures/sessionBrokerAdapterConformance.json", import.meta.url)), + "utf8", + ), +); + +const bundlePath = process.env.HUNK_NODE_ADAPTER_BUNDLE; +if (!bundlePath) throw new Error("HUNK_NODE_ADAPTER_BUNDLE must name the built Node adapter."); +const { serveSessionBrokerDaemon } = await import(pathToFileURL(bundlePath).href); + +async function reservePort() { + const server = createServer(); + await new Promise((resolve, reject) => { + server.once("error", reject); + server.listen(0, "127.0.0.1", resolve); + }); + const address = server.address(); + const port = typeof address === "object" && address ? address.port : 0; + await new Promise((resolve) => server.close(resolve)); + return port; +} + +function closeCode(socket) { + return new Promise((resolve, reject) => { + const timer = setTimeout( + () => reject(new Error("Timed out waiting for WebSocket close.")), + 2_000, + ); + socket.addEventListener( + "close", + (event) => { + clearTimeout(timer); + resolve(event.code); + }, + { once: true }, + ); + socket.addEventListener("error", () => {}, { once: true }); + }); +} + +async function openSocket(url) { + const socket = new WebSocket(url); + await new Promise((resolve, reject) => { + const timer = setTimeout( + () => reject(new Error("Timed out waiting for WebSocket open.")), + 2_000, + ); + socket.addEventListener( + "open", + () => { + clearTimeout(timer); + resolve(); + }, + { once: true }, + ); + socket.addEventListener( + "error", + () => { + clearTimeout(timer); + reject(new Error("WebSocket failed to open.")); + }, + { once: true }, + ); + }); + return socket; +} + +function fakeDaemon(overrides = {}, behavior = {}) { + const limits = { + maxWsMessageBytes: corpus.inbound.maxMessageBytes, + maxInFlightWsBytes: 64 * 1024 * 1024, + maxOutboundBytesPerPeer: 8 * 1024 * 1024, + maxOutboundBytesTotal: 64 * 1024 * 1024, + maxHttpResponseBytes: 8 * 1024 * 1024, + maxUnauthenticatedSockets: 64, + ...overrides, + }; + return { + limits, + stopped: new Promise(() => {}), + matchesSocketPath: (pathname) => pathname === "/session", + handleConnectionMessage: behavior.handleConnectionMessage ?? (() => {}), + handleConnectionClose() {}, + handleRequest: async (request) => + new URL(request.url).pathname === "/health" ? Response.json({ ok: true }) : null, + shutdown() {}, + }; +} + +test("Node WebCrypto Ed25519 and base64url work without Bun globals", async () => { + assert.equal(typeof globalThis.Bun, "undefined"); + const keys = await crypto.subtle.generateKey("Ed25519", false, ["sign", "verify"]); + const message = new TextEncoder().encode("session-broker-node"); + const signature = new Uint8Array(await crypto.subtle.sign("Ed25519", keys.privateKey, message)); + assert.equal(await crypto.subtle.verify("Ed25519", keys.publicKey, signature, message), true); + const encoded = Buffer.from(signature).toString("base64url"); + assert.deepEqual(Buffer.from(encoded, "base64url"), Buffer.from(signature)); +}); + +test("Node adapter consumes the shared text/binary/oversize/pressure corpus", async () => { + const port = await reservePort(); + const running = await serveSessionBrokerDaemon({ + daemon: fakeDaemon({ + maxWsMessageBytes: 8, + maxHttpResponseBytes: 8, + maxUnauthenticatedSockets: 1, + }), + hostname: "127.0.0.1", + port, + handleRequest: (request) => + new URL(request.url).pathname === "/large" + ? new Response("123456789", { headers: { "content-length": "9" } }) + : undefined, + }); + try { + const boundedResponse = await fetch(`http://127.0.0.1:${port}/large`); + assert.equal(boundedResponse.status, 503); + assert.equal(await boundedResponse.text(), ""); + const exact = await openSocket(`ws://127.0.0.1:${port}/session`); + exact.send("12345678"); + await new Promise((resolve) => setTimeout(resolve, 20)); + assert.equal(exact.readyState, WebSocket.OPEN); + await assert.rejects(openSocket(`ws://127.0.0.1:${port}/session`)); + const exactClosed = closeCode(exact); + exact.close(); + await exactClosed; + const afterRelease = await openSocket(`ws://127.0.0.1:${port}/session`); + const afterReleaseClosed = closeCode(afterRelease); + afterRelease.close(); + await afterReleaseClosed; + + const binary = await openSocket(`ws://127.0.0.1:${port}/session`); + const binaryClosed = closeCode(binary); + binary.send(new Uint8Array([1])); + assert.equal(await binaryClosed, corpus.textOnly.binaryCloseCode); + + const oversized = await openSocket(`ws://127.0.0.1:${port}/session`); + const oversizedClosed = closeCode(oversized); + oversized.send("123456789"); + assert.equal(await oversizedClosed, corpus.inbound.oversizedCloseCode); + + const malformed = await openSocket(`ws://127.0.0.1:${port}/session`); + const malformedClosed = closeCode(malformed); + malformed.send(Buffer.from([0xc0, 0xaf]), { binary: false }); + assert.equal(await malformedClosed, 1007); + } finally { + await running.stop(); + await running.stopped; + } + + const outboundPort = await reservePort(); + const outboundRunning = await serveSessionBrokerDaemon({ + daemon: fakeDaemon( + { maxOutboundBytesPerPeer: 1 }, + { handleConnectionMessage: (peer) => peer.send("too large") }, + ), + hostname: "127.0.0.1", + port: outboundPort, + }); + try { + const outbound = await openSocket(`ws://127.0.0.1:${outboundPort}/session`); + const outboundClosed = closeCode(outbound); + outbound.send("trigger"); + assert.equal(await outboundClosed, 1013); + } finally { + await outboundRunning.stop(); + await outboundRunning.stopped; + } + + const pressurePort = await reservePort(); + const pressureRunning = await serveSessionBrokerDaemon({ + daemon: fakeDaemon({ maxWsMessageBytes: 8, maxInFlightWsBytes: 0 }), + hostname: "127.0.0.1", + port: pressurePort, + }); + try { + const pressure = await openSocket(`ws://127.0.0.1:${pressurePort}/session`); + const pressureClosed = closeCode(pressure); + pressure.send("{}"); + assert.equal(await pressureClosed, corpus.inbound.pressureCloseCode); + } finally { + await pressureRunning.stop(); + await pressureRunning.stopped; + } +}); From 61dbde2d801df3967fc019e2e08936f03717f096 Mon Sep 17 00:00:00 2001 From: Ben Vinegar Date: Sat, 29 Aug 2026 15:35:09 -0400 Subject: [PATCH 5/5] feat(session): authenticate Hunk broker clients --- .../secure-session-broker-integration.md | 5 + docs/agent-workflows.md | 10 +- docs/session-broker-sdk.md | 13 + packages/session-broker-bun/src/serve.test.ts | 9 +- packages/session-broker-bun/src/serve.ts | 23 + .../src/brokerState.test.ts | 71 +++ .../session-broker-core/src/brokerState.ts | 28 +- .../session-broker-core/src/budgets.test.ts | 13 + packages/session-broker-core/src/budgets.ts | 46 ++ packages/session-broker-node/src/serve.ts | 24 + .../session-broker/src/authentication.test.ts | 26 + packages/session-broker/src/authentication.ts | 29 +- packages/session-broker/src/broker.ts | 15 +- .../src/clientAuthentication.test.ts | 208 +++++++ .../src/clientAuthentication.ts | 575 ++++++++++++++++++ .../session-broker/src/connection.test.ts | 58 +- packages/session-broker/src/connection.ts | 134 +++- packages/session-broker/src/daemon.test.ts | 104 +++- packages/session-broker/src/daemon.ts | 325 +++++++++- packages/session-broker/src/index.ts | 1 + src/main.tsx | 2 +- src/session/agent/cliClient.test.ts | 13 +- src/session/agent/cliClient.ts | 102 +++- src/session/agent/commands.ts | 96 +-- src/session/broker/appContract.ts | 15 + src/session/broker/brokerClient.test.ts | 31 +- src/session/broker/brokerClient.ts | 87 +-- src/session/broker/brokerConfig.test.ts | 15 + src/session/broker/brokerLauncher.ts | 10 +- .../broker/brokerServer.helpers.test.ts | 4 +- src/session/broker/brokerServer.test.ts | 223 +++++-- src/session/broker/brokerServer.ts | 162 ++++- src/session/broker/credentials.test.ts | 91 +++ src/session/broker/credentials.ts | 375 ++++++++++++ src/session/broker/state.ts | 3 +- src/session/client/capabilities.ts | 4 +- test/session-broker-node/adapter.test.mjs | 16 +- test/session/broker-e2e.test.ts | 25 +- test/session/cli.test.ts | 54 +- test/session/daemon.test.ts | 6 +- 40 files changed, 2689 insertions(+), 362 deletions(-) create mode 100644 .changeset/secure-session-broker-integration.md create mode 100644 packages/session-broker/src/clientAuthentication.test.ts create mode 100644 packages/session-broker/src/clientAuthentication.ts create mode 100644 src/session/broker/appContract.ts create mode 100644 src/session/broker/credentials.test.ts create mode 100644 src/session/broker/credentials.ts diff --git a/.changeset/secure-session-broker-integration.md b/.changeset/secure-session-broker-integration.md new file mode 100644 index 000000000..4c356d2fe --- /dev/null +++ b/.changeset/secure-session-broker-integration.md @@ -0,0 +1,5 @@ +--- +"hunkdiff": minor +--- + +Authenticate local session producers and CLI controls with automatically discovered owner-private credentials, signed responses, scoped reconnect replacement, and bounded handshakes. Expose only minimal public daemon health and refuse unsafe PID-based replacement of legacy listeners. diff --git a/docs/agent-workflows.md b/docs/agent-workflows.md index 390247dd1..09ce109f1 100644 --- a/docs/agent-workflows.md +++ b/docs/agent-workflows.md @@ -25,15 +25,7 @@ When a Hunk TUI starts, it registers with a local loopback daemon. `hunk session Most users only need `hunk session ...`. Use `hunk mcp serve` only for manual startup or debugging of the local daemon. -If `hunk session list` reports no sessions while Hunk is visibly running, the agent sandbox may be blocking loopback access. Probe the daemon directly: - -```bash -curl -s -X POST http://127.0.0.1:47657/session-api \ - -H 'content-type: application/json' \ - --data '{"action":"list"}' -``` - -If this shows sessions, rerun the command with the agent's network/sandbox escalation. If you run the daemon with a custom `HUNK_MCP_PORT`, use that port instead. +If `hunk session list` reports no sessions while Hunk is visibly running, the agent sandbox may be blocking loopback access. Rerun `hunk session list --json` with the agent's network/sandbox escalation. Do not probe `/session-api` with raw `curl`: session controls require an automatically discovered, owner-private caller credential and signed responses, and Hunk intentionally exposes no credential flags. ## The commands you will use most diff --git a/docs/session-broker-sdk.md b/docs/session-broker-sdk.md index 12a2c2fea..6250715dc 100644 --- a/docs/session-broker-sdk.md +++ b/docs/session-broker-sdk.md @@ -555,6 +555,19 @@ Security outranks wire compatibility: no migration accepts unauthenticated contr Preservation means paths, selectors, outputs, and automatic credential discovery for upgraded clients—not interoperability with pre-authentication binaries. +Hunk's fixed-endpoint Phase-1 credential store uses a home-local `.hunk` parent when +`XDG_RUNTIME_DIR` is unavailable, rather than a predictable name in a shared temporary directory. +It inherits the current user's ACL when it creates the `hunk-mcp/security-v1` directory +on Windows and rejects symbolic-link redirection. Node does not +provide a portable owner/DACL or general reparse-point inspection API, so this integration cannot +detect a pre-existing custom permissive DACL or every non-symlink reparse point; completing native +Windows ACL validation remains a release-gate item before the reusable package is published. + +The fixed-endpoint integration authenticates bootstrap reconnects, distinguishes `register` from +`reconnect` scope, atomically retires the previous socket, and rejects its uncertain work. It does +not yet claim the durable candidate-key `registered`/`registration-ack` rotation sequence above; +that sequence remains a publication gate rather than an unauthenticated compatibility fallback. + Before publishing even `initializing`, a Hunk candidate binds and retains the legacy guard endpoint; only its holder may enter coordinator election. A contender unable to bind waits a bounded startup interval for authenticated coordinator publication, then reports an unverifiable listener and launches diff --git a/packages/session-broker-bun/src/serve.test.ts b/packages/session-broker-bun/src/serve.test.ts index 16d5e6af1..591e198f2 100644 --- a/packages/session-broker-bun/src/serve.test.ts +++ b/packages/session-broker-bun/src/serve.test.ts @@ -236,11 +236,13 @@ describe("session broker bun adapter", () => { } }); - test("admits exactly the configured number of active websocket peers", async () => { + test("admits exactly the configured number of unauthenticated websocket peers", async () => { const broker = new SessionBroker({ protocolParsers }); const daemon = createSessionBrokerDaemon({ broker, - limits: { maxUnauthenticatedSockets: 1 }, + limits: { maxUnauthenticatedSockets: 1, maxHandshakeDurationMs: 50 }, + helloAuthenticator: {} as never, + producerEndpoint: "ws://127.0.0.1/session", }); const port = await reserveLoopbackPort(); const server = serveSessionBrokerDaemon({ daemon, hostname: "127.0.0.1", port }); @@ -248,8 +250,7 @@ describe("session broker bun adapter", () => { const first = await openTestSocket(`ws://127.0.0.1:${port}/session`); await expect(openTestSocket(`ws://127.0.0.1:${port}/session`)).rejects.toThrow(); const closed = testSocketCloseCode(first); - first.close(); - await closed; + expect(await closed).toBe(1008); const afterRelease = await openTestSocket(`ws://127.0.0.1:${port}/session`); afterRelease.close(); } finally { diff --git a/packages/session-broker-bun/src/serve.ts b/packages/session-broker-bun/src/serve.ts index 9bbfd883c..ba15f8448 100644 --- a/packages/session-broker-bun/src/serve.ts +++ b/packages/session-broker-bun/src/serve.ts @@ -10,6 +10,7 @@ import type { SessionBrokerDaemon, SessionBrokerPeer } from "@hunk/session-broke interface BrokerWebSocketData { admission: BudgetReservation; + handshakeTimer?: ReturnType; } export interface ServeSessionBrokerDaemonOptions< @@ -123,6 +124,16 @@ export function serveSessionBrokerDaemon< } }, close: (code, reason) => socket.close(code, reason), + markAuthenticated() { + const data = (socket as typeof socket & { data?: BrokerWebSocketData }).data; + if (!data) return; + if (data.handshakeTimer) { + clearTimeout(data.handshakeTimer); + data.handshakeTimer = undefined; + } + activeAdmissions.delete(data.admission); + data.admission.release(); + }, }; peers.set(key, peer); return peer; @@ -192,6 +203,17 @@ export function serveSessionBrokerDaemon< ); }, websocket: { + open: (socket) => { + if (!options.daemon.requiresProducerAuthentication) { + activeAdmissions.delete(socket.data.admission); + socket.data.admission.release(); + return; + } + socket.data.handshakeTimer = setTimeout(() => { + socket.close(1008, "Session broker authentication timed out."); + }, options.daemon.limits.maxHandshakeDurationMs); + socket.data.handshakeTimer.unref?.(); + }, // Bun cannot customize the close code of its native payload rejection. Keep the native cap // at the fixed aggregate ceiling so decoded messages above the per-message limit reach the // portable 1009 path while runtime buffering remains bounded. @@ -239,6 +261,7 @@ export function serveSessionBrokerDaemon< }, close: (socket) => { const key = socket as object; + if (socket.data.handshakeTimer) clearTimeout(socket.data.handshakeTimer); bufferedReservations.get(key)?.release(); bufferedReservations.delete(key); activeAdmissions.delete(socket.data.admission); diff --git a/packages/session-broker-core/src/brokerState.test.ts b/packages/session-broker-core/src/brokerState.test.ts index 6c2855d06..1edf642ce 100644 --- a/packages/session-broker-core/src/brokerState.test.ts +++ b/packages/session-broker-core/src/brokerState.test.ts @@ -527,6 +527,77 @@ describe("session broker state", () => { expect(state.listSessions()).toHaveLength(1); }); + test("atomically replaces a live owner without leaking the replacement socket's prior reservations", () => { + const registration = createRegistration(); + const snapshot = createSnapshot(); + const retainedBytes = + new TextEncoder().encode(JSON.stringify({ registration, snapshot })).byteLength + 256; + const expandedRegistration = createRegistration({ + info: { ...registration.info, title: "x".repeat(64) }, + }); + const expandedBytes = + new TextEncoder().encode( + JSON.stringify({ + registration: expandedRegistration, + snapshot: createSnapshot({ updatedAt: "2026-03-22T00:00:01.000Z" }), + }), + ).byteLength + 256; + const state = createState({ + limits: { + maxSessions: 2, + maxRetainedSessionBytes: expandedBytes, + maxRetainedBytes: retainedBytes * 2, + }, + }); + const originalSocket = { send() {} }; + const replacementSocket = { send() {} }; + state.registerSession(originalSocket, registration, snapshot); + state.registerSession( + replacementSocket, + createRegistration({ sessionId: "session-2" }), + snapshot, + ); + + expect( + state.registerSession( + replacementSocket, + expandedRegistration, + createSnapshot({ updatedAt: "2026-03-22T00:00:01.000Z" }), + { replaceOwner: true }, + ), + ).toBe("registered"); + expect(state.markSessionSeen(originalSocket, "session-1")).toBe("not-owner"); + expect(state.markSessionSeen(replacementSocket, "session-1")).toBe("seen"); + state.unregisterSocket(originalSocket); + expect(state.listSessions()).toHaveLength(1); + }); + + test("releases the replacement socket's prior session count reservation", () => { + const state = createState({ limits: { maxSessions: 2 } }); + const originalSocket = { send() {} }; + const replacementSocket = { send() {} }; + const thirdSocket = { send() {} }; + state.registerSession(originalSocket, createRegistration(), createSnapshot()); + state.registerSession( + replacementSocket, + createRegistration({ sessionId: "session-2" }), + createSnapshot(), + ); + expect( + state.registerSession(replacementSocket, createRegistration(), createSnapshot(), { + replaceOwner: true, + }), + ).toBe("registered"); + expect( + state.registerSession( + thirdSocket, + createRegistration({ sessionId: "session-3" }), + createSnapshot(), + ), + ).toBe("registered"); + expect(state.listSessions()).toHaveLength(2); + }); + test("rejects commands immediately when the live session socket cannot accept them", async () => { const state = createState(); const socket = { diff --git a/packages/session-broker-core/src/brokerState.ts b/packages/session-broker-core/src/brokerState.ts index 2f2f52bff..8d2f0f7c0 100644 --- a/packages/session-broker-core/src/brokerState.ts +++ b/packages/session-broker-core/src/brokerState.ts @@ -287,6 +287,7 @@ export class SessionBrokerState< socket: DaemonSessionSocket, registrationInput: unknown, snapshotInput: unknown, + options: { replaceOwner?: boolean } = {}, ): RegisterSessionResult { let registration: SessionRegistration | null; let snapshot: SessionSnapshot | null; @@ -312,7 +313,7 @@ export class SessionBrokerState< if (retainedBytes > this.limits.maxRetainedSessionBytes) return "capacity-exceeded"; const existing = this.sessions.get(registration.sessionId); - if (existing && existing.socket !== socket) return "already-connected"; + if (existing && existing.socket !== socket && !options.replaceOwner) return "already-connected"; const previousSessionId = this.sessionIdsBySocket.get(socket); const transferSessionId = existing ? registration.sessionId : previousSessionId; const previousRetained = transferSessionId @@ -321,13 +322,27 @@ export class SessionBrokerState< const previousCount = transferSessionId ? this.sessionReservations.get(transferSessionId) : undefined; + const abandonedRetained = + existing && previousSessionId && previousSessionId !== registration.sessionId + ? this.retainedReservations.get(previousSessionId) + : undefined; + const abandonedCount = + existing && previousSessionId && previousSessionId !== registration.sessionId + ? this.sessionReservations.get(previousSessionId) + : undefined; let retainedReservation: BudgetReservation | null = null; let sessionReservation: BudgetReservation | null = null; try { try { retainedReservation = previousRetained - ? this.retainedByteBudget.resize(previousRetained, retainedBytes) + ? abandonedRetained + ? this.retainedByteBudget.resizeWithCredit( + previousRetained, + retainedBytes, + abandonedRetained, + ) + : this.retainedByteBudget.resize(previousRetained, retainedBytes) : this.retainedByteBudget.reserve(retainedBytes); sessionReservation = previousCount ?? this.sessionBudget.reserve(); } catch { @@ -335,11 +350,20 @@ export class SessionBrokerState< } const now = new Date().toISOString(); + if (existing && existing.socket !== socket) { + this.sessionIdsBySocket.delete(existing.socket); + this.rejectPendingCommandsForSession( + registration.sessionId, + new Error("The session owner reconnected."), + ); + } if (previousSessionId && previousSessionId !== registration.sessionId) { // Detach the old identity without releasing the reservations transferred to its replacement. this.sessions.delete(previousSessionId); this.retainedReservations.delete(previousSessionId); this.sessionReservations.delete(previousSessionId); + abandonedRetained?.release(); + abandonedCount?.release(); this.rejectPendingCommandsForSession( previousSessionId, new Error("The session registration was replaced."), diff --git a/packages/session-broker-core/src/budgets.test.ts b/packages/session-broker-core/src/budgets.test.ts index 86f189d68..97705a0b8 100644 --- a/packages/session-broker-core/src/budgets.test.ts +++ b/packages/session-broker-core/src/budgets.test.ts @@ -27,6 +27,7 @@ describe("session broker limits", () => { maxHttpResponseBytes: 8 * 1024 * 1024, maxWsMessageBytes: 8 * 1024 * 1024, maxInFlightWsBytes: 64 * 1024 * 1024, + maxHandshakeDurationMs: 15_000, }); expect(Object.isFrozen(DEFAULT_SESSION_BROKER_LIMITS)).toBe(true); }); @@ -75,6 +76,18 @@ describe("resource reservations", () => { expect(budget.used).toBe(0); }); + test("combines a replacement and retired reservation without transient over-admission", () => { + const budget = new ResourceBudget(10, "bytes"); + const target = budget.reserve(6); + const credit = budget.reserve(4); + const replacement = budget.resizeWithCredit(target, 9, credit); + expect(budget.used).toBe(9); + expect(target.released).toBe(true); + expect(credit.released).toBe(true); + replacement.release(); + expect(budget.used).toBe(0); + }); + test("resizes retained records by their delta and transfers release ownership", () => { const budget = new ResourceBudget(4, "bytes"); const original = budget.reserve(4); diff --git a/packages/session-broker-core/src/budgets.ts b/packages/session-broker-core/src/budgets.ts index 51372b8bc..5bb1236c8 100644 --- a/packages/session-broker-core/src/budgets.ts +++ b/packages/session-broker-core/src/budgets.ts @@ -22,6 +22,7 @@ export interface SessionBrokerLimits { readonly maxOutboundBytesPerPeer: number; readonly maxOutboundBytesTotal: number; readonly maxUnauthenticatedSockets: number; + readonly maxHandshakeDurationMs: number; readonly maxIncompleteHandshakes: number; readonly maxIncompleteHandshakeBytes: number; readonly maxHandshakeProposalBytes: number; @@ -52,6 +53,7 @@ export const DEFAULT_SESSION_BROKER_LIMITS: Readonly = Obje maxOutboundBytesPerPeer: 8 * 1024 * 1024, maxOutboundBytesTotal: 64 * 1024 * 1024, maxUnauthenticatedSockets: 64, + maxHandshakeDurationMs: 15_000, maxIncompleteHandshakes: 128, maxIncompleteHandshakeBytes: 4 * 1024 * 1024, maxHandshakeProposalBytes: 64 * 1024, @@ -227,6 +229,50 @@ export class ResourceBudget { this.reservationStates.set(replacement, replacementState); return replacement; } + + /** Atomically resize one reservation while retiring a second reservation from this budget. */ + resizeWithCredit( + previous: BudgetReservation, + amount: number, + credit: BudgetReservation, + ): BudgetReservation { + assertLimit(amount, this.resource); + const previousState = this.reservationStates.get(previous); + const creditState = this.reservationStates.get(credit); + if ( + previous === credit || + !previousState || + previousState.released || + !creditState || + creditState.released + ) { + throw new TypeError(`Cannot combine inactive ${this.resource} reservations.`); + } + const delta = amount - previousState.amount - creditState.amount; + if (delta > this.capacity - this.reserved) { + throw new BrokerCapacityError(this.code, this.resource); + } + this.reserved += delta; + const replacementState = { amount, released: false }; + const replacement: BudgetReservation = { + amount, + get released() { + return replacementState.released; + }, + release: () => { + if (replacementState.released) return; + replacementState.released = true; + this.reservationStates.delete(replacement); + this.reserved -= replacementState.amount; + }, + }; + previousState.released = true; + creditState.released = true; + this.reservationStates.delete(previous); + this.reservationStates.delete(credit); + this.reservationStates.set(replacement, replacementState); + return replacement; + } } /** Own several incremental reservations and release all of them idempotently. */ diff --git a/packages/session-broker-node/src/serve.ts b/packages/session-broker-node/src/serve.ts index fc5b6a52a..c5d6595ca 100644 --- a/packages/session-broker-node/src/serve.ts +++ b/packages/session-broker-node/src/serve.ts @@ -49,6 +49,7 @@ function toNodeConnection( socket: WebSocket, outboundBudget: ResourceBudget, maxPeerBytes: number, + markAuthenticated: () => void, ): SessionBrokerPeer { return { send(data: string) { @@ -77,6 +78,7 @@ function toNodeConnection( close(code?: number, reason?: string) { socket.close(code, reason); }, + markAuthenticated, }; } @@ -215,6 +217,7 @@ export async function serveSessionBrokerDaemon< // connection object that registration and message handling used earlier. const peerBySocket = new WeakMap(); const admissionBySocket = new WeakMap(); + const handshakeTimers = new WeakMap>(); const activeWebSockets = new Set(); const activeSockets = new Set(); server.on("connection", (socket) => { @@ -238,11 +241,28 @@ export async function serveSessionBrokerDaemon< webSocketServer.on("connection", (socket: WebSocket) => { activeWebSockets.add(socket); + const markAuthenticated = () => { + admissionBySocket.get(socket)?.release(); + admissionBySocket.delete(socket); + const timer = handshakeTimers.get(socket); + if (timer) clearTimeout(timer); + handshakeTimers.delete(socket); + }; const peer = toNodeConnection( socket, outboundBudget, options.daemon.limits.maxOutboundBytesPerPeer, + markAuthenticated, ); + if (options.daemon.requiresProducerAuthentication) { + const timer = setTimeout(() => { + socket.close(1008, "Session broker authentication timed out."); + }, options.daemon.limits.maxHandshakeDurationMs); + timer.unref?.(); + handshakeTimers.set(socket, timer); + } else { + markAuthenticated(); + } peerBySocket.set(socket, peer); socket.on("message", (message: Buffer | ArrayBuffer | Buffer[], isBinary: boolean) => { if (isBinary) { @@ -288,7 +308,11 @@ export async function serveSessionBrokerDaemon< socket.on("error", () => {}); socket.on("close", (code: number, reason: Buffer) => { activeWebSockets.delete(socket); + const timer = handshakeTimers.get(socket); + if (timer) clearTimeout(timer); + handshakeTimers.delete(socket); admissionBySocket.get(socket)?.release(); + admissionBySocket.delete(socket); options.daemon.handleConnectionClose(peerBySocket.get(socket) ?? peer); // The runtime-neutral daemon only cares that the transport closed; Node-specific close data // stays ignored here instead of leaking into the shared broker API. diff --git a/packages/session-broker/src/authentication.test.ts b/packages/session-broker/src/authentication.test.ts index 190853f92..7082d07ad 100644 --- a/packages/session-broker/src/authentication.test.ts +++ b/packages/session-broker/src/authentication.test.ts @@ -212,6 +212,32 @@ describe("session broker signed authentication", () => { }); }); + test("rechecks producer expiry and revocation after hello completion", async () => { + let revoked = false; + const values = await setup({ revoked: () => revoked }); + const request = challengeRequest("producer"); + const challenge = await values.authenticator.issueChallenge(request, request.endpoint); + const transcript = challengeTranscriptForClient(request, challenge, "generation-1"); + const signature = encodeBase64Url( + await webSessionBrokerCrypto.sign(values.producer.privateKey, transcript), + ); + const hello = await values.authenticator.completeProducerHello( + { challengeId: challenge.challengeId, signature }, + "connection-1", + ); + + expect(() => values.authenticator.assertProducerActive(hello.principal)).not.toThrow(); + revoked = true; + expect(() => values.authenticator.assertProducerActive(hello.principal)).toThrow( + SessionBrokerAuthenticationError, + ); + revoked = false; + values.setNow(10_001); + expect(() => values.authenticator.assertProducerActive(hello.principal)).toThrow( + SessionBrokerAuthenticationError, + ); + }); + test("rejects missing, wrong, expired, revoked, and reused credentials with redacted errors", async () => { const values = await setup(); await expect( diff --git a/packages/session-broker/src/authentication.ts b/packages/session-broker/src/authentication.ts index 08e36cb01..afda09487 100644 --- a/packages/session-broker/src/authentication.ts +++ b/packages/session-broker/src/authentication.ts @@ -117,6 +117,7 @@ export interface SessionBrokerHelloChallengeRequest { export interface SessionBrokerHelloChallenge { readonly challengeId: string; + readonly generation: string; readonly responderNonce: string; readonly expiresAt: number; readonly daemonKeyId: string; @@ -541,8 +542,20 @@ export function canonicalHttpTarget(url: URL): string { return query ? `${path}?${query}` : path; } +export interface SessionBrokerHelloAuthenticator { + issueChallenge(request: unknown, listenerEndpoint: string): Promise; + completeCallerHello(proofInput: unknown): Promise; + completeProducerHello( + proofInput: unknown, + connectionId: unknown, + ): Promise; + assertProducerActive(principal: ProducerPrincipal): void; +} + /** Authenticate bounded producer hellos and generation-bound signed caller request sessions. */ -export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { +export class SessionBrokerAuthenticator + implements CallerRequestAuthenticator, SessionBrokerHelloAuthenticator +{ private readonly crypto: SessionBrokerCrypto; private readonly config: AuthenticatorSnapshot; private readonly credentials: Map; @@ -626,6 +639,7 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { ); return Object.freeze({ challengeId, + generation: this.config.generation, responderNonce, expiresAt, daemonKeyId: this.config.daemonIdentity.keyId, @@ -846,6 +860,19 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { }); } + /** Recheck the configured producer grant before every connection-owned mutation. */ + assertProducerActive(principal: ProducerPrincipal): void { + const credential = this.credentials.get(`producer:${principal.keyId}:${principal.grantId}`); + if ( + !credential || + credential.grant.kind !== "producer" || + JSON.stringify(principalFromGrant(credential.grant)) !== JSON.stringify(principal) + ) { + authenticationError("invalid-credential"); + } + this.requireActiveGrant(credential.grant); + } + /** Revoke one in-memory caller session without exposing whether it previously existed. */ revokeCallerSession(callerSessionId: string): void { this.deleteCallerSession(callerSessionId); diff --git a/packages/session-broker/src/broker.ts b/packages/session-broker/src/broker.ts index 44ff0752a..8a20b52bf 100644 --- a/packages/session-broker/src/broker.ts +++ b/packages/session-broker/src/broker.ts @@ -19,6 +19,7 @@ import type { SessionBrokerProtocolParsers } from "./protocolParsers"; export interface SessionBrokerPeer { send(data: string): unknown; close?(code?: number, reason?: string): unknown; + markAuthenticated?(): void; } /** One raw live session record with the original registration and snapshot payloads intact. */ @@ -64,12 +65,15 @@ export interface SessionBrokerController< readonly limits?: Readonly; listSessions(): SessionView[]; getSession(selector: SessionTargetSelector): SessionView; + resolveSessionId(selector: SessionTargetSelector): string; + getSessionIds(): string[]; getSessionCount(): number; getPendingCommandCount(): number; registerSession( connection: SessionBrokerPeer, registrationInput: unknown, snapshotInput: unknown, + options?: { replaceOwner?: boolean }, ): RegisterSessionResult; updateSnapshot( connection: SessionBrokerPeer, @@ -191,6 +195,14 @@ export class SessionBroker< return this.state.getSession(selector); } + resolveSessionId(selector: SessionTargetSelector) { + return this.state.getSession(selector).sessionId; + } + + getSessionIds() { + return this.state.listSessions().map((session) => session.sessionId); + } + getSessionCount() { return this.state.getSessionCount(); } @@ -203,8 +215,9 @@ export class SessionBroker< connection: SessionBrokerPeer, registrationInput: unknown, snapshotInput: unknown, + options?: { replaceOwner?: boolean }, ) { - return this.state.registerSession(connection, registrationInput, snapshotInput); + return this.state.registerSession(connection, registrationInput, snapshotInput, options); } updateSnapshot( diff --git a/packages/session-broker/src/clientAuthentication.test.ts b/packages/session-broker/src/clientAuthentication.test.ts new file mode 100644 index 000000000..b09ecc8d6 --- /dev/null +++ b/packages/session-broker/src/clientAuthentication.test.ts @@ -0,0 +1,208 @@ +import { describe, expect, test } from "bun:test"; +import { SESSION_BROKER_SIGNATURE_ALGORITHM, type CallerGrant } from "@hunk/session-broker-core"; +import { SessionBrokerAuthenticator } from "./authentication"; +import { SessionBrokerCallerClient } from "./clientAuthentication"; + +async function keyPair() { + const generated = (await crypto.subtle.generateKey("Ed25519", true, [ + "sign", + "verify", + ])) as CryptoKeyPair; + const privateBytes = await crypto.subtle.exportKey("pkcs8", generated.privateKey); + return { + publicKey: generated.publicKey, + privateKey: await crypto.subtle.importKey("pkcs8", privateBytes, "Ed25519", false, ["sign"]), + }; +} + +async function setup() { + const daemon = await keyPair(); + const caller = await keyPair(); + const grant: CallerGrant = { + kind: "caller", + appId: "dev.example", + principalId: "caller-1", + keyId: "caller-key-1", + grantId: "caller-grant-1", + algorithm: SESSION_BROKER_SIGNATURE_ALGORITHM, + issuedAt: Date.now() - 1_000, + expiresAt: Date.now() + 60_000, + revocationId: "caller-revocation-1", + mayDelegate: false, + operations: ["list"], + commands: [], + }; + const authenticator = new SessionBrokerAuthenticator({ + appId: "dev.example", + appRevision: 7, + generation: "generation-1", + daemonIdentity: { keyId: "daemon-key-1", privateKey: daemon.privateKey }, + credentials: [{ grant, publicKey: caller.publicKey }], + }); + return { daemon, caller, grant, authenticator }; +} + +/** Build an in-memory HTTP adapter exercising the exact generic challenge/proof/request bytes. */ +function createFetch( + authenticator: SessionBrokerAuthenticator, + proofCount: { value: number }, + targetSpecific = false, +) { + return (async (input: string | URL | Request, init?: RequestInit) => { + const request = input instanceof Request ? input : new Request(input, init); + const url = new URL(request.url); + if (url.pathname === "/session-auth/challenge") { + return Response.json(await authenticator.issueChallenge(await request.json(), request.url)); + } + if (url.pathname === "/session-auth/proof") { + proofCount.value += 1; + return Response.json(await authenticator.completeCallerHello(await request.json())); + } + const body = new Uint8Array(await request.arrayBuffer()); + try { + const authenticated = await authenticator.authenticate({ request, body }); + const responseBody = { sessions: [] }; + return Response.json({ + body: responseBody, + authentication: await authenticated.signResponse({ + httpStatus: 200, + body: responseBody, + ...(targetSpecific ? { appContract: { appRevision: 7, features: [] } } : {}), + }), + }); + } catch { + return Response.json({ error: "authentication-required" }, { status: 401 }); + } + }) as typeof fetch; +} + +describe("session broker caller client", () => { + test("negotiates once, allocates monotonic signed sequences, and verifies signed responses", async () => { + const values = await setup(); + const proofCount = { value: 0 }; + const client = new SessionBrokerCallerClient({ + appId: "dev.example", + appRevision: 7, + origin: "http://broker.test", + credential: { grant: values.grant, privateKey: values.caller.privateKey }, + daemon: { keyId: "daemon-key-1", publicKey: values.daemon.publicKey }, + fetch: createFetch(values.authenticator, proofCount), + }); + + await expect( + client + .request("/control", { method: "POST", body: "{}" }) + .then((response) => response.json()), + ).resolves.toEqual({ sessions: [] }); + await expect( + client + .request("/control", { method: "POST", body: "{}" }) + .then((response) => response.json()), + ).resolves.toEqual({ sessions: [] }); + expect(proofCount.value).toBe(1); + }); + + test("requires the exact Hunk-style application contract on target-specific responses", async () => { + const values = await setup(); + const proofCount = { value: 0 }; + const client = new SessionBrokerCallerClient({ + appId: "dev.example", + appRevision: 7, + origin: "http://broker.test", + credential: { grant: values.grant, privateKey: values.caller.privateKey }, + daemon: { keyId: "daemon-key-1", publicKey: values.daemon.publicKey }, + fetch: createFetch(values.authenticator, proofCount, true), + }); + + await expect( + client + .request("/control", { method: "POST", body: "{}" }, { targetSpecific: true }) + .then((response) => response.json()), + ).resolves.toEqual({ sessions: [] }); + }); + + test("rejects an unsigned second 401 after one fresh-session retry", async () => { + const values = await setup(); + const proofCount = { value: 0 }; + const authenticatedFetch = createFetch(values.authenticator, proofCount); + const fetchWithForgedControls = (async (input: string | URL | Request, init?: RequestInit) => { + const request = input instanceof Request ? input : new Request(input, init); + return new URL(request.url).pathname === "/control" + ? Response.json({ error: "forged" }, { status: 401 }) + : authenticatedFetch(request); + }) as typeof fetch; + const client = new SessionBrokerCallerClient({ + appId: "dev.example", + appRevision: 7, + origin: "http://broker.test", + credential: { grant: values.grant, privateKey: values.caller.privateKey }, + daemon: { keyId: "daemon-key-1", publicKey: values.daemon.publicKey }, + fetch: fetchWithForgedControls, + }); + + await expect(client.request("/control", { method: "POST", body: "{}" })).rejects.toThrow( + "daemon identity could not be verified", + ); + expect(proofCount.value).toBe(2); + }); + + test("propagates the control abort signal through challenge and proof negotiation", async () => { + const values = await setup(); + const proofCount = { value: 0 }; + const signals: Array = []; + const authenticatedFetch = createFetch(values.authenticator, proofCount); + const observingFetch = (async (input: string | URL | Request, init?: RequestInit) => { + signals.push(input instanceof Request ? input.signal : init?.signal); + return authenticatedFetch(input, init); + }) as typeof fetch; + const client = new SessionBrokerCallerClient({ + appId: "dev.example", + appRevision: 7, + origin: "http://broker.test", + credential: { grant: values.grant, privateKey: values.caller.privateKey }, + daemon: { keyId: "daemon-key-1", publicKey: values.daemon.publicKey }, + fetch: observingFetch, + }); + const controller = new AbortController(); + + await client.request("/control", { method: "POST", body: "{}", signal: controller.signal }); + expect(signals).toHaveLength(3); + expect(signals.every((signal) => signal === controller.signal)).toBe(true); + }); + + test("rejects oversized unauthenticated challenge responses before parsing", async () => { + const values = await setup(); + const client = new SessionBrokerCallerClient({ + appId: "dev.example", + appRevision: 7, + origin: "http://broker.test", + credential: { grant: values.grant, privateKey: values.caller.privateKey }, + daemon: { keyId: "daemon-key-1", publicKey: values.daemon.publicKey }, + maxResponseBytes: 32, + fetch: (async () => Response.json({ padding: "x".repeat(128) })) as unknown as typeof fetch, + }); + + await expect(client.request("/control")).rejects.toThrow( + "daemon identity could not be verified", + ); + }); + + test("verifies the daemon challenge before presenting caller proof", async () => { + const values = await setup(); + const wrongDaemon = await keyPair(); + const proofCount = { value: 0 }; + const client = new SessionBrokerCallerClient({ + appId: "dev.example", + appRevision: 7, + origin: "http://broker.test", + credential: { grant: values.grant, privateKey: values.caller.privateKey }, + daemon: { keyId: "daemon-key-1", publicKey: wrongDaemon.publicKey }, + fetch: createFetch(values.authenticator, proofCount), + }); + + await expect(client.request("/control", { method: "POST", body: "{}" })).rejects.toThrow( + "daemon identity could not be verified", + ); + expect(proofCount.value).toBe(0); + }); +}); diff --git a/packages/session-broker/src/clientAuthentication.ts b/packages/session-broker/src/clientAuthentication.ts new file mode 100644 index 000000000..40cbc58d3 --- /dev/null +++ b/packages/session-broker/src/clientAuthentication.ts @@ -0,0 +1,575 @@ +import { + CallerSequenceAllocator, + DEFAULT_SESSION_BROKER_LIMITS, + SESSION_BROKER_PROTOCOL_REVISION, + buildBrokerHelloAckTranscript, + buildBrokerResponseTranscript, + buildCallerRequestTranscript, + canonicalJsonBytes, + isValidBrokerIdentifier, + type BrokerGrant, + type BrokerHelloProposal, + type CallerGrant, + type CanonicalJsonValue, + type ProducerGrant, +} from "@hunk/session-broker-core"; +import { + canonicalHttpTarget, + challengeTranscriptForClient, + type AuthenticatedCallerSession, + type AuthenticatedProducerHello, + type SessionBrokerHelloChallenge, + type SessionBrokerHelloChallengeRequest, +} from "./authentication"; +import { + decodeBase64Url, + encodeBase64Url, + webSessionBrokerCrypto, + type SessionBrokerCrypto, +} from "./crypto"; +import type { SessionBrokerAuthenticatedResponse } from "./types"; + +export interface SessionBrokerClientCredential { + readonly grant: Grant; + readonly privateKey: CryptoKey; +} + +export interface SessionBrokerDaemonVerifier { + readonly keyId: string; + readonly publicKey: CryptoKey; +} + +export interface SessionBrokerHelloClientOptions { + readonly appId: string; + readonly appRevision: number; + readonly endpoint: string; + readonly credential: SessionBrokerClientCredential; + readonly daemon: SessionBrokerDaemonVerifier; + readonly crypto?: SessionBrokerCrypto; +} + +export interface PendingSessionBrokerHello { + readonly request: SessionBrokerHelloChallengeRequest; + readonly transcript: Uint8Array; + readonly transcriptHash: string; + readonly proof: { readonly challengeId: string; readonly signature: string }; + readonly challenge: SessionBrokerHelloChallenge; + readonly options: SessionBrokerHelloClientOptions; +} + +function clientAuthError(): never { + throw new Error( + "Session broker authentication failed or the daemon identity could not be verified.", + ); +} + +function exactRecord(value: unknown, keys: readonly string[]): Record { + if (!value || typeof value !== "object" || Array.isArray(value)) clientAuthError(); + const record = value as Record; + if ( + Object.keys(record).length !== keys.length || + keys.some((key) => !Object.hasOwn(record, key)) || + Object.keys(record).some((key) => !keys.includes(key)) + ) + clientAuthError(); + return record; +} + +function parseChallenge(value: unknown): SessionBrokerHelloChallenge { + const record = exactRecord(value, [ + "challengeId", + "generation", + "responderNonce", + "expiresAt", + "daemonKeyId", + "daemonSignature", + ]); + if ( + !isValidBrokerIdentifier(record.challengeId) || + !isValidBrokerIdentifier(record.generation) || + !isValidBrokerIdentifier(record.responderNonce) || + !Number.isFinite(record.expiresAt) || + typeof record.daemonKeyId !== "string" || + typeof record.daemonSignature !== "string" + ) + clientAuthError(); + return record as unknown as SessionBrokerHelloChallenge; +} + +function randomId(cryptoImpl: SessionBrokerCrypto) { + return `b_${encodeBase64Url(cryptoImpl.randomBytes(24))}_0`; +} + +function fixedProposal(appRevision: number): BrokerHelloProposal { + return { brokerRevision: SESSION_BROKER_PROTOCOL_REVISION, appRevision, features: [] }; +} + +/** Create the credential-free hello proposal that starts either producer or caller authentication. */ +export function createSessionBrokerHelloRequest( + options: SessionBrokerHelloClientOptions, +): SessionBrokerHelloChallengeRequest { + const cryptoImpl = options.crypto ?? webSessionBrokerCrypto; + return Object.freeze({ + role: options.credential.grant.kind, + appId: options.appId, + endpoint: options.endpoint, + keyId: options.credential.grant.keyId, + grantId: options.credential.grant.grantId, + initiatorNonce: randomId(cryptoImpl), + proposal: fixedProposal(options.appRevision), + }); +} + +/** Verify the daemon challenge before signing the same generation-bound transcript. */ +export async function answerSessionBrokerHelloChallenge( + options: SessionBrokerHelloClientOptions, + request: SessionBrokerHelloChallengeRequest, + challenge: SessionBrokerHelloChallenge, +): Promise> { + const cryptoImpl = options.crypto ?? webSessionBrokerCrypto; + if ( + challenge.daemonKeyId !== options.daemon.keyId || + !isValidBrokerIdentifier(challenge.challengeId) || + !isValidBrokerIdentifier(challenge.generation) || + !isValidBrokerIdentifier(challenge.responderNonce) || + !Number.isFinite(challenge.expiresAt) || + Date.now() >= challenge.expiresAt + ) + clientAuthError(); + const transcript = challengeTranscriptForClient(request, challenge, challenge.generation); + const daemonSignature = decodeBase64Url(challenge.daemonSignature); + if ( + !daemonSignature || + !(await cryptoImpl.verify(options.daemon.publicKey, daemonSignature, transcript)) + ) { + clientAuthError(); + } + const signature = encodeBase64Url( + await cryptoImpl.sign(options.credential.privateKey, transcript), + ); + return Object.freeze({ + request, + transcript, + transcriptHash: encodeBase64Url(await cryptoImpl.sha256(transcript)), + proof: Object.freeze({ challengeId: challenge.challengeId, signature }), + challenge, + options, + }); +} + +/** Verify a signed producer acknowledgement against the authenticated hello transcript. */ +export async function verifyProducerHelloAck( + pending: PendingSessionBrokerHello, + ack: AuthenticatedProducerHello, +): Promise { + exactRecord(ack, [ + "principal", + "connectionId", + "brokerRevision", + "appRevision", + "features", + "helloTranscriptHash", + "daemonKeyId", + "daemonSignature", + ]); + const grant = pending.options.credential.grant; + const principal = exactRecord(ack.principal, [ + "kind", + "appId", + "principalId", + "keyId", + "grantId", + "scopes", + ...(grant.sessionId ? ["sessionId"] : []), + ]); + const cryptoImpl = pending.options.crypto ?? webSessionBrokerCrypto; + if ( + principal.kind !== "producer" || + principal.appId !== grant.appId || + principal.principalId !== grant.principalId || + principal.keyId !== grant.keyId || + principal.grantId !== grant.grantId || + principal.sessionId !== grant.sessionId || + JSON.stringify(principal.scopes) !== JSON.stringify(grant.operations) || + ack.daemonKeyId !== pending.options.daemon.keyId || + ack.helloTranscriptHash !== pending.transcriptHash || + ack.brokerRevision !== SESSION_BROKER_PROTOCOL_REVISION || + ack.appRevision !== pending.options.appRevision || + !Array.isArray(ack.features) || + ack.features.length !== 0 || + !isValidBrokerIdentifier(ack.connectionId) + ) + clientAuthError(); + const signature = decodeBase64Url(ack.daemonSignature); + if ( + !signature || + !(await cryptoImpl.verify( + pending.options.daemon.publicKey, + signature, + buildBrokerHelloAckTranscript({ + role: "producer", + appId: pending.options.appId, + generation: pending.challenge.generation, + keyId: pending.options.credential.grant.keyId, + grantId: pending.options.credential.grant.grantId, + helloTranscriptHash: pending.transcriptHash, + selection: fixedProposal(pending.options.appRevision), + connectionId: ack.connectionId, + }), + )) + ) + clientAuthError(); +} + +export type SessionBrokerSignedRequestInit = Omit & { + readonly body?: string | null; +}; + +export interface SessionBrokerCallerClientOptions { + readonly appId: string; + readonly appRevision: number; + readonly origin: string; + readonly credential: SessionBrokerClientCredential; + readonly daemon: SessionBrokerDaemonVerifier; + readonly fetch?: typeof fetch; + readonly crypto?: SessionBrokerCrypto; + readonly challengePath?: string; + readonly proofPath?: string; + readonly maxResponseBytes?: number; +} + +/** Read one untrusted response through a strict byte ceiling before JSON decoding. */ +async function readBoundedResponseJson(response: Response, maxBytes: number): Promise { + const declared = response.headers.get("content-length"); + if (declared && (!/^(?:0|[1-9][0-9]*)$/.test(declared) || Number(declared) > maxBytes)) { + clientAuthError(); + } + if (!response.body) clientAuthError(); + const reader = response.body.getReader(); + const chunks: Uint8Array[] = []; + let total = 0; + try { + for (;;) { + const { done, value } = await reader.read(); + if (done) break; + total += value.byteLength; + if (total > maxBytes) { + await reader.cancel(); + clientAuthError(); + } + chunks.push(value); + } + } catch { + clientAuthError(); + } + const bytes = new Uint8Array(total); + let offset = 0; + for (const chunk of chunks) { + bytes.set(chunk, offset); + offset += chunk.byteLength; + } + try { + return JSON.parse(new TextDecoder("utf-8", { fatal: true }).decode(bytes)); + } catch { + clientAuthError(); + } +} + +/** Negotiates short-lived caller sessions and signs/verifies every exact HTTP control payload. */ +export class SessionBrokerCallerClient { + private session: AuthenticatedCallerSession | null = null; + private sequence: CallerSequenceAllocator | null = null; + private pending: PendingSessionBrokerHello | null = null; + private readonly fetchImpl: typeof fetch; + private readonly cryptoImpl: SessionBrokerCrypto; + + constructor(private readonly options: SessionBrokerCallerClientOptions) { + this.fetchImpl = options.fetch ?? fetch; + this.cryptoImpl = options.crypto ?? webSessionBrokerCrypto; + } + + /** Issue one signed request, renegotiating once after restart, expiry, or replay rejection. */ + async request( + path: string, + init: SessionBrokerSignedRequestInit = {}, + options: { readonly targetSpecific?: boolean } = {}, + ): Promise { + for (let attempt = 0; attempt < 2; attempt += 1) { + if (!this.session || Date.now() >= this.session.expiresAt) await this.negotiate(init.signal); + const response = await this.signedRequest(path, init, options.targetSpecific ?? false); + if (response === null) { + if (attempt === 0) { + this.clear(); + continue; + } + clientAuthError(); + } + return response; + } + clientAuthError(); + } + + clear() { + this.session = null; + this.sequence = null; + this.pending = null; + } + + private async negotiate(signal?: AbortSignal | null) { + const challengePath = this.options.challengePath ?? "/session-auth/challenge"; + const proofPath = this.options.proofPath ?? "/session-auth/proof"; + const endpoint = `${this.options.origin}${challengePath}`; + const helloOptions: SessionBrokerHelloClientOptions = { + appId: this.options.appId, + appRevision: this.options.appRevision, + endpoint, + credential: this.options.credential, + daemon: this.options.daemon, + crypto: this.cryptoImpl, + }; + const request = createSessionBrokerHelloRequest(helloOptions); + const challengeResponse = await this.fetchImpl(endpoint, { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify(request), + signal, + }); + if (!challengeResponse.ok) clientAuthError(); + const challenge = parseChallenge( + await readBoundedResponseJson( + challengeResponse, + this.options.maxResponseBytes ?? DEFAULT_SESSION_BROKER_LIMITS.maxHttpResponseBytes, + ), + ); + const pending = await answerSessionBrokerHelloChallenge(helloOptions, request, challenge); + const proofResponse = await this.fetchImpl(`${this.options.origin}${proofPath}`, { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify(pending.proof), + signal, + }); + if (!proofResponse.ok) clientAuthError(); + const sessionValue = await readBoundedResponseJson( + proofResponse, + this.options.maxResponseBytes ?? DEFAULT_SESSION_BROKER_LIMITS.maxHttpResponseBytes, + ); + const sessionRecord = exactRecord(sessionValue, [ + "callerSessionId", + "principal", + "expiresAt", + "initialSequence", + "brokerRevision", + "appRevision", + "features", + "helloTranscriptHash", + "daemonKeyId", + "daemonSignature", + ]); + const session = sessionRecord as unknown as AuthenticatedCallerSession; + await this.verifyCallerAck(pending, session); + this.pending = pending; + this.session = session; + this.sequence = new CallerSequenceAllocator(BigInt(session.initialSequence)); + } + + private async verifyCallerAck( + pending: PendingSessionBrokerHello, + session: AuthenticatedCallerSession, + ) { + const grant = this.options.credential.grant; + const principal = exactRecord(session.principal, [ + "kind", + "appId", + "principalId", + "keyId", + "grantId", + "operations", + "commands", + ...(grant.sessionId ? ["sessionId"] : []), + ]); + if ( + principal.kind !== "caller" || + principal.appId !== grant.appId || + principal.principalId !== grant.principalId || + principal.keyId !== grant.keyId || + principal.grantId !== grant.grantId || + principal.sessionId !== grant.sessionId || + JSON.stringify(principal.operations) !== JSON.stringify(grant.operations) || + JSON.stringify(principal.commands) !== JSON.stringify(grant.commands) || + session.daemonKeyId !== this.options.daemon.keyId || + session.helloTranscriptHash !== pending.transcriptHash || + session.brokerRevision !== SESSION_BROKER_PROTOCOL_REVISION || + session.appRevision !== this.options.appRevision || + !Array.isArray(session.features) || + session.features.length !== 0 || + !Number.isFinite(session.expiresAt) || + session.initialSequence !== "1" || + !isValidBrokerIdentifier(session.callerSessionId) + ) + clientAuthError(); + const signature = decodeBase64Url(session.daemonSignature); + if ( + !signature || + !(await this.cryptoImpl.verify( + this.options.daemon.publicKey, + signature, + buildBrokerHelloAckTranscript({ + role: "caller", + appId: this.options.appId, + generation: pending.challenge.generation, + keyId: this.options.credential.grant.keyId, + grantId: this.options.credential.grant.grantId, + helloTranscriptHash: pending.transcriptHash, + selection: fixedProposal(this.options.appRevision), + callerSessionId: session.callerSessionId, + initialSequence: session.initialSequence, + }), + )) + ) + clientAuthError(); + } + + private async signedRequest( + path: string, + init: SessionBrokerSignedRequestInit, + targetSpecific: boolean, + ) { + const session = this.session!; + const pending = this.pending!; + const sequence = this.sequence!.allocate(); + if (!sequence) clientAuthError(); + const method = (init.method ?? "GET").toUpperCase(); + const bodyBytes = + typeof init.body === "string" + ? new TextEncoder().encode(init.body) + : init.body == null + ? new Uint8Array() + : clientAuthError(); + const url = new URL(path, this.options.origin); + if ( + url.origin !== new URL(this.options.origin).origin || + url.username || + url.password || + url.hash + ) { + clientAuthError(); + } + const requestId = randomId(this.cryptoImpl); + const bodyDigest = encodeBase64Url(await this.cryptoImpl.sha256(bodyBytes)); + const signature = encodeBase64Url( + await this.cryptoImpl.sign( + this.options.credential.privateKey, + buildCallerRequestTranscript({ + appId: this.options.appId, + generation: pending.challenge.generation, + callerSessionId: session.callerSessionId, + keyId: this.options.credential.grant.keyId, + grantId: this.options.credential.grant.grantId, + helloTranscriptHash: pending.transcriptHash, + method, + target: canonicalHttpTarget(url), + bodyDigest, + requestId, + sequence, + }), + ), + ); + const headers = new Headers(init.headers); + headers.set("x-session-broker-caller-session", session.callerSessionId); + headers.set("x-session-broker-request-id", requestId); + headers.set("x-session-broker-sequence", sequence); + headers.set("x-session-broker-signature", signature); + const response = await this.fetchImpl(url, { ...init, method, headers }); + let envelope: SessionBrokerAuthenticatedResponse; + try { + envelope = (await readBoundedResponseJson( + response, + this.options.maxResponseBytes ?? DEFAULT_SESSION_BROKER_LIMITS.maxHttpResponseBytes, + )) as SessionBrokerAuthenticatedResponse; + await this.verifyResponse( + envelope, + response.status, + requestId, + pending.challenge.generation, + targetSpecific, + ); + } catch { + if (response.status === 401) return null; + throw new Error( + "Session broker authentication failed or the daemon identity could not be verified.", + ); + } + return new Response(JSON.stringify(envelope.body), { + status: response.status, + headers: { "content-type": "application/json" }, + }); + } + + private async verifyResponse( + envelope: SessionBrokerAuthenticatedResponse, + status: number, + requestId: string, + generation: string, + targetSpecific: boolean, + ) { + const envelopeRecord = exactRecord(envelope, ["body", "authentication"]); + const authenticationKeys = [ + "generation", + "brokerRevision", + "requestId", + "httpStatus", + "bodyDigest", + "daemonKeyId", + "daemonSignature", + ...(targetSpecific ? ["appContract"] : []), + ]; + const auth = exactRecord( + envelopeRecord.authentication, + authenticationKeys, + ) as unknown as SessionBrokerAuthenticatedResponse["authentication"]; + const appContract = auth.appContract + ? exactRecord(auth.appContract, ["appRevision", "features"]) + : undefined; + if ( + !auth || + typeof auth.bodyDigest !== "string" || + typeof auth.daemonSignature !== "string" || + auth.generation !== generation || + auth.requestId !== requestId || + auth.httpStatus !== status || + auth.brokerRevision !== SESSION_BROKER_PROTOCOL_REVISION || + auth.daemonKeyId !== this.options.daemon.keyId || + (targetSpecific ? !auth.appContract : auth.appContract !== undefined) + ) + clientAuthError(); + if ( + appContract && + (appContract.appRevision !== this.options.appRevision || + !Array.isArray(appContract.features) || + appContract.features.length !== 0) + ) + clientAuthError(); + const bodyDigest = encodeBase64Url( + await this.cryptoImpl.sha256(canonicalJsonBytes(envelopeRecord.body as CanonicalJsonValue)), + ); + if (bodyDigest !== auth.bodyDigest) clientAuthError(); + const signature = decodeBase64Url(auth.daemonSignature); + if ( + !signature || + !(await this.cryptoImpl.verify( + this.options.daemon.publicKey, + signature, + buildBrokerResponseTranscript({ + appId: this.options.appId, + generation, + brokerRevision: SESSION_BROKER_PROTOCOL_REVISION, + requestId, + httpStatus: status, + bodyDigest, + ...(auth.appContract ? { appContract: auth.appContract } : {}), + }), + )) + ) + clientAuthError(); + } +} diff --git a/packages/session-broker/src/connection.test.ts b/packages/session-broker/src/connection.test.ts index fe3dc41fb..b03a9bfae 100644 --- a/packages/session-broker/src/connection.test.ts +++ b/packages/session-broker/src/connection.test.ts @@ -1,10 +1,14 @@ import { describe, expect, test } from "bun:test"; import type { + ProducerGrant, SessionRegistration, SessionServerMessage, SessionSnapshot, } from "@hunk/session-broker-core"; -import { SESSION_BROKER_REGISTRATION_VERSION } from "@hunk/session-broker-core"; +import { + SESSION_BROKER_REGISTRATION_VERSION, + SESSION_BROKER_SIGNATURE_ALGORITHM, +} from "@hunk/session-broker-core"; import { createSessionBrokerConnection } from "./connection"; import { createSessionBrokerProtocolParsers } from "./protocolParsers"; import type { SessionBrokerSocketLike } from "./types"; @@ -149,6 +153,58 @@ describe("session broker connection", () => { }); }); + test("withholds registration and replacement updates until producer authentication completes", async () => { + const socket = new TestSocket(); + const pair = (await crypto.subtle.generateKey("Ed25519", false, [ + "sign", + "verify", + ])) as CryptoKeyPair; + const grant: ProducerGrant = { + kind: "producer", + appId: "dev.example", + principalId: "producer-1", + keyId: "producer-key-1", + grantId: "producer-grant-1", + algorithm: SESSION_BROKER_SIGNATURE_ALGORITHM, + issuedAt: Date.now() - 1_000, + expiresAt: Date.now() + 60_000, + revocationId: "producer-revocation-1", + mayDelegate: false, + operations: ["register", "reconnect"], + }; + const connection = createSessionBrokerConnection< + TestSessionInfo, + TestSessionState, + TestSocket, + TestServerMessage, + { ok: true } + >({ + url: "ws://broker.test/session", + createSocket: () => socket, + registration: createRegistration(), + snapshot: createSnapshot(), + protocolParsers, + producerAuthentication: { + appId: "dev.example", + appRevision: 1, + credential: { grant, privateKey: pair.privateKey }, + daemon: { keyId: "daemon-key-1", publicKey: pair.publicKey }, + }, + }); + + connection.start(); + socket.emitOpen(); + connection.updateSnapshot({ + ...createSnapshot(), + state: { selectedIndex: 2 }, + }); + connection.replaceSession(createRegistration(), createSnapshot()); + + expect(socket.sent).toHaveLength(1); + expect(JSON.parse(socket.sent[0]!)).toMatchObject({ type: "hello-init" }); + connection.stop(); + }); + test("keeps the previous registration when replacement send throws", () => { const socket = new TestSocket(); const registration = createRegistration(); diff --git a/packages/session-broker/src/connection.ts b/packages/session-broker/src/connection.ts index dbc72ac89..c23113c61 100644 --- a/packages/session-broker/src/connection.ts +++ b/packages/session-broker/src/connection.ts @@ -15,6 +15,20 @@ import { } from "@hunk/session-broker-core"; import type { SessionBrokerProtocolParsers } from "./protocolParsers"; import { parseSessionBrokerJsonText } from "./protocolParsers"; +import { + answerSessionBrokerHelloChallenge, + createSessionBrokerHelloRequest, + verifyProducerHelloAck, + type PendingSessionBrokerHello, + type SessionBrokerClientCredential, + type SessionBrokerDaemonVerifier, +} from "./clientAuthentication"; +import type { + AuthenticatedProducerHello, + SessionBrokerHelloChallenge, + SessionBrokerHelloChallengeRequest, +} from "./authentication"; +import type { ProducerGrant } from "@hunk/session-broker-core"; import type { SessionBrokerConnectionCloseDirective, SessionBrokerSocketCloseEvent, @@ -59,6 +73,12 @@ export interface SessionBrokerConnectionOptions< snapshot: SessionSnapshot; bridge?: SessionBrokerConnectionBridge | null; protocolParsers: SessionBrokerProtocolParsers; + producerAuthentication?: { + readonly appId: string; + readonly appRevision: number; + readonly credential: SessionBrokerClientCredential; + readonly daemon: SessionBrokerDaemonVerifier; + }; heartbeatIntervalMs?: number; reconnectDelayMs?: number; openState?: number; @@ -80,6 +100,7 @@ export class SessionBrokerConnection< Result = unknown, > { private socket: Socket | null = null; + private activeSocket: Socket | null = null; private bridge: SessionBrokerConnectionBridge | null; readonly limits: Readonly; @@ -93,6 +114,14 @@ export class SessionBrokerConnection< private stopped = false; private registration: SessionRegistration; private snapshot: SessionSnapshot; + private readonly handshakeTimers = new WeakMap>(); + private readonly producerHellos = new WeakMap< + Socket, + { + request: SessionBrokerHelloChallengeRequest; + pending?: PendingSessionBrokerHello | null; + } + >(); constructor( private readonly options: SessionBrokerConnectionOptions< @@ -141,8 +170,14 @@ export class SessionBrokerConnection< } this.stopHeartbeat(); - this.socket?.close(); + if (this.socket) { + const handshakeTimer = this.handshakeTimers.get(this.socket); + if (handshakeTimer) clearTimeout(handshakeTimer); + this.handshakeTimers.delete(this.socket); + this.socket.close(); + } this.socket = null; + this.activeSocket = null; } getRegistration() { @@ -155,6 +190,12 @@ export class SessionBrokerConnection< } replaceSession(registration: SessionRegistration, snapshot: SessionSnapshot) { + if ( + this.options.producerAuthentication && + registration.sessionId !== this.registration.sessionId + ) { + throw new BrokerProtocolError("invalid-app-payload"); + } // Re-register instead of sending only a snapshot because selectors like cwd, repoRoot, and the // session id itself live in the registration envelope. Send before committing local state so // a throwing socket keeps the previous registration and snapshot coherent. @@ -183,19 +224,33 @@ export class SessionBrokerConnection< const socket = this.options.createSocket(this.options.url); this.socket = socket; + if (this.options.producerAuthentication) { + const timer = setTimeout(() => { + socket.close(1008, "Session broker authentication timed out."); + }, this.limits.maxHandshakeDurationMs); + timer.unref?.(); + this.handshakeTimers.set(socket, timer); + } socket.onopen = () => { - this.startHeartbeat(); - // Register on every fresh socket after the prior close retired its broker-side ownership. - this.sendToSocket(socket, { - type: "register", - registration: this.registration, - snapshot: this.snapshot, - }); - void this.flushQueuedMessages(socket); + if (this.options.producerAuthentication) { + const authentication = this.options.producerAuthentication; + const request = createSessionBrokerHelloRequest({ + ...authentication, + endpoint: this.options.url, + }); + this.producerHellos.set(socket, { request }); + socket.send(JSON.stringify({ type: "hello-init", hello: request })); + return; + } + this.activateSocket(socket); }; socket.onmessage = (event) => { + if (this.options.producerAuthentication && this.activeSocket !== socket) { + void this.handleProducerHello(socket, event.data); + return; + } let parsed: ServerMessage; try { const raw = parseSessionBrokerJsonText(event.data) as { input?: unknown }; @@ -217,8 +272,12 @@ export class SessionBrokerConnection< }; socket.onclose = (event) => { + const handshakeTimer = this.handshakeTimers.get(socket); + if (handshakeTimer) clearTimeout(handshakeTimer); + this.handshakeTimers.delete(socket); if (this.socket === socket) { this.socket = null; + this.activeSocket = null; this.stopHeartbeat(); } @@ -250,6 +309,58 @@ export class SessionBrokerConnection< }; } + private activateSocket(socket: Socket) { + if (this.socket !== socket || this.activeSocket === socket) return; + this.activeSocket = socket; + const handshakeTimer = this.handshakeTimers.get(socket); + if (handshakeTimer) clearTimeout(handshakeTimer); + this.handshakeTimers.delete(socket); + this.startHeartbeat(); + this.sendToSocket(socket, { + type: "register", + registration: this.registration, + snapshot: this.snapshot, + }); + void this.flushQueuedMessages(socket); + } + + /** Verify the daemon challenge and acknowledgement before registration leaves this process. */ + private async handleProducerHello(socket: Socket, message: unknown) { + try { + if (typeof message === "string" && utf8ByteLength(message) > this.limits.maxWsMessageBytes) { + socket.close(1009, "Session broker authentication message exceeded its limit."); + return; + } + const value = parseSessionBrokerJsonText(message) as Record; + const authentication = this.options.producerAuthentication!; + const hello = this.producerHellos.get(socket); + if (!hello) throw new Error(); + if (hello.pending === undefined) { + if (value?.type !== "hello-challenge") throw new Error(); + hello.pending = null; + const pending = await answerSessionBrokerHelloChallenge( + { ...authentication, endpoint: this.options.url }, + hello.request, + value.challenge as SessionBrokerHelloChallenge, + ); + if ( + this.socket !== socket || + socket.readyState !== (this.options.openState ?? DEFAULT_SOCKET_OPEN_STATE) + ) { + return; + } + hello.pending = pending; + socket.send(JSON.stringify({ type: "hello-proof", proof: pending.proof })); + return; + } + if (!hello.pending || value?.type !== "hello-ack") throw new Error(); + await verifyProducerHelloAck(hello.pending, value.ack as AuthenticatedProducerHello); + this.activateSocket(socket); + } catch { + socket.close(1008, "Session broker authentication failed."); + } + } + private scheduleReconnect(delayMs = this.options.reconnectDelayMs ?? DEFAULT_RECONNECT_DELAY_MS) { if (this.reconnectTimer || this.stopped) { return; @@ -288,17 +399,18 @@ export class SessionBrokerConnection< } private send(message: SessionClientMessage) { - if (!this.socket) { + if (!this.activeSocket) { return; } - this.sendToSocket(this.socket, message); + this.sendToSocket(this.activeSocket, message); } /** Send a response only through the still-active socket that received its command. */ private sendToSocket(socket: Socket, message: SessionClientMessage) { if ( this.socket !== socket || + this.activeSocket !== socket || socket.readyState !== (this.options.openState ?? DEFAULT_SOCKET_OPEN_STATE) ) { return; diff --git a/packages/session-broker/src/daemon.test.ts b/packages/session-broker/src/daemon.test.ts index 11cfb03e3..2120b05c8 100644 --- a/packages/session-broker/src/daemon.test.ts +++ b/packages/session-broker/src/daemon.test.ts @@ -5,6 +5,7 @@ import { parseSessionRegistrationEnvelope, parseSessionSnapshotEnvelope, type CallerPrincipal, + type ProducerOperation, type SessionRegistration, type SessionServerMessage, type SessionSnapshot, @@ -12,7 +13,11 @@ import { import { SessionBroker } from "./broker"; import { createSessionBrokerDaemon } from "./daemon"; import { createSessionBrokerProtocolParsers } from "./protocolParsers"; -import type { AuthenticatedCallerRequest } from "./authentication"; +import type { + AuthenticatedCallerRequest, + AuthenticatedProducerHello, + SessionBrokerHelloChallenge, +} from "./authentication"; interface TestSessionInfo { title: string; @@ -170,9 +175,13 @@ async function authenticatedBody(response: Response | null) { function createConnection() { const sent: string[] = []; let closed: { code?: number; reason?: string } | null = null; + let authenticated = false; return { sent, + get authenticated() { + return authenticated; + }, get closed() { return closed; }, @@ -183,6 +192,9 @@ function createConnection() { close(code?: number, reason?: string) { closed = { code, reason }; }, + markAuthenticated() { + authenticated = true; + }, }, }; } @@ -605,6 +617,81 @@ describe("session broker daemon", () => { daemon.shutdown(); }); + test("requires reconnect scope before atomically replacing an authenticated owner", async () => { + let operations: readonly ProducerOperation[] = ["register"]; + const principal = () => ({ + kind: "producer" as const, + appId: "dev.example", + principalId: "producer-1", + keyId: "producer-key-1", + grantId: "producer-grant-1", + scopes: operations, + }); + const daemon = createSessionBrokerDaemon({ + broker: createBroker(), + appId: "dev.example", + appRevision: 1, + producerEndpoint: "ws://broker.test/session", + helloAuthenticator: { + async issueChallenge() { + return { challengeId: "challenge-1" } as SessionBrokerHelloChallenge; + }, + async completeCallerHello() { + throw new Error("not used"); + }, + async completeProducerHello(_proof, connectionId) { + return { + principal: principal(), + connectionId: String(connectionId), + brokerRevision: 1, + appRevision: 1, + features: [], + helloTranscriptHash: "transcript-1", + daemonKeyId: "daemon-key-1", + daemonSignature: "signature-1", + } satisfies AuthenticatedProducerHello; + }, + assertProducerActive() {}, + }, + }); + const first = createConnection(); + const denied = createConnection(); + const replacement = createConnection(); + const authenticate = async (connection: ReturnType["connection"]) => { + daemon.handleConnectionMessage(connection, JSON.stringify({ type: "hello-init", hello: {} })); + await Bun.sleep(0); + daemon.handleConnectionMessage( + connection, + JSON.stringify({ type: "hello-proof", proof: {} }), + ); + await Bun.sleep(0); + }; + const register = (connection: ReturnType["connection"]) => + daemon.handleConnectionMessage( + connection, + JSON.stringify({ + type: "register", + registration: createRegistration(), + snapshot: createSnapshot(), + }), + ); + + await authenticate(first.connection); + expect(first.authenticated).toBe(false); + register(first.connection); + expect(first.authenticated).toBe(true); + await authenticate(denied.connection); + register(denied.connection); + expect(denied.closed?.reason).toContain("scope rejected"); + expect(first.closed).toBeNull(); + + operations = ["reconnect"]; + await authenticate(replacement.connection); + register(replacement.connection); + expect(first.closed?.reason).toContain("owner reconnected"); + expect(daemon.listSessions()).toHaveLength(1); + }); + test("rejects duplicate live registration without retiring the owner", () => { const daemon = createSessionBrokerDaemon({ broker: createBroker(), @@ -845,6 +932,15 @@ describe("session broker daemon", () => { return true; }, }); + const owner = createConnection(); + daemon.handleConnectionMessage( + owner.connection, + JSON.stringify({ + type: "register", + registration: createRegistration(), + snapshot: createSnapshot(), + }), + ); const post = (body: unknown) => daemon.handleRequest( new Request("http://broker.test/broker", { @@ -854,13 +950,13 @@ describe("session broker daemon", () => { }), ); - expect((await post({ action: "get", selector: { sessionId: "missing" } }))?.status).toBe(403); + expect((await post({ action: "get", selector: { sessionId: "session-1" } }))?.status).toBe(403); expect(appAuthorizerCalls).toBe(0); expect( ( await post({ action: "dispatch", - selector: { sessionId: "missing" }, + selector: { sessionId: "session-1" }, command: "forbidden", input: {}, }) @@ -871,7 +967,7 @@ describe("session broker daemon", () => { ( await post({ action: "dispatch", - selector: { sessionId: "missing" }, + selector: { sessionId: "session-1" }, command: "allowed", commandVersion: 0, input: {}, diff --git a/packages/session-broker/src/daemon.ts b/packages/session-broker/src/daemon.ts index a555b3a6e..621c505fa 100644 --- a/packages/session-broker/src/daemon.ts +++ b/packages/session-broker/src/daemon.ts @@ -8,6 +8,7 @@ import { mergeSessionBrokerLimits, DEFAULT_SESSION_BROKER_LIMITS, callerPrincipalAllows, + producerPrincipalAllows, canonicalizeJson, isValidBrokerAppId, isValidBrokerIdentifier, @@ -17,6 +18,7 @@ import { type CallerOperation, type CallerPrincipal, type CanonicalJsonValue, + type ProducerPrincipal, type SessionBrokerLimitOptions, type SessionBrokerLimits, type SessionServerMessage, @@ -27,6 +29,7 @@ import { SessionBrokerAuthenticationError, type AuthenticatedCallerRequest, type CallerRequestAuthenticator, + type SessionBrokerHelloAuthenticator, } from "./authentication"; import { parseSessionBrokerJsonBytes, @@ -65,6 +68,19 @@ const BROKER_STATE_LIMITS = [ "maxCommandTimeoutMs", ] as const satisfies readonly (keyof SessionBrokerLimits)[]; +export interface SessionBrokerAuthenticatedControlFacts { + readonly operation: CallerOperation; + readonly sessionId?: string; + readonly command?: string; + readonly commandVersion?: number; + readonly targetSpecific?: boolean; +} + +export interface SessionBrokerAuthenticatedControlResult { + readonly body: CanonicalJsonValue; + readonly status?: number; +} + export interface SessionBrokerDaemonOptions< SessionView = unknown, ServerMessage extends SessionServerMessage = SessionServerMessage, @@ -75,6 +91,10 @@ export interface SessionBrokerDaemonOptions< paths?: Partial; exposeHttpApi?: boolean; callerAuthenticator?: CallerRequestAuthenticator; + helloAuthenticator?: SessionBrokerHelloAuthenticator; + /** @deprecated Use helloAuthenticator. */ + producerAuthenticator?: SessionBrokerHelloAuthenticator; + producerEndpoint?: string; authorizer?: SessionBrokerAuthorizer; audit?: SessionBrokerAuditHook; appId?: string; @@ -114,6 +134,17 @@ function defaultTimeoutMessage(command: string) { return `Timed out waiting for the session to handle ${command}.`; } +/** Match the immutable producer identity that is allowed to reclaim one session. */ +function sameProducerBinding(left: ProducerPrincipal, right: ProducerPrincipal) { + return ( + left.appId === right.appId && + left.principalId === right.principalId && + left.keyId === right.keyId && + left.grantId === right.grantId && + left.sessionId === right.sessionId + ); +} + /** * Runtime-neutral daemon engine that owns broker lifecycle, health, stale pruning, and raw HTTP * plus websocket message handling without choosing Bun, Node, or any other server implementation. @@ -142,7 +173,21 @@ export class SessionBrokerDaemon< private readonly appId: string; private readonly appRevision?: number; private readonly callerAuthenticator?: CallerRequestAuthenticator; + private readonly helloAuthenticator?: SessionBrokerHelloAuthenticator; + private readonly producerEndpoint?: string; private readonly authorizer?: SessionBrokerAuthorizer; + private readonly producerAuthentication = new WeakMap< + SessionBrokerPeer, + { state: "challenged" | "authenticated"; principal?: ProducerPrincipal; sessionId?: string } + >(); + private readonly producerOwners = new Map< + string, + { connection: SessionBrokerPeer; principal: ProducerPrincipal } + >(); + private readonly producerReconnects = new Map< + string, + { principal: ProducerPrincipal; disconnectedAt: number } + >(); private readonly audit?: SessionBrokerAuditHook; private readonly httpControlBudget: ResourceBudget; private readonly httpBodyBudget: ResourceBudget; @@ -202,6 +247,14 @@ export class SessionBrokerDaemon< this.appId = options.appId ?? "session-broker"; this.appRevision = this.protocolParsers.appRevision; this.callerAuthenticator = options.callerAuthenticator; + this.helloAuthenticator = options.helloAuthenticator ?? options.producerAuthenticator; + this.producerEndpoint = options.producerEndpoint; + if (options.producerAuthenticator && !this.producerEndpoint) { + throw new TypeError("Authenticated producer transport requires its listener endpoint."); + } + if (this.producerEndpoint && !this.helloAuthenticator) { + throw new TypeError("Authenticated producer transport requires a hello authenticator."); + } this.authorizer = options.authorizer; this.audit = options.audit; this.idleTimeoutMs = options.idleTimeoutMs ?? DEFAULT_IDLE_TIMEOUT_MS; @@ -240,6 +293,10 @@ export class SessionBrokerDaemon< return pathname === this.paths.socket; } + get requiresProducerAuthentication() { + return this.producerEndpoint !== undefined; + } + /** Run one app-specific finite HTTP control through the daemon's shared count/body budgets. */ async handleBoundedControl( request: Request, @@ -283,6 +340,27 @@ export class SessionBrokerDaemon< async handleRequest(request: Request) { const url = new URL(request.url); + if (url.pathname === "/session-auth/challenge" || url.pathname === "/session-auth/proof") { + if (request.method !== "POST" || !hasJsonContentType(request) || !this.helloAuthenticator) { + return jsonError("Session broker authentication requires an upgraded client.", 401); + } + return this.handleBoundedControl(request, async (body) => { + try { + const input = parseSessionBrokerJsonBytes(body); + const result = url.pathname.endsWith("/challenge") + ? await this.helloAuthenticator!.issueChallenge(input, request.url) + : await this.helloAuthenticator!.completeCallerHello(input); + return Response.json(result); + } catch (error) { + const code = + error instanceof SessionBrokerAuthenticationError + ? error.code + : "authentication-required"; + return Response.json({ error: code }, { status: 401 }); + } + }); + } + if (url.pathname === this.paths.health) { // Treat health checks as a cheap maintenance pulse so stale sessions disappear even when the // daemon is mostly idle and no websocket traffic is flowing. @@ -292,6 +370,7 @@ export class SessionBrokerDaemon< if (removed > 0) { this.noteActivity(); } + this.reconcileProducerOwners(); // Public health is deliberately liveness-only. Apps may expose authenticated diagnostics on // a separate route, but broker identity, paths, counts, and process facts stay private. @@ -311,6 +390,65 @@ export class SessionBrokerDaemon< } handleConnectionMessage(connection: SessionBrokerPeer, message: unknown) { + if (typeof message === "string" && utf8ByteLength(message) > this.limits.maxWsMessageBytes) { + connection.close?.(1009, "Session broker message exceeded its limit."); + return; + } + if (this.producerEndpoint && this.helloAuthenticator) { + const authentication = this.producerAuthentication.get(connection); + if (authentication?.state !== "authenticated") { + void this.handleProducerHelloMessage(connection, message, authentication); + return; + } + } + this.handleAuthenticatedConnectionMessage(connection, message); + } + + /** Complete the producer hello before allowing any registration-shaped message to reach state. */ + private async handleProducerHelloMessage( + connection: SessionBrokerPeer, + message: unknown, + current?: { + state: "challenged" | "authenticated"; + principal?: ProducerPrincipal; + sessionId?: string; + }, + ) { + try { + const value = parseSessionBrokerJsonText(message) as Record; + if (!current) { + if (value?.type !== "hello-init" || !("hello" in value)) throw new Error(); + const challenged = { state: "challenged" as const }; + this.producerAuthentication.set(connection, challenged); + const challenge = await this.helloAuthenticator!.issueChallenge( + value.hello, + this.producerEndpoint!, + ); + if (this.producerAuthentication.get(connection) !== challenged) return; + connection.send(JSON.stringify({ type: "hello-challenge", challenge })); + return; + } + if (current.state !== "challenged" || value?.type !== "hello-proof" || !("proof" in value)) { + throw new Error(); + } + const connectionId = `b_${crypto.randomUUID().replaceAll("-", "")}_0`; + const ack = await this.helloAuthenticator!.completeProducerHello(value.proof, connectionId); + if (this.producerAuthentication.get(connection) !== current) return; + this.producerAuthentication.set(connection, { + state: "authenticated", + principal: ack.principal, + }); + connection.send(JSON.stringify({ type: "hello-ack", ack })); + } catch { + this.producerAuthentication.delete(connection); + connection.close?.( + INCOMPATIBLE_PAYLOAD_CLOSE_CODE, + "Session broker authentication required; upgrade Hunk.", + ); + } + } + + private handleAuthenticatedConnectionMessage(connection: SessionBrokerPeer, message: unknown) { let parsed; try { parsed = this.protocolParsers.parseClientMessage(parseSessionBrokerJsonText(message)); @@ -319,12 +457,47 @@ export class SessionBrokerDaemon< return; } + const producerAuthentication = this.producerAuthentication.get(connection); + if (this.producerEndpoint && this.helloAuthenticator && producerAuthentication?.principal) { + try { + this.helloAuthenticator.assertProducerActive(producerAuthentication.principal); + } catch { + connection.close?.(INCOMPATIBLE_PAYLOAD_CLOSE_CODE, "Session producer authority expired."); + return; + } + } switch (parsed.type) { case "register": { + const sessionId = (parsed.registration as { sessionId: string }).sessionId; + this.pruneProducerReconnects(); + const owner = this.producerOwners.get(sessionId); + const reconnect = + owner && owner.connection !== connection ? owner : this.producerReconnects.get(sessionId); + const operation = reconnect ? "reconnect" : "register"; + if ( + this.producerEndpoint && + this.helloAuthenticator && + (!producerAuthentication?.principal || + (producerAuthentication.sessionId !== undefined && + producerAuthentication.sessionId !== sessionId) || + (reconnect && + !sameProducerBinding(producerAuthentication.principal, reconnect.principal)) || + !producerPrincipalAllows(producerAuthentication.principal, { + appId: this.appId, + operation, + sessionId, + })) + ) { + connection.close?.(INCOMPATIBLE_PAYLOAD_CLOSE_CODE, "Session producer scope rejected."); + return; + } + if (producerAuthentication) producerAuthentication.sessionId = sessionId; + const replacedConnection = owner?.connection !== connection ? owner?.connection : undefined; const registrationResult = this.broker.registerSession( connection, parsed.registration, parsed.snapshot, + { replaceOwner: replacedConnection !== undefined }, ); if (registrationResult === "invalid") { // Close immediately when the registration payload is incompatible so the session does not @@ -342,10 +515,23 @@ export class SessionBrokerDaemon< return; } + if (producerAuthentication?.principal) { + this.producerOwners.set(sessionId, { + connection, + principal: producerAuthentication.principal, + }); + this.producerReconnects.delete(sessionId); + } + connection.markAuthenticated?.(); + replacedConnection?.close?.(1000, "Session owner reconnected."); this.noteActivity(); break; } case "snapshot": { + if (this.producerEndpoint && producerAuthentication?.sessionId !== parsed.sessionId) { + connection.close?.(INCOMPATIBLE_PAYLOAD_CLOSE_CODE, "Session producer scope rejected."); + return; + } // Snapshot updates are only valid after registration. Closing missing or invalid sessions // keeps the broker state single-sourced instead of guessing how to recover. const updateResult = this.broker.updateSnapshot( @@ -371,6 +557,10 @@ export class SessionBrokerDaemon< break; } case "heartbeat": { + if (this.producerEndpoint && producerAuthentication?.sessionId !== parsed.sessionId) { + connection.close?.(INCOMPATIBLE_PAYLOAD_CLOSE_CODE, "Session producer scope rejected."); + return; + } const seenResult = this.broker.markSessionSeen(connection, parsed.sessionId); if (seenResult === "not-owner") { connection.close?.(INCOMPATIBLE_PAYLOAD_CLOSE_CODE, "Session ownership rejected."); @@ -401,10 +591,56 @@ export class SessionBrokerDaemon< } handleConnectionClose(connection: SessionBrokerPeer) { + const authentication = this.producerAuthentication.get(connection); + this.producerAuthentication.delete(connection); + const sessionId = authentication?.sessionId; + if (sessionId && authentication.principal) { + const owner = this.producerOwners.get(sessionId); + if (owner?.connection === connection) { + this.producerOwners.delete(sessionId); + this.pruneProducerReconnects(); + if (this.producerReconnects.size >= this.limits.maxSessions) { + const oldest = this.producerReconnects.keys().next().value as string | undefined; + if (oldest) this.producerReconnects.delete(oldest); + } + this.producerReconnects.set(sessionId, { + principal: authentication.principal, + disconnectedAt: Date.now(), + }); + } + } this.broker.unregisterConnection(connection); this.noteActivity(); } + /** Retire producer sockets whose session vanished or whose configured grant is no longer active. */ + private reconcileProducerOwners() { + const live = new Set(this.broker.getSessionIds()); + for (const [sessionId, owner] of this.producerOwners) { + let active = live.has(sessionId); + if (active && this.helloAuthenticator) { + try { + this.helloAuthenticator.assertProducerActive(owner.principal); + } catch { + active = false; + } + } + if (active) continue; + this.producerOwners.delete(sessionId); + this.broker.unregisterConnection(owner.connection); + owner.connection.close?.(1000, "Session producer authority retired."); + } + } + + /** Expire bounded reconnect authority on the same horizon as disconnected session state. */ + private pruneProducerReconnects(now = Date.now()) { + for (const [sessionId, reconnect] of this.producerReconnects) { + if (now - reconnect.disconnectedAt >= this.staleSessionTtlMs) { + this.producerReconnects.delete(sessionId); + } + } + } + shutdown(error = new Error("The session broker daemon shut down.")) { if (this.shuttingDown) { return; @@ -422,6 +658,8 @@ export class SessionBrokerDaemon< } this.broker.shutdown(error); + this.producerOwners.clear(); + this.producerReconnects.clear(); this.callerAuthenticator?.clear?.(); this.resolveStopped?.(); this.resolveStopped = null; @@ -435,6 +673,7 @@ export class SessionBrokerDaemon< if (removed > 0) { this.noteActivity(); } + this.reconcileProducerOwners(); }, this.staleSessionSweepIntervalMs); this.sweepTimer.unref?.(); @@ -483,6 +722,69 @@ export class SessionBrokerDaemon< }, remainingMs); } + /** Authenticate, authorize, execute, and sign one app-owned finite JSON control. */ + async handleAuthenticatedControl( + request: Request, + options: { + resolve: (body: Uint8Array) => SessionBrokerAuthenticatedControlFacts; + resolveFailureTargetSpecific?: (body: Uint8Array) => boolean; + handle: ( + body: Uint8Array, + facts: SessionBrokerAuthenticatedControlFacts, + ) => + | SessionBrokerAuthenticatedControlResult + | Promise; + }, + ): Promise { + return this.handleBoundedControl(request, async (body) => { + const authenticated = await this.authenticateRequest(request, body, "list"); + if (authenticated instanceof Response) return authenticated; + let facts: SessionBrokerAuthenticatedControlFacts; + try { + facts = options.resolve(body); + } catch { + let targetSpecific = false; + try { + targetSpecific = options.resolveFailureTargetSpecific?.(body) ?? false; + } catch { + // Malformed bodies have no trustworthy target contract. + } + return this.authenticatedResponse( + authenticated, + { error: "protocol-validation-failed" }, + 400, + targetSpecific, + ); + } + if (!(await this.authorize(request, authenticated, facts))) { + return this.authenticatedResponse( + authenticated, + { error: "authorization-denied" }, + 403, + facts.targetSpecific ?? facts.operation !== "list", + ); + } + const inactive = this.rejectInactiveRequest(authenticated); + if (inactive) return inactive; + try { + const result = await options.handle(body, facts); + return this.authenticatedResponse( + authenticated, + result.body, + result.status ?? 200, + facts.targetSpecific ?? facts.operation !== "list", + ); + } catch { + return this.authenticatedResponse( + authenticated, + { error: "session-control-failed" }, + 400, + facts.targetSpecific ?? facts.operation !== "list", + ); + } + }); + } + private async authenticateRequest( request: Request, body: Uint8Array, @@ -738,7 +1040,20 @@ export class SessionBrokerDaemon< const operation = input.action as CallerOperation; const selector = "selector" in input ? input.selector : undefined; - const sessionId = selector?.sessionId; + const targetSpecific = input.action !== "list"; + let sessionId: string | undefined; + if (selector) { + try { + sessionId = this.broker.resolveSessionId(selector); + } catch (error) { + return this.authenticatedResponse( + authenticated, + { error: error instanceof Error ? error.message : "Session target resolution failed." }, + 400, + true, + ); + } + } const command = input.action === "dispatch" ? input.command : undefined; const commandVersion = input.action === "dispatch" ? (input.commandVersion ?? 1) : undefined; const facts = { @@ -746,7 +1061,6 @@ export class SessionBrokerDaemon< ...(sessionId !== undefined ? { sessionId } : {}), ...(command !== undefined ? { command, commandVersion } : {}), }; - const targetSpecific = input.action !== "list"; if (!(await this.authorize(request, authenticated, facts))) { return this.authenticatedResponse( authenticated, @@ -765,15 +1079,12 @@ export class SessionBrokerDaemon< response = { sessions: this.broker.listSessions() }; break; case "get": - response = { session: this.broker.getSession(input.selector) }; + response = { session: this.broker.getSession({ sessionId: sessionId! }) }; break; case "dispatch": { - // Resolve the target before invoking app-owned parsing so the exact target contract is - // selected first. This read-only lookup happens only after authentication/authorization. - this.broker.getSession(input.selector); response = { result: await this.broker.dispatchCommand({ - selector: input.selector, + selector: { sessionId: sessionId! }, command: input.command, commandVersion: input.commandVersion ?? 1, input: input.input, diff --git a/packages/session-broker/src/index.ts b/packages/session-broker/src/index.ts index 7e0be3777..dddd0821f 100644 --- a/packages/session-broker/src/index.ts +++ b/packages/session-broker/src/index.ts @@ -5,4 +5,5 @@ export * from "./daemon"; export * from "./connection"; export * from "./crypto"; export * from "./authentication"; +export * from "./clientAuthentication"; export * from "./protocolParsers"; diff --git a/src/main.tsx b/src/main.tsx index 23ab6c3b7..69b784fda 100644 --- a/src/main.tsx +++ b/src/main.tsx @@ -16,7 +16,7 @@ async function main() { } if (startupPlan.kind === "daemon-serve") { - const server = serveSessionBrokerDaemon(); + const server = await serveSessionBrokerDaemon(); await server.stopped; return; } diff --git a/src/session/agent/cliClient.test.ts b/src/session/agent/cliClient.test.ts index ab4738856..b58b0a9f9 100644 --- a/src/session/agent/cliClient.test.ts +++ b/src/session/agent/cliClient.test.ts @@ -34,6 +34,9 @@ import { const selector = { sessionId: "session-1" } satisfies SessionSelectorInput; const originalFetch = globalThis.fetch; +const injectedCaller = { + request: (path: string, init?: RequestInit) => globalThis.fetch(path, init), +}; afterEach(() => { globalThis.fetch = originalFetch; @@ -134,7 +137,7 @@ describe("HTTP Hunk session CLI client", () => { return Response.json(responses[request.action as keyof typeof responses]); }) as typeof fetch; - const client = createHttpHunkSessionCliClient(); + const client = createHttpHunkSessionCliClient({ caller: injectedCaller }); expect(await client.getCapabilities()).toMatchObject({ version: HUNK_SESSION_API_VERSION }); expect(await client.listSessions()).toEqual([session]); @@ -327,7 +330,7 @@ describe("HTTP Hunk session CLI client", () => { }); }) as typeof fetch; - const client = createHttpHunkSessionCliClient({ timeoutMs: 10 }); + const client = createHttpHunkSessionCliClient({ timeoutMs: 10, caller: injectedCaller }); await expect(client.listSessions()).rejects.toThrow( "Timed out waiting for the Hunk session daemon to complete session list.", @@ -340,7 +343,7 @@ describe("HTTP Hunk session CLI client", () => { sessions: [{ sessionId: "partial", unknown: true }], })) as unknown as typeof fetch; - const client = createHttpHunkSessionCliClient(); + const client = createHttpHunkSessionCliClient({ caller: injectedCaller }); await expect(client.listSessions()).rejects.toThrow( "Invalid Hunk session daemon response for list.", ); @@ -356,7 +359,7 @@ describe("HTTP Hunk session CLI client", () => { globalThis.fetch = (async () => Response.json({ sessions: [session] })) as unknown as typeof fetch; - const client = createHttpHunkSessionCliClient(); + const client = createHttpHunkSessionCliClient({ caller: injectedCaller }); const result = await client.listSessions(); expect(result).toEqual([session]); expect(result[0]).not.toBe(session); @@ -369,7 +372,7 @@ describe("HTTP Hunk session CLI client", () => { { status: 404, statusText: "Not Found" }, )) as unknown as typeof fetch; - const client = createHttpHunkSessionCliClient(); + const client = createHttpHunkSessionCliClient({ caller: injectedCaller }); await expect(client.listSessions()).rejects.toThrow("No matching session."); globalThis.fetch = (async () => diff --git a/src/session/agent/cliClient.ts b/src/session/agent/cliClient.ts index ca22c0f56..0942e9fa3 100644 --- a/src/session/agent/cliClient.ts +++ b/src/session/agent/cliClient.ts @@ -1,19 +1,28 @@ import { sanitizeTerminalText } from "../../lib/terminalText"; import { resolveSessionBrokerConfig } from "../broker/brokerConfig"; +import { + SessionBrokerCallerClient, + type SessionBrokerSignedRequestInit, +} from "@hunk/session-broker"; import type { SessionTerminalLocation, SessionTerminalMetadata } from "@hunk/session-broker-core"; -import { readHunkSessionDaemonCapabilities } from "../client/capabilities"; import { HUNK_SESSION_DAEMON_HTTP_TIMEOUT_MS, - requestSessionDaemonHttp, + withSessionDaemonHttpTimeout, } from "../client/daemonHttp"; +import { loadOrCreateHunkSessionBrokerCredentials } from "../broker/credentials"; +import { + HUNK_SESSION_BROKER_APP_ID, + HUNK_SESSION_BROKER_APP_REVISION, +} from "../broker/appContract"; import { HUNK_SESSION_API_PATH, + HUNK_SESSION_CAPABILITIES_PATH, type SessionDaemonAction, type SessionDaemonCapabilities, type SessionDaemonRequest, type SessionDaemonResponses, } from "../protocol"; -import { parseSessionDaemonResponse } from "../protocolSchemas"; +import { parseSessionDaemonCapabilities, parseSessionDaemonResponse } from "../protocolSchemas"; import type { AppliedCommentBatchResult, AppliedCommentResult, @@ -76,31 +85,60 @@ async function extractResponseError(response: Response) { return response.statusText || "Unknown Hunk session daemon error."; } +interface HunkCallerTransport { + request( + path: string, + init?: SessionBrokerSignedRequestInit, + options?: { readonly targetSpecific?: boolean }, + ): Promise; +} + class HttpHunkSessionCliClient implements HunkSessionCliClient { private readonly config = resolveSessionBrokerConfig(); - - constructor(private readonly timeoutMs = HUNK_SESSION_DAEMON_HTTP_TIMEOUT_MS) {} + private callerPromise: Promise | null = null; + + constructor( + private readonly timeoutMs = HUNK_SESSION_DAEMON_HTTP_TIMEOUT_MS, + private readonly injectedCaller?: HunkCallerTransport, + ) {} + + private caller() { + if (this.injectedCaller) return Promise.resolve(this.injectedCaller); + this.callerPromise ??= loadOrCreateHunkSessionBrokerCredentials().then( + (credentials) => + new SessionBrokerCallerClient({ + appId: HUNK_SESSION_BROKER_APP_ID, + appRevision: HUNK_SESSION_BROKER_APP_REVISION, + origin: this.config.httpOrigin, + credential: credentials.caller, + daemon: { + keyId: credentials.daemonIdentity.keyId, + publicKey: credentials.daemonPublicKey, + }, + }), + ); + return this.callerPromise; + } private async request( input: Extract, ): Promise { - return requestSessionDaemonHttp({ - config: this.config, - path: HUNK_SESSION_API_PATH, + return withSessionDaemonHttpTimeout({ operation: `complete session ${input.action}`, timeoutMs: this.timeoutMs, - init: { - method: "POST", - headers: { - "content-type": "application/json", - }, - body: JSON.stringify(input), - }, - parse: async (response) => { - if (!response.ok) { - throw new Error(await extractResponseError(response)); - } - + task: async (signal) => { + const caller = await this.caller(); + const response = await caller.request( + HUNK_SESSION_API_PATH, + { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify(input), + signal, + }, + { targetSpecific: input.action !== "list" }, + ); + if (!response.ok) throw new Error(await extractResponseError(response)); let value: unknown; try { value = await response.json(); @@ -113,7 +151,24 @@ class HttpHunkSessionCliClient implements HunkSessionCliClient { } async getCapabilities() { - return readHunkSessionDaemonCapabilities(this.config, this.timeoutMs); + return withSessionDaemonHttpTimeout({ + operation: "report capabilities", + timeoutMs: this.timeoutMs, + task: async (signal) => { + try { + const response = await ( + await this.caller() + ).request(HUNK_SESSION_CAPABILITIES_PATH, { + method: "GET", + signal, + }); + if (!response.ok) return null; + return parseSessionDaemonCapabilities(await response.json()); + } catch { + return null; + } + }, + }); } async listSessions() { @@ -255,8 +310,9 @@ class HttpHunkSessionCliClient implements HunkSessionCliClient { /** Create the concrete Hunk session CLI client that speaks to the broker-backed HTTP API. */ export function createHttpHunkSessionCliClient({ timeoutMs, -}: { timeoutMs?: number } = {}): HunkSessionCliClient { - return new HttpHunkSessionCliClient(timeoutMs); + caller, +}: { timeoutMs?: number; caller?: HunkCallerTransport } = {}): HunkSessionCliClient { + return new HttpHunkSessionCliClient(timeoutMs, caller); } export function stringifyJson(value: unknown) { diff --git a/src/session/agent/commands.ts b/src/session/agent/commands.ts index e5cfca35f..afeb81448 100644 --- a/src/session/agent/commands.ts +++ b/src/session/agent/commands.ts @@ -5,15 +5,9 @@ import type { } from "../../core/run/commandInputs"; import type { SessionLiveCommentSummary, SessionReviewNoteSummary } from "../types"; import { NO_ACTIVE_SESSIONS_MESSAGE } from "./errors"; -import { - ensureSessionBrokerAvailable, - isSessionBrokerHealthy, - isLoopbackPortReachable, - readSessionBrokerHealth, - waitForSessionBrokerShutdown, -} from "../broker/brokerLauncher"; +import { isSessionBrokerHealthy, isLoopbackPortReachable } from "../broker/brokerLauncher"; import { resolveSessionBrokerConfig } from "../broker/brokerConfig"; -import { matchesSessionSelector, normalizeSessionSelector } from "@hunk/session-broker-core"; +import { normalizeSessionSelector } from "@hunk/session-broker-core"; import { createHttpHunkSessionCliClient, formatClearCommentsOutput, @@ -73,80 +67,32 @@ function createDaemonCliClient() { return sessionCommandTestHooks?.createClient?.() ?? createHttpHunkSessionCliClient(); } -async function waitForSessionRegistration(selector?: SessionSelectorInput, timeoutMs = 8_000) { - const deadline = Date.now() + timeoutMs; - - while (Date.now() < deadline) { - const client = createDaemonCliClient(); - - try { - const sessions = await client.listSessions(); - if (sessions.some((session) => matchesSessionSelector(session, selector))) { - return true; - } - } catch { - // Keep polling while the fresh daemon/session reconnects. - } - - await Bun.sleep(200); - } - - return false; -} - async function restartDaemonForMissingAction( action: SessionDaemonAction, - selector?: SessionSelectorInput, + _selector?: SessionSelectorInput, ) { - const health = await readSessionBrokerHealth(); - const pid = health?.pid; - const hadSessions = (health?.sessions ?? 0) > 0; - if (!pid || pid === process.pid) { - throw new Error( - `The running Hunk session daemon is missing required support for ${action}. ` + - `Restart Hunk so it can launch a fresh daemon from the current source tree.`, - ); - } - - process.kill(pid, "SIGTERM"); - - const shutDown = await waitForSessionBrokerShutdown(); - if (!shutDown) { - throw new Error( - `Stopped waiting for the old Hunk session daemon to exit after it was found missing ${action}.`, - ); - } - - const config = resolveSessionBrokerConfig(); - await ensureSessionBrokerAvailable({ - config, - timeoutMs: 3_000, - timeoutMessage: "Timed out waiting for the refreshed Hunk session daemon to start.", - }); - - // `hunk session list` can recover from a stale daemon even when the old process belonged to a - // sibling worktree that reports sessions which will never reconnect to this fresh daemon. - if (selector || (hadSessions && action !== "list")) { - const registered = await waitForSessionRegistration(selector); - if (!registered) { - throw new Error( - "Timed out waiting for the live Hunk session to reconnect after refreshing the session daemon. " + - "Restart that Hunk window if it was launched from an older build.", - ); - } - } + // Public health intentionally carries no PID or identity proof. Never signal an incumbent that + // has not completed the signed broker handshake. + throw new Error( + `The running Hunk session daemon is missing required support for ${action}. ` + + "Stop the conflicting or legacy daemon and restart Hunk; it cannot be replaced safely by PID.", + ); } -async function ensureRequiredAction(action: SessionDaemonAction, selector?: SessionSelectorInput) { - const client = createDaemonCliClient(); +async function ensureRequiredAction( + action: SessionDaemonAction, + selector?: SessionSelectorInput, + client = createDaemonCliClient(), +) { const capabilities = await client.getCapabilities(); if (capabilities?.version === HUNK_SESSION_API_VERSION && capabilities.actions.includes(action)) { - return; + return false; } reportHunkDaemonUpgradeRestart(); await (sessionCommandTestHooks?.restartDaemonForMissingAction?.(action, selector) ?? restartDaemonForMissingAction(action, selector)); + return true; } async function resolveDaemonAvailability(action: SessionCommandInput["action"]) { @@ -185,9 +131,13 @@ export async function runSessionCommand(input: SessionCommandInput) { const normalizedSelector = "selector" in input ? normalizeSessionSelector(input.selector) : null; const requiredAction = REQUIRED_ACTION_BY_COMMAND[input.action]; - await ensureRequiredAction(requiredAction, normalizedSelector ?? undefined); - - const client = createDaemonCliClient(); + let client = createDaemonCliClient(); + const refreshed = await ensureRequiredAction( + requiredAction, + normalizedSelector ?? undefined, + client, + ); + if (refreshed) client = createDaemonCliClient(); switch (input.action) { case "list": { diff --git a/src/session/broker/appContract.ts b/src/session/broker/appContract.ts new file mode 100644 index 000000000..fbc2612ae --- /dev/null +++ b/src/session/broker/appContract.ts @@ -0,0 +1,15 @@ +import { + SESSION_BROKER_PROTOCOL_REVISION, + type BrokerAppContract, +} from "@hunk/session-broker-core"; +import { HUNK_SESSION_DAEMON_VERSION } from "../protocol"; + +/** Defines Hunk's immutable Phase-1 broker and application wire contract. */ +export const HUNK_SESSION_BROKER_APP_ID = "dev.hunk" as const; +export const HUNK_SESSION_BROKER_REVISION = SESSION_BROKER_PROTOCOL_REVISION; +export const HUNK_SESSION_BROKER_APP_REVISION = HUNK_SESSION_DAEMON_VERSION; +export const HUNK_SESSION_BROKER_FEATURES = Object.freeze([]) as readonly []; +export const HUNK_SESSION_BROKER_APP_CONTRACT: Readonly = Object.freeze({ + appRevision: HUNK_SESSION_BROKER_APP_REVISION, + features: HUNK_SESSION_BROKER_FEATURES, +}); diff --git a/src/session/broker/brokerClient.test.ts b/src/session/broker/brokerClient.test.ts index 00c6ebc5c..f15bfd582 100644 --- a/src/session/broker/brokerClient.test.ts +++ b/src/session/broker/brokerClient.test.ts @@ -7,6 +7,7 @@ import { } from "../../../test/helpers/session-daemon-fixtures"; import { HUNK_SESSION_API_VERSION, HUNK_SESSION_DAEMON_VERSION } from "../protocol"; import { SessionBrokerClient } from "./brokerClient"; +import { loadOrCreateHunkSessionBrokerCredentials } from "./credentials"; const originalHost = process.env.HUNK_MCP_HOST; const originalPort = process.env.HUNK_MCP_PORT; @@ -120,33 +121,10 @@ describe("Hunk session daemon client", () => { } }, 10_000); - test("restartIncompatibleDaemon lets startup recover when the stale daemon already exited", async () => { - const server = createServer((_request, response) => { - response.writeHead(404, { "content-type": "text/plain" }); - response.end("gone"); - }); - await new Promise((resolve, reject) => { - server.once("error", reject); - server.listen(0, "127.0.0.1", () => resolve()); - }); - - const address = server.address(); - const port = typeof address === "object" && address ? address.port : 0; - const config = { - host: "127.0.0.1", - port, - httpOrigin: `http://127.0.0.1:${port}`, - wsOrigin: `ws://127.0.0.1:${port}`, - }; - + test("does not retain the legacy PID-based incompatible-daemon replacement path", () => { const client = new SessionBrokerClient(createRegistration(), createSnapshot()); - - try { - await expect((client as any).restartIncompatibleDaemon(config)).resolves.toBeUndefined(); - } finally { - client.stop(); - await new Promise((resolve) => server.close(() => resolve())); - } + expect((client as any).restartIncompatibleDaemon).toBeUndefined(); + client.stop(); }); test("logs one actionable warning when a refreshed daemon rejects an older Hunk window", async () => { @@ -224,6 +202,7 @@ describe("Hunk session daemon client", () => { await Bun.sleep(25); } + (client as any).credentials = await loadOrCreateHunkSessionBrokerCredentials(); await (client as any).connect({ host: "127.0.0.1", port, diff --git a/src/session/broker/brokerClient.ts b/src/session/broker/brokerClient.ts index 53ab54330..a3691abf8 100644 --- a/src/session/broker/brokerClient.ts +++ b/src/session/broker/brokerClient.ts @@ -10,16 +10,13 @@ import { resolveSessionBrokerConfig, type ResolvedSessionBrokerConfig, } from "./brokerConfig"; -import { - ensureSessionBrokerAvailable, - readSessionBrokerHealth, - waitForSessionBrokerShutdown, -} from "./brokerLauncher"; +import { ensureSessionBrokerAvailable } from "./brokerLauncher"; import { hunkSessionProtocolParsers } from "./protocolParsers"; import { - readHunkSessionDaemonCapabilities, - reportHunkDaemonUpgradeRestart, -} from "../client/capabilities"; + loadOrCreateHunkSessionBrokerCredentials, + type HunkSessionBrokerCredentials, +} from "./credentials"; +import { HUNK_SESSION_BROKER_APP_ID, HUNK_SESSION_BROKER_APP_REVISION } from "./appContract"; import type { HunkSessionCommandResult, HunkSessionInfo, @@ -34,6 +31,8 @@ const INCOMPATIBLE_SESSION_CLOSE_CODE = 1008; const INCOMPATIBLE_SESSION_CLOSE_REASON_PREFIX = "Incompatible session "; const INCOMPATIBLE_SESSION_CLOSE_MESSAGE = "This window is too old for the refreshed session broker daemon. Restart the window to reconnect."; +const AUTHENTICATION_REFUSAL_MESSAGE = + "The process on the Hunk session daemon endpoint could not complete signed authentication. Stop the conflicting or legacy daemon and restart Hunk; it will not be replaced by PID."; type SessionAppBridge = SessionBrokerConnectionBridge< HunkSessionServerMessage, @@ -62,6 +61,7 @@ export class SessionBrokerClient { private stopped = false; private startupPromise: Promise | null = null; private lastConnectionWarning: string | null = null; + private credentials: HunkSessionBrokerCredentials | null = null; constructor( private registration: SessionRegistration, @@ -127,6 +127,7 @@ export class SessionBrokerClient { private async ensureDaemonAndConnect() { const config = this.resolveConfig(); await this.ensureDaemonAvailable(config); + this.credentials = await loadOrCreateHunkSessionBrokerCredentials(); this.connect(config); } @@ -136,61 +137,11 @@ export class SessionBrokerClient { timeoutMs: this.timing.daemonStartupTimeoutMs ?? DAEMON_STARTUP_TIMEOUT_MS, }); - const capabilities = await readHunkSessionDaemonCapabilities(config); - if (!capabilities) { - await this.restartIncompatibleDaemon(config); - await ensureSessionBrokerAvailable({ - config, - timeoutMs: this.timing.daemonStartupTimeoutMs ?? DAEMON_STARTUP_TIMEOUT_MS, - }); - - if (!(await readHunkSessionDaemonCapabilities(config))) { - throw new Error( - "The running session broker daemon is incompatible with this build. " + - "Restart the app so it can launch a fresh daemon from the current source tree.", - ); - } - } - + // Minimal health proves only liveness. Compatibility and identity are established by the + // signed websocket hello; an unverifiable incumbent is never signalled or replaced by PID. this.lastConnectionWarning = null; } - private async restartIncompatibleDaemon(config: ResolvedSessionBrokerConfig) { - reportHunkDaemonUpgradeRestart(); - const health = await readSessionBrokerHealth(config); - const pid = health?.pid; - if (pid === process.pid) { - throw new Error( - "The running session broker daemon is incompatible with this build. " + - "Restart the app so it can launch a fresh daemon from the current source tree.", - ); - } - - // If the stale daemon already disappeared on its own, let the normal startup path launch a - // fresh one instead of turning that race into a manual restart error. - if (!pid) { - return; - } - - try { - process.kill(pid, "SIGTERM"); - } catch (error) { - if (!(error instanceof Error) || !("code" in error) || error.code !== "ESRCH") { - throw error; - } - } - - const shutDown = await waitForSessionBrokerShutdown({ - config, - timeoutMs: DAEMON_STARTUP_TIMEOUT_MS, - }); - if (!shutDown) { - throw new Error( - "Stopped waiting for the old session broker daemon to exit after it was found incompatible.", - ); - } - } - setBridge(bridge: SessionAppBridge | null) { this.bridge = bridge; this.connection?.setBridge(bridge); @@ -206,6 +157,7 @@ export class SessionBrokerClient { return; } + if (!this.credentials) return; this.connection = createSessionBrokerConnection< HunkSessionInfo, HunkSessionState, @@ -219,12 +171,25 @@ export class SessionBrokerClient { snapshot: this.snapshot, bridge: this.bridge, protocolParsers: hunkSessionProtocolParsers, + producerAuthentication: { + appId: HUNK_SESSION_BROKER_APP_ID, + appRevision: HUNK_SESSION_BROKER_APP_REVISION, + credential: this.credentials.producer, + daemon: { + keyId: this.credentials.daemonIdentity.keyId, + publicKey: this.credentials.daemonPublicKey, + }, + }, heartbeatIntervalMs: HEARTBEAT_INTERVAL_MS, reconnectDelayMs: this.timing.reconnectDelayMs ?? RECONNECT_DELAY_MS, resolveClose: (event) => this.isIncompatibleSessionClose(event) ? { reconnect: false, warning: INCOMPATIBLE_SESSION_CLOSE_MESSAGE } - : { reconnect: true }, + : event.code === INCOMPATIBLE_SESSION_CLOSE_CODE + ? { reconnect: false, warning: AUTHENTICATION_REFUSAL_MESSAGE } + : event.code === 1006 + ? { reconnect: true, warning: AUTHENTICATION_REFUSAL_MESSAGE } + : { reconnect: true }, onWarning: (message) => this.warnUnavailable(message), }); diff --git a/src/session/broker/brokerConfig.test.ts b/src/session/broker/brokerConfig.test.ts index e782f75f1..7eb925a4b 100644 --- a/src/session/broker/brokerConfig.test.ts +++ b/src/session/broker/brokerConfig.test.ts @@ -1,4 +1,11 @@ import { describe, expect, test } from "bun:test"; +import { HUNK_SESSION_DAEMON_VERSION } from "../protocol"; +import { + HUNK_SESSION_BROKER_APP_ID, + HUNK_SESSION_BROKER_APP_REVISION, + HUNK_SESSION_BROKER_FEATURES, + HUNK_SESSION_BROKER_REVISION, +} from "./appContract"; import { DEFAULT_SESSION_BROKER_HOST, DEFAULT_SESSION_BROKER_PORT, @@ -11,6 +18,14 @@ import { } from "./brokerConfig"; describe("Hunk session daemon config", () => { + test("exports one fixed Phase-1 Hunk contract", () => { + expect(HUNK_SESSION_BROKER_APP_ID).toBe("dev.hunk"); + expect(HUNK_SESSION_BROKER_REVISION).toBe(1); + expect(HUNK_SESSION_BROKER_APP_REVISION).toBe(HUNK_SESSION_DAEMON_VERSION); + expect(HUNK_SESSION_BROKER_FEATURES).toEqual([]); + expect(Object.isFrozen(HUNK_SESSION_BROKER_FEATURES)).toBe(true); + }); + test("resolves exported host and port metadata as runtime defaults", () => { expect(resolveSessionBrokerConfig({})).toMatchObject({ host: DEFAULT_SESSION_BROKER_HOST, diff --git a/src/session/broker/brokerLauncher.ts b/src/session/broker/brokerLauncher.ts index 3976484df..543900949 100644 --- a/src/session/broker/brokerLauncher.ts +++ b/src/session/broker/brokerLauncher.ts @@ -2,7 +2,7 @@ import { spawn } from "node:child_process"; import type { ChildProcess } from "node:child_process"; import { existsSync, mkdirSync, readFileSync, rmSync, statSync, writeFileSync } from "node:fs"; import { connect } from "node:net"; -import { tmpdir } from "node:os"; +import { homedir, tmpdir } from "node:os"; import { join } from "node:path"; import { parseBrokerSafeInteger, @@ -93,7 +93,11 @@ function safeRuntimeToken(value: string) { } function resolveRuntimeBaseDir(env: NodeJS.ProcessEnv = process.env) { - return env.XDG_RUNTIME_DIR?.trim() || tmpdir(); + const configured = env.XDG_RUNTIME_DIR?.trim(); + if (configured) return configured; + // Unix temporary directories are commonly shared across users. Keep the fallback beneath the + // current home directory instead of a predictable shared-/tmp name another account can pre-own. + return typeof process.getuid === "function" ? join(homedir(), ".hunk") : tmpdir(); } function isRunningPid(pid: number) { @@ -146,7 +150,7 @@ function tryAcquireDaemonLaunchLock({ staleAfterMs: number; }): SessionBrokerLaunchLock | null { const paths = resolveSessionBrokerRuntimePaths(config, env); - mkdirSync(paths.runtimeDir, { recursive: true }); + mkdirSync(paths.runtimeDir, { recursive: true, mode: 0o700 }); const payload: SessionBrokerLaunchLockFile = { ownerPid: process.pid, diff --git a/src/session/broker/brokerServer.helpers.test.ts b/src/session/broker/brokerServer.helpers.test.ts index c7040431a..0afebecce 100644 --- a/src/session/broker/brokerServer.helpers.test.ts +++ b/src/session/broker/brokerServer.helpers.test.ts @@ -78,8 +78,8 @@ describe("parseHostAndPort", () => { expect(parseHostAndPort("[::1]:0")).toBeNull(); }); - test("tolerates an unbracketed IPv6 literal by dropping the port", () => { - expect(parseHostAndPort("::1")).toEqual({ host: "::1", port: undefined }); + test("rejects ambiguous unbracketed IPv6 authorities", () => { + expect(parseHostAndPort("::1")).toBeNull(); }); }); diff --git a/src/session/broker/brokerServer.test.ts b/src/session/broker/brokerServer.test.ts index 7bdd791af..42c083af7 100644 --- a/src/session/broker/brokerServer.test.ts +++ b/src/session/broker/brokerServer.test.ts @@ -7,8 +7,19 @@ import { createTestSessionSnapshot, } from "../../../test/helpers/session-daemon-fixtures"; import { SessionBrokerState } from "@hunk/session-broker-core"; +import { + SessionBrokerCallerClient, + answerSessionBrokerHelloChallenge, + createSessionBrokerHelloRequest, + verifyProducerHelloAck, + type AuthenticatedProducerHello, + type SessionBrokerHelloChallenge, + type SessionBrokerSignedRequestInit, +} from "@hunk/session-broker"; import { HUNK_SESSION_API_VERSION, HUNK_SESSION_DAEMON_VERSION } from "../protocol"; import { serveSessionBrokerDaemon } from "./brokerServer"; +import { loadOrCreateHunkSessionBrokerCredentials } from "./credentials"; +import { HUNK_SESSION_BROKER_APP_ID, HUNK_SESSION_BROKER_APP_REVISION } from "./appContract"; const originalHost = process.env.HUNK_MCP_HOST; const originalPort = process.env.HUNK_MCP_PORT; @@ -16,9 +27,9 @@ const originalUnsafeRemote = process.env.HUNK_MCP_UNSAFE_ALLOW_REMOTE; interface HealthResponse { ok: boolean; - pid: number; - sessions: number; - pendingCommands: number; + pid?: number; + sessions?: number; + pendingCommands?: number; paths?: Record; sessionApi?: string; sessionCapabilities?: string; @@ -85,10 +96,41 @@ async function waitForShutdown(port: number, timeoutMs = 1_500) { ); } +async function authenticatedFetch( + port: number, + path: string, + init: SessionBrokerSignedRequestInit = {}, +) { + const credentials = await loadOrCreateHunkSessionBrokerCredentials(); + const caller = new SessionBrokerCallerClient({ + appId: HUNK_SESSION_BROKER_APP_ID, + appRevision: HUNK_SESSION_BROKER_APP_REVISION, + origin: `http://127.0.0.1:${port}`, + credential: credentials.caller, + daemon: { keyId: credentials.daemonIdentity.keyId, publicKey: credentials.daemonPublicKey }, + }); + const action = + typeof init.body === "string" + ? ((JSON.parse(init.body) as { action?: string }).action ?? "") + : ""; + return caller.request(path, init, { + targetSpecific: path === "/session-api" && action !== "list", + }); +} + async function waitForSessionCount(port: number, count: number) { await waitUntil("session registration", async () => { - const health = await readHealth(port); - return health?.sessions === count ? health : null; + try { + const response = await authenticatedFetch(port, "/session-api", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ action: "list" }), + }); + const body = (await response.json()) as { sessions?: unknown[] }; + return body.sessions?.length === count ? body : null; + } catch { + return null; + } }); } @@ -169,6 +211,49 @@ async function openRegisteredSession( snapshotOverrides: Parameters[0] = {}, ) { const socket = await openSessionSocket(port); + const credentials = await loadOrCreateHunkSessionBrokerCredentials(); + const options = { + appId: HUNK_SESSION_BROKER_APP_ID, + appRevision: HUNK_SESSION_BROKER_APP_REVISION, + endpoint: `ws://127.0.0.1:${port}/session`, + credential: credentials.producer, + daemon: { keyId: credentials.daemonIdentity.keyId, publicKey: credentials.daemonPublicKey }, + }; + const hello = createSessionBrokerHelloRequest(options); + socket.send(JSON.stringify({ type: "hello-init", hello })); + const challenge = await new Promise((resolve, reject) => { + const timeout = setTimeout( + () => reject(new Error("Timed out waiting for producer challenge.")), + 1_000, + ); + socket.addEventListener( + "message", + (event) => { + clearTimeout(timeout); + resolve( + (JSON.parse(String(event.data)) as { challenge: SessionBrokerHelloChallenge }).challenge, + ); + }, + { once: true }, + ); + }); + const pending = await answerSessionBrokerHelloChallenge(options, hello, challenge); + socket.send(JSON.stringify({ type: "hello-proof", proof: pending.proof })); + const ack = await new Promise((resolve, reject) => { + const timeout = setTimeout( + () => reject(new Error("Timed out waiting for producer acknowledgement.")), + 1_000, + ); + socket.addEventListener( + "message", + (event) => { + clearTimeout(timeout); + resolve((JSON.parse(String(event.data)) as { ack: AuthenticatedProducerHello }).ack); + }, + { once: true }, + ); + }); + await verifyProducerHelloAck(pending, ack); socket.send( JSON.stringify({ @@ -228,12 +313,12 @@ afterEach(() => { }); describe("Hunk session daemon server", () => { - test("refuses non-loopback binding unless explicitly allowed", () => { + test("refuses non-loopback binding unless explicitly allowed", async () => { process.env.HUNK_MCP_HOST = "0.0.0.0"; process.env.HUNK_MCP_PORT = "47657"; delete process.env.HUNK_MCP_UNSAFE_ALLOW_REMOTE; - expect(() => serveSessionBrokerDaemon()).toThrow("local-only by default"); + await expect(serveSessionBrokerDaemon()).rejects.toThrow("local-only by default"); }); test("reports a clear error when the daemon port is already in use", async () => { @@ -249,7 +334,7 @@ describe("Hunk session daemon server", () => { process.env.HUNK_MCP_PORT = String(port); try { - expect(() => serveSessionBrokerDaemon()).toThrow("port is already in use"); + await expect(serveSessionBrokerDaemon()).rejects.toThrow("port is already in use"); } finally { await new Promise((resolve) => listener.close(() => resolve())); } @@ -260,21 +345,13 @@ describe("Hunk session daemon server", () => { process.env.HUNK_MCP_HOST = "127.0.0.1"; process.env.HUNK_MCP_PORT = String(port); - const server = serveSessionBrokerDaemon(); + const server = await serveSessionBrokerDaemon(); try { const health = await fetch(`http://127.0.0.1:${port}/health`); expect(health.status).toBe(200); const healthPayload = (await health.json()) as HealthResponse; - expect(healthPayload.paths).toEqual({ - health: "/health", - socket: "/session", - }); - expect(healthPayload).toMatchObject({ - sessionApi: `http://127.0.0.1:${port}/session-api`, - sessionCapabilities: `http://127.0.0.1:${port}/session-api/capabilities`, - sessionSocket: `ws://127.0.0.1:${port}/session`, - }); + expect(healthPayload).toEqual({ ok: true }); const genericCapabilities = await fetch(`http://127.0.0.1:${port}/broker/capabilities`); expect(genericCapabilities.status).toBe(404); @@ -288,7 +365,7 @@ describe("Hunk session daemon server", () => { }); expect(genericBroker.status).toBe(404); - const capabilities = await fetch(`http://127.0.0.1:${port}/session-api/capabilities`); + const capabilities = await authenticatedFetch(port, "/session-api/capabilities"); expect(capabilities.status).toBe(200); await expect(capabilities.json()).resolves.toMatchObject({ version: HUNK_SESSION_API_VERSION, @@ -326,12 +403,41 @@ describe("Hunk session daemon server", () => { } }); + test("keeps generic caller and browser-review authority independent", async () => { + const port = await reserveLoopbackPort(); + process.env.HUNK_MCP_HOST = "127.0.0.1"; + process.env.HUNK_MCP_PORT = String(port); + const server = await serveSessionBrokerDaemon(); + try { + await expect(authenticatedFetch(port, "/review-api/missing/publication")).rejects.toThrow( + "daemon identity could not be verified", + ); + const genericHeadersWithoutReviewCapability = await fetch( + `http://127.0.0.1:${port}/review-api/missing/publication`, + { headers: { "x-session-broker-caller-session": "generic-only" } }, + ); + expect(genericHeadersWithoutReviewCapability.status).toBe(401); + + const reviewCapabilityOnSession = await fetch(`http://127.0.0.1:${port}/session-api`, { + method: "POST", + headers: { + "content-type": "application/json", + "hunk-review-capability": "review-only-capability", + }, + body: JSON.stringify({ action: "list" }), + }); + expect(reviewCapabilityOnSession.status).toBe(401); + } finally { + server.stop(true); + } + }); + test("rejects HTTP requests with non-loopback or wrong-port Host headers", async () => { const port = await reserveLoopbackPort(); process.env.HUNK_MCP_HOST = "127.0.0.1"; process.env.HUNK_MCP_PORT = String(port); - const server = serveSessionBrokerDaemon(); + const server = await serveSessionBrokerDaemon(); try { const attackerHostResponse = await fetch(`http://127.0.0.1:${port}/health`, { @@ -361,7 +467,7 @@ describe("Hunk session daemon server", () => { process.env.HUNK_MCP_HOST = "127.0.0.1"; process.env.HUNK_MCP_PORT = String(port); - const server = serveSessionBrokerDaemon(); + const server = await serveSessionBrokerDaemon(); try { const response = await fetch(`http://127.0.0.1:${port}/session-api/capabilities`, { @@ -379,15 +485,35 @@ describe("Hunk session daemon server", () => { } }); + test("requires GET with an empty body for authenticated Hunk capabilities", async () => { + const port = await reserveLoopbackPort(); + process.env.HUNK_MCP_HOST = "127.0.0.1"; + process.env.HUNK_MCP_PORT = String(port); + const server = await serveSessionBrokerDaemon(); + try { + const wrongMethod = await authenticatedFetch(port, "/session-api/capabilities", { + method: "POST", + headers: { "content-type": "application/json" }, + body: "{}", + }); + expect(wrongMethod.status).toBe(405); + await expect(wrongMethod.json()).resolves.toEqual({ + error: "Capabilities require GET with an empty body.", + }); + } finally { + server.stop(true); + } + }); + test("requires JSON content type for session API posts", async () => { const port = await reserveLoopbackPort(); process.env.HUNK_MCP_HOST = "127.0.0.1"; process.env.HUNK_MCP_PORT = String(port); - const server = serveSessionBrokerDaemon(); + const server = await serveSessionBrokerDaemon(); try { - const response = await fetch(`http://127.0.0.1:${port}/session-api`, { + const response = await authenticatedFetch(port, "/session-api", { method: "POST", headers: { "content-type": "text/plain" }, body: JSON.stringify({ action: "list" }), @@ -407,18 +533,22 @@ describe("Hunk session daemon server", () => { process.env.HUNK_MCP_HOST = "127.0.0.1"; process.env.HUNK_MCP_PORT = String(port); - const server = serveSessionBrokerDaemon(); + const server = await serveSessionBrokerDaemon(); try { const response = await fetch(`http://127.0.0.1:${port}/session-api`, { method: "POST", - headers: { "content-type": "application/json" }, + headers: { + "content-type": "application/json", + "x-session-broker-caller-session": "oversized-test-session", + }, body: JSON.stringify({ action: "list", filler: "x".repeat(5 * 1024 * 1024) }), }); expect(response.status).toBe(413); await expect(response.json()).resolves.toMatchObject({ - error: expect.stringContaining("session broker limit"), + error: "capacity-exceeded", + resource: "maxHttpBodyBytes", }); } finally { server.stop(true); @@ -436,7 +566,7 @@ describe("Hunk session daemon server", () => { process.env.HUNK_MCP_HOST = "127.0.0.1"; process.env.HUNK_MCP_PORT = String(port); - const server = serveSessionBrokerDaemon({ + const server = await serveSessionBrokerDaemon({ idleTimeoutMs: 250, staleSessionTtlMs: 500, staleSessionSweepIntervalMs: 25, @@ -455,7 +585,7 @@ describe("Hunk session daemon server", () => { await expect(closed).resolves.toEqual({ code: 1008, - reason: "Session ownership rejected.", + reason: "Session broker authentication required; upgrade Hunk.", }); } finally { socket.close(); @@ -468,7 +598,7 @@ describe("Hunk session daemon server", () => { process.env.HUNK_MCP_HOST = "127.0.0.1"; process.env.HUNK_MCP_PORT = String(port); - const server = serveSessionBrokerDaemon({ + const server = await serveSessionBrokerDaemon({ idleTimeoutMs: 250, staleSessionTtlMs: 500, staleSessionSweepIntervalMs: 25, @@ -498,7 +628,7 @@ describe("Hunk session daemon server", () => { 1_000, ); - const emptyList = await fetch(`http://127.0.0.1:${port}/session-api`, { + const emptyList = await authenticatedFetch(port, "/session-api", { method: "POST", headers: { "content-type": "application/json", @@ -510,7 +640,7 @@ describe("Hunk session daemon server", () => { const goodSocket = await openRegisteredSession(port, "session-good"); try { - const response = await fetch(`http://127.0.0.1:${port}/session-api`, { + const response = await authenticatedFetch(port, "/session-api", { method: "POST", headers: { "content-type": "application/json", @@ -536,7 +666,7 @@ describe("Hunk session daemon server", () => { process.env.HUNK_MCP_HOST = "127.0.0.1"; process.env.HUNK_MCP_PORT = String(port); - const server = serveSessionBrokerDaemon({ + const server = await serveSessionBrokerDaemon({ idleTimeoutMs: 60, staleSessionTtlMs: 500, staleSessionSweepIntervalMs: 25, @@ -545,10 +675,7 @@ describe("Hunk session daemon server", () => { try { await Bun.sleep(150); - await expect(waitForHealth(port)).resolves.toMatchObject({ - ok: true, - sessions: 1, - }); + await expect(waitForHealth(port)).resolves.toEqual({ ok: true }); } finally { socket.close(); server.stop(true); @@ -560,7 +687,7 @@ describe("Hunk session daemon server", () => { process.env.HUNK_MCP_HOST = "127.0.0.1"; process.env.HUNK_MCP_PORT = String(port); - const server = serveSessionBrokerDaemon({ + const server = await serveSessionBrokerDaemon({ idleTimeoutMs: 75, staleSessionTtlMs: 500, staleSessionSweepIntervalMs: 25, @@ -582,7 +709,7 @@ describe("Hunk session daemon server", () => { process.env.HUNK_MCP_HOST = "127.0.0.1"; process.env.HUNK_MCP_PORT = String(port); - const server = serveSessionBrokerDaemon({ + const server = await serveSessionBrokerDaemon({ idleTimeoutMs: 75, staleSessionTtlMs: 80, staleSessionSweepIntervalMs: 20, @@ -658,10 +785,10 @@ describe("Hunk session daemon server", () => { }; }; - const server = serveSessionBrokerDaemon(); + const server = await serveSessionBrokerDaemon(); try { - const response = await fetch(`http://127.0.0.1:${port}/session-api`, { + const response = await authenticatedFetch(port, "/session-api", { method: "POST", headers: { "content-type": "application/json", @@ -700,7 +827,7 @@ describe("Hunk session daemon server", () => { SessionBrokerState.prototype.dispatchCommand = (({ command, input }: any) => { expect(command).toBe("reload_session"); expect(input).toMatchObject({ - sessionPath: "/tmp/live-session", + sessionId: "session-1", sourcePath: "/tmp/source-repo", nextInput: { kind: "vcs", @@ -719,17 +846,17 @@ describe("Hunk session daemon server", () => { }); }) as SessionBrokerState["dispatchCommand"]; - const server = serveSessionBrokerDaemon(); + const server = await serveSessionBrokerDaemon(); try { - const response = await fetch(`http://127.0.0.1:${port}/session-api`, { + const response = await authenticatedFetch(port, "/session-api", { method: "POST", headers: { "content-type": "application/json", }, body: JSON.stringify({ action: "reload", - selector: { sessionPath: "/tmp/live-session" }, + selector: { sessionId: "session-1" }, sourcePath: "/tmp/source-repo", nextInput: { kind: "vcs", @@ -758,7 +885,7 @@ describe("Hunk session daemon server", () => { process.env.HUNK_MCP_HOST = "127.0.0.1"; process.env.HUNK_MCP_PORT = String(port); - const server = serveSessionBrokerDaemon(); + const server = await serveSessionBrokerDaemon(); const socket = await openRegisteredSession(port, "session-1", { reviewNoteCount: 2, reviewNotes: [ @@ -783,7 +910,7 @@ describe("Hunk session daemon server", () => { }); try { - const listResponse = await fetch(`http://127.0.0.1:${port}/session-api`, { + const listResponse = await authenticatedFetch(port, "/session-api", { method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify({ @@ -853,10 +980,10 @@ describe("Hunk session daemon server", () => { }); }) as SessionBrokerState["dispatchCommand"]; - const server = serveSessionBrokerDaemon(); + const server = await serveSessionBrokerDaemon(); try { - const response = await fetch(`http://127.0.0.1:${port}/session-api`, { + const response = await authenticatedFetch(port, "/session-api", { method: "POST", headers: { "content-type": "application/json", diff --git a/src/session/broker/brokerServer.ts b/src/session/broker/brokerServer.ts index c9ee84f2c..2a0d3eb43 100644 --- a/src/session/broker/brokerServer.ts +++ b/src/session/broker/brokerServer.ts @@ -1,4 +1,9 @@ -import { createSessionBrokerDaemon, type SessionBrokerController } from "@hunk/session-broker"; +import { + SessionBrokerAuthenticator, + createSessionBrokerDaemon, + type SessionBrokerAuthenticatedControlFacts, + type SessionBrokerController, +} from "@hunk/session-broker"; import { serveSessionBrokerDaemon as serveSessionBrokerDaemonWithBun, type RunningSessionBrokerDaemon as RunningBunSessionBrokerDaemon, @@ -38,10 +43,13 @@ import { HUNK_SESSION_DAEMON_VERSION, type SessionDaemonAction, type SessionDaemonCapabilities, + type SessionDaemonRequest, type SessionDaemonResponse, } from "../protocol"; import { parseSessionDaemonRequest } from "../protocolSchemas"; import { hunkSessionProtocolParsers } from "./protocolParsers"; +import { loadOrCreateHunkSessionBrokerCredentials } from "./credentials"; +import { HUNK_SESSION_BROKER_APP_ID, HUNK_SESSION_BROKER_APP_REVISION } from "./appContract"; const DEFAULT_STALE_SESSION_TTL_MS = 45_000; const DEFAULT_STALE_SESSION_SWEEP_INTERVAL_MS = 15_000; @@ -116,7 +124,7 @@ function hasJsonContentType(request: Request) { /** Parse a Host-style value into hostname and optional port pieces. */ export function parseHostAndPort(value: string) { const trimmed = value.trim(); - if (!trimmed) { + if (!trimmed || trimmed.includes(",")) { return null; } @@ -136,8 +144,10 @@ export function parseHostAndPort(value: string) { return null; } - const port = Number.parseInt(rest.slice(1), 10); - return Number.isInteger(port) && port > 0 ? { host, port } : null; + const rawPort = rest.slice(1); + if (!/^[0-9]+$/.test(rawPort)) return null; + const port = Number(rawPort); + return Number.isInteger(port) && port > 0 && port <= 65_535 ? { host, port } : null; } const colonCount = [...trimmed].filter((character) => character === ":").length; @@ -147,13 +157,14 @@ export function parseHostAndPort(value: string) { if (colonCount === 1) { const [host, rawPort] = trimmed.split(":"); - const port = Number.parseInt(rawPort ?? "", 10); - return host && Number.isInteger(port) && port > 0 ? { host, port } : null; + if (!host || !/^[0-9]+$/.test(rawPort ?? "")) return null; + const port = Number(rawPort); + return Number.isInteger(port) && port > 0 && port <= 65_535 ? { host, port } : null; } - // Unbracketed IPv6 literals are invalid in Host headers, but accepting the address without a - // port keeps validation strict enough for DNS-rebinding while tolerating unusual native clients. - return { host: trimmed, port: undefined }; + // URL authorities require brackets around IPv6 literals; accepting another spelling would make + // listener-derived authority comparison ambiguous. + return null; } /** Return whether a parsed authority targets an accepted broker host and port. */ @@ -189,6 +200,9 @@ export function validateOriginHeader(request: Request, expectedPort: number, all if (!origin) { return null; } + if (origin === "null" || origin.includes(",")) { + return jsonError("Origin is not allowed for the local session broker.", 403); + } let url: URL; try { @@ -197,7 +211,15 @@ export function validateOriginHeader(request: Request, expectedPort: number, all return jsonError("Origin is not allowed for the local session broker.", 403); } - if (url.protocol !== "http:" && url.protocol !== "https:") { + if ( + (url.protocol !== "http:" && url.protocol !== "https:") || + url.username || + url.password || + url.pathname !== "/" || + url.search || + url.hash || + url.origin !== origin + ) { return jsonError("Origin is not allowed for the local session broker.", 403); } @@ -221,10 +243,43 @@ function parseJsonRequestBytes(bytes: Uint8Array) { return parseSessionDaemonRequest(raw); } +/** Map each Hunk action to the generic operation and exact producer command scope it requires. */ +function sessionApiAuthorizationFacts( + state: HunkSessionBrokerState, + bytes: Uint8Array, +): SessionBrokerAuthenticatedControlFacts { + const input = parseJsonRequestBytes(bytes); + if (input.action === "list") return { operation: "list", targetSpecific: false }; + const sessionId = input.selector.sessionId ?? state.getSession(input.selector).sessionId; + if (["get", "context", "review", "comment-list"].includes(input.action)) { + return { operation: "get", sessionId, targetSpecific: true }; + } + const commandByAction = { + navigate: "navigate_to_hunk", + reload: "reload_session", + "comment-add": "comment", + "comment-apply": "comment_batch", + "comment-rm": "remove_comment", + "comment-clear": "clear_comments", + "highlight-add": "highlight", + "highlight-clear": "clear_highlights", + } as const; + const command = commandByAction[input.action as keyof typeof commandByAction]; + if (!command) throw new Error("Unknown session API action."); + return { + operation: "dispatch", + sessionId, + command, + commandVersion: 1, + targetSpecific: true, + }; +} + export async function handleSessionApiRequest( state: HunkSessionBrokerState, request: Request, bodyBytes?: Uint8Array, + resolvedSessionId?: string, ) { if (request.method !== "POST") { return jsonError("Session API requests must use POST.", 405); @@ -235,9 +290,13 @@ export async function handleSessionApiRequest( } try { - const input = parseJsonRequestBytes( + const parsedInput = parseJsonRequestBytes( bodyBytes ?? (await readRequestBytesWithLimit(request, MAX_HTTP_BODY_BYTES)), ); + const input: SessionDaemonRequest = + resolvedSessionId && parsedInput.action !== "list" + ? { ...parsedInput, selector: { sessionId: resolvedSessionId } } + : parsedInput; let response: SessionDaemonResponse; switch (input.action) { @@ -474,10 +533,12 @@ function createHunkBrokerController( limits: state.limits, listSessions: () => state.listSessions(), getSession: (selector) => state.getSession(selector), + resolveSessionId: (selector) => state.getSession(selector).sessionId, + getSessionIds: () => state.listSessions().map((session) => session.sessionId), getSessionCount: () => state.getSessionCount(), getPendingCommandCount: () => state.getPendingCommandCount(), - registerSession: (connection, registrationInput, snapshotInput) => - state.registerSession(connection, registrationInput, snapshotInput), + registerSession: (connection, registrationInput, snapshotInput, options) => + state.registerSession(connection, registrationInput, snapshotInput, options), updateSnapshot: (connection, sessionId, snapshotInput) => state.updateSnapshot(connection, sessionId, snapshotInput), markSessionSeen: (connection, sessionId) => state.markSessionSeen(connection, sessionId), @@ -493,9 +554,9 @@ function createHunkBrokerController( } /** Serve the local session broker daemon and websocket broker transport. */ -export function serveSessionBrokerDaemon( +export async function serveSessionBrokerDaemon( options: ServeSessionBrokerDaemonOptions = {}, -): RunningSessionBrokerDaemon { +): Promise { const config = resolveSessionBrokerConfig(); const allowRemote = allowsUnsafeRemoteSessionBroker(); const idleTimeoutMs = options.idleTimeoutMs ?? DEFAULT_IDLE_TIMEOUT_MS; @@ -503,6 +564,18 @@ export function serveSessionBrokerDaemon( const staleSessionSweepIntervalMs = options.staleSessionSweepIntervalMs ?? DEFAULT_STALE_SESSION_SWEEP_INTERVAL_MS; const state = createHunkSessionBrokerState(); + const credentials = await loadOrCreateHunkSessionBrokerCredentials(); + const generation = `h_${crypto.randomUUID().replaceAll("-", "")}_0`; + const authenticator = new SessionBrokerAuthenticator({ + appId: HUNK_SESSION_BROKER_APP_ID, + appRevision: HUNK_SESSION_BROKER_APP_REVISION, + generation, + daemonIdentity: credentials.daemonIdentity, + credentials: [credentials.producer, credentials.caller], + // A CLI process normally performs capabilities plus one action, then exits. Retire its caller + // session quickly so repeated short-lived commands cannot fill the generic retained-session cap. + callerSessionTtlMs: 30_000, + }); // One loopback process serves every attached review, rather than a port per terminal. const browserReview = new BrowserReviewServer(state); const daemon = createSessionBrokerDaemon({ @@ -515,6 +588,15 @@ export function serveSessionBrokerDaemon( idleTimeoutMs, staleSessionTtlMs, staleSessionSweepIntervalMs, + appId: HUNK_SESSION_BROKER_APP_ID, + appRevision: HUNK_SESSION_BROKER_APP_REVISION, + callerAuthenticator: authenticator, + helloAuthenticator: authenticator, + producerEndpoint: `${config.wsOrigin}${SESSION_BROKER_SOCKET_PATH}`, + authorizer: () => true, + // Hunk currently keeps audit decisions in-process; the generic hook guarantees only redacted + // principal/operation metadata can be wired to a future diagnostic sink. + audit: () => undefined, paths: { socket: SESSION_BROKER_SOCKET_PATH, }, @@ -538,29 +620,45 @@ export function serveSessionBrokerDaemon( const url = new URL(request.url); - if (url.pathname === "/health") { - // Extend the generic health payload with the Hunk-specific companion endpoints that older - // CLI clients and debugging workflows still expect to discover from one place. - return Response.json({ - ...daemon.getHealth(), - sessionApi: `${config.httpOrigin}${HUNK_SESSION_API_PATH}`, - sessionCapabilities: `${config.httpOrigin}${HUNK_SESSION_CAPABILITIES_PATH}`, - sessionSocket: `${config.wsOrigin}${SESSION_BROKER_SOCKET_PATH}`, - }); + if ( + (url.pathname === HUNK_SESSION_CAPABILITIES_PATH || + url.pathname === HUNK_SESSION_API_PATH) && + !request.headers.has("x-session-broker-caller-session") + ) { + return Response.json( + { + error: "authentication-required", + message: + "This Hunk session client must be upgraded to use automatic signed authentication.", + }, + { status: 401 }, + ); } if (url.pathname === HUNK_SESSION_CAPABILITIES_PATH) { - return Response.json(sessionCapabilities()); + return daemon.handleAuthenticatedControl(request, { + resolve: () => ({ operation: "diagnostics", targetSpecific: false }), + handle: (body) => + request.method === "GET" && body.byteLength === 0 + ? { body: sessionCapabilities() as never } + : { + body: { error: "Capabilities require GET with an empty body." }, + status: request.method === "GET" ? 400 : 405, + }, + }); } - // Keep the richer Hunk session API here rather than in the shared package so commands like - // review, reload, and comment flows stay app-specific. + // Keep Hunk action parsing and lowering app-owned while the generic hook authenticates, + // authorizes, budgets, and signs the exact transport body and response. if (url.pathname === HUNK_SESSION_API_PATH) { - return daemon.handleBoundedControl( - request, - (body) => handleSessionApiRequest(state, request, body), - { payloadTooLarge: (error) => jsonError(error.message, 413) }, - ); + return daemon.handleAuthenticatedControl(request, { + resolve: (body) => sessionApiAuthorizationFacts(state, body), + resolveFailureTargetSpecific: (body) => parseJsonRequestBytes(body).action !== "list", + handle: async (body, facts) => { + const response = await handleSessionApiRequest(state, request, body, facts.sessionId); + return { body: (await response.json()) as never, status: response.status }; + }, + }); } // The review surface authorizes every one of its own routes with a per-session diff --git a/src/session/broker/credentials.test.ts b/src/session/broker/credentials.test.ts new file mode 100644 index 000000000..62b48cfb9 --- /dev/null +++ b/src/session/broker/credentials.test.ts @@ -0,0 +1,91 @@ +import { afterEach, describe, expect, test } from "bun:test"; +import { + chmodSync, + lstatSync, + mkdtempSync, + readFileSync, + rmSync, + symlinkSync, + writeFileSync, +} from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { loadOrCreateHunkSessionBrokerCredentials } from "./credentials"; + +const roots: string[] = []; + +function isolatedEnv() { + const root = mkdtempSync(join(tmpdir(), "hunk-credentials-test-")); + roots.push(root); + return { ...process.env, XDG_RUNTIME_DIR: root }; +} + +afterEach(() => { + for (const root of roots.splice(0)) rmSync(root, { recursive: true, force: true }); +}); + +describe("Hunk session broker credential store", () => { + test("creates stable independent Ed25519 material with owner-private Unix permissions", async () => { + const env = isolatedEnv(); + const first = await loadOrCreateHunkSessionBrokerCredentials({ env }); + const second = await loadOrCreateHunkSessionBrokerCredentials({ env }); + + expect(second.daemonIdentity.keyId).toBe(first.daemonIdentity.keyId); + expect(second.producer.grant.keyId).toBe(first.producer.grant.keyId); + expect(second.caller.grant.keyId).toBe(first.caller.grant.keyId); + expect(first.producer.grant.keyId).not.toBe(first.caller.grant.keyId); + + const securityDir = join(env.XDG_RUNTIME_DIR!, "hunk-mcp", "security-v1"); + if (process.platform !== "win32") { + expect(lstatSync(securityDir).mode & 0o777).toBe(0o700); + for (const name of ["daemon.json", "producer.json", "caller.json"]) { + expect(lstatSync(join(securityDir, name)).mode & 0o777).toBe(0o600); + } + } + const callerFile = readFileSync(join(securityDir, "caller.json"), "utf8"); + expect(callerFile).not.toContain("hunk-review-capability"); + }); + + test("adopts one complete winner under concurrent first use", async () => { + const env = isolatedEnv(); + const results = await Promise.all( + Array.from({ length: 12 }, () => loadOrCreateHunkSessionBrokerCredentials({ env })), + ); + expect(new Set(results.map((value) => value.daemonIdentity.keyId)).size).toBe(1); + expect(new Set(results.map((value) => value.producer.grant.keyId)).size).toBe(1); + expect(new Set(results.map((value) => value.caller.grant.keyId)).size).toBe(1); + }); + + test("rejects malformed and overly permissive credential files without leaking private bytes", async () => { + const env = isolatedEnv(); + await loadOrCreateHunkSessionBrokerCredentials({ env }); + const callerPath = join(env.XDG_RUNTIME_DIR!, "hunk-mcp", "security-v1", "caller.json"); + const secret = "private-secret-sentinel"; + writeFileSync(callerPath, `{"privateKey":"${secret}"}`); + if (process.platform !== "win32") chmodSync(callerPath, 0o644); + + let message = ""; + try { + await loadOrCreateHunkSessionBrokerCredentials({ env }); + } catch (error) { + message = error instanceof Error ? error.message : String(error); + } + expect(message).toContain("unsafe or malformed"); + expect(message).not.toContain(secret); + }); + + test("rejects a symlinked security directory", async () => { + if (process.platform === "win32") return; + const env = isolatedEnv(); + const runtimeDir = join(env.XDG_RUNTIME_DIR!, "hunk-mcp"); + const target = join(env.XDG_RUNTIME_DIR!, "redirect"); + const { mkdirSync } = await import("node:fs"); + mkdirSync(runtimeDir, { mode: 0o700 }); + mkdirSync(target, { mode: 0o700 }); + symlinkSync(target, join(runtimeDir, "security-v1"), "dir"); + + await expect(loadOrCreateHunkSessionBrokerCredentials({ env })).rejects.toThrow( + "unsafe or malformed", + ); + }); +}); diff --git a/src/session/broker/credentials.ts b/src/session/broker/credentials.ts new file mode 100644 index 000000000..c3d715d73 --- /dev/null +++ b/src/session/broker/credentials.ts @@ -0,0 +1,375 @@ +import { + closeSync, + constants, + fsyncSync, + fstatSync, + lstatSync, + linkSync, + mkdirSync, + openSync, + readFileSync, + rmSync, + writeFileSync, +} from "node:fs"; +import { dirname, join } from "node:path"; +import { + SESSION_BROKER_SIGNATURE_ALGORITHM, + type CallerGrant, + type ProducerGrant, +} from "@hunk/session-broker-core"; +import { + importEd25519PrivateKey, + importEd25519PublicKey, + type SessionBrokerCredential, + type SessionBrokerDaemonIdentity, +} from "@hunk/session-broker"; +import { resolveSessionBrokerRuntimePaths } from "./brokerLauncher"; +import { HUNK_SESSION_BROKER_APP_ID } from "./appContract"; + +const CREDENTIAL_VERSION = 1; +const CREDENTIAL_LIFETIME_MS = 10 * 365 * 24 * 60 * 60 * 1_000; +const PRIVATE_MODE = 0o600; +const DIRECTORY_MODE = 0o700; + +const HUNK_COMMAND_SCOPES = [ + "navigate_to_hunk", + "reload_session", + "comment", + "comment_batch", + "remove_comment", + "clear_comments", + "highlight", + "clear_highlights", +].map((name) => ({ name, version: 1 })) as readonly { name: string; version: number }[]; + +interface StoredCredentialFile { + version: 1; + role: "daemon" | "producer" | "caller"; + keyId: string; + publicKey: string; + privateKey: string; + grant?: ProducerGrant | CallerGrant; +} + +export interface HunkSessionBrokerCredentials { + readonly daemonIdentity: SessionBrokerDaemonIdentity; + readonly daemonPublicKey: CryptoKey; + readonly producer: SessionBrokerCredential & { readonly privateKey: CryptoKey }; + readonly caller: SessionBrokerCredential & { readonly privateKey: CryptoKey }; +} + +export interface HunkCredentialStoreOptions { + readonly env?: NodeJS.ProcessEnv; + readonly now?: () => number; + readonly randomBytes?: (length: number) => Uint8Array; +} + +function securityError(): never { + throw new Error( + "Hunk session credentials are unavailable because their owner-private runtime state is unsafe or malformed.", + ); +} + +function encode(bytes: ArrayBuffer | Uint8Array) { + return Buffer.from(bytes instanceof Uint8Array ? bytes : new Uint8Array(bytes)).toString( + "base64url", + ); +} + +function decode(value: unknown): Uint8Array { + if (typeof value !== "string" || !/^[A-Za-z0-9_-]+$/.test(value)) securityError(); + const bytes = Buffer.from(value, "base64url"); + if (bytes.length === 0 || bytes.toString("base64url") !== value) securityError(); + return bytes; +} + +function randomId(randomBytes: (length: number) => Uint8Array) { + return `h_${Buffer.from(randomBytes(18)).toString("base64url")}_0`; +} + +/** Reject credential directories and files that can redirect reads or expose owner material. */ +function validateOwnerPrivatePath(path: string, kind: "directory" | "file") { + let stat; + try { + stat = lstatSync(path); + } catch { + securityError(); + } + if (stat.isSymbolicLink() || (kind === "directory" ? !stat.isDirectory() : !stat.isFile())) { + securityError(); + } + if (process.platform !== "win32") { + if (typeof process.getuid === "function" && stat.uid !== process.getuid()) securityError(); + const unsafeBits = kind === "directory" ? stat.mode & 0o077 : stat.mode & 0o177; + if (unsafeBits !== 0) securityError(); + } +} + +/** Validate the legacy namespace parent while allowing its historical read/execute mode. */ +function ensureRuntimeNamespace(path: string) { + mkdirSync(path, { recursive: true, mode: DIRECTORY_MODE }); + const stat = lstatSync(path); + if (stat.isSymbolicLink() || !stat.isDirectory()) securityError(); + if (process.platform !== "win32") { + if (typeof process.getuid === "function" && stat.uid !== process.getuid()) securityError(); + if ((stat.mode & 0o022) !== 0) securityError(); + } +} + +/** Create and validate the stable hunk-mcp owner-private security directory. */ +function ensureSecurityDirectory(path: string) { + mkdirSync(path, { recursive: true, mode: DIRECTORY_MODE }); + if (process.platform !== "win32") { + // mkdir honors umask by making permissions narrower, which is safe; never broaden an existing dir. + validateOwnerPrivatePath(path, "directory"); + } else { + validateOwnerPrivatePath(path, "directory"); + } +} + +/** Read a regular owner-private file through a no-follow descriptor where the runtime supports it. */ +function readPrivateFile(path: string): unknown { + validateOwnerPrivatePath(path, "file"); + let descriptor: number | null = null; + try { + const noFollow = (constants as typeof constants & { O_NOFOLLOW?: number }).O_NOFOLLOW ?? 0; + descriptor = openSync(path, constants.O_RDONLY | noFollow); + const stat = fstatSync(descriptor); + if (!stat.isFile() || stat.size <= 0 || stat.size > 64 * 1024) securityError(); + if (process.platform !== "win32") { + if (typeof process.getuid === "function" && stat.uid !== process.getuid()) securityError(); + if ((stat.mode & 0o177) !== 0) securityError(); + } + return JSON.parse(readFileSync(descriptor, "utf8")); + } catch { + securityError(); + } finally { + if (descriptor !== null) closeSync(descriptor); + } +} + +function parseStored(value: unknown, role: StoredCredentialFile["role"]): StoredCredentialFile { + if (!value || typeof value !== "object" || Array.isArray(value)) securityError(); + const record = value as Record; + const expected = new Set([ + "version", + "role", + "keyId", + "publicKey", + "privateKey", + ...(role === "daemon" ? [] : ["grant"]), + ]); + if ( + Object.keys(record).some((key) => !expected.has(key)) || + Object.keys(record).length !== expected.size + ) + securityError(); + if (record.version !== CREDENTIAL_VERSION || record.role !== role) securityError(); + if (typeof record.keyId !== "string" || !/^h_[A-Za-z0-9_-]+_0$/.test(record.keyId)) + securityError(); + decode(record.publicKey); + decode(record.privateKey); + if (role !== "daemon") { + const grant = record.grant as Record | undefined; + const grantKeys = new Set([ + "kind", + "appId", + "principalId", + "keyId", + "grantId", + "algorithm", + "issuedAt", + "expiresAt", + "revocationId", + "mayDelegate", + "operations", + ...(role === "caller" ? ["commands"] : []), + ]); + const expectedOperations = + role === "producer" ? ["register", "reconnect"] : ["list", "get", "dispatch", "diagnostics"]; + if ( + !grant || + Object.keys(grant).length !== grantKeys.size || + Object.keys(grant).some((key) => !grantKeys.has(key)) || + grant.kind !== role || + grant.appId !== HUNK_SESSION_BROKER_APP_ID || + grant.principalId !== `hunk-${role}` || + grant.keyId !== record.keyId || + grant.grantId !== `hunk-${role}-bootstrap-v1` || + grant.algorithm !== SESSION_BROKER_SIGNATURE_ALGORITHM || + !Number.isFinite(grant.issuedAt) || + !Number.isFinite(grant.expiresAt) || + (grant.issuedAt as number) >= (grant.expiresAt as number) || + grant.revocationId !== `hunk-${role}-bootstrap-v1` || + grant.mayDelegate !== false || + JSON.stringify(grant.operations) !== JSON.stringify(expectedOperations) || + (role === "caller" && JSON.stringify(grant.commands) !== JSON.stringify(HUNK_COMMAND_SCOPES)) + ) + securityError(); + } + return record as unknown as StoredCredentialFile; +} + +/** Atomically adopts a complete credential file without ever replacing a live winner. */ +function adoptPrivateFile( + path: string, + contents: string, + randomBytes: (length: number) => Uint8Array, +) { + const temp = `${path}.tmp-${process.pid}-${Buffer.from(randomBytes(9)).toString("hex")}`; + let descriptor: number | null = null; + try { + descriptor = openSync( + temp, + constants.O_WRONLY | constants.O_CREAT | constants.O_EXCL, + PRIVATE_MODE, + ); + writeFileSync(descriptor, contents, "utf8"); + fsyncSync(descriptor); + closeSync(descriptor); + descriptor = null; + try { + // A hard link publishes the already-complete inode and fails rather than replacing a winner. + requireLink(temp, path); + if (process.platform !== "win32") { + const directory = openSync(dirname(path), constants.O_RDONLY); + try { + fsyncSync(directory); + } catch (error) { + if (!["EINVAL", "ENOTSUP"].includes((error as NodeJS.ErrnoException).code ?? "")) { + throw error; + } + } finally { + closeSync(directory); + } + } + } catch (error) { + if ((error as NodeJS.ErrnoException).code !== "EEXIST") throw error; + } + } finally { + if (descriptor !== null) closeSync(descriptor); + rmSync(temp, { force: true }); + } +} + +function requireLink(source: string, destination: string) { + linkSync(source, destination); +} + +async function createStored( + role: StoredCredentialFile["role"], + now: number, + randomBytes: (length: number) => Uint8Array, +): Promise { + const pair = (await crypto.subtle.generateKey("Ed25519", true, [ + "sign", + "verify", + ])) as CryptoKeyPair; + const keyId = randomId(randomBytes); + const base = { + version: CREDENTIAL_VERSION, + role, + keyId, + publicKey: encode(await crypto.subtle.exportKey("spki", pair.publicKey)), + privateKey: encode(await crypto.subtle.exportKey("pkcs8", pair.privateKey)), + } as const; + if (role === "daemon") return base; + const common = { + kind: role, + appId: HUNK_SESSION_BROKER_APP_ID, + principalId: `hunk-${role}`, + keyId, + grantId: `hunk-${role}-bootstrap-v1`, + algorithm: SESSION_BROKER_SIGNATURE_ALGORITHM, + issuedAt: now, + expiresAt: now + CREDENTIAL_LIFETIME_MS, + revocationId: `hunk-${role}-bootstrap-v1`, + mayDelegate: false, + } as const; + const grant = + role === "producer" + ? ({ + ...common, + kind: "producer", + operations: ["register", "reconnect"], + } satisfies ProducerGrant) + : ({ + ...common, + kind: "caller", + operations: ["list", "get", "dispatch", "diagnostics"], + commands: HUNK_COMMAND_SCOPES, + } satisfies CallerGrant); + return { ...base, grant }; +} + +async function loadOrCreate( + path: string, + role: StoredCredentialFile["role"], + now: number, + randomBytes: (length: number) => Uint8Array, +) { + try { + return parseStored(readPrivateFile(path), role); + } catch (error) { + const code = (() => { + try { + lstatSync(path); + return "exists"; + } catch (cause) { + return (cause as NodeJS.ErrnoException).code; + } + })(); + if (code !== "ENOENT") throw error; + } + const generated = await createStored(role, now, randomBytes); + adoptPrivateFile(path, `${JSON.stringify(generated)}\n`, randomBytes); + return parseStored(readPrivateFile(path), role); +} + +/** Load or safely create Hunk's daemon, producer, and caller Ed25519 bootstrap material. */ +export async function loadOrCreateHunkSessionBrokerCredentials( + options: HunkCredentialStoreOptions = {}, +): Promise { + const env = options.env ?? process.env; + const randomBytes = + options.randomBytes ?? ((length) => crypto.getRandomValues(new Uint8Array(length))); + const runtimeDir = resolveSessionBrokerRuntimePaths(undefined, env).runtimeDir; + const securityDir = join(runtimeDir, "security-v1"); + ensureRuntimeNamespace(runtimeDir); + ensureSecurityDirectory(securityDir); + const now = (options.now ?? Date.now)(); + const [daemon, producer, caller] = await Promise.all([ + loadOrCreate(join(securityDir, "daemon.json"), "daemon", now, randomBytes), + loadOrCreate(join(securityDir, "producer.json"), "producer", now, randomBytes), + loadOrCreate(join(securityDir, "caller.json"), "caller", now, randomBytes), + ]); + const [ + daemonPublicKey, + daemonPrivateKey, + producerPublicKey, + producerPrivateKey, + callerPublicKey, + callerPrivateKey, + ] = await Promise.all([ + importEd25519PublicKey(decode(daemon.publicKey)), + importEd25519PrivateKey(decode(daemon.privateKey)), + importEd25519PublicKey(decode(producer.publicKey)), + importEd25519PrivateKey(decode(producer.privateKey)), + importEd25519PublicKey(decode(caller.publicKey)), + importEd25519PrivateKey(decode(caller.privateKey)), + ]); + return Object.freeze({ + daemonIdentity: Object.freeze({ keyId: daemon.keyId, privateKey: daemonPrivateKey }), + daemonPublicKey, + producer: Object.freeze({ + grant: producer.grant as ProducerGrant, + publicKey: producerPublicKey, + privateKey: producerPrivateKey, + }), + caller: Object.freeze({ + grant: caller.grant as CallerGrant, + publicKey: callerPublicKey, + privateKey: callerPrivateKey, + }), + }); +} diff --git a/src/session/broker/state.ts b/src/session/broker/state.ts index 48c8321c6..0a28fb64a 100644 --- a/src/session/broker/state.ts +++ b/src/session/broker/state.ts @@ -206,8 +206,9 @@ export class HunkSessionBrokerState extends SessionBrokerState< socket: HunkBrokerConnection, registrationInput: unknown, snapshotInput: unknown, + options?: { replaceOwner?: boolean }, ) { - const registered = super.registerSession(socket, registrationInput, snapshotInput); + const registered = super.registerSession(socket, registrationInput, snapshotInput, options); this.reconcileMirroredSessions(); if (registered !== "registered") { return registered; diff --git a/src/session/client/capabilities.ts b/src/session/client/capabilities.ts index 2892b5e98..464154b81 100644 --- a/src/session/client/capabilities.ts +++ b/src/session/client/capabilities.ts @@ -7,9 +7,9 @@ import { parseSessionDaemonCapabilities } from "../protocolSchemas"; import { HUNK_SESSION_DAEMON_HTTP_TIMEOUT_MS, requestSessionDaemonHttp } from "./daemonHttp"; export const HUNK_DAEMON_UPGRADE_RESTART_NOTICE = - "[hunk:session] Restarting stale session daemon after upgrade."; + "[hunk:session] The legacy session daemon requires a manual restart; Hunk will not signal an unverifiable PID."; -/** Tell the user that Hunk is refreshing an old daemon left running across an upgrade. */ +/** Tell the user that an unverifiable legacy daemon must be stopped manually. */ export function reportHunkDaemonUpgradeRestart(log: (message: string) => void = console.error) { log(HUNK_DAEMON_UPGRADE_RESTART_NOTICE); } diff --git a/test/session-broker-node/adapter.test.mjs b/test/session-broker-node/adapter.test.mjs index 9d9a27480..c05716d6e 100644 --- a/test/session-broker-node/adapter.test.mjs +++ b/test/session-broker-node/adapter.test.mjs @@ -85,11 +85,13 @@ function fakeDaemon(overrides = {}, behavior = {}) { maxOutboundBytesTotal: 64 * 1024 * 1024, maxHttpResponseBytes: 8 * 1024 * 1024, maxUnauthenticatedSockets: 64, + maxHandshakeDurationMs: 15_000, ...overrides, }; return { limits, stopped: new Promise(() => {}), + requiresProducerAuthentication: behavior.requiresProducerAuthentication ?? false, matchesSocketPath: (pathname) => pathname === "/session", handleConnectionMessage: behavior.handleConnectionMessage ?? (() => {}), handleConnectionClose() {}, @@ -112,11 +114,15 @@ test("Node WebCrypto Ed25519 and base64url work without Bun globals", async () = test("Node adapter consumes the shared text/binary/oversize/pressure corpus", async () => { const port = await reservePort(); const running = await serveSessionBrokerDaemon({ - daemon: fakeDaemon({ - maxWsMessageBytes: 8, - maxHttpResponseBytes: 8, - maxUnauthenticatedSockets: 1, - }), + daemon: fakeDaemon( + { + maxWsMessageBytes: 8, + maxHttpResponseBytes: 8, + maxUnauthenticatedSockets: 1, + maxHandshakeDurationMs: 1_000, + }, + { requiresProducerAuthentication: true }, + ), hostname: "127.0.0.1", port, handleRequest: (request) => diff --git a/test/session/broker-e2e.test.ts b/test/session/broker-e2e.test.ts index 52ac4a208..51c2033f1 100644 --- a/test/session/broker-e2e.test.ts +++ b/test/session/broker-e2e.test.ts @@ -1,5 +1,5 @@ import { afterAll, afterEach, describe, expect, test } from "bun:test"; -import { mkdtempSync, rmSync, writeFileSync } from "node:fs"; +import { mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; import { createServer } from "node:http"; import { tmpdir } from "node:os"; import { join } from "node:path"; @@ -32,8 +32,6 @@ const ttyToolsAvailable = supportsControllableScript(); interface HealthResponse { ok: boolean; - pid: number; - sessions: number; } interface SessionListJson { @@ -224,6 +222,19 @@ async function waitUntil( } } +/** Read the PID only from this test's launch metadata for teardown, never from public health. */ +function readLaunchedDaemonPid(port: number) { + try { + const runtimeBase = process.env.XDG_RUNTIME_DIR?.trim() || tmpdir(); + const metadata = JSON.parse( + readFileSync(join(runtimeBase, "hunk-mcp", `daemon-127-0-0-1-${port}.json`), "utf8"), + ) as { pid?: unknown }; + return typeof metadata.pid === "number" && metadata.pid > 0 ? metadata.pid : null; + } catch { + return null; + } +} + async function waitForHealth(port: number, timeoutMs = 15_000) { return waitUntil( "session daemon health endpoint", @@ -288,7 +299,7 @@ describe("session broker end-to-end", () => { try { const health = await waitForHealth(port); - daemonPid = health.pid; + daemonPid = readLaunchedDaemonPid(port); expect(health.ok).toBe(true); const listed = await waitUntil("registered Hunk session", async () => { @@ -394,7 +405,7 @@ describe("session broker end-to-end", () => { try { const health = await waitForHealth(port); - daemonPid = health.pid; + daemonPid = readLaunchedDaemonPid(port); expect(health.ok).toBe(true); const listed = await waitUntil("registered Hunk session", async () => { @@ -489,7 +500,7 @@ describe("session broker end-to-end", () => { try { const health = await waitForHealth(port); - daemonPid = health.pid; + daemonPid = readLaunchedDaemonPid(port); expect(health.ok).toBe(true); const listed = await waitUntil("registered Hunk session", async () => { @@ -621,7 +632,7 @@ describe("session broker end-to-end", () => { try { const health = await waitForHealth(port, 20_000); - daemonPid = health.pid; + daemonPid = readLaunchedDaemonPid(port); expect(health.ok).toBe(true); const sessions = await waitUntil("two registered Hunk sessions", async () => { diff --git a/test/session/cli.test.ts b/test/session/cli.test.ts index f24058d59..e9bfd22ff 100644 --- a/test/session/cli.test.ts +++ b/test/session/cli.test.ts @@ -1,5 +1,5 @@ import { afterAll, afterEach, describe, expect, test } from "bun:test"; -import { existsSync, mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs"; +import { existsSync, mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; import { createServer } from "node:net"; import { tmpdir } from "node:os"; import { join } from "node:path"; @@ -214,20 +214,24 @@ async function quitHunkSession( const ownedDaemonPids = new Map(); -/** Poll daemon health directly before exercising the CLI boundary once. */ +/** Poll through the authenticated CLI because public health intentionally exposes no session facts. */ async function waitForRegisteredSessions(port: number) { - await waitUntil("registered live session", async () => { - const health = await readDaemonHealth(port); - if (!health || (health.sessions ?? 0) === 0) return null; - ownedDaemonPids.set(port, health.pid); - return true; + return waitUntil("registered live session", () => { + const { proc, stdout } = runSessionCli(["list", "--json"], port); + if (proc.exitCode !== 0) return null; + const sessions = (JSON.parse(stdout) as SessionListJson).sessions; + if (sessions.length === 0) return null; + try { + const metadata = JSON.parse( + readFileSync(join(testRuntimeDir, "hunk-mcp", `daemon-127-0-0-1-${port}.json`), "utf8"), + ) as { pid?: unknown }; + if (typeof metadata.pid === "number" && metadata.pid > 0) + ownedDaemonPids.set(port, metadata.pid); + } catch { + // Teardown can still rely on daemon idleness if metadata publication raced this read. + } + return sessions; }); - - const { proc, stdout, stderr } = runSessionCli(["list", "--json"], port); - if (proc.exitCode !== 0) { - throw new Error(stderr.trim() || "Failed to list the registered Hunk session."); - } - return (JSON.parse(stdout) as SessionListJson).sessions; } /** Read one test daemon's health without leaking connection failures into teardown. */ @@ -235,7 +239,7 @@ async function readDaemonHealth(port: number) { try { const response = await fetch(`http://127.0.0.1:${port}/health`); if (!response.ok) return null; - return (await response.json()) as { pid: number; sessions?: number }; + return (await response.json()) as { ok: boolean }; } catch { return null; } @@ -267,9 +271,6 @@ async function waitForDaemonExit(port: number, pid: number, label: string) { label, async () => { const health = await readDaemonHealth(port); - if (health && health.pid !== pid) { - throw new Error(`Refusing to manage unexpected daemon ${health.pid} on port ${port}.`); - } return !isProcessRunning(pid) && health === null ? true : null; }, 1_500, @@ -283,17 +284,10 @@ async function stopTestDaemon(port: number) { ownedDaemonPids.delete(port); if (pid === undefined) return; - const health = await readDaemonHealth(port); - if (health && health.pid !== pid) { - throw new Error(`Refusing to stop unexpected daemon ${health.pid} on port ${port}.`); - } - signalProcess(pid, "SIGTERM"); try { await waitForDaemonExit(port, pid, "session daemon exit"); - } catch (error) { - const remaining = await readDaemonHealth(port); - if (remaining && remaining.pid !== pid) throw error; + } catch { signalProcess(pid, "SIGKILL"); await waitForDaemonExit(port, pid, "killed session daemon exit"); } @@ -491,7 +485,7 @@ sessionDescribe("session CLI integration", () => { } }, 20_000); - test("reload refuses option-like VCS ranges sent directly to the session API", async () => { + test("raw session API callers cannot present option-like VCS ranges", async () => { const port = await reserveLoopbackPort(); const fixture = createFixtureFiles( "reload-injection", @@ -505,8 +499,7 @@ sessionDescribe("session CLI integration", () => { const listed = await waitForRegisteredSessions(port); const sessionId = listed[0]!.sessionId; - // Bypass the CLI parser on purpose: the raw daemon surface is the attacker-controlled - // path, so reproduce the injected flag exactly as a hostile /session-api caller would. + // Raw callers never reach app parsing without the owner-private signed caller session. const sentinel = join(fixture.dir, "hunk-poc"); const response = await fetch(`http://127.0.0.1:${port}/session-api`, { method: "POST", @@ -523,9 +516,10 @@ sessionDescribe("session CLI integration", () => { }), }); - expect(response.status).toBe(400); + expect(response.status).toBe(401); await expect(response.json()).resolves.toMatchObject({ - error: expect.stringContaining("looks like a VCS option"), + error: "authentication-required", + message: expect.stringContaining("upgraded"), }); expect(existsSync(sentinel)).toBe(false); diff --git a/test/session/daemon.test.ts b/test/session/daemon.test.ts index c068dc7a2..b8c155daa 100644 --- a/test/session/daemon.test.ts +++ b/test/session/daemon.test.ts @@ -52,7 +52,7 @@ async function readHealth(port: number) { return null; } - return (await response.json()) as { ok: boolean; pid: number }; + return (await response.json()) as { ok: boolean }; } catch { return null; } @@ -96,8 +96,8 @@ describe("session daemon lifecycle", () => { exited = true; }); - // Windows may keep the `bun run` launcher separate from the child serving the daemon. - process.kill(health.pid, "SIGTERM"); + // This test owns the spawned process handle; public health intentionally exposes no PID. + proc.kill("SIGTERM"); await waitUntil("daemon serve process exit", () => (exited ? true : null), 1_500, 25); await waitUntil("daemon port close", async () =>