diff --git a/hub/docs/README.md b/hub/docs/README.md index 5d471ed3..79034d43 100644 --- a/hub/docs/README.md +++ b/hub/docs/README.md @@ -9,6 +9,7 @@ Hub is being rebuilt. The [Hub overview](../README.md) describes what the baseli - [Database](database.md): Prisma schema conventions, migrations, and the shared client. - [Logging](logging.md): application events, automatic request records, and payload safety. - [Background jobs](background-jobs.md): PostgreSQL queues, workers, schedules, setup, and test isolation. +- [WebSocket updates](websockets.md): local broadcasts, subscriptions, access hooks, and connection limits. - [Testing and operations](testing-and-operations.md): baseline checks and their limits. - [Storybook](storybook.md): retained presentation, fixtures, and screenshot evidence. - [Prototype canvas](prototype-canvas.md): the separately retained extension. diff --git a/hub/docs/architecture.md b/hub/docs/architecture.md index bf72e320..d5bf7066 100644 --- a/hub/docs/architecture.md +++ b/hub/docs/architecture.md @@ -22,6 +22,8 @@ The [Pino integration](logging.md) writes structured application and HTTP logs t The [background-job foundation](background-jobs.md) uses BullMQ with PostgreSQL and separate on-demand and scheduled worker processes. Setup migrates BullMQ's own schema independently of Prisma. Product processors and schedules are not registered yet. +The [WebSocket foundation](websockets.md) attaches to the API's HTTP server and broadcasts to local channel subscribers without persistence. The application rejects upgrades until access checks are supplied. Product channels, client recovery, and communication between processes are not implemented. + ## Frontend ```mermaid diff --git a/hub/docs/development.md b/hub/docs/development.md index 452682d3..3d344861 100644 --- a/hub/docs/development.md +++ b/hub/docs/development.md @@ -38,6 +38,8 @@ npm run dev Open . Vite serves the frontend and forwards `/api` requests to Express at . `GET /api/health` returns JSON; other API routes have not been built. Both servers accept connections only from this machine. Stop each with Ctrl+C. +Vite also forwards WebSocket upgrades under `/api`. The `/api/realtime` endpoint rejects connections until application access checks are connected; see [WebSocket updates](websockets.md). Stopping the API closes its WebSocket connections as well. + The API reads an optional `PORT` environment variable, defaulting to 3001. If you change it, update the local Vite proxy target to match. The editor forwards ports 3001, 5173, and 6006 from the devcontainer. API startup reads `.env` when present; Compose's environment takes precedence. ## Background workers diff --git a/hub/docs/websockets.md b/hub/docs/websockets.md new file mode 100644 index 00000000..5b2530ae --- /dev/null +++ b/hub/docs/websockets.md @@ -0,0 +1,95 @@ +# WebSocket updates + +Hub can send live updates to the client app and Copilot plugin through the `ws` integration under `server/integrations/ws/`. The integration owns connections and channel subscriptions in one API process. It does not store messages, replay missed events, or deliver between processes. + +The application attaches `/api/realtime` to the same HTTP server as Express and passes its broadcast function to `createApp()`. Vite forwards WebSocket upgrades under `/api` during development. The endpoint rejects connections until application access checks are supplied. There are no comment channels or product events yet. + +## Server operations + +`attachRealtimeServer({ server, path, logger, limits?, authenticate?, authorizeSubscription? })` receives an existing Node HTTP server and an injected Pino logger. It does not start listening or open another port. It returns an instance with two operations: + +- `broadcast(channel, { type, payload })` sends a JSON event to that channel's connected subscribers in this process. It validates the event even when there are no subscribers. Invalid input throws; an oversized event throws `RangeError`. Individual connection failures are logged and do not stop delivery to other subscribers. Sending is not an acknowledgment that a client processed the event. +- `close()` returns a promise, stops upgrades, cancels pending admission, clears timers and subscriptions, and closes sockets. It terminates sockets that do not close before the shutdown deadline. Repeated calls return the same promise. The caller still owns the HTTP server. Broadcasting after close throws. + +Application startup creates the shared server, attaches WebSockets, and installs Express as its HTTP request handler: + +```ts +const server = createServer() +const realtime = attachRealtimeServer({ server, path: "/api/realtime", logger }) +server.on("request", createApp({ broadcast: realtime.broadcast })) +server.listen(port) +``` + +`createApp({ logger?, broadcast? })` and route registration accept an optional broadcast function. Features that receive it can call `broadcast?.(channel, event)` after a successful database commit. Without it, HTTP features work without live updates. Pass the broadcaster explicitly rather than using a global server. If a supplied broadcaster fails, report that failure rather than pretending it sent the update. + +Keep comment storage and channel selection in the feature's domain. A feature must not report a saved comment as unsaved merely because a later notification failed. + +## Client messages + +Browser clients use native `WebSocket`. Other clients use a WebSocket client supported by their runtime. Clients may subscribe and unsubscribe, but cannot publish application events. + +| Client message | Server reply | +| --- | --- | +| `{"type":"subscribe","channel":"example:123"}` | `{"type":"subscribed","channel":"example:123"}` | +| `{"type":"unsubscribe","channel":"example:123"}` | `{"type":"unsubscribed","channel":"example:123"}` | + +Channel names are case-sensitive, contain 1 to 200 ASCII characters, and start with a letter or digit. Later characters may also contain `:`, `.`, `_`, `/`, or `-`. Repeating a subscription does not add another membership or consume another slot. Unsubscribing from a channel without membership is harmless. + +Broadcasts use this shape: + +```json +{ + "type": "event", + "channel": "example:123", + "event": { "type": "changed", "payload": { "id": "123" } } +} +``` + +Event types contain 1 to 100 characters; payloads must be JSON values. Feature code owns the payload contract. + +Rejected client messages receive `{"type":"error","code":"INVALID_MESSAGE"}` or the same shape with code `FORBIDDEN` or `SUBSCRIPTION_LIMIT`. Invalid JSON, binary messages, unknown operations, extra fields, and invalid channel names are rejected. Logs contain static reasons and connection IDs, not payloads, channel names, credentials, or raw request URLs. The first rejected upgrade per status is logged immediately; repeats are summarized at each heartbeat and shutdown instead of producing an unrestricted stream of rejection logs. + +The server installs membership before confirming a subscription. Clients should subscribe, wait for confirmation, then fetch current state. Events received during that fetch must trigger another refresh rather than being overwritten by an older response. Clients must resubscribe and refetch after reconnecting; this integration does not implement client recovery. + +## Access checks + +The integration does not implement sign-in or choose sessions, tokens, or user records. Application startup may supply: + +- `authenticate(request, signal)`: returns the caller's identity, or `null` to reject admission. +- `authorizeSubscription(identity, channel, signal)`: returns whether that caller may subscribe. + +These callbacks can be asynchronous. If either callback is missing, upgrades return HTTP 403. Authentication returning `null` produces 401, a failed check produces 500, and an admission timeout produces 408. A denied subscription returns `FORBIDDEN`; a failed subscription check closes the socket with code 1011. + +Callbacks must honor their abort signal and finish promptly. Admission timeout, disconnect, and shutdown cancel outstanding checks. Subscription checks also have a deadline. The integration bounds unresolved access checks even if a callback ignores cancellation; capacity remains occupied until those checks settle. Only one subscription check may be outstanding per socket; sending more messages while it is pending closes that socket with code 1008 instead of building a queue. + +When authentication is connected, its application policy must also check browser origins, credential expiry, and access revocation for existing subscriptions. Express middleware does not run automatically for WebSocket upgrades. Test-only callbacks that allow synthetic callers must not be used by application startup. + +## Resource limits + +The defaults in `server/integrations/ws/constants.ts` are initial safety limits, not measured production capacity. Supply positive integer overrides through `limits`. Tune connection limits against measured process memory with headroom before deployment. + +| Limit | Default | +| --- | --- | +| Pending and established connections per instance | 1,000 | +| Connections per peer address | 50 | +| Tracked peer addresses | 4,096 | +| Upgrade attempts per peer address per minute | 60 | +| Incoming message and serialized broadcast size | 16 KiB | +| Subscriptions per connection | 32 | +| Client messages, pings, and pongs per connection per second | 20 | +| Queued outgoing bytes per connection | 256 KiB | +| Heartbeat interval | 30 seconds | +| Admission and subscription check deadline | 5 seconds | +| Shutdown deadline | 5 seconds | + +Capacity is reserved before asynchronous admission and released on failure or disconnect. Full connection or address-tracking capacity returns HTTP 503; address connection and attempt limits return 429. Idle address counters expire after their attempt window. Compression is disabled. The transport also limits each message to 128 fragments and its receiver to 128 buffered chunks, so tiny fragments cannot bypass the payload byte limit through excessive bookkeeping. + +The integration uses the direct peer address and ignores forwarded headers. Behind a proxy, callers may therefore share one address. Operators must configure proxy-side caller limits and tune the peer cap for that topology; do not trust arbitrary forwarded headers. The proxy must support upgrades and limit handshakes and connections before they reach Node. Application limits do not protect against every form of network exhaustion. + +Oversized incoming messages close with code 1009. Excess message rates close with 1008. Dead or slow connections are terminated. Heartbeat checks, address tracking, subscriptions, and pending access checks all have bounded state. + +## Delivery limits + +Any route using the injected broadcaster reaches only this API process. A worker or another API process cannot reach these sockets without a shared signaling layer. Add that layer before porting a feature that publishes across processes. + +Notifications are best effort. A client can miss an event while disconnected, and a process can fail between saving a record and publishing. Feature clients should reconcile on reconnect and focus, with periodic refresh if a bounded delay is required. Durable records remain the source of truth. diff --git a/hub/package-lock.json b/hub/package-lock.json index 25ceab39..4b6b9cd1 100644 --- a/hub/package-lock.json +++ b/hub/package-lock.json @@ -21,6 +21,7 @@ "pino-http": "11.0.0", "react": "19.3.0", "react-dom": "19.3.0", + "ws": "8.22.0", "zod": "4.6.5" }, "devDependencies": { @@ -37,6 +38,7 @@ "@types/pg": "8.23.1", "@types/react": "19.3.0", "@types/react-dom": "19.3.0", + "@types/ws": "8.18.2", "@typescript/native": "npm:typescript@7.0.2", "@vitejs/plugin-react": "6.1.2", "ajv": "8.20.0", @@ -4048,6 +4050,16 @@ "dev": true, "license": "MIT" }, + "node_modules/@types/ws": { + "version": "8.18.2", + "resolved": "https://registry.npmjs.org/@types/ws/-/ws-8.18.2.tgz", + "integrity": "sha512-67MQl+fpWKVTT1NYdnmo3U4sc/xPo/zQBncVnI74qmQa0z/b+1g6iYqNmGCPbxO+zz2aklb08a0oHfegiVd0/w==", + "dev": true, + "license": "MIT", + "dependencies": { + "@types/node": "*" + } + }, "node_modules/@typescript-eslint/eslint-plugin": { "version": "8.71.2", "resolved": "https://registry.npmjs.org/@typescript-eslint/eslint-plugin/-/eslint-plugin-8.71.2.tgz", @@ -10038,10 +10050,9 @@ "license": "ISC" }, "node_modules/ws": { - "version": "8.21.3", - "resolved": "https://registry.npmjs.org/ws/-/ws-8.21.3.tgz", - "integrity": "sha512-201TZ/kPWxoPr/OKWjquZR1SWKXcvxdH+e1xrx89b3YbmzLMFCLfnaG1HFIgWzJOEWZ7MvpK++odZufgYR50Rw==", - "dev": true, + "version": "8.22.0", + "resolved": "https://registry.npmjs.org/ws/-/ws-8.22.0.tgz", + "integrity": "sha512-Ydggc987+RO0AnWtZ/7Wq9FtNvcrL1b/RO0ud9mWjUPgDrsAAwQSF51sm2hm1XofbU/4jkpGEsLFsZZxU+1DOg==", "license": "MIT", "engines": { "node": ">=10.0.0" diff --git a/hub/package.json b/hub/package.json index 259f8e00..1334554c 100644 --- a/hub/package.json +++ b/hub/package.json @@ -58,6 +58,7 @@ "pg": "8.23.1", "react": "19.3.0", "react-dom": "19.3.0", + "ws": "8.22.0", "zod": "4.6.5" }, "devDependencies": { @@ -74,6 +75,7 @@ "@types/pg": "8.23.1", "@types/react": "19.3.0", "@types/react-dom": "19.3.0", + "@types/ws": "8.18.2", "@typescript/native": "npm:typescript@7.0.2", "@vitejs/plugin-react": "6.1.2", "ajv": "8.20.0", diff --git a/hub/server/create-app/index.test.ts b/hub/server/create-app/index.test.ts index c9453fb5..a050819f 100644 --- a/hub/server/create-app/index.test.ts +++ b/hub/server/create-app/index.test.ts @@ -18,6 +18,28 @@ function parseRecord(line: string) { } describe("createApp()", () => { + it("accepts a broadcast dependency without changing HTTP routes", async () => { + const broadcast = vi.fn() + const server = createServer(createApp({ broadcast })) + + try { + server.listen(0, "127.0.0.1") + await once(server, "listening") + const { port } = server.address() as AddressInfo + const health = await fetch(`http://127.0.0.1:${port}/api/health`) + const missing = await fetch(`http://127.0.0.1:${port}/api/realtime`) + + expect(health.status).toBe(200) + expect(await health.json()).toEqual({ status: "ok" }) + expect(missing.status).toBe(404) + expect(broadcast).not.toHaveBeenCalled() + } finally { + await new Promise((resolve, reject) => { + server.close((error) => (error ? reject(error) : resolve())) + }) + } + }) + it("automatically logs health and not-found responses with the injected logger", async () => { const records: Record[] = [] const logger = createLogger({ diff --git a/hub/server/create-app/index.ts b/hub/server/create-app/index.ts index a3954deb..6ea4dbb4 100644 --- a/hub/server/create-app/index.ts +++ b/hub/server/create-app/index.ts @@ -1,19 +1,20 @@ import express from "express" -import type { Logger } from "pino" import { handleError } from "../integrations/express/handle-error/index.ts" import { respondNotFound } from "../integrations/express/respond-not-found/index.ts" import { createRequestLogger } from "../integrations/pino/create-request-logger/index.ts" import { routes } from "../routes/index.ts" +import type { CreateAppOptions } from "./types.ts" + /** Give every endpoint the same request logging and safe error responses. */ -export function createApp({ logger }: { logger?: Logger } = {}) { +export function createApp({ logger, broadcast }: CreateAppOptions = {}) { const app = express() app.disable("x-powered-by") app.use(createRequestLogger(logger)) app.use(express.json()) - routes(app) + routes(app, { broadcast }) app.use(respondNotFound) app.use(handleError) diff --git a/hub/server/create-app/types.ts b/hub/server/create-app/types.ts new file mode 100644 index 00000000..0d07eda1 --- /dev/null +++ b/hub/server/create-app/types.ts @@ -0,0 +1,8 @@ +import type { Logger } from "pino" + +import type { Broadcast } from "../integrations/ws/types.ts" + +export interface CreateAppOptions { + logger?: Logger + broadcast?: Broadcast +} diff --git a/hub/server/index.ts b/hub/server/index.ts index 64bafe13..5136dded 100644 --- a/hub/server/index.ts +++ b/hub/server/index.ts @@ -1,13 +1,31 @@ import "dotenv/config" +import { createServer } from "node:http" + import { createApp } from "./create-app/index.ts" import { log } from "./integrations/pino/log/index.ts" +import { logger } from "./integrations/pino/index.ts" +import { attachRealtimeServer } from "./integrations/ws/attach-realtime-server/index.ts" +import { shutdownServer } from "./shutdown-server/index.ts" const port = Number(process.env.PORT ?? 3001) if (!Number.isInteger(port) || port < 1 || port > 65535) { throw new Error("PORT must be an integer between 1 and 65535.") } -createApp().listen(port, "127.0.0.1", () => { +const server = createServer() +const realtime = attachRealtimeServer({ server, path: "/api/realtime", logger }) +server.on("request", createApp({ broadcast: realtime.broadcast })) +server.listen(port, "127.0.0.1", () => { log("info", "Hub API listening", { host: "127.0.0.1", port }) }) + +let shuttingDown = false +function shutdown() { + if (shuttingDown) return + shuttingDown = true + void shutdownServer(server, realtime) +} + +process.once("SIGINT", shutdown) +process.once("SIGTERM", shutdown) diff --git a/hub/server/integrations/ws/attach-realtime-server/index.test.ts b/hub/server/integrations/ws/attach-realtime-server/index.test.ts new file mode 100644 index 00000000..085bb85a --- /dev/null +++ b/hub/server/integrations/ws/attach-realtime-server/index.test.ts @@ -0,0 +1,319 @@ +import { once } from "node:events" +import { describe, expect, it, vi } from "vitest" +import WebSocket from "ws" + +import { deferred } from "../../../../test/ws/deferred" +import { openServer } from "../../../../test/ws/open-server" +import { parseMessage } from "../../../../test/ws/parse-message" +import { rejected } from "../../../../test/ws/rejected" +import { sendMessage } from "../../../../test/ws/send-message" + +import type { RealtimeOptions } from "../types" + +describe("attachRealtimeServer()", () => { + it("rejects upgrades when no access policy is configured", async () => { + const hub = await openServer() + + expect(await rejected(hub.client())).toBe(403) + }) + + it("rejects unauthenticated callers before upgrading", async () => { + const hub = await openServer({ + authenticate: () => null, + authorizeSubscription: () => true, + }) + + expect(await rejected(hub.client())).toBe(401) + }) + + it("passes the caller's identity to authorization and never subscribes denied channels", async () => { + const authorizeSubscription = vi.fn((identity: string, channel: string) => + Promise.resolve(identity === "reader" && channel === "allowed"), + ) + const hub = await openServer({ + authenticate: () => "reader", + authorizeSubscription, + }) + const socket = await hub.connect() + const received: unknown[] = [] + socket.on("message", (data) => received.push(parseMessage(data))) + + expect(await sendMessage(socket, { type: "subscribe", channel: "private" })).toEqual({ + type: "error", + code: "FORBIDDEN", + }) + hub.realtime.broadcast("private", { type: "secret", payload: "not delivered" }) + await sendMessage(socket, { type: "subscribe", channel: "allowed" }) + + expect(authorizeSubscription.mock.calls[0]?.slice(0, 2)).toEqual(["reader", "private"]) + expect(received).toEqual([ + { type: "error", code: "FORBIDDEN" }, + { type: "subscribed", channel: "allowed" }, + ]) + }) + + it("delivers only to current subscribers and releases their subscription capacity", async () => { + const hub = await openServer({ + authenticate: () => "reader", + authorizeSubscription: () => true, + limits: { maxSubscriptionsPerConnection: 1 }, + }) + const member = await hub.connect() + const outsider = await hub.connect() + const received: unknown[] = [] + outsider.on("message", (data) => received.push(parseMessage(data))) + expect(await sendMessage(member, { type: "subscribe", channel: "one" })).toEqual({ + type: "subscribed", + channel: "one", + }) + expect(await sendMessage(member, { type: "subscribe", channel: "two" })).toEqual({ + type: "error", + code: "SUBSCRIPTION_LIMIT", + }) + const delivery = once(member, "message") as Promise<[Buffer]> + const event = { type: "updated", payload: { count: 1 } } + + hub.realtime.broadcast("one", event) + + const [data] = await delivery + expect(JSON.parse(String(data))).toEqual({ type: "event", channel: "one", event }) + expect(await sendMessage(member, { type: "unsubscribe", channel: "one" })).toEqual({ + type: "unsubscribed", + channel: "one", + }) + expect(await sendMessage(member, { type: "subscribe", channel: "two" })).toEqual({ + type: "subscribed", + channel: "two", + }) + const remaining: unknown[] = [] + member.on("message", (message) => remaining.push(parseMessage(message))) + const nextDelivery = once(member, "message") + hub.realtime.broadcast("one", { type: "removed", payload: null }) + hub.realtime.broadcast("two", event) + await nextDelivery + await sendMessage(outsider, { type: "unsubscribe", channel: "one" }) + + expect(remaining).toEqual([{ type: "event", channel: "two", event }]) + expect(received).toEqual([{ type: "unsubscribed", channel: "one" }]) + }) + + it("rejects malformed JSON without making the connection unusable", async () => { + const hub = await openServer({ + authenticate: () => "reader", + authorizeSubscription: () => true, + }) + const socket = await hub.connect() + const received = once(socket, "message") as Promise<[Buffer]> + + socket.send("{") + + const [data] = await received + expect(JSON.parse(String(data))).toEqual({ type: "error", code: "INVALID_MESSAGE" }) + expect(await sendMessage(socket, { type: "subscribe", channel: "one" })).toEqual({ + type: "subscribed", + channel: "one", + }) + }) + + it("reserves capacity during authentication and retains it after timeout until the check settles", async () => { + const identity = deferred() + const authenticate = vi.fn["authenticate"]>>( + () => identity.promise, + ) + const hub = await openServer({ + authenticate, + authorizeSubscription: () => true, + limits: { maxConnections: 1, upgradeTimeoutMs: 100 }, + }) + const pending = rejected(hub.client()) + await vi.waitFor(() => expect(authenticate).toHaveBeenCalledTimes(1)) + + expect(await rejected(hub.client())).toBe(503) + expect(await pending).toBe(408) + expect(authenticate.mock.calls[0]?.[1].aborted).toBe(true) + expect(await rejected(hub.client())).toBe(503) + identity.resolve("reader") + const replacement = await hub.connect() + expect(replacement.readyState).toBe(WebSocket.OPEN) + }) + + it("times out subscription checks and retains their capacity until they settle", async () => { + const decision = deferred() + const authorizeSubscription = vi + .fn["authorizeSubscription"]>>() + .mockImplementationOnce(() => decision.promise) + .mockResolvedValue(true) + const hub = await openServer({ + authenticate: () => "reader", + authorizeSubscription, + limits: { maxConnections: 1, upgradeTimeoutMs: 100 }, + }) + const socket = await hub.connect() + const closed = once(socket, "close") as Promise<[number]> + socket.send(JSON.stringify({ type: "subscribe", channel: "one" })) + + const [code] = await closed + expect(code).toBe(1008) + expect(authorizeSubscription.mock.calls[0]?.[2].aborted).toBe(true) + expect(await rejected(hub.client())).toBe(503) + decision.resolve(true) + const replacement = await hub.connect() + expect(await sendMessage(replacement, { type: "subscribe", channel: "one" })).toEqual({ + type: "subscribed", + channel: "one", + }) + }) + + it("limits the direct peer address without trusting forwarded addresses", async () => { + const hub = await openServer({ + authenticate: () => "reader", + authorizeSubscription: () => true, + limits: { maxConnectionsPerAddress: 1 }, + }) + const first = await hub.connect({ headers: { "x-forwarded-for": "192.0.2.1" } }) + + expect(await rejected(hub.client({ headers: { "x-forwarded-for": "192.0.2.2" } }))).toBe(429) + const closed = once(first, "close") + first.close() + await closed + expect((await hub.connect()).readyState).toBe(WebSocket.OPEN) + }) + + it("bounds upgrade attempts even when authentication denies every request", async () => { + const authenticate = vi.fn(() => null) + const hub = await openServer({ + authenticate, + authorizeSubscription: () => true, + limits: { maxUpgradeAttemptsPerMinute: 2 }, + }) + + expect(await rejected(hub.client())).toBe(401) + expect(await rejected(hub.client())).toBe(401) + expect(await rejected(hub.client())).toBe(429) + expect(authenticate).toHaveBeenCalledTimes(2) + }) + + it("closes oversized inbound messages with the protocol size code", async () => { + const hub = await openServer({ + authenticate: () => "reader", + authorizeSubscription: () => true, + limits: { maxMessageBytes: 64 }, + }) + const socket = await hub.connect() + const closed = once(socket, "close") as Promise<[number]> + + socket.send("a".repeat(65)) + + const [code] = await closed + expect(code).toBe(1009) + }) + + it("rejects oversized broadcasts without disconnecting a healthy member", async () => { + const hub = await openServer({ + authenticate: () => "reader", + authorizeSubscription: () => true, + limits: { maxMessageBytes: 128 }, + }) + const socket = await hub.connect() + await sendMessage(socket, { type: "subscribe", channel: "one" }) + + expect(() => + hub.realtime.broadcast("one", { type: "large", payload: "a".repeat(128) }), + ).toThrow(RangeError) + const delivery = once(socket, "message") as Promise<[Buffer]> + hub.realtime.broadcast("one", { type: "small", payload: null }) + const [data] = await delivery + expect(JSON.parse(String(data))).toMatchObject({ event: { type: "small" } }) + }) + + it("closes a connection that exceeds its message rate", async () => { + const hub = await openServer({ + authenticate: () => "reader", + authorizeSubscription: () => true, + limits: { maxMessagesPerSecond: 1 }, + }) + const socket = await hub.connect() + await sendMessage(socket, { type: "subscribe", channel: "one" }) + const closed = once(socket, "close") as Promise<[number]> + + socket.send(JSON.stringify({ type: "unsubscribe", channel: "one" })) + + const [code] = await closed + expect(code).toBe(1008) + }) + + it("terminates a member before delivery and framing exceed its send buffer budget", async () => { + const event = { type: "updated", payload: "a".repeat(60) } + const bytes = Buffer.byteLength(JSON.stringify({ type: "event", channel: "one", event })) + const hub = await openServer({ + authenticate: () => "reader", + authorizeSubscription: () => true, + limits: { maxBufferedBytes: bytes }, + }) + const socket = await hub.connect() + await sendMessage(socket, { type: "subscribe", channel: "one" }) + const closed = once(socket, "close") as Promise<[number]> + + hub.realtime.broadcast("one", event) + + const [code] = await closed + expect(code).toBe(1006) + }) + + it("terminates clients that do not answer heartbeats but keeps responsive clients", async () => { + const hub = await openServer({ + authenticate: () => "reader", + authorizeSubscription: () => true, + limits: { heartbeatIntervalMs: 40 }, + }) + const responsive = await hub.connect() + const unresponsive = await hub.connect({ autoPong: false }) + const closed = once(unresponsive, "close") as Promise<[number]> + + const [code] = await closed + + expect(code).toBe(1006) + expect(responsive.readyState).toBe(WebSocket.OPEN) + expect(await sendMessage(responsive, { type: "subscribe", channel: "one" })).toEqual({ + type: "subscribed", + channel: "one", + }) + }) + + it("closes clients, removes its listener, and leaves HTTP serving", async () => { + const hub = await openServer({ + authenticate: () => "reader", + authorizeSubscription: () => true, + }) + const socket = await hub.connect() + const closed = once(socket, "close") as Promise<[number]> + + const shutdown = hub.realtime.close() + expect(hub.realtime.close()).toBe(shutdown) + const [code] = await closed + await shutdown + + expect(code).toBe(1001) + expect(hub.server.listenerCount("upgrade")).toBe(0) + expect(() => hub.realtime.broadcast("one", { type: "updated", payload: null })).toThrow() + const response = await fetch(hub.url.replace("ws:", "http:")) + expect(await response.text()).toBe("healthy") + }) + + it("finishes shutdown when a client never processes the closing handshake", async () => { + const hub = await openServer({ + authenticate: () => "reader", + authorizeSubscription: () => true, + limits: { shutdownTimeoutMs: 40 }, + }) + const socket = await hub.connect() + socket.pause() + const closed = once(socket, "close") + + await hub.realtime.close() + socket.resume() + await closed + + expect(socket.readyState).toBe(WebSocket.CLOSED) + }) +}) diff --git a/hub/server/integrations/ws/attach-realtime-server/index.ts b/hub/server/integrations/ws/attach-realtime-server/index.ts new file mode 100644 index 00000000..3ff51081 --- /dev/null +++ b/hub/server/integrations/ws/attach-realtime-server/index.ts @@ -0,0 +1,59 @@ +import type { IncomingMessage } from "node:http" +import type { Duplex } from "node:stream" +import { WebSocketServer } from "ws" + +import { broadcast } from "../broadcast/index.ts" +import { checkConnections } from "../check-connections/index.ts" +import { closeRealtimeServer } from "../close-realtime-server/index.ts" +import { DEFAULT_LIMITS } from "../constants.ts" +import type { RealtimeOptions, RealtimeServer, RealtimeState } from "../types.ts" +import { upgradeConnection } from "../upgrade-connection/index.ts" + +/** Own only local connections; application callbacks decide who may connect and subscribe. */ +export function attachRealtimeServer({ + ...options +}: RealtimeOptions): RealtimeServer { + const { server, path, limits: overrides } = options + if (!/^\/[^?#]*$/.test(path)) throw new Error("WebSocket path must be an absolute URL path.") + const limits = { ...DEFAULT_LIMITS, ...overrides } + for (const [name, value] of Object.entries(limits)) { + if (!Number.isSafeInteger(value) || value < 1 || value > 2_147_483_647) { + throw new Error( + `WebSocket limit ${name} must be a positive integer no greater than 2147483647.`, + ) + } + } + + const state: RealtimeState = { + logger: options.logger, + limits, + wss: new WebSocketServer({ + noServer: true, + maxPayload: limits.maxMessageBytes, + maxFragments: 128, + maxBufferedChunks: 128, + perMessageDeflate: false, + autoPong: false, + }), + connections: new Map(), + channels: new Map(), + addresses: new Map(), + pending: new Map(), + rejectedUpgrades: new Map(), + reserved: 0, + accessChecks: 0, + closing: false, + } + function upgrade(request: IncomingMessage, socket: Duplex, head: Buffer) { + void upgradeConnection(state, options, request, socket, head) + } + const heartbeat = setInterval(() => checkConnections(state), limits.heartbeatIntervalMs) + heartbeat.unref() + server.on("upgrade", upgrade) + state.wss.on("error", () => state.logger.error({}, "WebSocket server failed")) + + return { + broadcast: (channel, event) => broadcast(state, channel, event), + close: () => closeRealtimeServer(state, server, upgrade, heartbeat), + } +} diff --git a/hub/server/integrations/ws/broadcast/index.ts b/hub/server/integrations/ws/broadcast/index.ts new file mode 100644 index 00000000..f01326c6 --- /dev/null +++ b/hub/server/integrations/ws/broadcast/index.ts @@ -0,0 +1,17 @@ +import { CHANNEL_SCHEMA, EVENT_SCHEMA } from "../constants.ts" +import { sendMessage } from "../send-message/index.ts" +import type { Broadcast, RealtimeState } from "../types.ts" + +export function broadcast(state: RealtimeState, channel: string, event: Parameters[1]) { + if (state.closing) throw new Error("WebSocket server is closed.") + CHANNEL_SCHEMA.parse(channel) + EVENT_SCHEMA.parse(event) + const message = JSON.stringify({ type: "event", channel, event }) + if (Buffer.byteLength(message) > state.limits.maxMessageBytes) { + throw new RangeError("WebSocket event exceeds the message size limit.") + } + for (const socket of state.channels.get(channel) ?? []) { + const connection = state.connections.get(socket) + if (connection) sendMessage(state, socket, connection, message) + } +} diff --git a/hub/server/integrations/ws/can-send/index.ts b/hub/server/integrations/ws/can-send/index.ts new file mode 100644 index 00000000..f8c01a23 --- /dev/null +++ b/hub/server/integrations/ws/can-send/index.ts @@ -0,0 +1,18 @@ +import { WebSocket } from "ws" + +import { disconnect } from "../disconnect/index.ts" +import type { Connection, RealtimeState } from "../types.ts" + +export function canSend( + state: RealtimeState, + socket: WebSocket, + connection: Connection, + bytes: number, +) { + if (socket.readyState !== WebSocket.OPEN) return false + if (socket.bufferedAmount + bytes > state.limits.maxBufferedBytes) { + disconnect(state, socket, connection, "send_buffer_limit") + return false + } + return true +} diff --git a/hub/server/integrations/ws/check-connections/index.ts b/hub/server/integrations/ws/check-connections/index.ts new file mode 100644 index 00000000..218069a9 --- /dev/null +++ b/hub/server/integrations/ws/check-connections/index.ts @@ -0,0 +1,24 @@ +import { WebSocket } from "ws" + +import { canSend } from "../can-send/index.ts" +import { disconnect } from "../disconnect/index.ts" +import { pruneAddresses } from "../prune-addresses/index.ts" +import { reportRejections } from "../report-rejections/index.ts" +import type { RealtimeState } from "../types.ts" + +export function checkConnections(state: RealtimeState) { + reportRejections(state) + pruneAddresses(state) + for (const [socket, connection] of state.connections) { + if (!connection.alive) { + disconnect(state, socket, connection, "heartbeat_timeout") + } else if (socket.readyState === WebSocket.OPEN) { + connection.alive = false + if (canSend(state, socket, connection, 2)) { + socket.ping(undefined, false, (error) => { + if (error) disconnect(state, socket, connection, "send_failed") + }) + } + } + } +} diff --git a/hub/server/integrations/ws/close-realtime-server/index.ts b/hub/server/integrations/ws/close-realtime-server/index.ts new file mode 100644 index 00000000..80755b0d --- /dev/null +++ b/hub/server/integrations/ws/close-realtime-server/index.ts @@ -0,0 +1,33 @@ +import type { Server } from "node:http" + +import { reportRejections } from "../report-rejections/index.ts" +import type { RealtimeState, UpgradeListener } from "../types.ts" + +export function closeRealtimeServer( + state: RealtimeState, + server: Server, + upgrade: UpgradeListener, + heartbeat: NodeJS.Timeout, +): Promise { + if (state.closePromise) return state.closePromise + state.closing = true + server.off("upgrade", upgrade) + clearInterval(heartbeat) + reportRejections(state) + for (const cancel of state.pending.values()) cancel() + state.addresses.clear() + state.channels.clear() + state.closePromise = new Promise((resolve, reject) => { + const deadline = setTimeout(() => { + for (const socket of state.connections.keys()) socket.terminate() + }, state.limits.shutdownTimeoutMs) + deadline.unref() + for (const socket of state.connections.keys()) socket.close(1001, "Server shutting down") + state.wss.close((error) => { + clearTimeout(deadline) + if (error) reject(error) + else resolve() + }) + }) + return state.closePromise +} diff --git a/hub/server/integrations/ws/constants.ts b/hub/server/integrations/ws/constants.ts new file mode 100644 index 00000000..0d4d66a9 --- /dev/null +++ b/hub/server/integrations/ws/constants.ts @@ -0,0 +1,28 @@ +import { z } from "zod" + +import type { RealtimeLimits } from "./types.ts" + +export const CHANNEL_SCHEMA = z + .string() + .min(1) + .max(200) + .regex(/^[a-zA-Z0-9][a-zA-Z0-9:._/-]*$/) +export const MESSAGE_SCHEMA = z.strictObject({ + type: z.enum(["subscribe", "unsubscribe"]), + channel: CHANNEL_SCHEMA, +}) +export const EVENT_SCHEMA = z.strictObject({ type: z.string().min(1).max(100), payload: z.json() }) + +export const DEFAULT_LIMITS: Readonly = { + maxConnections: 1000, + maxConnectionsPerAddress: 50, + maxTrackedAddresses: 4096, + maxUpgradeAttemptsPerMinute: 60, + maxMessageBytes: 16 * 1024, + maxSubscriptionsPerConnection: 32, + maxMessagesPerSecond: 20, + maxBufferedBytes: 256 * 1024, + heartbeatIntervalMs: 30_000, + upgradeTimeoutMs: 5000, + shutdownTimeoutMs: 5000, +} diff --git a/hub/server/integrations/ws/count-message/index.ts b/hub/server/integrations/ws/count-message/index.ts new file mode 100644 index 00000000..5db69547 --- /dev/null +++ b/hub/server/integrations/ws/count-message/index.ts @@ -0,0 +1,18 @@ +import { WebSocket } from "ws" + +import type { Connection, RealtimeState } from "../types.ts" + +export function countMessage(state: RealtimeState, socket: WebSocket, connection: Connection) { + if (state.closing || socket.readyState !== WebSocket.OPEN) return false + if (Date.now() >= connection.windowEndsAt) { + connection.messages = 0 + connection.windowEndsAt = Date.now() + 1000 + } + connection.messages += 1 + if (connection.messages > state.limits.maxMessagesPerSecond) { + state.logger.warn({ connectionId: connection.id }, "WebSocket message rate exceeded") + socket.close(1008, "Message rate exceeded") + return false + } + return true +} diff --git a/hub/server/integrations/ws/disconnect/index.ts b/hub/server/integrations/ws/disconnect/index.ts new file mode 100644 index 00000000..7ed02b0f --- /dev/null +++ b/hub/server/integrations/ws/disconnect/index.ts @@ -0,0 +1,13 @@ +import type { WebSocket } from "ws" + +import type { Connection, RealtimeState } from "../types.ts" + +export function disconnect( + { logger }: RealtimeState, + socket: WebSocket, + connection: Connection, + reason: string, +) { + logger.warn({ connectionId: connection.id, reason }, "WebSocket connection terminated") + socket.terminate() +} diff --git a/hub/server/integrations/ws/handle-message/index.ts b/hub/server/integrations/ws/handle-message/index.ts new file mode 100644 index 00000000..2b562abf --- /dev/null +++ b/hub/server/integrations/ws/handle-message/index.ts @@ -0,0 +1,93 @@ +import { WebSocket } from "ws" + +import { MESSAGE_SCHEMA } from "../constants.ts" +import { countMessage } from "../count-message/index.ts" +import { removeSubscription } from "../remove-subscription/index.ts" +import { sendError } from "../send-error/index.ts" +import { sendMessage } from "../send-message/index.ts" +import type { Connection, RealtimeOptions, RealtimeState } from "../types.ts" + +export async function handleMessage( + state: RealtimeState, + socket: WebSocket, + connection: Connection, + identity: Identity, + data: Buffer, + isBinary: boolean, + authorizeSubscription: RealtimeOptions["authorizeSubscription"], +) { + const { logger, limits } = state + if (!countMessage(state, socket, connection)) return + if (connection.busy) { + logger.warn({ connectionId: connection.id }, "WebSocket subscription check already pending") + socket.close(1008, "Subscription check already pending") + return + } + let input: unknown + try { + if (isBinary) throw new Error("Binary messages are not supported.") + input = JSON.parse(data.toString()) + } catch { + sendError(state, socket, connection, "INVALID_MESSAGE") + return + } + const parsed = MESSAGE_SCHEMA.safeParse(input) + if (!parsed.success) { + sendError(state, socket, connection, "INVALID_MESSAGE") + return + } + const { type, channel } = parsed.data + if (type === "unsubscribe") { + removeSubscription(state, socket, connection, channel) + sendMessage(state, socket, connection, JSON.stringify({ type: "unsubscribed", channel })) + return + } + if ( + !connection.subscriptions.has(channel) && + connection.subscriptions.size >= limits.maxSubscriptionsPerConnection + ) { + sendError(state, socket, connection, "SUBSCRIPTION_LIMIT") + return + } + if (state.accessChecks >= limits.maxConnections) { + logger.warn({ connectionId: connection.id }, "WebSocket access check capacity exceeded") + socket.close(1013, "Server busy") + return + } + connection.busy = true + state.accessChecks += 1 + const controller = new AbortController() + connection.check = controller + const timeout = setTimeout(() => { + controller.abort() + logger.warn({ connectionId: connection.id }, "WebSocket subscription check timed out") + socket.close(1008, "Subscription check timed out") + }, limits.upgradeTimeoutMs) + timeout.unref() + try { + const allowed = await authorizeSubscription?.(identity, channel, controller.signal) + if (state.closing || socket.readyState !== WebSocket.OPEN) return + if (!allowed) { + removeSubscription(state, socket, connection, channel) + sendError(state, socket, connection, "FORBIDDEN") + return + } + connection.subscriptions.add(channel) + let subscribers = state.channels.get(channel) + if (!subscribers) { + subscribers = new Set() + state.channels.set(channel, subscribers) + } + subscribers.add(socket) + sendMessage(state, socket, connection, JSON.stringify({ type: "subscribed", channel })) + } catch { + if (controller.signal.aborted) return + logger.error({ connectionId: connection.id }, "WebSocket subscription check failed") + socket.close(1011, "Subscription check failed") + } finally { + clearTimeout(timeout) + state.accessChecks -= 1 + connection.check = undefined + connection.busy = false + } +} diff --git a/hub/server/integrations/ws/prune-addresses/index.ts b/hub/server/integrations/ws/prune-addresses/index.ts new file mode 100644 index 00000000..4c752967 --- /dev/null +++ b/hub/server/integrations/ws/prune-addresses/index.ts @@ -0,0 +1,8 @@ +import type { RealtimeState } from "../types.ts" + +export function pruneAddresses({ addresses }: RealtimeState) { + const now = Date.now() + for (const [address, state] of addresses) { + if (state.connections === 0 && state.windowEndsAt <= now) addresses.delete(address) + } +} diff --git a/hub/server/integrations/ws/register-connection/index.ts b/hub/server/integrations/ws/register-connection/index.ts new file mode 100644 index 00000000..d839cb08 --- /dev/null +++ b/hub/server/integrations/ws/register-connection/index.ts @@ -0,0 +1,66 @@ +import { randomUUID } from "node:crypto" +import type { WebSocket } from "ws" + +import { canSend } from "../can-send/index.ts" +import { countMessage } from "../count-message/index.ts" +import { disconnect } from "../disconnect/index.ts" +import { handleMessage } from "../handle-message/index.ts" +import { removeSubscription } from "../remove-subscription/index.ts" +import type { Connection, RealtimeOptions, RealtimeState } from "../types.ts" + +export function registerConnection( + state: RealtimeState, + client: WebSocket, + identity: Identity, + authorizeSubscription: RealtimeOptions["authorizeSubscription"], + release: () => void, +) { + const connection: Connection = { + id: randomUUID(), + subscriptions: new Set(), + alive: true, + busy: false, + messages: 0, + windowEndsAt: Date.now() + 1000, + } + state.connections.set(client, connection) + state.logger.info( + { connectionId: connection.id, activeConnections: state.connections.size }, + "WebSocket connected", + ) + client.on("pong", () => { + if (countMessage(state, client, connection)) connection.alive = true + }) + client.on("ping", (data) => { + if ( + !countMessage(state, client, connection) || + !canSend(state, client, connection, data.length + 2) + ) { + return + } + client.pong(data, false, (error) => { + if (error) disconnect(state, client, connection, "send_failed") + }) + }) + client.on("message", (data, isBinary) => { + const buffer = Array.isArray(data) + ? Buffer.concat(data) + : Buffer.isBuffer(data) + ? data + : Buffer.from(data) + void handleMessage(state, client, connection, identity, buffer, isBinary, authorizeSubscription) + }) + client.on("error", () => disconnect(state, client, connection, "transport_failed")) + client.once("close", () => { + connection.check?.abort() + for (const channel of connection.subscriptions) { + removeSubscription(state, client, connection, channel) + } + state.connections.delete(client) + release() + state.logger.info( + { connectionId: connection.id, activeConnections: state.connections.size }, + "WebSocket disconnected", + ) + }) +} diff --git a/hub/server/integrations/ws/reject-upgrade/index.ts b/hub/server/integrations/ws/reject-upgrade/index.ts new file mode 100644 index 00000000..1c06e4dc --- /dev/null +++ b/hub/server/integrations/ws/reject-upgrade/index.ts @@ -0,0 +1,19 @@ +import type { Duplex } from "node:stream" + +import type { RealtimeState } from "../types.ts" + +export function rejectUpgrade( + { logger, rejectedUpgrades }: RealtimeState, + socket: Duplex, + status: number, +) { + const count = rejectedUpgrades.get(status) ?? 0 + if (count === 0) logger.warn({ statusCode: status }, "WebSocket upgrade rejected") + rejectedUpgrades.set(status, count + 1) + if (!socket.destroyed) { + socket.end( + `HTTP/1.1 ${status} Rejected\r\nConnection: close\r\nContent-Length: 0\r\n\r\n`, + () => socket.destroy(), + ) + } +} diff --git a/hub/server/integrations/ws/remove-subscription/index.ts b/hub/server/integrations/ws/remove-subscription/index.ts new file mode 100644 index 00000000..e77a2fbf --- /dev/null +++ b/hub/server/integrations/ws/remove-subscription/index.ts @@ -0,0 +1,15 @@ +import type { WebSocket } from "ws" + +import type { Connection, RealtimeState } from "../types.ts" + +export function removeSubscription( + { channels }: RealtimeState, + socket: WebSocket, + connection: Connection, + channel: string, +) { + connection.subscriptions.delete(channel) + const subscribers = channels.get(channel) + subscribers?.delete(socket) + if (subscribers?.size === 0) channels.delete(channel) +} diff --git a/hub/server/integrations/ws/report-rejections/index.ts b/hub/server/integrations/ws/report-rejections/index.ts new file mode 100644 index 00000000..0250db57 --- /dev/null +++ b/hub/server/integrations/ws/report-rejections/index.ts @@ -0,0 +1,10 @@ +import type { RealtimeState } from "../types.ts" + +export function reportRejections({ logger, rejectedUpgrades }: RealtimeState) { + for (const [statusCode, count] of rejectedUpgrades) { + if (count > 1) { + logger.warn({ statusCode, count: count - 1 }, "Additional WebSocket upgrades rejected") + } + } + rejectedUpgrades.clear() +} diff --git a/hub/server/integrations/ws/send-error/index.ts b/hub/server/integrations/ws/send-error/index.ts new file mode 100644 index 00000000..0dc8fa7a --- /dev/null +++ b/hub/server/integrations/ws/send-error/index.ts @@ -0,0 +1,14 @@ +import type { WebSocket } from "ws" + +import { sendMessage } from "../send-message/index.ts" +import type { Connection, RealtimeState } from "../types.ts" + +export function sendError( + state: RealtimeState, + socket: WebSocket, + connection: Connection, + code: string, +) { + state.logger.warn({ connectionId: connection.id, code }, "WebSocket message rejected") + sendMessage(state, socket, connection, JSON.stringify({ type: "error", code })) +} diff --git a/hub/server/integrations/ws/send-message/index.ts b/hub/server/integrations/ws/send-message/index.ts new file mode 100644 index 00000000..60a10ab0 --- /dev/null +++ b/hub/server/integrations/ws/send-message/index.ts @@ -0,0 +1,19 @@ +import type { WebSocket } from "ws" + +import { canSend } from "../can-send/index.ts" +import { disconnect } from "../disconnect/index.ts" +import type { Connection, RealtimeState } from "../types.ts" + +export function sendMessage( + state: RealtimeState, + socket: WebSocket, + connection: Connection, + message: string, +) { + const bytes = Buffer.byteLength(message) + const frameBytes = bytes < 126 ? 2 : bytes <= 65535 ? 4 : 10 + if (!canSend(state, socket, connection, bytes + frameBytes)) return + socket.send(message, (error) => { + if (error) disconnect(state, socket, connection, "send_failed") + }) +} diff --git a/hub/server/integrations/ws/types.ts b/hub/server/integrations/ws/types.ts new file mode 100644 index 00000000..c4861399 --- /dev/null +++ b/hub/server/integrations/ws/types.ts @@ -0,0 +1,77 @@ +import type { IncomingMessage, Server } from "node:http" +import type { Duplex } from "node:stream" +import type { Logger } from "pino" +import type { WebSocket, WebSocketServer } from "ws" + +export type JsonValue = + null | boolean | number | string | JsonValue[] | { [key: string]: JsonValue } + +export type Broadcast = (channel: string, event: { type: string; payload: JsonValue }) => void + +export interface RealtimeLimits { + maxConnections: number + maxConnectionsPerAddress: number + maxTrackedAddresses: number + maxUpgradeAttemptsPerMinute: number + maxMessageBytes: number + maxSubscriptionsPerConnection: number + maxMessagesPerSecond: number + maxBufferedBytes: number + heartbeatIntervalMs: number + upgradeTimeoutMs: number + shutdownTimeoutMs: number +} + +export interface RealtimeOptions { + server: Server + path: string + logger: Logger + limits?: Partial + authenticate?: ( + request: IncomingMessage, + signal: AbortSignal, + ) => Identity | null | Promise + authorizeSubscription?: ( + identity: Identity, + channel: string, + signal: AbortSignal, + ) => boolean | Promise +} + +export interface RealtimeServer { + broadcast: Broadcast + close: () => Promise +} + +export interface AddressState { + connections: number + attempts: number + windowEndsAt: number +} + +export interface Connection { + id: string + subscriptions: Set + alive: boolean + busy: boolean + messages: number + windowEndsAt: number + check?: AbortController +} + +export interface RealtimeState { + logger: Logger + limits: RealtimeLimits + wss: WebSocketServer + connections: Map + channels: Map> + addresses: Map + pending: Map void> + rejectedUpgrades: Map + reserved: number + accessChecks: number + closing: boolean + closePromise?: Promise +} + +export type UpgradeListener = (request: IncomingMessage, socket: Duplex, head: Buffer) => void diff --git a/hub/server/integrations/ws/upgrade-connection/index.ts b/hub/server/integrations/ws/upgrade-connection/index.ts new file mode 100644 index 00000000..b42dd9f3 --- /dev/null +++ b/hub/server/integrations/ws/upgrade-connection/index.ts @@ -0,0 +1,126 @@ +import type { IncomingMessage } from "node:http" +import type { Duplex } from "node:stream" + +import { pruneAddresses } from "../prune-addresses/index.ts" +import { registerConnection } from "../register-connection/index.ts" +import { rejectUpgrade } from "../reject-upgrade/index.ts" +import type { RealtimeOptions, RealtimeState } from "../types.ts" + +export async function upgradeConnection( + state: RealtimeState, + { path, authenticate, authorizeSubscription }: RealtimeOptions, + request: IncomingMessage, + socket: Duplex, + head: Buffer, +) { + const { logger, limits } = state + function socketError() { + logger.warn({}, "WebSocket transport failed") + socket.destroy() + } + socket.on("error", socketError) + socket.once("close", () => socket.off("error", socketError)) + if (request.url?.split("?")[0] !== path) { + rejectUpgrade(state, socket, 404) + return + } + if (state.closing) { + rejectUpgrade(state, socket, 503) + return + } + + // Use the peer address, never an untrusted forwarded header. + const address = request.socket.remoteAddress + if (!address) { + rejectUpgrade(state, socket, 400) + return + } + let addressState = state.addresses.get(address) + const now = Date.now() + if (!addressState) { + if (state.addresses.size >= limits.maxTrackedAddresses) pruneAddresses(state) + if (state.addresses.size >= limits.maxTrackedAddresses) { + rejectUpgrade(state, socket, 503) + return + } + addressState = { connections: 0, attempts: 0, windowEndsAt: now + 60_000 } + state.addresses.set(address, addressState) + } + if (addressState.windowEndsAt <= now) { + addressState.attempts = 0 + addressState.windowEndsAt = now + 60_000 + } + addressState.attempts += 1 + if ( + addressState.attempts > limits.maxUpgradeAttemptsPerMinute || + addressState.connections >= limits.maxConnectionsPerAddress + ) { + rejectUpgrade(state, socket, 429) + return + } + if (state.reserved >= limits.maxConnections || state.accessChecks >= limits.maxConnections) { + rejectUpgrade(state, socket, 503) + return + } + if (!authenticate || !authorizeSubscription) { + rejectUpgrade(state, socket, 403) + return + } + + state.reserved += 1 + addressState.connections += 1 + const reservation = addressState + let released = false + let completed = false + const controller = new AbortController() + function release() { + if (released) return + released = true + state.reserved -= 1 + reservation.connections -= 1 + } + function cancel() { + if (completed) return + completed = true + clearTimeout(timeout) + state.pending.delete(socket) + controller.abort() + release() + } + const timeout = setTimeout(() => { + cancel() + rejectUpgrade(state, socket, 408) + }, limits.upgradeTimeoutMs) + timeout.unref() + state.pending.set(socket, () => { + cancel() + socket.destroy() + }) + socket.once("close", cancel) + + state.accessChecks += 1 + try { + const identity = await authenticate(request, controller.signal) + if (completed || socket.destroyed || state.closing) return + if (identity === null || identity === undefined) { + cancel() + rejectUpgrade(state, socket, 401) + return + } + state.wss.handleUpgrade(request, socket, head, (client) => { + completed = true + clearTimeout(timeout) + state.pending.delete(socket) + socket.off("close", cancel) + socket.off("error", socketError) + registerConnection(state, client, identity, authorizeSubscription, release) + }) + } catch { + if (completed || socket.destroyed) return + cancel() + logger.error({}, "WebSocket admission failed") + rejectUpgrade(state, socket, 500) + } finally { + state.accessChecks -= 1 + } +} diff --git a/hub/server/routes/index.ts b/hub/server/routes/index.ts index e5216064..d7b4f2f8 100644 --- a/hub/server/routes/index.ts +++ b/hub/server/routes/index.ts @@ -1,8 +1,10 @@ import type { Express } from "express" import { get as handleApiHealthGet } from "./api/health/get/index.ts" +import type { RouteDependencies } from "./types.ts" /** Keep endpoint registration separate from the app's shared middleware. */ -export function routes(app: Express): void { +// eslint-disable-next-line @typescript-eslint/no-unused-vars -- Product routes will consume these dependencies when ported. +export function routes(app: Express, _dependencies: RouteDependencies = {}): void { app.get("/api/health", handleApiHealthGet) } diff --git a/hub/server/routes/types.ts b/hub/server/routes/types.ts new file mode 100644 index 00000000..bab536ec --- /dev/null +++ b/hub/server/routes/types.ts @@ -0,0 +1,5 @@ +import type { Broadcast } from "../integrations/ws/types.ts" + +export interface RouteDependencies { + broadcast?: Broadcast +} diff --git a/hub/server/shutdown-server/index.ts b/hub/server/shutdown-server/index.ts new file mode 100644 index 00000000..f24d441c --- /dev/null +++ b/hub/server/shutdown-server/index.ts @@ -0,0 +1,16 @@ +import type { Server } from "node:http" + +import { log } from "../integrations/pino/log/index.ts" +import type { RealtimeServer } from "../integrations/ws/types.ts" + +export async function shutdownServer(server: Server, realtime: RealtimeServer) { + const httpClosed = new Promise((resolve, reject) => { + server.close((error) => (error ? reject(error) : resolve())) + }) + try { + await Promise.all([realtime.close(), httpClosed]) + } catch { + log("error", "Hub API shutdown failed") + process.exitCode = 1 + } +} diff --git a/hub/test/ws/deferred/index.ts b/hub/test/ws/deferred/index.ts new file mode 100644 index 00000000..7455b080 --- /dev/null +++ b/hub/test/ws/deferred/index.ts @@ -0,0 +1,7 @@ +export function deferred() { + let resolve!: (value: T) => void + const promise = new Promise((settle) => { + resolve = settle + }) + return { promise, resolve } +} diff --git a/hub/test/ws/open-server/index.ts b/hub/test/ws/open-server/index.ts new file mode 100644 index 00000000..b4203e2a --- /dev/null +++ b/hub/test/ws/open-server/index.ts @@ -0,0 +1,50 @@ +import { once } from "node:events" +import { createServer } from "node:http" +import type { AddressInfo } from "node:net" +import pino from "pino" +import { onTestFinished } from "vitest" +import WebSocket from "ws" + +import { attachRealtimeServer } from "../../../server/integrations/ws/attach-realtime-server" + +import type { Options } from "./types" + +export async function openServer(options: Options = {}) { + const server = createServer((_request, response) => response.end("healthy")) + const realtime = attachRealtimeServer({ + ...options, + server, + path: "/realtime", + logger: options.logger ?? pino({ level: "silent" }), + }) + const sockets: WebSocket[] = [] + + onTestFinished(async () => { + for (const socket of sockets) socket.terminate() + await realtime.close() + await new Promise((resolve, reject) => { + server.close((error) => (error ? reject(error) : resolve())) + }) + }) + + server.listen(0, "127.0.0.1") + await once(server, "listening") + const { port } = server.address() as AddressInfo + const url = `ws://127.0.0.1:${port}/realtime` + + function client(clientOptions: WebSocket.ClientOptions = {}, target = url) { + const socket = new WebSocket(target, clientOptions) + // Rejected handshakes emit an expected error after their HTTP response. + socket.on("error", () => {}) + sockets.push(socket) + return socket + } + + async function connect(clientOptions?: WebSocket.ClientOptions) { + const socket = client(clientOptions) + await once(socket, "open") + return socket + } + + return { server, realtime, url, client, connect } +} diff --git a/hub/test/ws/open-server/types.ts b/hub/test/ws/open-server/types.ts new file mode 100644 index 00000000..f08d59ae --- /dev/null +++ b/hub/test/ws/open-server/types.ts @@ -0,0 +1,5 @@ +import type { RealtimeOptions } from "../../../server/integrations/ws/types" + +export type Options = Omit, "server" | "path" | "logger"> & { + logger?: RealtimeOptions["logger"] +} diff --git a/hub/test/ws/parse-message/index.ts b/hub/test/ws/parse-message/index.ts new file mode 100644 index 00000000..56270e11 --- /dev/null +++ b/hub/test/ws/parse-message/index.ts @@ -0,0 +1,10 @@ +import type WebSocket from "ws" + +export function parseMessage(data: WebSocket.RawData): unknown { + const buffer = Array.isArray(data) + ? Buffer.concat(data) + : Buffer.isBuffer(data) + ? data + : Buffer.from(data) + return JSON.parse(buffer.toString()) as unknown +} diff --git a/hub/test/ws/rejected/index.ts b/hub/test/ws/rejected/index.ts new file mode 100644 index 00000000..c87dbe6f --- /dev/null +++ b/hub/test/ws/rejected/index.ts @@ -0,0 +1,10 @@ +import { once } from "node:events" +import type { IncomingMessage } from "node:http" +import type WebSocket from "ws" + +export async function rejected(socket: WebSocket) { + const [, response] = (await once(socket, "unexpected-response")) as [unknown, IncomingMessage] + response.resume() + socket.terminate() + return response.statusCode +} diff --git a/hub/test/ws/send-message/index.ts b/hub/test/ws/send-message/index.ts new file mode 100644 index 00000000..6a738ec0 --- /dev/null +++ b/hub/test/ws/send-message/index.ts @@ -0,0 +1,9 @@ +import { once } from "node:events" +import type WebSocket from "ws" + +export async function sendMessage(socket: WebSocket, value: unknown) { + const received = once(socket, "message") as Promise<[Buffer]> + socket.send(JSON.stringify(value)) + const [data] = await received + return JSON.parse(String(data)) as unknown +} diff --git a/hub/vite.config.ts b/hub/vite.config.ts index 6d2147d6..f4c163e6 100644 --- a/hub/vite.config.ts +++ b/hub/vite.config.ts @@ -14,7 +14,7 @@ export default defineConfig({ server: { port: 5173, strictPort: true, - proxy: { "/api": "http://127.0.0.1:3001" }, + proxy: { "/api": { target: "http://127.0.0.1:3001", ws: true } }, }, resolve: { alias: { ...viteAliases, [controlsValidatorModule]: controlsValidatorPath } }, optimizeDeps: { include: [controlsValidatorModule] },