Add a selection of chainable websocket endpoints

Author: pimterry

src/endpoint-chain.ts

@@ -0,0 +1,36 @@
+import { StatusError } from '@httptoolkit/util';
+
+const MAX_CHAIN_DEPTH = 10;
+
+export interface ChainableEndpoint {
+    matchPath: (path: string, hostnamePrefix?: string) => boolean;
+    getRemainingPath?: (path: string) => string | undefined;
+}
+
+export function resolveEndpointChain<T extends ChainableEndpoint>(
+    endpoints: Array<T & { name: string }>,
+    initialPath: string,
+    hostnamePrefix?: string
+): Array<{ endpoint: T & { name: string }; path: string }> {
+    const entries: Array<{ endpoint: T & { name: string }; path: string }> = [];
+    let path: string | undefined = initialPath;
+
+    while (path && entries.length <= MAX_CHAIN_DEPTH) {
+        // matchPath may throw StatusError for invalid parameters
+        const endpoint = endpoints.find(ep => ep.matchPath(path!, hostnamePrefix));
+        if (!endpoint) {
+            throw new StatusError(404, `Could not match endpoint for ${initialPath}${
+                hostnamePrefix ? ` (${hostnamePrefix})` : ''
+            }`);
+        }
+
+        entries.push({ endpoint, path });
+        path = endpoint.getRemainingPath?.(path);
+    }
+
+    if (path) {
+        throw new StatusError(400, `Endpoint chain exceeded maximum depth for ${initialPath}`);
+    }
+
+    return entries;
+}

src/endpoints/http/delay.ts

@@ -1,9 +1,7 @@
-import { delay } from '@httptoolkit/util';
+import { delay, StatusError } from '@httptoolkit/util';
 import { HttpEndpoint, HttpHandler } from '../http-index.js';
 import { buildHttpBinAnythingEndpoint } from '../../httpbin-compat.js';
 
-const matchPath = (path: string) => path.startsWith('/delay/');
-
 const getRemainingPath = (path: string): string | undefined => {
     const idx = path.indexOf('/', '/delay/'.length);
     return idx !== -1 ? path.slice(idx) : undefined;
@@ -21,13 +19,6 @@ const defaultAnythingEndpoint = buildHttpBinAnythingEndpoint({
 
 const handle: HttpHandler = async (req, res, { path }) => {
     const delaySeconds = parseDelaySeconds(path);
-
-    if (isNaN(delaySeconds)) {
-        res.writeHead(400);
-        res.end('Invalid delay duration');
-        return;
-    }
-
     const cappedDelayMs = Math.min(delaySeconds, 10) * 1000;
     await delay(cappedDelayMs);
 
@@ -39,7 +30,14 @@ const handle: HttpHandler = async (req, res, { path }) => {
 };
 
 export const delayEndpoint: HttpEndpoint = {
-    matchPath,
+    matchPath: (path) => {
+        if (!path.startsWith('/delay/')) return false;
+        const delaySeconds = parseDelaySeconds(path);
+        if (isNaN(delaySeconds)) {
+            throw new StatusError(400, `Invalid delay duration in ${path}`);
+        }
+        return true;
+    },
     handle,
     getRemainingPath
 };

src/endpoints/http/status.ts

@@ -1,19 +1,24 @@
+import { StatusError } from '@httptoolkit/util';
 import { HttpEndpoint, HttpHandler } from '../http-index.js';
 
-const matchPath = (path: string) => path.startsWith('/status/');
+const parseStatusCode = (path: string): number => {
+    return parseInt(path.slice('/status/'.length), 10);
+};
 
 const handle: HttpHandler = (_req, res, { path }) => {
-    const statusCode = parseInt(path.slice('/status/'.length), 10);
-    if (isNaN(statusCode)) {
-        res.writeHead(400);
-        res.end('Invalid status code');
-    } else {
-        res.writeHead(statusCode);
-        res.end();
-    }
+    const statusCode = parseStatusCode(path);
+    res.writeHead(statusCode);
+    res.end();
 }
 
 export const status: HttpEndpoint = {
-    matchPath,
+    matchPath: (path) => {
+        if (!path.startsWith('/status/')) return false;
+        const statusCode = parseStatusCode(path);
+        if (isNaN(statusCode)) {
+            throw new StatusError(400, `Invalid status code in ${path}`);
+        }
+        return true;
+    },
     handle
 }

src/endpoints/ws-index.ts

@@ -2,11 +2,16 @@ import { WebSocket } from 'ws';
 import { IncomingMessage } from 'http';
 
 export interface WebSocketEndpoint {
+    /** Return true to match, false to skip, or throw StatusError for invalid params. */
     matchPath: (path: string, hostnamePrefix?: string) => boolean;
+    getRemainingPath?: (path: string) => string | undefined;
     handle: (ws: WebSocket, req: IncomingMessage, options: {
         path: string;
         query: URLSearchParams;
-    }) => void;
+    }) => void | Promise<void>;
 }
 
 export * from './ws/echo.js';
+export * from './ws/delay.js';
+export * from './ws/close.js';
+export * from './ws/message.js';

src/endpoints/ws/close.ts

@@ -0,0 +1,23 @@
+import { StatusError } from '@httptoolkit/util';
+import { WebSocketEndpoint } from '../ws-index.js';
+
+const parseCloseCode = (path: string): number => {
+    if (path === '/ws/close') return 1000;
+    return parseInt(path.slice('/ws/close/'.length), 10);
+};
+
+export const wsCloseEndpoint: WebSocketEndpoint = {
+    matchPath: (path) => {
+        if (!path.startsWith('/ws/close/') && path !== '/ws/close') return false;
+        const code = parseCloseCode(path);
+        if (isNaN(code) || code < 1000 || code > 4999) {
+            throw new StatusError(400, `Invalid WebSocket close code: ${code}`);
+        }
+        return true;
+    },
+    handle: (ws, req, { path, query }) => {
+        const code = parseCloseCode(path);
+        const reason = query.get('reason') ?? undefined;
+        ws.close(code, reason);
+    }
+};

src/endpoints/ws/delay.ts

@@ -0,0 +1,30 @@
+import { delay, StatusError } from '@httptoolkit/util';
+import { WebSocketEndpoint } from '../ws-index.js';
+
+const getRemainingPath = (path: string): string | undefined => {
+    const idx = path.indexOf('/', '/ws/delay/'.length);
+    return idx !== -1 ? '/ws' + path.slice(idx) : undefined;
+};
+
+const parseDelaySeconds = (path: string): number => {
+    const idx = path.indexOf('/', '/ws/delay/'.length);
+    const end = idx !== -1 ? idx : path.length;
+    return parseFloat(path.slice('/ws/delay/'.length, end));
+};
+
+export const wsDelayEndpoint: WebSocketEndpoint = {
+    matchPath: (path) => {
+        if (!path.startsWith('/ws/delay/')) return false;
+        const delaySeconds = parseDelaySeconds(path);
+        if (isNaN(delaySeconds)) {
+            throw new StatusError(400, `Invalid delay duration in ${path}`);
+        }
+        return true;
+    },
+    getRemainingPath,
+    handle: async (ws, req, { path }) => {
+        const delaySeconds = parseDelaySeconds(path);
+        const cappedDelayMs = Math.min(delaySeconds, 10) * 1000;
+        await delay(cappedDelayMs);
+    }
+};

src/endpoints/ws/message.ts

@@ -0,0 +1,25 @@
+import { WebSocketEndpoint } from '../ws-index.js';
+
+const matchPath = (path: string) => path.startsWith('/ws/message/');
+
+const getRemainingPath = (path: string): string | undefined => {
+    const idx = path.indexOf('/', '/ws/message/'.length);
+    return idx !== -1 ? '/ws' + path.slice(idx) : undefined;
+};
+
+const parseMessage = (path: string): string => {
+    const idx = path.indexOf('/', '/ws/message/'.length);
+    const end = idx !== -1 ? idx : path.length;
+    return decodeURIComponent(path.slice('/ws/message/'.length, end));
+};
+
+export const wsMessageEndpoint: WebSocketEndpoint = {
+    matchPath,
+    getRemainingPath,
+    handle: (ws, req, { path }) => {
+        const message = parseMessage(path);
+        if (ws.readyState === ws.OPEN) {
+            ws.send(message);
+        }
+    }
+};

src/http-handler.ts

@@ -6,35 +6,7 @@ import { MaybePromise, StatusError } from '@httptoolkit/util';
 import { httpEndpoints } from './endpoints/endpoint-index.js';
 import { HttpRequest, HttpResponse } from './endpoints/http-index.js';
 import { handleWebSocketUpgrade } from './ws-handler.js';
-
-const MAX_CHAIN_DEPTH = 10;
-
-function resolveEndpointChain(initialPath: string, hostnamePrefix?: string) {
-    const entries: Array<{ endpoint: typeof httpEndpoints[number]; path: string }> = [];
-    let needsRawData = false;
-    let path: string | undefined = initialPath;
-
-    while (path && entries.length <= MAX_CHAIN_DEPTH) {
-        const endpoint = httpEndpoints.find(ep => ep.matchPath(path!, hostnamePrefix));
-        if (!endpoint) {
-            throw new StatusError(404, `Could not match endpoint for ${initialPath}${
-                hostnamePrefix
-                ? ` (${hostnamePrefix})`
-                : ''
-            }`);
-        }
-
-        entries.push({ endpoint, path });
-        needsRawData ||= !!endpoint.needsRawData;
-        path = endpoint.getRemainingPath?.(path);
-    }
-
-    if (path) {
-        throw new StatusError(400, `Endpoint chain exceeded maximum depth with path: ${initialPath}`);
-    }
-
-    return { entries, needsRawData };
-}
+import { resolveEndpointChain } from './endpoint-chain.js';
 
 function stopRawDataCapture(req: HttpRequest): void {
     if (req.httpVersion === '2.0') {
@@ -144,19 +116,13 @@ function createHttpRequestHandler(options: {
             return;
         }
 
-        const { entries, needsRawData } = resolveEndpointChain(path, hostnamePrefix);
+        const entries = resolveEndpointChain(httpEndpoints, path, hostnamePrefix);
+        const needsRawData = entries.some(e => e.endpoint.needsRawData);
 
         if (!needsRawData) {
             stopRawDataCapture(req);
         }
 
-        if (entries.length === 0) {
-            console.log(`Request to ${path} matched no endpoints`);
-            res.writeHead(404);
-            res.end(`No handler for ${req.url}`);
-            return;
-        }
-
         const endpointNames = entries.map(e => e.endpoint.name).join(' → ');
         console.log(`Request to ${path}${
             hostnamePrefix ? ` ('${hostnamePrefix}' prefix)` : ` (${options.rootDomain})`

src/ws-handler.ts

@@ -1,8 +1,10 @@
 import { IncomingMessage } from 'http';
 import { Duplex } from 'stream';
 import { WebSocketServer } from 'ws';
+import { StatusError } from '@httptoolkit/util';
 
 import { wsEndpoints } from './endpoints/endpoint-index.js';
+import { resolveEndpointChain } from './endpoint-chain.js';
 
 const wss = new WebSocketServer({ noServer: true });
 
@@ -19,36 +21,48 @@ export function handleWebSocketUpgrade(
         ? url.hostname.slice(0, -options.rootDomain.length - 1)
         : undefined;
 
-    const endpoint = wsEndpoints.find(ep => ep.matchPath(path, hostnamePrefix));
-
-    if (!endpoint) {
-        console.log(`WebSocket upgrade to ${path} matched no endpoints`);
-        socket.write('HTTP/1.1 404 Not Found\r\n\r\n');
+    let entries: typeof wsEndpoints extends Array<infer T> ? Array<{ endpoint: T; path: string }> : never;
+    try {
+        entries = resolveEndpointChain(wsEndpoints, path, hostnamePrefix);
+    } catch (err) {
+        if (err instanceof StatusError) {
+            console.log(`WebSocket upgrade to ${path}: ${err.message}`);
+            socket.write(`HTTP/1.1 ${err.statusCode} ${err.statusCode === 404 ? 'Not Found' : 'Bad Request'}\r\n\r\n`);
+        } else {
+            console.log(`WebSocket upgrade to ${path}: unexpected error`, err);
+            socket.write('HTTP/1.1 500 Internal Server Error\r\n\r\n');
+        }
         socket.destroy();
         return;
     }
 
+    const endpointNames = entries.map(e => e.endpoint.name).join(' → ');
     console.log(`WebSocket upgrade to ${path}${
         hostnamePrefix ? ` ('${hostnamePrefix}' prefix)` : ''
-    } matched: ${endpoint.name}`);
+    } matched: ${endpointNames}`);
 
     socket.on('error', (err) => {
         console.log('WebSocket upgrade socket error:', err.message);
     });
 
-    wss.handleUpgrade(req, socket, head, (ws) => {
+    wss.handleUpgrade(req, socket, head, async (ws) => {
         ws.on('error', (err) => {
             console.log(`WebSocket error on ${path}:`, err.message);
         });
 
         try {
-            endpoint.handle(ws, req, {
-                path,
-                query: url.searchParams
-            });
+            for (const { endpoint, path: entryPath } of entries) {
+                if (ws.readyState !== ws.OPEN) return;
+                await endpoint.handle(ws, req, {
+                    path: entryPath,
+                    query: url.searchParams
+                });
+            }
         } catch (err) {
             console.log(`WebSocket handler error on ${path}:`, err);
-            ws.close(1011, 'Internal error');
+            if (ws.readyState === ws.OPEN) {
+                ws.close(1011, 'Internal error');
+            }
         }
     });
 }