sync-server-v2/source/services/transport/ws-transport.ts

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 type { ApplicationRequest, ApplicationRouter } from '../router.ts';

import { ApplicationError, normalizePublicError } from '../../errors/index.ts';
import { HTTP_STATUS_CODE_BAD_REQUEST } from '../../constants.ts';
import { WSStream } from '../stream/ws-stream.ts';
import { type AppEnv, BaseTransport } from './base-transport.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 extends BaseTransport {
    /**
     * 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;

    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,
    ) {
        super();

        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,
        });
    }

    /**
     * Register the connection endpoint and translate socket lifecycle events.
     *
     * @param server - Hono instance to attach the upgrade handler to.
     */
    register(server: Hono<AppEnv>): 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<void> {
        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<void>((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<void> {
        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<ApplicationRequest> {
        const payload = await WsTransportRouter.decodeWebSocketMessage(message);
        let decoded: unknown;

        try {
            decoded = fromExtendedJson(payload);
        } catch {
            throw new ApplicationError(HTTP_STATUS_CODE_BAD_REQUEST, '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<string> {
        if (typeof message === 'string') {
            return message;
        }

        if (message instanceof Blob) {
            return Buffer.from(await message.arrayBuffer()).toString('utf8');
        }

        return Buffer.from(message).toString('utf8');
    }
}