diff --git a/src/app.ts b/src/app.ts index 522af6a..d14b9bc 100644 --- a/src/app.ts +++ b/src/app.ts @@ -3,6 +3,7 @@ import { Database, MigrationService } from './services/storage/index.ts'; import { Broadcaster } from './services/broadcaster.ts'; import { ApplicationRouter } from './services/router.ts'; import { HttpTransportRouter } from './services/transport/http-transport.ts'; +import { WsTransportRouter } from './services/transport/ws-transport.ts'; import { ServerHost } from './services/server-host.ts'; import { Logger } from './utils/logger.ts'; @@ -31,7 +32,8 @@ export class App { const router = await ApplicationRouter.create(routes); const http = new HttpTransportRouter(router, debug); - const host = new ServerHost(config, debug, [http]); + const ws = new WsTransportRouter(router, debug, config.server.maxRequestBodyBytes); + const host = new ServerHost(config, debug, [http, ws]); return new App(host, database); } diff --git a/src/services/stream/ws-stream.ts b/src/services/stream/ws-stream.ts new file mode 100644 index 0000000..08731d1 --- /dev/null +++ b/src/services/stream/ws-stream.ts @@ -0,0 +1,94 @@ +import { randomUUID } from 'node:crypto'; +import type { WSContext } from 'hono/ws'; +import { toExtendedJson } from '@xo-cash/utils'; + +import { BaseStream, type StreamMessage } from './base-stream.ts'; + +/** WebSocket readyState value indicating an open socket. */ +const OPEN = 1; + +/** Adapts one WebSocket connection to the shared connection-stream contract. */ +export class WSStream extends BaseStream { + readonly streaming = true; + readonly bidirectional = true; + + /** Stable server-side identity used only for diagnostics. */ + readonly id: string; + + /** Observers notified when the socket closes. */ + private closeCallbacks: Array<() => void> = []; + + /** Guards against sends after local or remote closure. */ + private closed = false; + + /** + * @param ws - Minimal WebSocket surface required for send/close operations. + */ + constructor(private readonly ws: Pick) { + super(); + this.id = randomUUID(); + } + + /** + * Encode and send one Extended JSON message frame. + * + * @param message - Neutral response or event envelope to send. + */ + async send(message: StreamMessage): Promise { + if (this.closed || this.ws.readyState !== OPEN) { + throw new Error('Cannot send to a closed WebSocket stream'); + } + + // WebSocket carries the complete neutral response or event envelope. + this.ws.send(toExtendedJson(message)); + } + + /** Close the socket from the server side and notify observers. */ + close(): void { + if (this.closed) { + return; + } + + this.closed = true; + this.ws.close(); + this.emitClose(); + } + + /** + * Register a callback invoked when the socket closes. + * + * @param callback - Called immediately if the stream is already closed. + */ + onClose(callback: () => void): void { + if (this.closed) { + callback(); + + return; + } + + this.closeCallbacks.push(callback); + } + + /** + * Record closure reported by WebSocket callbacks without closing the socket + * again. This shares the same observer notification path as local closure. + */ + markClosed(): void { + if (this.closed) { + return; + } + + this.closed = true; + this.emitClose(); + } + + /** Drain and invoke all registered close observers. */ + private emitClose(): void { + const callbacks = this.closeCallbacks; + this.closeCallbacks = []; + + for (const callback of callbacks) { + callback(); + } + } +} diff --git a/src/services/transport/ws-transport.ts b/src/services/transport/ws-transport.ts new file mode 100644 index 0000000..71a60f7 --- /dev/null +++ b/src/services/transport/ws-transport.ts @@ -0,0 +1,230 @@ +import { upgradeWebSocket } from '@hono/node-server'; +import type { WebSocketServerLike } from '@hono/node-server'; +import type { Hono } from 'hono'; +import type { WSContext, WSMessageReceive } from 'hono/ws'; +import { WebSocketServer } from 'ws'; +import { toExtendedJson, fromExtendedJson } from '@xo-cash/utils'; +import { z } from 'zod'; + +import type { Logger } from '../../utils/logger.ts'; +import { ApplicationError, normalizePublicError } from '../../errors/index.ts'; +import type { ApplicationRequest, ApplicationRouter } from '../router.ts'; +import { WSStream } from '../stream/ws-stream.ts'; +import type { AppEnv, UpgradeTransportRouter } from './transport-router.ts'; + +/** Default WebSocket upgrade path for application messages. */ +const WS_ROUTE = '/ws'; + +// Strict validation prevents legacy or protocol-specific fields reaching routes. +const wsRequestSchema = z + .object({ + id: z.string().min(1).optional(), + path: z.string().min(1), + body: z.unknown().optional(), + }) + .strict(); + +/** + * WebSocket protocol adapter. + * + * One WSStream is shared by every message on a socket, allowing subscribe and + * unsubscribe requests to operate on the same broadcaster registration. + */ +export class WsTransportRouter implements UpgradeTransportRouter { + /** + * Native Node WebSocket server used by Hono's Node adapter. + * + * Hono handles the HTTP upgrade and lifecycle callbacks, while the `ws` + * server owns frame parsing and therefore enforces the payload limit. + */ + private readonly wsServer: WebSocketServer; + + /** Upgrade server wired into the Node HTTP listener by ServerHost. */ + readonly websocketServer: WebSocketServerLike; + + private readonly debug: Logger; + + /** + * @param router - Shared application router for route dispatch. + * @param debug - Root logger extended with a ws-transport namespace. + * @param maxRequestBodyBytes - Maximum encoded WebSocket message size. + * @param url - WebSocket upgrade path exposed by this adapter. + */ + constructor( + private readonly router: ApplicationRouter, + debug: Logger, + private readonly maxRequestBodyBytes: number, + private readonly url: string = WS_ROUTE, + ) { + this.debug = debug.extend('ws-transport'); + + // Hono's Node WebSocket helper delegates frame handling to `ws`; unlike + // HTTP requests, WebSocket messages do not pass through Hono middleware. + // maxPayload rejects oversized messages before onMessage receives them. + this.wsServer = new WebSocketServer({ + noServer: true, + maxPayload: this.maxRequestBodyBytes, + }); + this.websocketServer = this.wsServer as unknown as WebSocketServerLike; + } + + /** + * Register the connection endpoint and translate socket lifecycle events. + * + * @param server - Hono instance to attach the upgrade handler to. + */ + register(server: Hono): void { + server.get( + this.url, + upgradeWebSocket( + () => { + // One WSStream per socket, not per message. Broadcaster topic + // subscriptions live on the connection stream so a later + // /data/unsubscribe message can remove topics registered by an + // earlier /data/subscribe on the same socket. + let stream: WSStream | undefined; + + return { + onOpen: (_event, ws): void => { + stream = new WSStream(ws); + this.debug('socket opened: %s', stream.id); + }, + onMessage: (event, ws): void => { + // onOpen may not run before the first message on some stacks. + stream ??= new WSStream(ws); + + // Each WebSocket frame is an independent application request + // (path + body + optional id). Dispatch is fire-and-forget so + // multiple in-flight requests can share one connection. + this.handleMessage(ws, stream, event.data).catch((error: unknown) => { + this.debug('unexpected message failure: %O', error); + }); + }, + onClose: (): void => { + this.debug('socket closed: %s', stream?.id ?? 'unknown'); + // markClosed, not close — the socket is already gone. This + // notifies the broadcaster's onClose hook so all topic + // registrations for this connection are torn down. + stream?.markClosed(); + }, + onError: (event): void => { + this.debug('socket error: %O', event); + stream?.markClosed(); + }, + }; + }, + { onError: (error) => this.debug('upgrade failed: %O', error) }, + ), + ); + } + + /** Stop accepting upgrades and close active sockets during shutdown. */ + async stop(): Promise { + this.debug('closing %d active WebSocket(s)', this.wsServer.clients.size); + for (const client of this.wsServer.clients) { + client.close(1001, 'Server shutting down'); + } + + await new Promise((resolve, reject) => { + this.wsServer.close((error) => (error ? reject(error) : resolve())); + }); + } + + /** + * Dispatch one independently correlated message. + * + * Route failures become error envelopes and intentionally leave the socket open. + * + * @param ws - Active WebSocket context for error replies. + * @param stream - Shared connection stream for this socket. + * @param message - Raw WebSocket payload from the client. + */ + private async handleMessage(ws: WSContext, stream: WSStream, message: WSMessageReceive): Promise { + let request: ApplicationRequest; + + try { + // Decode the transport envelope { id?, path, body? } and revive any + // Extended JSON values before the message reaches application routes. + request = await WsTransportRouter.decodeWebSocketRequest(message); + } catch (error) { + this.debug('invalid message: %O', error); + // Malformed envelopes have no request id to echo back. + this.sendError(ws, undefined, error); + + return; + } + + try { + // ApplicationRouter wraps the shared connection stream in a per-request + // RouteStream (body + requestId). A subscription dispatch intentionally + // remains pending until a later request unsubscribes or the socket closes. + await this.router.dispatch(request, stream); + } catch (error) { + this.debug('dispatch failed for %s on socket %s: %O', request.path, stream.id, error); + // Reply with a correlated error envelope but keep the socket open. + // The client may have other in-flight requests or active subscriptions + // on this connection that must survive a single route failure. + this.sendError(ws, request.requestId, error); + } + } + + /** + * Serialize the shared public error contract without exposing internals. + * + * @param ws - Active WebSocket context to send the error on. + * @param requestId - Optional correlation ID echoed back to the client. + * @param error - Failure to normalize into the public error shape. + */ + private sendError(ws: WSContext, requestId: string | undefined, error: unknown): void { + ws.send( + toExtendedJson({ + ...(requestId === undefined ? {} : { id: requestId }), + type: 'error', + ...normalizePublicError(error), + }), + ); + } + + /** + * Decode the WebSocket message into an application request. + * + * @param message - Raw WebSocket payload from the client. + * @returns The decoded WebSocket message as an application request. + */ + static async decodeWebSocketRequest(message: WSMessageReceive): Promise { + const payload = await WsTransportRouter.decodeWebSocketMessage(message); + let decoded: unknown; + + try { + decoded = fromExtendedJson(payload); + } catch { + throw new ApplicationError(400, 'Invalid JSON in WebSocket message'); + } + + const envelope = wsRequestSchema.parse(decoded); + + return { + path: envelope.path, + ...(envelope.id === undefined ? {} : { requestId: envelope.id }), + ...(envelope.body === undefined ? {} : { body: envelope.body }), + }; + } + + /** + * Decode the WebSocket message into a string. + * + * @param message - Raw WebSocket payload from the client. + * @returns The decoded WebSocket message as a string. + */ + static async decodeWebSocketMessage(message: WSMessageReceive): Promise { + if (typeof message === 'string') { + return message; + } + + if (message instanceof Blob) { + return Buffer.from(await message.arrayBuffer()).toString('utf8'); + } + + return Buffer.from(message).toString('utf8'); + } +} diff --git a/tests/services/stream/stream.test.ts b/tests/services/stream/stream.test.ts new file mode 100644 index 0000000..5dd9571 --- /dev/null +++ b/tests/services/stream/stream.test.ts @@ -0,0 +1,62 @@ +import { describe, expect, it, vi } from "vitest"; + +import { HonoSSEStream } from "../../../src/services/stream/hono-sse-stream.js"; +import { HttpRequestStream } from "../../../src/services/stream/http-request-stream.js"; +import { WSStream } from "../../../src/services/stream/ws-stream.js"; + +describe("stream lifecycle observers", () => { + it("buffers exactly one normal HTTP response", async () => { + const stream = new HttpRequestStream(); + + await stream.send({ + type: "response", + statusCode: 200, + body: { ok: true }, + }); + + expect(stream.getResponse()).toEqual({ + type: "response", + statusCode: 200, + body: { ok: true }, + }); + await expect( + stream.send({ + type: "response", + statusCode: 200, + body: { second: true }, + }), + ).rejects.toThrow("only send one response"); + }); + + it("notifies WebSocket observers registered after remote closure", () => { + const stream = new WSStream({ + send: vi.fn(), + close: vi.fn(), + readyState: 1, + }); + const onClose = vi.fn(); + + stream.markClosed(); + stream.onClose(onClose); + + expect(onClose).toHaveBeenCalledOnce(); + }); + + it("notifies SSE observers registered after local closure", () => { + const streamApi = { + writeSSE: vi.fn(), + close: vi.fn(), + }; + const stream = new HonoSSEStream( + streamApi as unknown as ConstructorParameters[0], + ); + const onClose = vi.fn(); + + stream.close(); + stream.onClose(onClose); + stream.close(); + + expect(streamApi.close).toHaveBeenCalledOnce(); + expect(onClose).toHaveBeenCalledOnce(); + }); +}); diff --git a/tests/services/transports/ws-transport.test.ts b/tests/services/transports/ws-transport.test.ts new file mode 100644 index 0000000..eeb8f40 --- /dev/null +++ b/tests/services/transports/ws-transport.test.ts @@ -0,0 +1,58 @@ +import { describe, expect, it } from "vitest"; +import { z } from "zod"; +import { WebSocketServer } from "ws"; + +import { ApplicationRouter } from "../../../src/services/router.js"; +import { + WsTransportRouter, +} from "../../../src/services/transport/ws-transport.js"; +import { Logger } from "../../../src/utils/logger.js"; +import { toExtendedJson } from "@xo-cash/utils"; + +describe("WebSocket request decoding", () => { + it("decodes the minimal route-agnostic envelope and Extended JSON body", async () => { + await expect( + WsTransportRouter.decodeWebSocketRequest( + toExtendedJson({ + id: "request-1", + path: "/data/write", + body: { value: new Uint8Array([1, 2, 3]) }, + }), + ), + ).resolves.toEqual({ + requestId: "request-1", + path: "/data/write", + body: { value: new Uint8Array([1, 2, 3]) }, + }); + }); + + it.each([ + "{}", + '{"path":42}', + '{"path":"/data/get","id":1}', + '{"path":"/data/get","method":"POST"}', + ])("rejects an invalid envelope: %s", async (payload) => { + await expect(WsTransportRouter.decodeWebSocketRequest(payload)).rejects.toBeInstanceOf( + z.ZodError, + ); + }); + + it("rejects malformed JSON", async () => { + await expect(WsTransportRouter.decodeWebSocketRequest("{")).rejects.toMatchObject({ + statusCode: 400, + message: "Invalid JSON in WebSocket message", + }); + }); +}); + +describe("WsTransportRouter payload limits", () => { + it("configures Hono's ws server with the requested maxPayload", async () => { + const debug = new Logger("ws-transport-test"); + const router = await ApplicationRouter.create([]); + const transport = new WsTransportRouter(router, debug, 1024); + const wsServer = transport.websocketServer as unknown as WebSocketServer; + + expect(wsServer.options.maxPayload).toBe(1024); + await transport.stop(); + }); +});