From e4b283fbd248e9e8eb83fbd457f78cd8d43e0d0a Mon Sep 17 00:00:00 2001 From: Ben Vinegar Date: Sun, 30 Aug 2026 17:48:51 -0400 Subject: [PATCH] feat(session-broker): authenticate local peers --- .changeset/fair-brokers-rest.md | 2 + .github/workflows/ci.yml | 8 + .github/workflows/pr-ci.yml | 3 + docs/session-broker-sdk.md | 74 ++- packages/session-broker-bun/package.json | 2 +- packages/session-broker-bun/src/serve.test.ts | 184 ++++- packages/session-broker-bun/src/serve.ts | 52 +- packages/session-broker-core/package.json | 2 +- packages/session-broker-core/src/auth.test.ts | 4 +- packages/session-broker-core/src/auth.ts | 4 + .../src/brokerState.test.ts | 91 +++ .../session-broker-core/src/brokerState.ts | 38 +- .../session-broker-core/src/budgets.test.ts | 38 ++ packages/session-broker-core/src/budgets.ts | 61 ++ .../session-broker-core/src/limits.test.ts | 127 +++- packages/session-broker-core/src/limits.ts | 86 ++- packages/session-broker-node/package.json | 2 +- .../session-broker-node/src/serve.test.ts | 10 +- packages/session-broker-node/src/serve.ts | 57 +- packages/session-broker/package.json | 2 +- .../session-broker/src/authentication.test.ts | 104 ++- packages/session-broker/src/authentication.ts | 42 +- packages/session-broker/src/broker.ts | 15 +- .../src/clientAuthentication.test.ts | 433 ++++++++++++ .../src/clientAuthentication.ts | 629 ++++++++++++++++++ .../session-broker/src/connection.test.ts | 159 ++++- packages/session-broker/src/connection.ts | 177 ++++- packages/session-broker/src/daemon.test.ts | 419 +++++++++++- packages/session-broker/src/daemon.ts | 454 ++++++++++++- packages/session-broker/src/index.ts | 1 + packages/session-broker/src/types.ts | 4 +- src/session/broker/brokerServer.ts | 2 + .../sessionBrokerAdapterConformance.json | 3 +- test/session-broker-node/adapter.test.mjs | 68 +- 34 files changed, 3165 insertions(+), 192 deletions(-) create mode 100644 .changeset/fair-brokers-rest.md create mode 100644 packages/session-broker/src/clientAuthentication.test.ts create mode 100644 packages/session-broker/src/clientAuthentication.ts diff --git a/.changeset/fair-brokers-rest.md b/.changeset/fair-brokers-rest.md new file mode 100644 index 000000000..a845151cc --- /dev/null +++ b/.changeset/fair-brokers-rest.md @@ -0,0 +1,2 @@ +--- +--- diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 802573b0f..eabc5cf5d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -43,6 +43,11 @@ jobs: with: bun-version: 1.3.14 + - name: Set up Node + uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7.0.0 + with: + node-version: 22 + - name: Install Jujutsu uses: taiki-e/install-action@07b4745e0c39a41822af610387492e3e53aa222b # v2.83.4 with: @@ -81,6 +86,9 @@ jobs: - name: Test suite run: bun run test + - name: Real Node broker adapter tests + run: bun run test:session-broker-node + - name: PTY integration tests run: bun run test:integration diff --git a/.github/workflows/pr-ci.yml b/.github/workflows/pr-ci.yml index 9b53ad6f2..86b125a71 100644 --- a/.github/workflows/pr-ci.yml +++ b/.github/workflows/pr-ci.yml @@ -187,6 +187,9 @@ jobs: - name: Test suite run: bun run test + - name: Real Node broker adapter tests + run: bun run test:session-broker-node + - name: PTY integration tests run: bun run test:integration diff --git a/docs/session-broker-sdk.md b/docs/session-broker-sdk.md index 12a2c2fea..14d273013 100644 --- a/docs/session-broker-sdk.md +++ b/docs/session-broker-sdk.md @@ -207,8 +207,9 @@ Each request carries caller session, caller request ID, and a canonical uint64 d `(?:0|[1-9][0-9]{0,19})`, at most `18446744073709551615`, parsed with integer/BigInt rather than JSON number. The signature binds generation, caller session, grant/key ID, hello transcript hash, HTTP method, canonical path and sorted/encoded query, canonical body digest, request ID, and sequence. -Authorization precedes execution and cache lookup. Signed target command envelopes and signed -responses repeat the exact selected application revision/features; producers reject mismatches +Responses bind the same caller session, request ID, and sequence so signatures cannot cross caller +sessions. Authorization precedes execution and cache lookup. Signed target command envelopes and +signed responses repeat the exact selected application revision/features; producers reject mismatches before parsing input. Target incompatibility returns a structured error without mutating the caller session. @@ -286,9 +287,10 @@ input/result for the exact contract, and cross-process health/capabilities. App revision. Failures return stable codes without stacks/parser internals, commit no partial mutation or result, -close producers when stale assumptions would remain, and behave identically on Node/Bun. Reject -binary WebSocket frames with `1003`, oversized frames with `1009`, and malformed discriminants, -numbers, and versions rather than casting them. +and close producers when stale assumptions would remain. Reject binary WebSocket frames with `1003` +and malformed discriminants, numbers, and versions rather than casting them. Native per-message caps +reject oversized frames before application decoding: Node's `ws` reports `1009`, while Bun 1.3 may +surface its non-configurable abnormal `1006` because it does not invoke the application callback. ## Local security contract @@ -398,10 +400,10 @@ follow their descriptor capacities and per-group FIFO. | Resource | Initial default | | ---------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| Sessions and commands | 256 sessions/daemon; 64 queued+in-flight/session; 1,024/daemon; 32 waiting for a missing bridge; one active/session by default | +| Sessions and commands | 256 sessions/daemon; 64 queued+in-flight/session; 1,024/daemon; 32 queued+executing through one producer bridge; one active/session by default | | Handshakes | Per daemon: 64 unauthenticated/challenged sockets, 128 incomplete records, 4 MiB incomplete bytes; 64 KiB/proposal | | Callers and HTTP | Per daemon: 256 caller sessions at 8 KiB each/2 MiB total, 32 concurrent controls, 64 MiB in-flight body bytes; 4 MiB maximum decoded request; 8 MiB maximum decoded/aggregated response | -| WebSocket | 8 MiB maximum inbound message; 64 MiB decoded/in-flight inbound/daemon; buffered outbound: 8 MiB/peer and 64 MiB/daemon | +| WebSocket | 8 MiB native inbound message cap; 64 MiB broker-owned delivered-message processing/daemon; 64 socket admissions; buffered outbound: 8 MiB/peer and 64 MiB/daemon | | Retained session state | 4 MiB metadata+snapshot/session; 256 MiB/daemon | | Command data/time | 1 MiB validated input/entry; 64 MiB queued-command bytes/daemon; 15 s default timeout; 5 min caller maximum | | Idempotency | 1,024 entries/session; 65,536 entries/daemon; 10 min TTL; 1 MiB result/entry; 64 MiB/daemon | @@ -410,10 +412,15 @@ Hosts may lower limits. Raising a public network/body/frame/buffer ceiling requi unsafe-limits configuration and is outside supported defaults. App validators add tighter collection, string, nesting, state, and command limits. -Overflow returns structured `busy`, `queue-full`, or `capacity-exceeded`. Reserve aggregate inbound -bytes before read/decode/parse and release in `finally`; unavailable WebSocket capacity or slow-peer -outbound overflow closes with retryable `1013`. Unwritten work is `not-delivered`; written work -follows the delivery matrix. Reserve outbound serialized bytes per peer/daemon until flush/close. HTTP readers reserve declared/incremental chunks and return `503` when unavailable. +Overflow returns structured `busy`, `queue-full`, or `capacity-exceeded`. Native WebSocket frame +assembly and runtime-owned queues are outside broker byte accounting; adapters bound them per message +with the native cap and bound peer count at socket admission. Once a message is delivered, reserve its +bytes before broker decode/parse/handling and release in `finally`. Unavailable socket-count admission +returns `503`; unavailable delivered-message capacity or slow-peer outbound overflow closes with +retryable `1013`. Unwritten work is `not-delivered`; written work follows the delivery matrix. Reserve +outbound serialized bytes per peer/daemon until flush/close. HTTP readers reserve the permitted source +maximum before the first pull, resize to actual bytes, and require aggregate capacity for the +source-plus-copy peak before retaining the merged body; unavailable capacity returns `503`. Unauthenticated overload allocates no handshake state. Response overflow fails before unbounded aggregation. New-session failure never evicts existing state. @@ -540,20 +547,41 @@ resources/cache/assembly, comments/highlights/navigation/reload/intents; HTTP ac discovery; and executable launch/upgrade copy. Generic state removes Hunk-shaped projections; Hunk builds them from read-only generic entries/events, and the package never imports Hunk types. -| Existing contract | Migration requirement | -| ---------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------ | -| `HUNK_MCP_HOST`, `HUNK_MCP_PORT`, `HUNK_MCP_DISABLE` | Preserve in Hunk adapter for at least one documented minor window. | -| `HUNK_MCP_UNSAFE_ALLOW_REMOTE` | Temporary unsupported Hunk-only escape hatch; never generic. | -| runtime directory `hunk-mcp` | Preserve or dual-read so old/new binaries cannot race separate namespaces. | -| default `127.0.0.1:47657` | Winning candidate reserves/retains it as guard before coordinator publication; explicit port reserves that endpoint in the same namespace. | -| `/session`, `/session-api` | Preserve CLI semantics but require new authentication; old clients get actionable upgrade refusal. | -| `/mcp` returns `410` | Retain until separately deprecated. | -| exact-version restart | Never signal from unverifiable PID; legacy daemon is an actionable manual stop/restart conflict; only authenticated generation may stop automatically. | -| missing `repoBoundary` | Preserve containment fallback; new clients may provide VCS-aware boundary. | +| Existing contract | Migration requirement | +| ---------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `HUNK_MCP_HOST`, `HUNK_MCP_PORT`, `HUNK_MCP_DISABLE` | Preserve in Hunk adapter for at least one documented minor window. | +| `HUNK_MCP_UNSAFE_ALLOW_REMOTE` | Temporary unsupported Hunk-only escape hatch; never generic. | +| runtime directory `hunk-mcp` | Preserve or dual-read so old/new binaries cannot race separate namespaces. | +| default `127.0.0.1:47657` | Winning candidate reserves/retains it as guard before coordinator publication; explicit port reserves that endpoint in the same namespace. | +| `/session`, `/session-api` | Preserve CLI semantics but require new authentication; old clients get actionable upgrade refusal. | +| `/mcp` returns `410` | Retain until separately deprecated. | +| incompatible daemon | Never signal from unverifiable PID. Interactive Hunk waits and retries until the incumbent exits while idle; forced replacement requires an authenticated generation. | +| missing `repoBoundary` | Preserve containment fallback; new clients may provide VCS-aware boundary. | Security outranks wire compatibility: no migration accepts unauthenticated control/registration. -Preservation means paths, selectors, outputs, and automatic credential discovery for upgraded -clients—not interoperability with pre-authentication binaries. +An interactive Hunk window keeps reviewing locally while an incompatible or pre-authentication +incumbent owns the endpoint. After the first signed-handshake refusal it polls only minimal health, +so repeated WebSocket closes cannot postpone the incumbent's idle timeout; once health disappears, +the same bounded connection object reruns discovery, signed negotiation, and registration. One-shot +session commands fail promptly with instructions instead +of becoming a second restart owner. Daemons too old to retire while idle, or hung incumbents, still +require manual termination. 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. Retained +producer authority is rechecked before inbound mutations and before any queued command bytes leave +the daemon. 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 diff --git a/packages/session-broker-bun/package.json b/packages/session-broker-bun/package.json index 1873c83c4..569d65c1d 100644 --- a/packages/session-broker-bun/package.json +++ b/packages/session-broker-bun/package.json @@ -19,7 +19,7 @@ "@hunk/session-broker": "workspace:*" }, "engines": { - "bun": ">=1.0.0", + "bun": ">=1.3.14", "node": ">=22" } } diff --git a/packages/session-broker-bun/src/serve.test.ts b/packages/session-broker-bun/src/serve.test.ts index dc8ed76b4..20fede6a6 100644 --- a/packages/session-broker-bun/src/serve.test.ts +++ b/packages/session-broker-bun/src/serve.test.ts @@ -1,11 +1,13 @@ import { afterEach, describe, expect, test } from "bun:test"; -import { createServer } from "node:net"; +import { connect, createServer } from "node:net"; import { + BrokerCapacityError, SESSION_BROKER_REGISTRATION_VERSION, brokerWireParsers, parseSessionRegistrationEnvelope, parseSessionSnapshotEnvelope, type SessionRegistration, + type SessionServerMessage, type SessionSnapshot, } from "@hunk/session-broker-core"; import { @@ -115,6 +117,36 @@ async function waitUntil( } } +/** Return the HTTP status from one raw WebSocket upgrade attempt. */ +async function rawWebSocketUpgradeStatus(port: number) { + return await new Promise((resolve, reject) => { + const socket = connect({ host: "127.0.0.1", port }, () => { + socket.write( + [ + "GET /session HTTP/1.1", + `Host: 127.0.0.1:${port}`, + "Connection: Upgrade", + "Upgrade: websocket", + "Sec-WebSocket-Version: 13", + "Sec-WebSocket-Key: dGVzdC1zZXNzaW9uLWtleQ==", + "", + "", + ].join("\r\n"), + ); + }); + let response = ""; + socket.on("data", (chunk) => { + response += chunk.toString("utf8"); + if (!response.includes("\r\n\r\n")) return; + const status = Number(response.match(/^HTTP\/1\.1 (\d{3})/)?.[1]); + socket.destroy(); + if (Number.isInteger(status)) resolve(status); + else reject(new Error("WebSocket upgrade returned an invalid HTTP response.")); + }); + socket.on("error", reject); + }); +} + async function openTestSocket(url: string) { const socket = new WebSocket(url); await new Promise((resolve, reject) => { @@ -196,18 +228,11 @@ 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 () => { + test("closes binary and oversized messages per the shared corpus", async () => { const broker = new SessionBroker({ protocolParsers }); const daemon = createSessionBrokerDaemon({ broker, - limits: { maxWsMessageBytes: 8, maxInFlightWsBytes: 0 }, + limits: { maxWsMessageBytes: 8 }, }); const port = await reserveLoopbackPort(); const server = serveSessionBrokerDaemon({ daemon, hostname: "127.0.0.1", port }); @@ -220,16 +245,60 @@ describe("session broker bun adapter", () => { 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, + expect(SESSION_BROKER_ADAPTER_CONFORMANCE.inbound.bunNativeOversizedCloseCodes).toContain( + await oversizedClosed, ); + } finally { + server.stop(true); + await server.stopped; + } + }); - 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, + test("returns the shared HTTP status when socket admission is full and releases on close", async () => { + const broker = new SessionBroker({ protocolParsers }); + const daemon = createSessionBrokerDaemon({ + broker, + limits: { maxUnauthenticatedSockets: 1, maxHandshakeDurationMs: 1_000 }, + helloAuthenticator: {} as never, + producerEndpoint: "ws://127.0.0.1/session", + }); + 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`); + expect(await rawWebSocketUpgradeStatus(port)).toBe( + SESSION_BROKER_ADAPTER_CONFORMANCE.inbound.admissionHttpStatus, ); + 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("contains handler failures and closes the affected peer", async () => { + const broker = new SessionBroker({ protocolParsers }); + const daemon = createSessionBrokerDaemon({ broker }); + daemon.handleConnectionMessage = (_peer, message) => { + if (message === "capacity") throw new BrokerCapacityError("busy", "test"); + throw new Error("unexpected handler failure"); + }; + const port = await reserveLoopbackPort(); + const server = serveSessionBrokerDaemon({ daemon, hostname: "127.0.0.1", port }); + try { + const capacity = await openTestSocket(`ws://127.0.0.1:${port}/session`); + const capacityClosed = testSocketCloseCode(capacity); + capacity.send("capacity"); + expect(await capacityClosed).toBe(1013); + + const unexpected = await openTestSocket(`ws://127.0.0.1:${port}/session`); + const unexpectedClosed = testSocketCloseCode(unexpected); + unexpected.send("unexpected"); + expect(await unexpectedClosed).toBe(1011); } finally { server.stop(true); await server.stopped; @@ -262,6 +331,76 @@ describe("session broker bun adapter", () => { } }); + test("closes outbound aggregate pressure and releases socket capacity for reconnect", async () => { + type PressureMessage = SessionServerMessage<"annotate", { summary: string }>; + const pressureParsers = createSessionBrokerProtocolParsers< + TestSessionInfo, + TestSessionState, + PressureMessage, + { applied: true } + >({ + appRevision: 1, + features: [], + parseRegistration: (value) => parseSessionRegistrationEnvelope(value, parseInfo), + parseSnapshot: (value) => parseSessionSnapshotEnvelope(value, parseState), + commands: [ + { + command: "annotate", + version: 1, + parseInput: (value) => + typeof (value as { summary?: unknown })?.summary === "string" + ? { summary: (value as { summary: string }).summary } + : null, + parseResult: (value) => + (value as { applied?: unknown })?.applied === true ? { applied: true } : null, + }, + ], + }); + const broker = new SessionBroker({ protocolParsers: pressureParsers }); + const daemon = createSessionBrokerDaemon({ + broker, + limits: { maxOutboundBytesTotal: 8 }, + }); + const port = await reserveLoopbackPort(); + const server = serveSessionBrokerDaemon({ daemon, hostname: "127.0.0.1", port }); + try { + const socket = await openTestSocket(`ws://127.0.0.1:${port}/session`); + socket.send( + JSON.stringify({ + type: "register", + registration: createRegistration(), + snapshot: createSnapshot(), + }), + ); + await waitUntil("pressure registration", () => + broker.listSessions().length === 1 ? true : null, + ); + const closed = testSocketCloseCode(socket); + const dispatchResult = broker + .dispatchCommand({ + selector: { sessionId: "session-1" }, + command: "annotate", + input: { summary: "pressure" }, + timeoutMessage: "timeout", + }) + .then( + () => null, + (error: unknown) => error, + ); + expect(await closed).toBe(SESSION_BROKER_ADAPTER_CONFORMANCE.outbound.pressureCloseCode); + expect(await dispatchResult).toBeInstanceOf(Error); + await waitUntil("pressure disconnect cleanup", () => + broker.listSessions().length === 0 ? true : null, + ); + + const afterRelease = await openTestSocket(`ws://127.0.0.1:${port}/session`); + afterRelease.close(); + } finally { + server.stop(true); + await server.stopped; + } + }); + test("manual stop retires peer admission and rejects late message delivery", async () => { let snapshotCalls = 0; const countingParsers = createSessionBrokerProtocolParsers({ @@ -366,11 +505,13 @@ describe("session broker bun adapter", () => { expect(stoppedSettled).toBe(true); }); - 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 }); @@ -378,8 +519,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 { @@ -413,7 +553,9 @@ describe("session broker bun adapter", () => { generation: "generation-1", brokerRevision: 1 as const, ...(appContract ? { appContract } : {}), + callerSessionId: "caller-session-1", requestId: "request-1", + sequence: "1", httpStatus, bodyDigest: "test-digest", daemonKeyId: "daemon-key-1", diff --git a/packages/session-broker-bun/src/serve.ts b/packages/session-broker-bun/src/serve.ts index 468c7d5d9..6ac39ee64 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< @@ -158,6 +159,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; @@ -212,15 +223,20 @@ export function serveSessionBrokerDaemon< const admission = unauthenticatedSocketBudget.tryReserve(); if (!admission) return new Response(null, { status: 503 }); activeAdmissions.add(admission); - if (bunServer.upgrade(request, { data: { admission } })) { - return undefined; + try { + if (bunServer.upgrade(request, { data: { admission } })) { + return undefined; + } + } catch (error) { + activeAdmissions.delete(admission); + admission.release(); + throw error; } 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. - return new Response("Expected websocket upgrade.", { status: 426 }); } @@ -236,16 +252,20 @@ export function serveSessionBrokerDaemon< } }, websocket: { - // 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, - ), - ), + 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 bounds native frame assembly per message but exposes no accounting hook before it + // delivers the decoded string. The broker budget below covers only application processing. + maxPayloadLength: Math.max(1, options.daemon.limits.maxWsMessageBytes), message: (socket, message) => { const peer = peerFor(socket); if (typeof message !== "string") { @@ -265,6 +285,11 @@ export function serveSessionBrokerDaemon< } try { options.daemon.handleConnectionMessage(peer, message); + } catch (error) { + socket.close( + error instanceof BrokerCapacityError ? 1013 : 1011, + "Session broker message handling failed.", + ); } finally { reservation.release(); } @@ -283,6 +308,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/package.json b/packages/session-broker-core/package.json index 36e39e37e..afc6294aa 100644 --- a/packages/session-broker-core/package.json +++ b/packages/session-broker-core/package.json @@ -16,7 +16,7 @@ } }, "engines": { - "bun": ">=1.0.0", + "bun": ">=1.3.14", "node": ">=22" } } diff --git a/packages/session-broker-core/src/auth.test.ts b/packages/session-broker-core/src/auth.test.ts index 7e4fb3f70..ad1809937 100644 --- a/packages/session-broker-core/src/auth.test.ts +++ b/packages/session-broker-core/src/auth.test.ts @@ -75,13 +75,15 @@ describe("session broker authentication core", () => { generation: "generation-1", brokerRevision: 1, appContract: { appRevision: 7, features: [] }, + callerSessionId: "caller-session-1", requestId: "request-1", + sequence: "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"}', + '{"appContract":{"appRevision":7,"features":[]},"appId":"dev.example","bodyDigest":"body-hash","brokerRevision":1,"callerSessionId":"caller-session-1","domain":"dev.hunk.session-broker.v1/caller-response","generation":"generation-1","httpStatus":200,"requestId":"request-1","sequence":"1"}', ); expect( new TextDecoder().decode( diff --git a/packages/session-broker-core/src/auth.ts b/packages/session-broker-core/src/auth.ts index e75f3f038..3199942f9 100644 --- a/packages/session-broker-core/src/auth.ts +++ b/packages/session-broker-core/src/auth.ts @@ -134,7 +134,9 @@ export interface BrokerResponseTranscriptInput { readonly appId: string; readonly generation: string; readonly brokerRevision: typeof SESSION_BROKER_PROTOCOL_REVISION; + readonly callerSessionId: string; readonly requestId: string; + readonly sequence: string; readonly httpStatus: number; readonly bodyDigest: string; readonly appContract?: BrokerAppContract; @@ -349,10 +351,12 @@ export function buildBrokerResponseTranscript(input: BrokerResponseTranscriptInp : {}), bodyDigest: input.bodyDigest, brokerRevision: input.brokerRevision, + callerSessionId: input.callerSessionId, domain: `${SESSION_BROKER_AUTH_DOMAIN}/caller-response`, generation: input.generation, httpStatus: input.httpStatus, requestId: input.requestId, + sequence: input.sequence, }); } diff --git a/packages/session-broker-core/src/brokerState.test.ts b/packages/session-broker-core/src/brokerState.test.ts index 4a91691d0..40a7a2ea2 100644 --- a/packages/session-broker-core/src/brokerState.test.ts +++ b/packages/session-broker-core/src/brokerState.test.ts @@ -199,6 +199,26 @@ function createListedSession(overrides: Partial = {}): TestLi } describe("session broker state", () => { + test("keeps shutdown terminal against registration and command re-admission", () => { + const state = createState(); + const socket = { send() {} }; + const shutdownError = new Error("terminal shutdown"); + state.shutdown(shutdownError); + state.shutdown(new Error("ignored second shutdown")); + + expect(state.registerSession(socket, createRegistration(), createSnapshot())).toBe("shutdown"); + expect(state.getSessionCount()).toBe(0); + expect(() => + state.dispatchCommand({ + selector: { sessionId: "session-1" }, + command: "annotate", + input: { filePath: "a.ts", summary: "late" }, + timeoutMessage: "timeout", + }), + ).toThrow(shutdownError); + expect(state.getPendingCommandCount()).toBe(0); + }); + test("resolves one target session by session id, session path, repo root, or sole-session fallback", () => { const one = [createListedSession()]; const two = [ @@ -527,6 +547,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..34abbd9fd 100644 --- a/packages/session-broker-core/src/brokerState.ts +++ b/packages/session-broker-core/src/brokerState.ts @@ -93,7 +93,8 @@ export type RegisterSessionResult = | "registered" | "invalid" | "already-connected" - | "capacity-exceeded"; + | "capacity-exceeded" + | "shutdown"; export type UpdateSnapshotResult = "updated" | "invalid" | "not-owner" | "capacity-exceeded"; export type MarkSessionSeenResult = "seen" | "not-owner"; export type HandleCommandResult = "handled" | "not-found" | "not-owner" | "invalid"; @@ -220,6 +221,7 @@ export class SessionBrokerState< private readonly queuedCommandByteBudget: ResourceBudget; private readonly retainedByteBudget: ResourceBudget; private lastPruneAt: number | null = null; + private shutdownError: Error | null = null; constructor( private view: SessionBrokerViewAdapter< @@ -287,7 +289,10 @@ export class SessionBrokerState< socket: DaemonSessionSocket, registrationInput: unknown, snapshotInput: unknown, + options: { replaceOwner?: boolean } = {}, ): RegisterSessionResult { + if (this.shutdownError) return "shutdown"; + let registration: SessionRegistration | null; let snapshot: SessionSnapshot | null; try { @@ -312,7 +317,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 +326,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 +354,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."), @@ -495,6 +523,7 @@ export class SessionBrokerState< timeoutMessage: string; timeoutMs?: number; }) { + if (this.shutdownError) throw this.shutdownError; if (!isValidBrokerRevision(commandVersion)) { throw new TypeError("Command version must be a positive safe integer."); } @@ -625,6 +654,9 @@ export class SessionBrokerState< } shutdown(error = new Error("The session broker daemon shut down.")) { + if (this.shutdownError) return; + this.shutdownError = error; + for (const pending of this.pendingCommands.values()) { this.finishPending(pending, () => pending.reject(error), false); } diff --git a/packages/session-broker-core/src/budgets.test.ts b/packages/session-broker-core/src/budgets.test.ts index aeb0a5300..3d1e5d9fe 100644 --- a/packages/session-broker-core/src/budgets.test.ts +++ b/packages/session-broker-core/src/budgets.test.ts @@ -30,6 +30,7 @@ describe("session broker limits", () => { maxInFlightWsBytes: 64 * 1024 * 1024, challengeTtlMs: 15_000, callerSessionTtlMs: 5 * 60_000, + maxHandshakeDurationMs: 15_000, }); expect(Object.isFrozen(DEFAULT_SESSION_BROKER_LIMITS)).toBe(true); }); @@ -64,6 +65,31 @@ describe("session broker limits", () => { expect(() => resolveSessionBrokerLimits({ limits: { unknown: 1 } as never })).toThrow( "Unknown session broker limit", ); + expect(() => + resolveSessionBrokerLimits({ + limits: { maxWsMessageBytes: 8, maxInFlightWsBytes: 7 }, + }), + ).toThrow("WebSocket message bytes must not exceed in-flight bytes"); + expect( + resolveSessionBrokerLimits({ + limits: { maxHttpBodyBytes: 4, maxInFlightHttpBodyBytes: 8 }, + }).maxHttpBodyBytes, + ).toBe(4); + expect(() => + resolveSessionBrokerLimits({ + limits: { maxHttpBodyBytes: 4, maxInFlightHttpBodyBytes: 7 }, + }), + ).toThrow("source-plus-copy peak"); + expect( + resolveSessionBrokerLimits({ + limits: { maxHttpResponseBytes: 4, maxInFlightHttpResponseBytes: 8 }, + }).maxHttpResponseBytes, + ).toBe(4); + expect(() => + resolveSessionBrokerLimits({ + limits: { maxHttpResponseBytes: 4, maxInFlightHttpResponseBytes: 7 }, + }), + ).toThrow("source-plus-copy peak"); }); }); @@ -78,6 +104,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 f2dc20c68..12cd6ef62 100644 --- a/packages/session-broker-core/src/budgets.ts +++ b/packages/session-broker-core/src/budgets.ts @@ -5,6 +5,7 @@ export interface SessionBrokerLimits { readonly maxSessions: number; readonly maxCommandsPerSession: number; readonly maxCommandsTotal: number; + /** Bound producer commands retained while queued or executing through one bridge. */ readonly maxPreBridgeCommands: number; readonly maxCommandInputBytes: number; readonly maxCommandResultBytes: number; @@ -19,10 +20,12 @@ export interface SessionBrokerLimits { readonly maxHttpResponseBytes: number; readonly maxInFlightHttpResponseBytes: number; readonly maxWsMessageBytes: number; + /** Bound broker-owned processing after native WebSocket delivery, excluding runtime queues. */ readonly maxInFlightWsBytes: number; readonly maxOutboundBytesPerPeer: number; readonly maxOutboundBytesTotal: number; readonly maxUnauthenticatedSockets: number; + readonly maxHandshakeDurationMs: number; readonly maxIncompleteHandshakes: number; readonly maxIncompleteHandshakeBytes: number; readonly maxHandshakeProposalBytes: number; @@ -56,6 +59,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, @@ -139,6 +143,19 @@ export function mergeSessionBrokerLimits( if (resolved.maxCallerSessionBytes > resolved.maxCallerSessionsBytes) { throw new TypeError("Session broker per-caller retained bytes must not exceed daemon bytes."); } + if (resolved.maxHttpBodyBytes > Math.floor(resolved.maxInFlightHttpBodyBytes / 2)) { + throw new TypeError( + "Session broker in-flight HTTP body bytes must cover the source-plus-copy peak.", + ); + } + if (resolved.maxHttpResponseBytes > Math.floor(resolved.maxInFlightHttpResponseBytes / 2)) { + throw new TypeError( + "Session broker in-flight HTTP response bytes must cover the source-plus-copy peak.", + ); + } + if (resolved.maxWsMessageBytes > resolved.maxInFlightWsBytes) { + throw new TypeError("Session broker WebSocket message bytes must not exceed in-flight bytes."); + } return Object.freeze(resolved); } @@ -233,6 +250,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-core/src/limits.test.ts b/packages/session-broker-core/src/limits.test.ts index e8e65f448..b6ea78026 100644 --- a/packages/session-broker-core/src/limits.test.ts +++ b/packages/session-broker-core/src/limits.test.ts @@ -36,16 +36,32 @@ function streamingRequest(byteLength: number, chunkSize = 64 * 1024) { } describe("readRequestTextWithLimit", () => { - test("rejects an oversized declared Content-Length before reading the body", async () => { + test("cancels an oversized declared body without pulling it", async () => { + let pulls = 0; + let cancelled = false; + const body = new ReadableStream( + { + pull(controller) { + pulls += 1; + controller.enqueue(new Uint8Array([1])); + }, + cancel() { + cancelled = true; + }, + }, + { highWaterMark: 0 }, + ); const request = new Request("http://broker.test/api", { method: "POST", - headers: { "content-type": "application/json", "content-length": String(10 * 1024 * 1024) }, - body: "ignored", - }); + headers: { "content-length": String(10 * 1024 * 1024) }, + body, + duplex: "half", + } as RequestInit); await expect(readRequestTextWithLimit(request, 1024)).rejects.toBeInstanceOf( PayloadTooLargeError, ); + expect({ pulls, cancelled }).toEqual({ pulls: 0, cancelled: true }); }); test("aborts the stream when a missing Content-Length hides an oversized body", async () => { @@ -56,15 +72,80 @@ describe("readRequestTextWithLimit", () => { ); }); - test("rejects malformed Content-Length instead of treating it as undeclared", async () => { + test("cancels a malformed declared body without pulling it", async () => { + let pulls = 0; + let cancelled = false; + const body = new ReadableStream( + { + pull(controller) { + pulls += 1; + controller.enqueue(new Uint8Array([1])); + }, + cancel() { + cancelled = true; + }, + }, + { highWaterMark: 0 }, + ); const request = new Request("http://broker.test/api", { method: "POST", headers: { "content-length": "01" }, - body: "x", - }); + body, + duplex: "half", + } as RequestInit); await expect(readRequestBytesWithLimit(request, 1024)).rejects.toBeInstanceOf( InvalidContentLengthError, ); + expect({ pulls, cancelled }).toEqual({ pulls: 0, cancelled: true }); + }); + + test("rejects a full aggregate budget before pulling and cancels the request body", async () => { + let pulls = 0; + let cancelled = false; + const body = new ReadableStream( + { + pull(controller) { + pulls += 1; + controller.enqueue(new Uint8Array([1])); + controller.close(); + }, + cancel() { + cancelled = true; + }, + }, + { highWaterMark: 0 }, + ); + const request = new Request("http://broker.test/api", { + method: "POST", + body, + duplex: "half", + } as RequestInit); + const budget = new ResourceBudget(4, "http"); + const occupied = budget.reserve(4); + + await expect(readRequestBytesWithReservation(request, 4, budget)).rejects.toBeInstanceOf( + BrokerCapacityError, + ); + expect({ pulls, cancelled, used: budget.used }).toEqual({ + pulls: 0, + cancelled: true, + used: 4, + }); + occupied.release(); + }); + + test("charges the maximum before reading a body whose declared length is dishonest", async () => { + const budget = new ResourceBudget(8, "http"); + const request = new Request("http://broker.test/api", { + method: "POST", + headers: { "content-length": "1" }, + body: "1234", + }); + const read = await readRequestBytesWithReservation(request, 4, budget); + expect(new TextDecoder().decode(read.bytes)).toBe("1234"); + expect(budget.used).toBe(4); + read.reservation.release(); + expect(budget.used).toBe(0); }); test("accounts source-plus-merged peak and transfers retained body capacity", async () => { @@ -119,6 +200,34 @@ describe("readRequestTextWithLimit", () => { }); describe("boundHttpResponse", () => { + test("rejects a full aggregate budget before pulling and cancels the response body", async () => { + let pulls = 0; + let cancelled = false; + const body = new ReadableStream( + { + pull(controller) { + pulls += 1; + controller.enqueue(new Uint8Array([1])); + controller.close(); + }, + cancel() { + cancelled = true; + }, + }, + { highWaterMark: 0 }, + ); + const budget = new ResourceBudget(4, "response"); + const occupied = budget.reserve(4); + const response = await boundHttpResponse(new Response(body), 4, budget); + expect(response.status).toBe(503); + expect({ pulls, cancelled, used: budget.used }).toEqual({ + pulls: 0, + cancelled: true, + used: 4, + }); + occupied.release(); + }); + test("accepts an exact-ceiling response and charges it until pull", async () => { const budget = new ResourceBudget(8, "response"); const response = await boundHttpResponse(new Response("1234"), 4, budget); @@ -162,9 +271,11 @@ describe("boundHttpResponse", () => { }); describe("utf8ByteLength", () => { - test("counts multi-byte characters by their encoded size", () => { + test("counts multi-byte characters and replacement sequences without allocating", () => { expect(utf8ByteLength("abc")).toBe(3); expect(utf8ByteLength("é")).toBe(2); expect(utf8ByteLength("😀")).toBe(4); + expect(utf8ByteLength("\ud800")).toBe(3); + expect(utf8ByteLength("\udc00")).toBe(3); }); }); diff --git a/packages/session-broker-core/src/limits.ts b/packages/session-broker-core/src/limits.ts index 7e12f9eb5..c379621c9 100644 --- a/packages/session-broker-core/src/limits.ts +++ b/packages/session-broker-core/src/limits.ts @@ -52,13 +52,32 @@ export class InvalidContentLengthError 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. */ +/** Count UTF-8 bytes without allocating an encoded copy before resource admission. */ export function utf8ByteLength(value: string): number { - return sharedTextEncoder.encode(value).length; + let bytes = 0; + for (let index = 0; index < value.length; index += 1) { + const codeUnit = value.charCodeAt(index); + if (codeUnit <= 0x7f) { + bytes += 1; + } else if (codeUnit <= 0x7ff) { + bytes += 2; + } else if ( + codeUnit >= 0xd800 && + codeUnit <= 0xdbff && + index + 1 < value.length && + value.charCodeAt(index + 1) >= 0xdc00 && + value.charCodeAt(index + 1) <= 0xdfff + ) { + bytes += 4; + index += 1; + } else { + // TextEncoder replaces every unpaired surrogate with the three-byte U+FFFD sequence. + bytes += 3; + } + } + return bytes; } /** @@ -73,24 +92,26 @@ export async function readRequestBytesWithReservation( maxBytes: number, aggregateBudget?: ResourceBudget, ): Promise<{ bytes: Uint8Array; reservation: BudgetReservation }> { + const body = request.body; const declaredHeader = request.headers.get("content-length"); if (declaredHeader !== null && !/^(?:0|[1-9][0-9]*)$/.test(declaredHeader)) { + await body?.cancel().catch(() => {}); throw new InvalidContentLengthError(); } const declared = declaredHeader === null ? null : Number(declaredHeader); if (declared !== null && (!Number.isSafeInteger(declared) || declared > maxBytes)) { + await body?.cancel().catch(() => {}); throw new PayloadTooLargeError(maxBytes); } - const sourceReservations = new ReservationGroup(); const retainedReservation = new ReservationGroup(); - try { - if (aggregateBudget && declared !== null) { - sourceReservations.add(aggregateBudget.reserve(declared)); - } - const body = request.body; - if (!body) return { bytes: new Uint8Array(), reservation: sourceReservations }; + if (!body) return { bytes: new Uint8Array(), reservation: retainedReservation }; + let sourceReservation: BudgetReservation | null = null; + try { + // Stream APIs expose bytes only after pulling them. Reserve the full permitted source before + // the first pull so an unknown or dishonest length cannot create an uncharged transient chunk. + if (aggregateBudget) sourceReservation = aggregateBudget.reserve(maxBytes); const reader = body.getReader(); const chunks: Uint8Array[] = []; let total = 0; @@ -105,11 +126,6 @@ export async function readRequestBytesWithReservation( await reader.cancel().catch(() => {}); throw new PayloadTooLargeError(maxBytes); } - if (aggregateBudget && nextTotal > (declared ?? 0)) { - sourceReservations.add( - aggregateBudget.reserve(nextTotal - Math.max(total, declared ?? 0)), - ); - } total = nextTotal; chunks.push(value); } @@ -117,20 +133,23 @@ export async function readRequestBytesWithReservation( reader.releaseLock(); } - // Retain capacity for the merged copy before allocating it; source chunks stay charged until - // copying completes so aggregate accounting covers the real peak. - if (aggregateBudget) retainedReservation.add(aggregateBudget.reserve(total)); + if (aggregateBudget && sourceReservation) { + sourceReservation = aggregateBudget.resize(sourceReservation, total); + retainedReservation.add(aggregateBudget.reserve(total)); + } const merged = new Uint8Array(total); let offset = 0; for (const chunk of chunks) { merged.set(chunk, offset); offset += chunk.byteLength; } - sourceReservations.release(); + sourceReservation?.release(); + sourceReservation = null; return { bytes: merged, reservation: retainedReservation }; } catch (error) { - sourceReservations.release(); + sourceReservation?.release(); retainedReservation.release(); + if (!body.locked) await body.cancel().catch(() => {}); throw error; } } @@ -157,13 +176,16 @@ export async function boundHttpResponse( } if (!response.body) return response; - const reader = response.body.getReader(); - const sourceReservations = new ReservationGroup(); + let sourceReservation: BudgetReservation | null = null; const retainedReservation = new ReservationGroup(); + let reader: ReadableStreamDefaultReader | null = null; const chunks: Uint8Array[] = []; let total = 0; let transferred = false; try { + // Reserve before the first pull because WHATWG streams do not expose chunk size beforehand. + if (aggregateBudget) sourceReservation = aggregateBudget.reserve(maxBytes); + reader = response.body.getReader(); for (;;) { const { done, value } = await reader.read(); if (done) break; @@ -173,19 +195,20 @@ export async function boundHttpResponse( await reader.cancel().catch(() => {}); return new Response(null, { status: 503 }); } - if (aggregateBudget) sourceReservations.add(aggregateBudget.reserve(value.byteLength)); chunks.push(value); } - // Charge both the source chunks and their merged replacement during the copy peak, then retain - // only the replacement body until the downstream transport first pulls or cancels it. - if (aggregateBudget) retainedReservation.add(aggregateBudget.reserve(total)); + if (aggregateBudget && sourceReservation) { + sourceReservation = aggregateBudget.resize(sourceReservation, total); + retainedReservation.add(aggregateBudget.reserve(total)); + } const body = new Uint8Array(total); let offset = 0; for (const chunk of chunks) { body.set(chunk, offset); offset += chunk.byteLength; } - sourceReservations.release(); + sourceReservation?.release(); + sourceReservation = null; const headers = new Headers(response.headers); headers.set("content-length", String(total)); let delivered = false; @@ -213,14 +236,15 @@ export async function boundHttpResponse( }); } catch (error) { if (error instanceof BrokerCapacityError) { - await reader.cancel().catch(() => {}); + if (reader) await reader.cancel().catch(() => {}); + else await response.body.cancel().catch(() => {}); return new Response(null, { status: 503 }); } throw error; } finally { - sourceReservations.release(); + sourceReservation?.release(); if (!transferred) retainedReservation.release(); - reader.releaseLock(); + reader?.releaseLock(); } } diff --git a/packages/session-broker-node/package.json b/packages/session-broker-node/package.json index c417693e6..aa7631355 100644 --- a/packages/session-broker-node/package.json +++ b/packages/session-broker-node/package.json @@ -20,7 +20,7 @@ "ws": "^8.18.3" }, "engines": { - "bun": ">=1.0.0", + "bun": ">=1.3.14", "node": ">=22" } } diff --git a/packages/session-broker-node/src/serve.test.ts b/packages/session-broker-node/src/serve.test.ts index 876c568dd..29c099523 100644 --- a/packages/session-broker-node/src/serve.test.ts +++ b/packages/session-broker-node/src/serve.test.ts @@ -13,7 +13,6 @@ 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 { @@ -116,13 +115,6 @@ 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({ @@ -148,7 +140,9 @@ describe("session broker node adapter", () => { generation: "generation-1", brokerRevision: 1 as const, ...(appContract ? { appContract } : {}), + callerSessionId: "caller-session-1", requestId: "request-1", + sequence: "1", httpStatus, bodyDigest: "test-digest", daemonKeyId: "daemon-key-1", diff --git a/packages/session-broker-node/src/serve.ts b/packages/session-broker-node/src/serve.ts index f3c680bd9..1fe75fa79 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, }; } @@ -231,6 +233,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) => { @@ -254,11 +257,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 (stopping) { @@ -269,21 +289,26 @@ export async function serveSessionBrokerDaemon< 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) { + const byteLength = Array.isArray(message) + ? message.reduce((total, chunk) => total + chunk.byteLength, 0) + : message.byteLength; + if (byteLength > options.daemon.limits.maxWsMessageBytes) { socket.close(1009, "Message exceeds the session broker size limit."); return; } - const reservation = inboundBudget.tryReserve(bytes.byteLength); + const reservation = inboundBudget.tryReserve(byteLength); if (!reservation) { socket.close(1013, "Session broker inbound pressure exceeded."); return; } try { + // ws usually supplies one Buffer. Concatenate only fragmented array variants, and create a + // zero-copy Buffer view for ArrayBuffer so broker accounting covers any required copy. + const bytes = Array.isArray(message) + ? Buffer.concat(message, byteLength) + : message instanceof ArrayBuffer + ? Buffer.from(message) + : message; let text: string; try { text = new TextDecoder("utf-8", { fatal: true }).decode(bytes); @@ -308,7 +333,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. @@ -323,8 +352,18 @@ export async function serveSessionBrokerDaemon< socket.destroy(); return; } - const pathname = new URL(`http://${options.hostname}:${options.port}${request.url ?? "/"}`) - .pathname; + let pathname: string; + try { + const target = request.url; + if (!target?.startsWith("/") || target.startsWith("//")) { + throw new TypeError("Expected an origin-form WebSocket request target."); + } + pathname = new URL(target, `http://${options.hostname}:${options.port}`).pathname; + } catch { + socket.write("HTTP/1.1 400 Bad Request\r\nContent-Length: 0\r\n\r\n"); + socket.destroy(); + return; + } if (!options.daemon.matchesSocketPath(pathname)) { socket.destroy(); return; diff --git a/packages/session-broker/package.json b/packages/session-broker/package.json index 27dcc8bbc..693d321ee 100644 --- a/packages/session-broker/package.json +++ b/packages/session-broker/package.json @@ -19,7 +19,7 @@ "@hunk/session-broker-core": "workspace:*" }, "engines": { - "bun": ">=1.0.0", + "bun": ">=1.3.14", "node": ">=22" } } diff --git a/packages/session-broker/src/authentication.test.ts b/packages/session-broker/src/authentication.test.ts index 40b5b79dd..ab5d52fec 100644 --- a/packages/session-broker/src/authentication.test.ts +++ b/packages/session-broker/src/authentication.test.ts @@ -269,10 +269,89 @@ describe("session broker signed authentication", () => { "connection-1", ), ).resolves.toMatchObject({ + ack: { + connectionId: "connection-1", + daemonKeyId: "daemon-key-1", + principal: { kind: "producer", scopes: ["register"] }, + }, + assertActive: expect.any(Function), + }); + }); + + test("producer authority rechecks revocation and credential-clear epochs", async () => { + let revoked = false; + const values = await setup({ revoked: () => revoked }); + const proof = await signedHelloProof(values, "producer"); + const authority = await values.authenticator.completeProducerHello( + { challengeId: proof.challengeId, signature: proof.signature }, + "connection-1", + ); + expect(() => authority.assertActive()).not.toThrow(); + expect(JSON.parse(JSON.stringify(authority.ack))).toMatchObject({ connectionId: "connection-1", - daemonKeyId: "daemon-key-1", - principal: { kind: "producer", scopes: ["register"] }, + principal: { kind: "producer" }, }); + expect("assertActive" in authority.ack).toBe(false); + revoked = true; + expect(() => authority.assertActive()).toThrow( + expect.objectContaining({ code: "credential-revoked" }), + ); + + const cleared = await setup(); + const clearedProof = await signedHelloProof(cleared, "producer"); + const clearedAuthority = await cleared.authenticator.completeProducerHello( + { challengeId: clearedProof.challengeId, signature: clearedProof.signature }, + "connection-2", + ); + cleared.authenticator.clear(); + expect(() => clearedAuthority.assertActive()).toThrow( + expect.objectContaining({ code: "invalid-credential" }), + ); + }); + + test("rejects producer acknowledgement when revocation races asynchronous signing", async () => { + let revoked = false; + let revokeAfterSign = false; + const cryptoWithRevocation: SessionBrokerCrypto = { + ...webSessionBrokerCrypto, + async sign(privateKey, value) { + const signature = await webSessionBrokerCrypto.sign(privateKey, value); + if (revokeAfterSign) revoked = true; + return signature; + }, + }; + const values = await setup({ revoked: () => revoked, crypto: cryptoWithRevocation }); + const proof = await signedHelloProof(values, "producer"); + revokeAfterSign = true; + await expect( + values.authenticator.completeProducerHello( + { challengeId: proof.challengeId, signature: proof.signature }, + "connection-1", + ), + ).rejects.toMatchObject({ code: "credential-revoked" }); + }); + + 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(hello.ack.principal.kind).toBe("producer"); + expect(hello.assertActive).not.toThrow(); + revoked = true; + expect(hello.assertActive).toThrow(SessionBrokerAuthenticationError); + revoked = false; + values.setNow(10_001); + expect(hello.assertActive).toThrow(SessionBrokerAuthenticationError); }); test("rejects missing, wrong, expired, revoked, and reused credentials with redacted errors", async () => { @@ -443,7 +522,7 @@ describe("session broker signed authentication", () => { { challengeId: next.challengeId, signature: next.signature }, "connection-2", ), - ).resolves.toMatchObject({ connectionId: "connection-2" }); + ).resolves.toMatchObject({ ack: { connectionId: "connection-2" } }); }); test("clear invalidates response signing across deferred crypto", async () => { @@ -535,7 +614,9 @@ describe("session broker signed authentication", () => { generation: response.generation, brokerRevision: 1 as const, appContract: { appRevision: 1, features: [] }, + callerSessionId: response.callerSessionId, requestId: response.requestId, + sequence: response.sequence, httpStatus: response.httpStatus, bodyDigest: response.bodyDigest, }; @@ -567,6 +648,23 @@ describe("session broker signed authentication", () => { buildBrokerResponseTranscript({ ...transcriptInput, bodyDigest: "tampered" }), ), ).toBe(false); + expect( + await webSessionBrokerCrypto.verify( + values.daemon.publicKey, + signature, + buildBrokerResponseTranscript({ + ...transcriptInput, + callerSessionId: "caller-session-2", + }), + ), + ).toBe(false); + expect( + await webSessionBrokerCrypto.verify( + values.daemon.publicKey, + signature, + buildBrokerResponseTranscript({ ...transcriptInput, sequence: "2" }), + ), + ).toBe(false); }); test("uses bounded collision-safe IDs with a deterministic custom random source", async () => { diff --git a/packages/session-broker/src/authentication.ts b/packages/session-broker/src/authentication.ts index afcceb984..b4dd46f2e 100644 --- a/packages/session-broker/src/authentication.ts +++ b/packages/session-broker/src/authentication.ts @@ -115,6 +115,7 @@ export interface SessionBrokerHelloChallengeRequest { export interface SessionBrokerHelloChallenge { readonly challengeId: string; + readonly generation: string; readonly responderNonce: string; readonly expiresAt: number; readonly daemonKeyId: string; @@ -139,7 +140,7 @@ export interface AuthenticatedCallerSession { readonly daemonSignature: string; } -export interface AuthenticatedProducerHello { +export interface SessionBrokerProducerHelloAck { readonly principal: ProducerPrincipal; readonly connectionId: string; readonly brokerRevision: typeof SESSION_BROKER_PROTOCOL_REVISION; @@ -150,6 +151,13 @@ export interface AuthenticatedProducerHello { readonly daemonSignature: string; } +/** Keep signed wire data separate from server-only authority retained for the live peer. */ +export interface AuthenticatedProducerHello { + readonly ack: SessionBrokerProducerHelloAck; + /** Reject producer work after credential revocation, expiry, or a clear epoch. */ + assertActive(): void; +} + export interface CallerRequestAuthenticationInput { readonly request: Request; readonly body: Uint8Array; @@ -159,7 +167,9 @@ export interface SessionBrokerResponseAuthentication { readonly generation: string; readonly brokerRevision: typeof SESSION_BROKER_PROTOCOL_REVISION; readonly appContract?: BrokerAppContract; + readonly callerSessionId: string; readonly requestId: string; + readonly sequence: string; readonly httpStatus: number; readonly bodyDigest: string; readonly daemonKeyId: string; @@ -541,8 +551,19 @@ 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; +} + /** 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; @@ -635,6 +656,7 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { committed = true; return Object.freeze({ challengeId, + generation: this.config.generation, responderNonce, expiresAt, daemonKeyId: this.config.daemonIdentity.keyId, @@ -774,7 +796,12 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { ), ); this.assertClearEpoch(epoch); - return Object.freeze({ + this.requireActiveGrant(grant); + const assertActive = () => { + this.assertClearEpoch(epoch); + this.requireActiveGrant(grant); + }; + const ack: SessionBrokerProducerHelloAck = Object.freeze({ principal: principalFromGrant(grant), connectionId, brokerRevision: SESSION_BROKER_PROTOCOL_REVISION, @@ -784,6 +811,7 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { daemonKeyId: this.config.daemonIdentity.keyId, daemonSignature, }); + return Object.freeze({ ack, assertActive }); } finally { pending.reservation.release(); } @@ -868,7 +896,7 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { requestId, assertActive, signResponse: (input: CallerResponseSigningInput) => - this.signResponse(requestId, input, assertActive), + this.signResponse(callerSessionId, requestId, sequence, input, assertActive), }); } @@ -910,7 +938,9 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { } private async signResponse( + callerSessionId: string, requestId: string, + sequence: string, input: CallerResponseSigningInput, assertActive: () => void, ): Promise { @@ -935,7 +965,9 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { appId: this.config.appId, generation: this.config.generation, brokerRevision: SESSION_BROKER_PROTOCOL_REVISION, + callerSessionId, requestId, + sequence, httpStatus: input.httpStatus, bodyDigest, ...(appContract ? { appContract } : {}), @@ -948,7 +980,9 @@ export class SessionBrokerAuthenticator implements CallerRequestAuthenticator { generation: this.config.generation, brokerRevision: SESSION_BROKER_PROTOCOL_REVISION, ...(appContract ? { appContract } : {}), + callerSessionId, requestId, + sequence, httpStatus: input.httpStatus, bodyDigest, daemonKeyId: this.config.daemonIdentity.keyId, 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..ac2e29078 --- /dev/null +++ b/packages/session-broker/src/clientAuthentication.test.ts @@ -0,0 +1,433 @@ +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, + SessionBrokerClientAuthenticationError, +} 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("rejects responses replayed across caller sessions or request sequences", async () => { + for (const [field, replacement] of [ + ["callerSessionId", "caller-session-replayed"], + ["sequence", "2"], + ] as const) { + const values = await setup(); + const proofCount = { value: 0 }; + const authenticatedFetch = createFetch(values.authenticator, proofCount); + const tamperingFetch = (async (input: string | URL | Request, init?: RequestInit) => { + const request = input instanceof Request ? input : new Request(input, init); + const response = await authenticatedFetch(request); + if (new URL(request.url).pathname !== "/control") return response; + const envelope = (await response.json()) as { + body: unknown; + authentication: Record; + }; + envelope.authentication[field] = replacement; + return Response.json(envelope, { status: response.status }); + }) 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: tamperingFetch, + }); + + await expect( + client.request("/control", { method: "POST", body: "{}" }), + ).rejects.toBeInstanceOf(SessionBrokerClientAuthenticationError); + } + }); + + 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("delayed stale 401s do not clear an overlapping shared recovery negotiation", async () => { + const values = await setup(); + const proofCount = { value: 0 }; + const authenticatedFetch = createFetch(values.authenticator, proofCount); + let staleControls = 0; + let staleMode = false; + let releaseFirst401!: () => void; + let releaseSecond401!: () => void; + let releaseRecovery!: () => void; + let signalBothStale!: () => void; + let signalRecovery!: () => void; + const first401 = new Promise((resolve) => (releaseFirst401 = resolve)); + const second401 = new Promise((resolve) => (releaseSecond401 = resolve)); + const recoveryGate = new Promise((resolve) => (releaseRecovery = resolve)); + const bothStale = new Promise((resolve) => (signalBothStale = resolve)); + const recoveryStarted = new Promise((resolve) => (signalRecovery = resolve)); + const fetchImpl = (async (input: string | URL | Request, init?: RequestInit) => { + const request = input instanceof Request ? input : new Request(input, init); + const pathname = new URL(request.url).pathname; + if (staleMode && pathname === "/session-auth/challenge" && proofCount.value === 1) { + signalRecovery(); + await recoveryGate; + } + if (staleMode && pathname === "/control" && proofCount.value === 1 && staleControls < 2) { + staleControls += 1; + if (staleControls === 2) signalBothStale(); + await (staleControls === 1 ? first401 : second401); + return Response.json({ error: "stale-session" }, { status: 401 }); + } + return 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: fetchImpl, + }); + + await expect(client.request("/control")).resolves.toBeInstanceOf(Response); + staleMode = true; + const first = client.request("/control"); + const second = client.request("/control"); + await bothStale; + releaseFirst401(); + await recoveryStarted; + releaseSecond401(); + releaseRecovery(); + + await expect(Promise.all([first, second])).resolves.toHaveLength(2); + expect(proofCount.value).toBe(2); + }); + + test("coalesces concurrent negotiations while retaining unique request sequences", 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), + }); + + const responses = await Promise.all( + Array.from({ length: 32 }, () => client.request("/control", { method: "POST", body: "{}" })), + ); + expect(responses).toHaveLength(32); + expect(proofCount.value).toBe(1); + }); + + test("aborting one negotiation waiter does not cancel another", async () => { + const values = await setup(); + const proofCount = { value: 0 }; + const authenticatedFetch = createFetch(values.authenticator, proofCount); + let release!: () => void; + const gate = new Promise((resolve) => { + release = resolve; + }); + const gatedFetch = (async (input: string | URL | Request, init?: RequestInit) => { + const request = input instanceof Request ? input : new Request(input, init); + if (new URL(request.url).pathname === "/session-auth/challenge") await gate; + return 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: gatedFetch, + }); + const controller = new AbortController(); + const aborted = client.request("/control", { signal: controller.signal }); + const surviving = client.request("/control"); + controller.abort(new Error("caller stopped")); + release(); + + await expect(aborted).rejects.toThrow("caller stopped"); + await expect(surviving).resolves.toBeInstanceOf(Response); + expect(proofCount.value).toBe(1); + }); + + test("clear invalidates stale negotiation installation and failures allow retry", async () => { + const values = await setup(); + const authenticatedFetch = createFetch(values.authenticator, { value: 0 }); + let challengeCount = 0; + let releaseFirst!: () => void; + const firstGate = new Promise((resolve) => { + releaseFirst = resolve; + }); + const fetchImpl = (async (input: string | URL | Request, init?: RequestInit) => { + const request = input instanceof Request ? input : new Request(input, init); + if (new URL(request.url).pathname === "/session-auth/challenge") { + challengeCount += 1; + if (challengeCount === 1) await firstGate; + if (challengeCount === 2) return new Response("no", { status: 503 }); + } + return 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: fetchImpl, + }); + + const stale = client.request("/control"); + client.clear(); + await expect(client.request("/control")).rejects.toBeInstanceOf( + SessionBrokerClientAuthenticationError, + ); + releaseFirst(); + await expect(stale).rejects.toBeInstanceOf(SessionBrokerClientAuthenticationError); + await expect(client.request("/control")).resolves.toBeInstanceOf(Response); + expect(challengeCount).toBe(3); + }); + + 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("cancels malformed or oversized declared response bodies before rejecting", async () => { + for (const declared of ["invalid", "33"]) { + const values = await setup(); + let cancelled = false; + 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 () => + new Response( + new ReadableStream({ + cancel() { + cancelled = true; + }, + }), + { headers: { "content-length": declared } }, + )) as unknown as typeof fetch, + }); + + await expect(client.request("/control")).rejects.toThrow( + "daemon identity could not be verified", + ); + expect(cancelled).toBe(true); + } + }); + + test("rejects challenge records with unknown or dangerous own keys", async () => { + for (const extra of ["extra", "__proto__"]) { + const values = await setup(); + const authenticatedFetch = createFetch(values.authenticator, { + value: 0, + }); + const fetchWithExtra = (async (input: string | URL | Request, init?: RequestInit) => { + const request = input instanceof Request ? input : new Request(input, init); + const response = await authenticatedFetch(request); + if (new URL(request.url).pathname !== "/session-auth/challenge") return response; + const challenge = (await response.json()) as Record; + Object.defineProperty(challenge, extra, { + value: true, + enumerable: true, + }); + return Response.json(challenge); + }) 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: fetchWithExtra, + }); + + 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..304d93718 --- /dev/null +++ b/packages/session-broker/src/clientAuthentication.ts @@ -0,0 +1,629 @@ +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 SessionBrokerProducerHelloAck, + 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; +} + +export class SessionBrokerClientAuthenticationError extends Error { + constructor() { + super("Session broker authentication failed or the daemon identity could not be verified."); + this.name = "SessionBrokerClientAuthenticationError"; + } +} + +function clientAuthError(): never { + throw new SessionBrokerClientAuthenticationError(); +} + +function exactRecord(value: unknown, keys: readonly string[]): Record { + if (!value || typeof value !== "object" || Array.isArray(value)) clientAuthError(); + const prototype = Object.getPrototypeOf(value); + if (prototype !== Object.prototype && prototype !== null) clientAuthError(); + const record = value as Record; + const ownKeys = Object.keys(record); + if ( + ownKeys.some((key) => ["__proto__", "prototype", "constructor"].includes(key)) || + ownKeys.length !== keys.length || + keys.some((key) => !Object.hasOwn(record, key)) || + ownKeys.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: SessionBrokerProducerHelloAck, +): 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)) { + await response.body?.cancel().catch(() => undefined); + 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 negotiation: { epoch: number; promise: Promise } | null = null; + private authenticationEpoch = 0; + 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.ensureNegotiated(init.signal); + } + const attemptedSession = this.session; + const attemptedEpoch = this.authenticationEpoch; + const response = await this.signedRequest(path, init, options.targetSpecific ?? false); + if (response === null) { + if (attempt === 0) { + // A delayed 401 from an older session must not invalidate recovery another request already + // completed. Only the request that still owns the current authentication epoch clears it. + if (this.session === attemptedSession && this.authenticationEpoch === attemptedEpoch) { + this.clear(); + } + continue; + } + clientAuthError(); + } + return response; + } + clientAuthError(); + } + + clear() { + this.authenticationEpoch += 1; + this.session = null; + this.sequence = null; + this.pending = null; + this.negotiation = null; + } + + /** Share one negotiation while allowing each waiting request to abort independently. */ + private async ensureNegotiated(signal?: AbortSignal | null) { + if (!this.negotiation) { + const epoch = this.authenticationEpoch; + const promise = this.negotiate(epoch).finally(() => { + if (this.negotiation?.promise === promise) this.negotiation = null; + }); + this.negotiation = { epoch, promise }; + } + const promise = this.negotiation.promise; + if (!signal) return promise; + if (signal.aborted) throw signal.reason ?? new DOMException("Aborted", "AbortError"); + return new Promise((resolve, reject) => { + const onAbort = () => reject(signal.reason ?? new DOMException("Aborted", "AbortError")); + signal.addEventListener("abort", onAbort, { once: true }); + void promise + .then(resolve, reject) + .finally(() => signal.removeEventListener("abort", onAbort)); + }); + } + + private async negotiate(epoch: number) { + 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), + }); + 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), + }); + 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); + if (epoch !== this.authenticationEpoch) clientAuthError(); + 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, + session.callerSessionId, + requestId, + sequence, + pending.challenge.generation, + targetSpecific, + ); + } catch { + if (response.status === 401) return null; + clientAuthError(); + } + return new Response(JSON.stringify(envelope.body), { + status: response.status, + headers: { "content-type": "application/json" }, + }); + } + + private async verifyResponse( + envelope: SessionBrokerAuthenticatedResponse, + status: number, + callerSessionId: string, + requestId: string, + sequence: string, + generation: string, + targetSpecific: boolean, + ) { + const envelopeRecord = exactRecord(envelope, ["body", "authentication"]); + const authenticationKeys = [ + "generation", + "brokerRevision", + "callerSessionId", + "requestId", + "sequence", + "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.callerSessionId !== callerSessionId || + auth.requestId !== requestId || + auth.sequence !== sequence || + 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, + callerSessionId, + requestId, + sequence, + 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 df9765ecc..2eed8032c 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(); @@ -853,6 +909,107 @@ describe("session broker connection", () => { connection.stop(); }); + test("rejects producer hello wrappers with unknown or dangerous keys", 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"], + }; + const connection = createSessionBrokerConnection({ + 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 }, + }, + reconnectDelayMs: 10_000, + }); + connection.start(); + socket.emitOpen(); + + for (const message of [ + '{"type":"hello-challenge","challenge":{},"extra":true}', + '{"type":"hello-challenge","challenge":{},"__proto__":{}}', + ]) { + socket.readyState = 1; + socket.emitMessage(message); + expect(socket.lastClose).toEqual({ + code: 1008, + reason: "Session broker authentication failed.", + }); + } + connection.stop(); + }); + + test("prepares reconnect once per attempt and stops after awaited preparation", async () => { + const sockets: TestSocket[] = []; + const warnings: string[] = []; + let attempts = 0; + let markSecondSocketCreated!: () => void; + const secondSocketCreated = new Promise((resolve) => (markSecondSocketCreated = resolve)); + let prepare = async () => { + attempts += 1; + if (attempts === 1) throw new Error("incumbent still alive"); + }; + const connection = createSessionBrokerConnection({ + url: "ws://broker.test/session", + createSocket: () => { + const socket = new TestSocket(); + sockets.push(socket); + if (sockets.length === 2) markSecondSocketCreated(); + return socket; + }, + registration: createRegistration(), + snapshot: createSnapshot(), + protocolParsers, + reconnectDelayMs: 1, + prepareReconnect: () => prepare(), + onWarning: (message) => warnings.push(message), + }); + connection.start(); + sockets[0]!.emitOpen(); + sockets[0]!.emitClose(); + await secondSocketCreated; + + expect(attempts).toBe(2); + expect(warnings).toEqual(["incumbent still alive"]); + expect(sockets).toHaveLength(2); + + let release!: () => void; + let markPreparationStarted!: () => void; + const gate = new Promise((resolve) => (release = resolve)); + const preparationStarted = new Promise((resolve) => (markPreparationStarted = resolve)); + sockets[1]!.emitOpen(); + prepare = () => { + markPreparationStarted(); + return gate; + }; + sockets[1]!.emitClose(); + await preparationStarted; + connection.stop(); + release(); + await Bun.sleep(0); + expect(sockets).toHaveLength(2); + }); + 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 dd2406c72..6d17d331b 100644 --- a/packages/session-broker/src/connection.ts +++ b/packages/session-broker/src/connection.ts @@ -3,6 +3,7 @@ import { BrokerProtocolError, ReservationGroup, ResourceBudget, + parseExactBrokerRecord, resolveSessionBrokerLimits, utf8ByteLength, type BudgetReservation, @@ -15,6 +16,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 { + SessionBrokerProducerHelloAck, + SessionBrokerHelloChallenge, + SessionBrokerHelloChallengeRequest, +} from "./authentication"; +import type { ProducerGrant } from "@hunk/session-broker-core"; import type { SessionBrokerConnectionCloseDirective, SessionBrokerSocketCloseEvent, @@ -39,6 +54,13 @@ function commandValueBytes(value: unknown): number { return utf8ByteLength(serialized); } +/** Parse one exact handshake wrapper before its payload reaches the authentication parser. */ +function exactHelloEnvelope(value: unknown, type: string, payloadKey: "challenge" | "ack") { + const record = parseExactBrokerRecord(value, ["type", payloadKey] as const, [] as const); + if (record.type !== type) throw new BrokerProtocolError("invalid-discriminant"); + return record; +} + export interface SessionBrokerConnectionBridge< ServerMessage extends SessionServerMessage = SessionServerMessage, Result = unknown, @@ -59,10 +81,19 @@ 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; resolveClose?: (event: SessionBrokerSocketCloseEvent) => SessionBrokerConnectionCloseDirective; + /** Prepare application-owned discovery before one reconnect attempt opens a new socket. */ + prepareReconnect?: () => Promise; + onConnected?: () => void; onWarning?: (message: string) => void; limits?: SessionBrokerLimitOptions["limits"]; unsafeLimits?: SessionBrokerLimitOptions["unsafeLimits"]; @@ -80,6 +111,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 +125,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 +181,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 +201,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,16 +235,26 @@ 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) => { @@ -205,7 +267,10 @@ export class SessionBrokerConnection< socket.close(1009, "Message exceeds the session broker size limit."); return; } - + 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 }; @@ -227,8 +292,13 @@ export class SessionBrokerConnection< }; socket.onclose = (event) => { + const wasAuthenticated = this.activeSocket === socket; + 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(); } @@ -243,7 +313,11 @@ export class SessionBrokerConnection< return; } - const directive = this.options.resolveClose?.(event) ?? { reconnect: true }; + const directive = this.options.resolveClose?.({ + code: event.code, + reason: event.reason, + authenticated: wasAuthenticated, + }) ?? { reconnect: true }; if (directive.warning) { this.options.onWarning?.(directive.warning); } @@ -260,19 +334,85 @@ export class SessionBrokerConnection< }; } - private scheduleReconnect(delayMs = this.options.reconnectDelayMs ?? DEFAULT_RECONNECT_DELAY_MS) { - if (this.reconnectTimer || this.stopped) { - return; + 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.options.onConnected?.(); + 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); + const authentication = this.options.producerAuthentication!; + const hello = this.producerHellos.get(socket); + if (!hello) throw new Error(); + if (hello.pending === undefined) { + const envelope = exactHelloEnvelope(value, "hello-challenge", "challenge"); + hello.pending = null; + const pending = await answerSessionBrokerHelloChallenge( + { ...authentication, endpoint: this.options.url }, + hello.request, + envelope.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) throw new Error(); + const envelope = exactHelloEnvelope(value, "hello-ack", "ack"); + await verifyProducerHelloAck(hello.pending, envelope.ack as SessionBrokerProducerHelloAck); + 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; this.reconnectTimer = setTimeout(() => { this.reconnectTimer = null; - this.connect(); + void this.prepareAndReconnect(); }, delayMs); - this.reconnectTimer.unref?.(); } + /** Run app discovery once per retry while retaining this connection's aggregate budgets. */ + private async prepareAndReconnect() { + try { + await this.options.prepareReconnect?.(); + } catch (error) { + if (this.stopped) return; + this.options.onWarning?.( + error instanceof Error ? error.message : "Session broker reconnect preparation failed.", + ); + this.scheduleReconnect(); + return; + } + if (!this.stopped) this.connect(); + } + private startHeartbeat() { if (this.heartbeatTimer) { return; @@ -298,17 +438,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 4371990f8..9322a3041 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; @@ -131,7 +136,9 @@ function authenticatedRequest(principal: CallerPrincipal): AuthenticatedCallerRe generation: "generation-1", brokerRevision: 1, ...(input.appContract ? { appContract: input.appContract } : {}), + callerSessionId: "caller-session-1", requestId: "request-1", + sequence: "1", httpStatus: input.httpStatus, bodyDigest: "test-body-digest", daemonKeyId: "daemon-key-1", @@ -170,9 +177,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,11 +194,84 @@ function createConnection() { close(code?: number, reason?: string) { closed = { code, reason }; }, + markAuthenticated() { + authenticated = true; + }, }, }; } describe("session broker daemon", () => { + test("closes late producer messages without parsing or recreating state after shutdown", () => { + let registrationParses = 0; + const parsers = createSessionBrokerProtocolParsers({ + appRevision: 1, + features: [], + parseRegistration: (value) => { + registrationParses += 1; + return parseSessionRegistrationEnvelope(value, parseInfo); + }, + parseSnapshot: (value) => parseSessionSnapshotEnvelope(value, parseState), + commands: [], + }); + const broker = new SessionBroker({ protocolParsers: parsers }); + const daemon = createSessionBrokerDaemon({ broker }); + const peer = createConnection(); + daemon.shutdown(); + daemon.handleConnectionMessage( + peer.connection, + JSON.stringify({ + type: "register", + registration: createRegistration(), + snapshot: createSnapshot(), + }), + ); + + expect(registrationParses).toBe(0); + expect(broker.listSessions()).toHaveLength(0); + expect(peer.closed).toEqual({ + code: 1001, + reason: "Session broker shutting down.", + }); + }); + + test("does not send a deferred producer challenge after shutdown", async () => { + let release!: () => void; + const gate = new Promise((resolve) => { + release = resolve; + }); + const daemon = createSessionBrokerDaemon({ + broker: createBroker(), + appId: "dev.example", + appRevision: 1, + producerEndpoint: "ws://broker.test/session", + helloAuthenticator: { + async issueChallenge() { + await gate; + return { challengeId: "challenge-1" } as SessionBrokerHelloChallenge; + }, + async completeCallerHello() { + throw new Error("not used"); + }, + async completeProducerHello() { + throw new Error("not used"); + }, + }, + }); + const peer = createConnection(); + daemon.handleConnectionMessage( + peer.connection, + JSON.stringify({ type: "hello-init", hello: {} }), + ); + await Bun.sleep(0); + daemon.shutdown(); + release(); + await Bun.sleep(0); + + expect(peer.sent).toEqual([]); + expect(daemon.listSessions()).toEqual([]); + }); + test("serves health and raw list/get requests when the HTTP API is enabled", async () => { const daemon = createSessionBrokerDaemon({ broker: createBroker(), @@ -315,7 +399,10 @@ describe("session broker daemon", () => { ["HEAD", new Headers({ "content-length": "1" })], ] as const) { const response = await daemon.handleRequest( - new Request("http://broker.test/broker/capabilities", { method, headers }), + new Request("http://broker.test/broker/capabilities", { + method, + headers, + }), ); expect(response?.status).toBe(400); await expect(response?.json()).resolves.toEqual({ @@ -679,6 +766,299 @@ describe("session broker daemon", () => { daemon.shutdown(); }); + test("rejects producer hello wrappers with unknown or dangerous keys", async () => { + let challengeCalls = 0; + const daemon = createSessionBrokerDaemon({ + broker: createBroker(), + appId: "dev.example", + appRevision: 1, + producerEndpoint: "ws://broker.test/session", + helloAuthenticator: { + async issueChallenge() { + challengeCalls += 1; + return { challengeId: "challenge-1" } as SessionBrokerHelloChallenge; + }, + async completeCallerHello() { + throw new Error("not used"); + }, + async completeProducerHello() { + throw new Error("not used"); + }, + }, + }); + + for (const message of [ + '{"type":"hello-init","hello":{},"extra":true}', + '{"type":"hello-init","hello":{},"__proto__":{}}', + ]) { + const peer = createConnection(); + daemon.handleConnectionMessage(peer.connection, message); + await Bun.sleep(0); + expect(peer.closed?.reason).toContain("authentication required"); + } + expect(challengeCalls).toBe(0); + daemon.shutdown(); + }); + + test("pre-registration authentication failures do not postpone idle shutdown", async () => { + const daemon = createSessionBrokerDaemon({ + broker: createBroker(), + appId: "dev.example", + appRevision: 1, + producerEndpoint: "ws://broker.test/session", + idleTimeoutMs: 50, + helloAuthenticator: { + async issueChallenge() { + throw new Error("incompatible"); + }, + async completeCallerHello() { + throw new Error("not used"); + }, + async completeProducerHello() { + throw new Error("not used"); + }, + }, + }); + const activityBeforeRefusal = (daemon as any).lastActivityAt; + await Bun.sleep(10); + const peer = createConnection(); + daemon.handleConnectionMessage( + peer.connection, + JSON.stringify({ type: "hello-init", hello: {} }), + ); + await Bun.sleep(0); + daemon.handleConnectionClose(peer.connection); + expect((daemon as any).lastActivityAt).toBe(activityBeforeRefusal); + + const outcome = await Promise.race([ + daemon.stopped.then(() => "stopped"), + Bun.sleep(80).then(() => "timed-out"), + ]); + expect(outcome).toBe("stopped"); + }); + + test("requires reconnect scope and rechecks retained producer authority", async () => { + let operations: readonly ProducerOperation[] = ["register"]; + let active = true; + let activeChecks = 0; + const principal = () => ({ + kind: "producer" as const, + appId: "dev.example", + principalId: "producer-1", + keyId: "producer-key-1", + grantId: "producer-grant-1", + scopes: operations, + }); + const broker = createBroker(); + const daemon = createSessionBrokerDaemon({ + broker, + 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 { + ack: { + principal: principal(), + connectionId: String(connectionId), + brokerRevision: 1, + appRevision: 1, + features: [], + helloTranscriptHash: "transcript-1", + daemonKeyId: "daemon-key-1", + daemonSignature: "signature-1", + }, + assertActive() { + activeChecks += 1; + if (!active) throw new Error("revoked"); + }, + } satisfies AuthenticatedProducerHello; + }, + }, + }); + 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); + + // Deliver work that was already queued on the displaced transport after replacement. Its + // retired authentication state must not let it reclaim the session. + register(first.connection); + expect(first.closed?.reason).toContain("authentication required"); + expect(daemon.listSessions()).toHaveLength(1); + + const checksBeforeRevocation = activeChecks; + const sentBeforeRevocation = replacement.sent.length; + active = false; + await expect( + broker.dispatchCommand({ + selector: { sessionId: "session-1" }, + command: "annotate", + input: { summary: "must stay private" }, + timeoutMessage: "timed out", + }), + ).rejects.toThrow("revoked"); + expect(replacement.sent).toHaveLength(sentBeforeRevocation); + expect(replacement.closed?.reason).toContain("authority expired"); + + daemon.handleConnectionMessage( + replacement.connection, + JSON.stringify({ type: "heartbeat", sessionId: "session-1" }), + ); + daemon.handleConnectionMessage( + replacement.connection, + JSON.stringify({ + type: "snapshot", + sessionId: "session-1", + snapshot: createSnapshot({ selectedIndex: 1 }), + }), + ); + expect(activeChecks).toBe(checksBeforeRevocation + 3); + expect(replacement.closed?.reason).toContain("authority expired"); + expect(daemon.getSession({ sessionId: "session-1" })).toMatchObject({ + snapshot: { state: { selectedIndex: 0 } }, + }); + }); + + test("retains active reconnect ownership when reconciliation prunes a stale session", async () => { + let operations: readonly ProducerOperation[] = ["register", "reconnect"]; + let principalId = "producer-1"; + let active = true; + const broker = createBroker(); + const daemon = createSessionBrokerDaemon({ + broker, + 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 { + ack: { + principal: { + kind: "producer" as const, + appId: "dev.example", + principalId, + keyId: "producer-key-1", + grantId: "producer-grant-1", + scopes: operations, + }, + connectionId: String(connectionId), + brokerRevision: 1, + appRevision: 1, + features: [], + helloTranscriptHash: "transcript-1", + daemonKeyId: "daemon-key-1", + daemonSignature: "signature-1", + }, + assertActive() { + if (!active) throw new Error("revoked"); + }, + } satisfies AuthenticatedProducerHello; + }, + }, + }); + const authenticate = async (peer: ReturnType) => { + daemon.handleConnectionMessage( + peer.connection, + JSON.stringify({ type: "hello-init", hello: {} }), + ); + await Bun.sleep(0); + daemon.handleConnectionMessage( + peer.connection, + JSON.stringify({ type: "hello-proof", proof: {} }), + ); + await Bun.sleep(0); + }; + const register = (peer: ReturnType) => + daemon.handleConnectionMessage( + peer.connection, + JSON.stringify({ + type: "register", + registration: createRegistration(), + snapshot: createSnapshot(), + }), + ); + + const owner = createConnection(); + await authenticate(owner); + register(owner); + const staleAt = Date.now() + 2_000; + expect(broker.pruneStaleSessions({ ttlMs: 1_000, now: staleAt })).toBe(1); + (daemon as any).reconcileProducerOwners(); + expect(owner.closed?.reason).toContain("authority retired"); + + principalId = "different-producer"; + operations = ["reconnect"]; + const differentBinding = createConnection(); + await authenticate(differentBinding); + register(differentBinding); + expect(differentBinding.closed?.reason).toContain("scope rejected"); + + principalId = "producer-1"; + operations = ["register"]; + const registerOnly = createConnection(); + await authenticate(registerOnly); + register(registerOnly); + expect(registerOnly.closed?.reason).toContain("scope rejected"); + + operations = ["reconnect"]; + const replacement = createConnection(); + await authenticate(replacement); + register(replacement); + expect(replacement.closed).toBeNull(); + expect(daemon.listSessions()).toHaveLength(1); + + active = false; + expect(broker.pruneStaleSessions({ ttlMs: 1_000, now: staleAt + 1 })).toBe(1); + (daemon as any).reconcileProducerOwners(); + expect((daemon as any).producerReconnects.size).toBe(0); + daemon.shutdown(); + }); + test("rejects duplicate live registration without retiring the owner", () => { const daemon = createSessionBrokerDaemon({ broker: createBroker(), @@ -919,6 +1299,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", { @@ -928,13 +1317,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: {}, }) @@ -945,7 +1334,7 @@ describe("session broker daemon", () => { ( await post({ action: "dispatch", - selector: { sessionId: "missing" }, + selector: { sessionId: "session-1" }, command: "allowed", commandVersion: 0, input: {}, @@ -960,6 +1349,7 @@ describe("session broker daemon", () => { test("returns stable redacted authentication failures without invoking app authorization", async () => { let authorized = false; + const auditedOperations: string[] = []; const daemon = createSessionBrokerDaemon({ broker: createBroker(), exposeHttpApi: true, @@ -975,6 +1365,9 @@ describe("session broker daemon", () => { authorized = true; return true; }, + audit: (event) => { + auditedOperations.push(event.operation); + }, }); const response = await daemon.handleRequest( new Request("http://broker.test/broker", { @@ -989,8 +1382,22 @@ describe("session broker daemon", () => { error: "authentication-failed", code: "invalid-signature", }); + const capabilitiesResponse = await daemon.handleRequest( + new Request("http://broker.test/broker/capabilities"), + ); + expect(capabilitiesResponse?.status).toBe(401); + const customResponse = await daemon.handleAuthenticatedControl( + new Request("http://broker.test/custom"), + { + authenticationFailureOperation: "shutdown", + resolve: () => ({ operation: "shutdown" }), + handle: () => ({ body: { ok: true } }), + }, + ); + expect(customResponse.status).toBe(401); expect(authorized).toBe(false); expect(responseText).not.toContain("private"); + expect(auditedOperations).toEqual(["unknown", "diagnostics", "shutdown"]); daemon.shutdown(); }); @@ -1042,7 +1449,7 @@ describe("session broker daemon", () => { test("supports a lower route-specific body ceiling and releases its reservation", async () => { const daemon = createSessionBrokerDaemon({ broker: createBroker(), - limits: { maxHttpBodyBytes: 8, maxInFlightHttpBodyBytes: 8 }, + limits: { maxHttpBodyBytes: 4, maxInFlightHttpBodyBytes: 8 }, }); let handled = 0; const invoke = (body: string) => diff --git a/packages/session-broker/src/daemon.ts b/packages/session-broker/src/daemon.ts index 90719ddbe..def4278c6 100644 --- a/packages/session-broker/src/daemon.ts +++ b/packages/session-broker/src/daemon.ts @@ -8,15 +8,18 @@ import { mergeSessionBrokerLimits, DEFAULT_SESSION_BROKER_LIMITS, callerPrincipalAllows, + producerPrincipalAllows, canonicalizeJson, isValidBrokerAppId, isValidBrokerIdentifier, + parseExactBrokerRecord, isValidBrokerRevision, utf8ByteLength, type BudgetReservation, type CallerOperation, type CallerPrincipal, type CanonicalJsonValue, + type ProducerPrincipal, type SessionBrokerLimitOptions, type SessionBrokerLimits, type SessionServerMessage, @@ -27,6 +30,7 @@ import { SessionBrokerAuthenticationError, type AuthenticatedCallerRequest, type CallerRequestAuthenticator, + type SessionBrokerHelloAuthenticator, } from "./authentication"; import { parseSessionBrokerJsonBytes, @@ -65,6 +69,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 +92,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; @@ -129,6 +150,38 @@ function defaultTimeoutMessage(command: string) { return `Timed out waiting for the session to handle ${command}.`; } +interface ProducerAuthenticationState { + state: "challenged" | "authenticated"; + principal?: ProducerPrincipal; + sessionId?: string; + assertActive?: () => void; + brokerPeer?: SessionBrokerPeer; +} + +interface ProducerOwner { + connection: SessionBrokerPeer; + brokerPeer: SessionBrokerPeer; + principal: ProducerPrincipal; +} + +/** Parse one exact producer handshake wrapper before forwarding its opaque payload. */ +function exactProducerHelloEnvelope(value: unknown, type: string, payloadKey: "hello" | "proof") { + const record = parseExactBrokerRecord(value, ["type", payloadKey] as const, [] as const); + if (record.type !== type) throw new BrokerProtocolError("invalid-discriminant"); + return record; +} + +/** 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. @@ -157,7 +210,18 @@ 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, + ProducerAuthenticationState + >(); + private readonly producerOwners = new Map(); + private readonly producerReconnects = new Map< + string, + { principal: ProducerPrincipal; disconnectedAt: number } + >(); private readonly audit?: SessionBrokerAuditHook; private readonly httpControlBudget: ResourceBudget; private readonly httpBodyBudget: ResourceBudget; @@ -217,6 +281,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; @@ -255,6 +327,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, @@ -304,6 +380,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. @@ -313,6 +410,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. @@ -332,6 +430,101 @@ export class SessionBrokerDaemon< } handleConnectionMessage(connection: SessionBrokerPeer, message: unknown) { + if (this.shuttingDown) { + connection.close?.(1001, "Session broker shutting down."); + return; + } + 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; + } + try { + authentication.assertActive?.(); + } catch { + connection.close?.(INCOMPATIBLE_PAYLOAD_CLOSE_CODE, "Session producer authority expired."); + 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?: ProducerAuthenticationState, + ) { + try { + const value = parseSessionBrokerJsonText(message); + if (!current) { + const envelope = exactProducerHelloEnvelope(value, "hello-init", "hello"); + const challenged = { state: "challenged" as const }; + this.producerAuthentication.set(connection, challenged); + const challenge = await this.helloAuthenticator!.issueChallenge( + envelope.hello, + this.producerEndpoint!, + ); + if (this.shuttingDown || this.producerAuthentication.get(connection) !== challenged) { + return; + } + connection.send(JSON.stringify({ type: "hello-challenge", challenge })); + return; + } + if (current.state !== "challenged") throw new Error(); + const envelope = exactProducerHelloEnvelope(value, "hello-proof", "proof"); + const connectionId = `b_${crypto.randomUUID().replaceAll("-", "")}_0`; + const authority = await this.helloAuthenticator!.completeProducerHello( + envelope.proof, + connectionId, + ); + if (this.shuttingDown || this.producerAuthentication.get(connection) !== current) { + return; + } + authority.assertActive(); + const brokerPeer: SessionBrokerPeer = { + send: (data) => { + try { + authority.assertActive(); + } catch (error) { + connection.close?.( + INCOMPATIBLE_PAYLOAD_CLOSE_CODE, + "Session producer authority expired.", + ); + throw error; + } + return connection.send(data); + }, + close: (code, reason) => connection.close?.(code, reason), + markAuthenticated: () => connection.markAuthenticated?.(), + }; + this.producerAuthentication.set(connection, { + state: "authenticated", + principal: authority.ack.principal, + assertActive: authority.assertActive, + brokerPeer, + }); + connection.send(JSON.stringify({ type: "hello-ack", ack: authority.ack })); + } catch { + this.producerAuthentication.delete(connection); + connection.close?.( + INCOMPATIBLE_PAYLOAD_CLOSE_CODE, + "Session broker authentication required; upgrade Hunk.", + ); + } + } + + private handleAuthenticatedConnectionMessage(connection: SessionBrokerPeer, message: unknown) { + if (this.shuttingDown) { + connection.close?.(1001, "Session broker shutting down."); + return; + } + let parsed; try { parsed = this.protocolParsers.parseClientMessage(parseSessionBrokerJsonText(message)); @@ -340,12 +533,39 @@ export class SessionBrokerDaemon< return; } + const producerAuthentication = this.producerAuthentication.get(connection); + const brokerPeer = producerAuthentication?.brokerPeer ?? connection; 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; + } + const replacedConnection = owner?.connection !== connection ? owner?.connection : undefined; const registrationResult = this.broker.registerSession( - connection, + brokerPeer, parsed.registration, parsed.snapshot, + { replaceOwner: replacedConnection !== undefined }, ); if (registrationResult === "invalid") { // Close immediately when the registration payload is incompatible so the session does not @@ -362,15 +582,37 @@ export class SessionBrokerDaemon< connection.close?.(1013, "Session broker capacity exceeded."); return; } + if (registrationResult === "shutdown") { + connection.close?.(1001, "Session broker shutting down."); + return; + } + if (producerAuthentication?.principal) { + // Retire the displaced transport before publishing the new owner. A queued message from + // the old socket must re-enter as unauthenticated and can never reclaim the session. + if (replacedConnection) this.producerAuthentication.delete(replacedConnection); + producerAuthentication.sessionId = sessionId; + this.producerOwners.set(sessionId, { + connection, + brokerPeer, + 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( - connection, + brokerPeer, parsed.sessionId, parsed.snapshot, ); @@ -392,7 +634,11 @@ export class SessionBrokerDaemon< break; } case "heartbeat": { - const seenResult = this.broker.markSessionSeen(connection, parsed.sessionId); + if (this.producerEndpoint && producerAuthentication?.sessionId !== parsed.sessionId) { + connection.close?.(INCOMPATIBLE_PAYLOAD_CLOSE_CODE, "Session producer scope rejected."); + return; + } + const seenResult = this.broker.markSessionSeen(brokerPeer, parsed.sessionId); if (seenResult === "not-owner") { connection.close?.(INCOMPATIBLE_PAYLOAD_CLOSE_CODE, "Session ownership rejected."); return; @@ -402,7 +648,7 @@ export class SessionBrokerDaemon< break; } case "command-result": { - const result = this.broker.handleCommandResult(connection, parsed); + const result = this.broker.handleCommandResult(brokerPeer, parsed); if (result === "not-owner") { connection.close?.(INCOMPATIBLE_PAYLOAD_CLOSE_CODE, "Command ownership rejected."); return; @@ -422,8 +668,79 @@ export class SessionBrokerDaemon< } handleConnectionClose(connection: SessionBrokerPeer) { - this.broker.unregisterConnection(connection); - this.noteActivity(); + 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); + try { + if (!authentication.assertActive) throw new Error("Producer authority is unavailable."); + authentication.assertActive(); + this.rememberProducerReconnect(sessionId, authentication.principal); + } catch { + // Revoked grants must not leave behind reconnect ownership. + } + } + } + this.broker.unregisterConnection(authentication?.brokerPeer ?? connection); + // Pre-registration authentication failures must not postpone quiescent shutdown. This is also + // what lets a newer client wait out an incompatible incumbent without keeping it alive. + if (!this.producerEndpoint || sessionId !== undefined) 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) { + const sessionIsLive = live.has(sessionId); + let authorityIsActive = !this.helloAuthenticator; + if (this.helloAuthenticator) { + try { + const authentication = this.producerAuthentication.get(owner.connection); + if (authentication?.state === "authenticated" && authentication.assertActive) { + authentication.assertActive(); + authorityIsActive = true; + } + } catch { + authorityIsActive = false; + } + } + if (sessionIsLive && authorityIsActive) continue; + + // Clear both ownership maps before closing so queued messages and the close callback cannot + // reuse the retired transport. A stale session keeps only its still-active binding. + this.producerOwners.delete(sessionId); + this.producerAuthentication.delete(owner.connection); + this.broker.unregisterConnection(owner.brokerPeer); + if (!sessionIsLive && authorityIsActive) { + this.rememberProducerReconnect(sessionId, owner.principal); + } + owner.connection.close?.(1000, "Session producer authority retired."); + } + } + + /** Retain one bounded producer binding after its authenticated transport disconnects. */ + private rememberProducerReconnect(sessionId: string, principal: ProducerPrincipal) { + 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, + disconnectedAt: Date.now(), + }); + } + + /** 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.")) { @@ -443,6 +760,8 @@ export class SessionBrokerDaemon< } this.broker.shutdown(error); + this.producerOwners.clear(); + this.producerReconnects.clear(); this.callerAuthenticator?.clear?.(); this.resolveStopped?.(); this.resolveStopped = null; @@ -456,6 +775,7 @@ export class SessionBrokerDaemon< if (removed > 0) { this.noteActivity(); } + this.reconcileProducerOwners(); }, this.staleSessionSweepIntervalMs); this.sweepTimer.unref?.(); @@ -504,10 +824,78 @@ export class SessionBrokerDaemon< }, remainingMs); } + /** Authenticate, authorize, execute, and sign one app-owned finite JSON control. */ + async handleAuthenticatedControl( + request: Request, + options: { + resolve: (body: Uint8Array) => SessionBrokerAuthenticatedControlFacts; + authenticationFailureOperation?: CallerOperation; + 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, + options.authenticationFailureOperation ?? "unknown", + ); + 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, - operation: CallerOperation, + operation: CallerOperation | "unknown", ): Promise { const requestId = request.headers.get("x-session-broker-request-id") ?? undefined; try { @@ -612,10 +1000,18 @@ export class SessionBrokerDaemon< let responseStatus = status; const targetContract = targetSpecific && this.appRevision !== undefined - ? { appContract: { appRevision: this.appRevision, features: [] as const } } + ? { + appContract: { + appRevision: this.appRevision, + features: [] as const, + }, + } : {}; if (utf8ByteLength(canonicalizeJson(structuredBody)) > this.limits.maxHttpResponseBytes) { - structuredBody = { error: "capacity-exceeded", resource: "maxHttpResponseBytes" }; + structuredBody = { + error: "capacity-exceeded", + resource: "maxHttpResponseBytes", + }; responseStatus = 503; } const authentication = await authenticated.signResponse({ @@ -629,7 +1025,10 @@ export class SessionBrokerDaemon< }; let serializedEnvelope = canonicalizeJson(envelope as unknown as CanonicalJsonValue); if (utf8ByteLength(serializedEnvelope) > this.limits.maxHttpResponseBytes) { - structuredBody = { error: "capacity-exceeded", resource: "maxHttpResponseBytes" }; + structuredBody = { + error: "capacity-exceeded", + resource: "maxHttpResponseBytes", + }; responseStatus = 503; envelope = { body: structuredBody, @@ -698,7 +1097,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); @@ -749,7 +1152,7 @@ export class SessionBrokerDaemon< } // Authenticate the exact transport bytes before decoding or interpreting attacker-controlled JSON. - const authenticated = await this.authenticateRequest(request, body, "list"); + const authenticated = await this.authenticateRequest(request, body, "unknown"); if (authenticated instanceof Response) return authenticated; let input; @@ -761,7 +1164,22 @@ 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 = { @@ -769,7 +1187,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, @@ -788,15 +1205,14 @@ 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/packages/session-broker/src/types.ts b/packages/session-broker/src/types.ts index 0c35c5fa2..700e613d2 100644 --- a/packages/session-broker/src/types.ts +++ b/packages/session-broker/src/types.ts @@ -83,6 +83,8 @@ export interface SessionBrokerHealth { export interface SessionBrokerSocketCloseEvent { code: number; reason: string; + /** Whether this socket completed authentication and became the active producer transport. */ + authenticated?: boolean; } export interface SessionBrokerSocketMessageEvent { @@ -126,7 +128,7 @@ export interface SessionBrokerAuditEvent { readonly principalId?: string; readonly keyId?: string; readonly sessionId?: string; - readonly operation: CallerOperation; + readonly operation: CallerOperation | "unknown"; readonly command?: string; readonly commandVersion?: number; readonly requestId?: string; diff --git a/src/session/broker/brokerServer.ts b/src/session/broker/brokerServer.ts index 5be7a5ac8..fb3c6972a 100644 --- a/src/session/broker/brokerServer.ts +++ b/src/session/broker/brokerServer.ts @@ -495,6 +495,8 @@ 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) => diff --git a/test/fixtures/sessionBrokerAdapterConformance.json b/test/fixtures/sessionBrokerAdapterConformance.json index 381186040..a763579f9 100644 --- a/test/fixtures/sessionBrokerAdapterConformance.json +++ b/test/fixtures/sessionBrokerAdapterConformance.json @@ -4,7 +4,8 @@ }, "inbound": { "oversizedCloseCode": 1009, - "pressureCloseCode": 1013, + "bunNativeOversizedCloseCodes": [1006, 1009], + "admissionHttpStatus": 503, "maxMessageBytes": 8388608 }, "outbound": { diff --git a/test/session-broker-node/adapter.test.mjs b/test/session-broker-node/adapter.test.mjs index aa583c85b..90f8d9439 100644 --- a/test/session-broker-node/adapter.test.mjs +++ b/test/session-broker-node/adapter.test.mjs @@ -99,11 +99,13 @@ function fakeDaemon(overrides = {}, behavior = {}) { maxHttpResponseBytes: 8 * 1024 * 1024, maxInFlightHttpResponseBytes: 64 * 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() {}, @@ -203,11 +205,15 @@ test("Node adapter waits for active HTTP handlers and preserves bodyless framing 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) => @@ -216,6 +222,17 @@ test("Node adapter consumes the shared text/binary/oversize/pressure corpus", as : undefined, }); try { + const malformedUpgrade = await rawHttp(port, [ + "GET * HTTP/1.1", + `Host: 127.0.0.1:${port}`, + "Connection: Upgrade", + "Upgrade: websocket", + "", + "", + ]); + assert.match(malformedUpgrade, /^HTTP\/1\.1 400/); + + // A successful follow-up request proves the malformed upgrade did not escape the listener. const boundedResponse = await fetch(`http://127.0.0.1:${port}/large`); assert.equal(boundedResponse.status, 503); assert.equal(await boundedResponse.text(), ""); @@ -223,7 +240,17 @@ test("Node adapter consumes the shared text/binary/oversize/pressure corpus", as 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 fullAdmission = await rawHttp(port, [ + "GET /session HTTP/1.1", + `Host: 127.0.0.1:${port}`, + "Connection: Upgrade", + "Upgrade: websocket", + "Sec-WebSocket-Version: 13", + "Sec-WebSocket-Key: dGVzdC1zZXNzaW9uLWtleQ==", + "", + "", + ]); + assert.match(fullAdmission, new RegExp(`^HTTP/1.1 ${corpus.inbound.admissionHttpStatus}`)); const exactClosed = closeCode(exact); exact.close(); await exactClosed; @@ -264,25 +291,32 @@ test("Node adapter consumes the shared text/binary/oversize/pressure corpus", as const outbound = await openSocket(`ws://127.0.0.1:${outboundPort}/session`); const outboundClosed = closeCode(outbound); outbound.send("trigger"); - assert.equal(await outboundClosed, 1013); + assert.equal(await outboundClosed, corpus.outbound.pressureCloseCode); } finally { await outboundRunning.stop(); await outboundRunning.stopped; } - const pressurePort = await reservePort(); - const pressureRunning = await serveSessionBrokerDaemon({ - daemon: fakeDaemon({ maxWsMessageBytes: 8, maxInFlightWsBytes: 0 }), + const handlerPort = await reservePort(); + const handlerRunning = await serveSessionBrokerDaemon({ + daemon: fakeDaemon( + {}, + { + handleConnectionMessage: () => { + throw new Error("unexpected handler failure"); + }, + }, + ), hostname: "127.0.0.1", - port: pressurePort, + port: handlerPort, }); 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); + const handlerFailure = await openSocket(`ws://127.0.0.1:${handlerPort}/session`); + const handlerFailureClosed = closeCode(handlerFailure); + handlerFailure.send("trigger"); + assert.equal(await handlerFailureClosed, 1011); } finally { - await pressureRunning.stop(); - await pressureRunning.stopped; + await handlerRunning.stop(); + await handlerRunning.stopped; } });