Add support for websockets

Author: pimterry

package-lock.json

@@ -40,13 +40,15 @@
     "lodash": "^4.17.23",
     "parse-multipart-data": "^1.5.0",
     "read-tls-client-hello": "^1.1.0",
-    "tsx": "^4.19.3"
+    "tsx": "^4.19.3",
+    "ws": "^8.19.0"
   },
   "devDependencies": {
     "@types/chai": "^4.3.14",
     "@types/lodash": "^4.17.0",
     "@types/mocha": "^10.0.6",
     "@types/node": "^22.15.30",
+    "@types/ws": "^8.18.1",
     "chai": "^5.1.0",
     "destroyable-server": "^1.0.1",
     "mocha": "^10.8.2",

package.json

@@ -1,8 +1,12 @@
 import * as httpIndex from './http-index.js';
 import * as tlsIndex from './tls-index.js';
+import * as wsIndex from './ws-index.js';
 
 export const httpEndpoints: Array<httpIndex.HttpEndpoint & { name: string }> = Object.entries(httpIndex)
     .map(([key, value]) => ({ ...value, name: key }));
 
 export const tlsEndpoints: Array<tlsIndex.TlsEndpoint & { name: string }> = Object.entries(tlsIndex)
-    .map(([key, value]) => ({ ...value, name: key }));
+    .map(([key, value]) => ({ ...value, name: key }));
+
+export const wsEndpoints: Array<wsIndex.WebSocketEndpoint & { name: string }> = Object.entries(wsIndex)
+    .map(([key, value]) => ({ ...(value as wsIndex.WebSocketEndpoint), name: key }));

src/endpoints/endpoint-index.ts

@@ -0,0 +1,12 @@
+import { WebSocket } from 'ws';
+import { IncomingMessage } from 'http';
+
+export interface WebSocketEndpoint {
+    matchPath: (path: string, hostnamePrefix?: string) => boolean;
+    handle: (ws: WebSocket, req: IncomingMessage, options: {
+        path: string;
+        query: URLSearchParams;
+    }) => void;
+}
+
+export * from './ws/echo.js';

src/endpoints/ws-index.ts

@@ -0,0 +1,12 @@
+import { WebSocketEndpoint } from '../ws-index.js';
+
+export const wsEchoEndpoint: WebSocketEndpoint = {
+    matchPath: (path) => path === '/ws/echo',
+    handle: (ws) => {
+        ws.on('message', (data, isBinary) => {
+            if (ws.readyState === ws.OPEN) {
+                ws.send(data, { binary: isBinary });
+            }
+        });
+    }
+};

src/endpoints/ws/echo.ts

@@ -5,6 +5,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;
 
@@ -209,6 +210,10 @@ export function createHttp1Handler(options: {
 
     handler.on('error', (err) => console.error('HTTP handler error', err));
 
+    handler.on('upgrade', (req, socket, head) => {
+        handleWebSocketUpgrade(req, socket, head, options);
+    });
+
     return handler;
 }
 

src/http-handler.ts

@@ -0,0 +1,54 @@
+import { IncomingMessage } from 'http';
+import { Duplex } from 'stream';
+import { WebSocketServer } from 'ws';
+
+import { wsEndpoints } from './endpoints/endpoint-index.js';
+
+const wss = new WebSocketServer({ noServer: true });
+
+export function handleWebSocketUpgrade(
+    req: IncomingMessage,
+    socket: Duplex,
+    head: Buffer,
+    options: { rootDomain: string }
+) {
+    const url = new URL(req.url!, `http://${req.headers.host}`);
+    const path = url.pathname;
+
+    const hostnamePrefix = url.hostname.endsWith(options.rootDomain)
+        ? 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');
+        socket.destroy();
+        return;
+    }
+
+    console.log(`WebSocket upgrade to ${path}${
+        hostnamePrefix ? ` ('${hostnamePrefix}' prefix)` : ''
+    } matched: ${endpoint.name}`);
+
+    socket.on('error', (err) => {
+        console.log('WebSocket upgrade socket error:', err.message);
+    });
+
+    wss.handleUpgrade(req, socket, head, (ws) => {
+        ws.on('error', (err) => {
+            console.log(`WebSocket error on ${path}:`, err.message);
+        });
+
+        try {
+            endpoint.handle(ws, req, {
+                path,
+                query: url.searchParams
+            });
+        } catch (err) {
+            console.log(`WebSocket handler error on ${path}:`, err);
+            ws.close(1011, 'Internal error');
+        }
+    });
+}

src/ws-handler.ts

@@ -0,0 +1,198 @@
+import * as net from 'net';
+import { expect } from 'chai';
+import { WebSocket } from 'ws';
+import { DestroyableServer, makeDestroyable } from 'destroyable-server';
+
+import { createServer } from '../src/server.js';
+
+describe("WebSocket Echo endpoint", () => {
+
+    let server: DestroyableServer;
+    let serverPort: number;
+
+    beforeEach(async () => {
+        server = makeDestroyable(await createServer({
+            domain: 'localhost'
+        }));
+        await new Promise<void>((resolve) => server.listen(resolve));
+        serverPort = (server.address() as net.AddressInfo).port;
+    });
+
+    afterEach(async () => {
+        await server.destroy();
+    });
+
+    it("echoes text messages", async () => {
+        const ws = new WebSocket(`ws://localhost:${serverPort}/ws/echo`);
+
+        await new Promise<void>((resolve, reject) => {
+            ws.on('open', resolve);
+            ws.on('error', reject);
+        });
+
+        ws.send('Hello, WebSocket!');
+
+        const receivedMessage = await new Promise<string>((resolve, reject) => {
+            ws.on('message', (data, isBinary) => {
+                expect(isBinary).to.equal(false);
+                resolve(data.toString());
+            });
+            ws.on('error', reject);
+        });
+
+        expect(receivedMessage).to.equal('Hello, WebSocket!');
+        ws.close();
+    });
+
+    it("echoes binary messages", async () => {
+        const ws = new WebSocket(`ws://localhost:${serverPort}/ws/echo`);
+
+        await new Promise<void>((resolve, reject) => {
+            ws.on('open', resolve);
+            ws.on('error', reject);
+        });
+
+        const testData = Buffer.from([0x01, 0x02, 0x03, 0xFF, 0xFE, 0x00]);
+        ws.send(testData);
+
+        const receivedData = await new Promise<{ data: Buffer; isBinary: boolean }>((resolve, reject) => {
+            ws.on('message', (data, isBinary) => {
+                resolve({ data: data as Buffer, isBinary });
+            });
+            ws.on('error', reject);
+        });
+
+        expect(receivedData.isBinary).to.equal(true);
+        expect(Buffer.compare(receivedData.data, testData)).to.equal(0);
+        ws.close();
+    });
+
+    it("echoes multiple messages in sequence", async () => {
+        const ws = new WebSocket(`ws://localhost:${serverPort}/ws/echo`);
+
+        await new Promise<void>((resolve, reject) => {
+            ws.on('open', resolve);
+            ws.on('error', reject);
+        });
+
+        const messages = ['first', 'second', 'third'];
+        const received: string[] = [];
+
+        for (const msg of messages) {
+            ws.send(msg);
+        }
+
+        await new Promise<void>((resolve, reject) => {
+            ws.on('message', (data) => {
+                received.push(data.toString());
+                if (received.length === messages.length) {
+                    resolve();
+                }
+            });
+            ws.on('error', reject);
+        });
+
+        expect(received).to.deep.equal(messages);
+        ws.close();
+    });
+
+    it("returns 404 for unknown WebSocket paths", async () => {
+        const ws = new WebSocket(`ws://localhost:${serverPort}/ws/unknown`);
+
+        const result = await new Promise<{ statusCode: number }>((resolve, reject) => {
+            ws.on('open', () => {
+                reject(new Error('Should not have connected'));
+            });
+            ws.on('unexpected-response', (req, res) => {
+                resolve({ statusCode: res.statusCode! });
+            });
+            ws.on('error', reject);
+        });
+
+        expect(result.statusCode).to.equal(404);
+    });
+
+    it("returns 404 for non-WebSocket paths", async () => {
+        const ws = new WebSocket(`ws://localhost:${serverPort}/status/200`);
+
+        const result = await new Promise<{ statusCode: number }>((resolve, reject) => {
+            ws.on('open', () => {
+                reject(new Error('Should not have connected'));
+            });
+            ws.on('unexpected-response', (req, res) => {
+                resolve({ statusCode: res.statusCode! });
+            });
+            ws.on('error', reject);
+        });
+
+        expect(result.statusCode).to.equal(404);
+    });
+
+    it("works over HTTPS (wss://)", async () => {
+        const ws = new WebSocket(`wss://localhost:${serverPort}/ws/echo`, {
+            rejectUnauthorized: false
+        });
+
+        await new Promise<void>((resolve, reject) => {
+            ws.on('open', resolve);
+            ws.on('error', reject);
+        });
+
+        ws.send('Secure WebSocket!');
+
+        const receivedMessage = await new Promise<string>((resolve, reject) => {
+            ws.on('message', (data) => {
+                resolve(data.toString());
+            });
+            ws.on('error', reject);
+        });
+
+        expect(receivedMessage).to.equal('Secure WebSocket!');
+        ws.close();
+    });
+
+    it("handles empty messages", async () => {
+        const ws = new WebSocket(`ws://localhost:${serverPort}/ws/echo`);
+
+        await new Promise<void>((resolve, reject) => {
+            ws.on('open', resolve);
+            ws.on('error', reject);
+        });
+
+        ws.send('');
+
+        const receivedMessage = await new Promise<string>((resolve, reject) => {
+            ws.on('message', (data) => {
+                resolve(data.toString());
+            });
+            ws.on('error', reject);
+        });
+
+        expect(receivedMessage).to.equal('');
+        ws.close();
+    });
+
+    it("handles large messages", async () => {
+        const ws = new WebSocket(`ws://localhost:${serverPort}/ws/echo`);
+
+        await new Promise<void>((resolve, reject) => {
+            ws.on('open', resolve);
+            ws.on('error', reject);
+        });
+
+        const largeMessage = 'x'.repeat(100000);
+        ws.send(largeMessage);
+
+        const receivedMessage = await new Promise<string>((resolve, reject) => {
+            ws.on('message', (data) => {
+                resolve(data.toString());
+            });
+            ws.on('error', reject);
+        });
+
+        expect(receivedMessage).to.equal(largeMessage);
+        expect(receivedMessage.length).to.equal(100000);
+        ws.close();
+    });
+
+});