From 5022eebde1da8ef60dc0b0fe766793d284924d09 Mon Sep 17 00:00:00 2001 From: nicoske Date: Wed, 22 Jul 2026 20:36:54 +0200 Subject: [PATCH] fix(websocket): capture upgrade request and pass wrapped socket to lifecycle hooks The adapter never installed an upgrade handler, so the HTTP upgrade request (path, query string, headers) was thrown away and there was no way for a gateway to read auth tokens or client metadata from the connection. handleConnection/handleDisconnect also received the raw uWS socket while message handlers get the wrapped UwsSocket, so data attached during handleConnection was invisible to handlers and the documented UwsSocket API (emit, join, data) crashed inside lifecycle hooks. Capture the upgrade request in an upgrade handler, expose it as socket.handshake (url, parsed query, headers, remote address), and hand the same wrapped socket instance to lifecycle hooks and message handlers. Matches the Lifecycle.md examples, which already type the hook parameter as UwsSocket and call client.emit() in handleConnection. --- src/websocket/adapter/uws.adapter.ts | 87 +++++++++- src/websocket/core/socket.ts | 14 +- .../interfaces/uws-socket.interface.ts | 40 +++++ test/websocket/gateway-lifecycle.e2e.spec.ts | 162 ++++++++++++++++++ 4 files changed, 298 insertions(+), 5 deletions(-) create mode 100644 test/websocket/gateway-lifecycle.e2e.spec.ts diff --git a/src/websocket/adapter/uws.adapter.ts b/src/websocket/adapter/uws.adapter.ts index fae8fff..4079ec4 100644 --- a/src/websocket/adapter/uws.adapter.ts +++ b/src/websocket/adapter/uws.adapter.ts @@ -9,6 +9,7 @@ import type { ResolvedUwsAdapterOptions, } from '../interfaces/uws-options.interface'; import type { WebSocketClient } from '../interfaces/websocket-client.interface'; +import type { SocketHandshake } from '../interfaces/uws-socket.interface'; import { MetadataScanner } from '../routing/metadata-scanner'; import { MessageRouter } from '../routing/message-router'; import { HandlerExecutor } from '../routing/handler-executor'; @@ -22,6 +23,18 @@ import { DefaultModuleRef, NestJsModuleRef, type ModuleRef } from '../../shared/ */ interface ExtendedWebSocket extends uWS.WebSocket { id?: string; + upgradeRequest?: UpgradeRequestData; +} + +/** + * Raw upgrade request data captured in the `upgrade` handler. + * uWebSockets.js invalidates the HttpRequest as soon as the upgrade callback + * returns, so everything needed later has to be copied out here. + */ +interface UpgradeRequestData { + url: string; + queryString: string; + headers: Record; } /** @@ -252,6 +265,33 @@ export class UwsAdapter implements WebSocketAdapter { closeOnBackpressureLimit: this.options.closeOnBackpressureLimit, sendPingsAutomatically: this.options.sendPingsAutomatically, + upgrade: ( + res: uWS.HttpResponse, + req: uWS.HttpRequest, + context: uWS.us_socket_context_t + ) => { + // Copy the upgrade request out now: uWebSockets.js invalidates + // `req` as soon as this callback returns. + const headers: Record = {}; + req.forEach((key, value) => { + headers[key] = value; + }); + + const upgradeRequest: UpgradeRequestData = { + url: req.getUrl(), + queryString: req.getQuery() || '', + headers, + }; + + res.upgrade( + { upgradeRequest } as unknown as WebSocketClient, + req.getHeader('sec-websocket-key'), + req.getHeader('sec-websocket-protocol'), + req.getHeader('sec-websocket-extensions'), + context + ); + }, + open: (ws: uWS.WebSocket) => { const extWs = ws as ExtendedWebSocket; const id = this.generateId(); @@ -265,13 +305,17 @@ export class UwsAdapter implements WebSocketAdapter { this.roomManager, this.broadcastToRooms.bind(this) ); + socket.handshake = this.buildHandshake(extWs); this.sockets.set(id, socket); try { - // Call lifecycle hooks for all registered gateways + // Lifecycle hooks receive the wrapped socket (the same instance + // message handlers get via @ConnectedSocket), so data attached in + // handleConnection stays visible to handlers and the documented + // UwsSocket API (emit/join/handshake) works there. this.gateways.forEach((_name, gateway) => { try { - this.lifecycleHooksManager.callConnectionHook(gateway, extWs); + this.lifecycleHooksManager.callConnectionHook(gateway, socket); } catch (error) { this.logger.error( `Connection hook error for ${gateway.constructor?.name}: ${this.formatError(error)}` @@ -315,6 +359,7 @@ export class UwsAdapter implements WebSocketAdapter { close: (ws: uWS.WebSocket, _code: number, _message: ArrayBuffer) => { const extWs = ws as ExtendedWebSocket; const id = extWs.id; + const socket = id ? this.sockets.get(id) : undefined; if (id) { // Remove client from all rooms @@ -325,10 +370,11 @@ export class UwsAdapter implements WebSocketAdapter { } try { - // Call lifecycle hooks for all registered gateways + // Call lifecycle hooks for all registered gateways with the same + // wrapped socket handleConnection received this.gateways.forEach((_name, gateway) => { try { - this.lifecycleHooksManager.callDisconnectHook(gateway, extWs); + this.lifecycleHooksManager.callDisconnectHook(gateway, socket ?? extWs); } catch (error) { this.logger.error( `Disconnect hook error for ${gateway.constructor?.name}: ${this.formatError(error)}` @@ -850,6 +896,39 @@ export class UwsAdapter implements WebSocketAdapter { return randomBytes(8).toString('hex'); } + /** + * Build the socket handshake from the upgrade request captured in the + * `upgrade` handler + * @param extWs - Native socket carrying the captured upgrade request + * @returns Handshake with url, parsed query, headers and remote address + */ + private buildHandshake(extWs: ExtendedWebSocket): SocketHandshake { + const upgradeRequest = extWs.upgradeRequest; + const query: Record = {}; + + if (upgradeRequest?.queryString) { + for (const [key, value] of new URLSearchParams(upgradeRequest.queryString)) { + if (!(key in query)) { + query[key] = value; + } + } + } + + let address = ''; + try { + address = Buffer.from(extWs.getRemoteAddressAsText()).toString('utf-8'); + } catch { + // Socket already closed - leave the address empty + } + + return { + url: upgradeRequest?.url ?? '', + query, + headers: upgradeRequest?.headers ?? {}, + address, + }; + } + /** * Serialize data to JSON string * @param data - Data to serialize diff --git a/src/websocket/core/socket.ts b/src/websocket/core/socket.ts index 90c65dd..dc95bec 100644 --- a/src/websocket/core/socket.ts +++ b/src/websocket/core/socket.ts @@ -1,5 +1,10 @@ import * as uWS from 'uWebSockets.js'; -import { UwsSocket, BroadcastOperator as IBroadcastOperator, WebSocketClient } from '../interfaces'; +import { + UwsSocket, + SocketHandshake, + BroadcastOperator as IBroadcastOperator, + WebSocketClient, +} from '../interfaces'; import { RoomManager } from '../rooms'; import { BroadcastOperator } from './broadcast-operator'; @@ -22,6 +27,13 @@ export class UwsSocketImpl implements UwsS > { private _id: string; private _data: TData; + + /** + * Data captured from the HTTP upgrade request (path, query, headers, + * remote address). Set by the adapter when the connection is opened. + */ + handshake?: SocketHandshake; + private nativeSocket: uWS.WebSocket; private roomManager: RoomManager; private broadcastFn: ( diff --git a/src/websocket/interfaces/uws-socket.interface.ts b/src/websocket/interfaces/uws-socket.interface.ts index 1cf081d..0d2cd1c 100644 --- a/src/websocket/interfaces/uws-socket.interface.ts +++ b/src/websocket/interfaces/uws-socket.interface.ts @@ -1,3 +1,30 @@ +/** + * Data captured from the HTTP upgrade request when the WebSocket connection + * was established + */ +export interface SocketHandshake { + /** + * Request path of the upgrade request (e.g. '/chat') + */ + url: string; + + /** + * Query parameters parsed from the upgrade request URL. + * When a key appears multiple times, the first value wins. + */ + query: Record; + + /** + * Headers sent with the upgrade request (header names are lowercase) + */ + headers: Record; + + /** + * Remote address of the connection as reported by uWebSockets.js + */ + address: string; +} + /** * Enhanced WebSocket interface with uWestJS features * This wraps the native uWebSockets.js socket with a Socket.IO-like API @@ -10,6 +37,19 @@ export interface UwsSocket { */ readonly id: string; + /** + * Data captured from the HTTP upgrade request (path, query, headers, + * remote address). Always set by the adapter for sockets it creates. + * @example + * ```typescript + * handleConnection(client: UwsSocket) { + * const token = client.handshake?.query.token; + * const userAgent = client.handshake?.headers['user-agent']; + * } + * ``` + */ + handshake?: SocketHandshake; + /** * Custom data attached to the socket * Use this to store user information, session data, etc. diff --git a/test/websocket/gateway-lifecycle.e2e.spec.ts b/test/websocket/gateway-lifecycle.e2e.spec.ts new file mode 100644 index 0000000..12d436e --- /dev/null +++ b/test/websocket/gateway-lifecycle.e2e.spec.ts @@ -0,0 +1,162 @@ +import 'reflect-metadata'; +import { UwsAdapter } from '../../src/websocket/adapter/uws.adapter'; +import { UwsSocket, SocketHandshake } from '../../src/websocket/interfaces'; +import { ParamType, PARAM_ARGS_METADATA } from '../../src/websocket/decorators'; + +const MESSAGE_MAPPING_METADATA = 'websockets:message_mapping'; +const MESSAGE_METADATA = 'message'; + +function addMessageMetadata(target: object, event: string): void { + Reflect.defineMetadata(MESSAGE_MAPPING_METADATA, true, target); + Reflect.defineMetadata(MESSAGE_METADATA, event, target); +} + +function addParamMetadata( + target: object, + methodName: string, + params: Array<{ index: number; type: ParamType }> +): void { + Reflect.defineMetadata(PARAM_ARGS_METADATA, params, target, methodName); +} + +type TestSocketData = { label?: string }; + +/** + * Gateway following the Lifecycle.md examples: handleConnection uses the + * UwsSocket API (emit, data, handshake) and message handlers receive the + * connected socket via parameter metadata. + */ +class LifecycleGateway { + connectionClients: UwsSocket[] = []; + disconnectClients: UwsSocket[] = []; + handlerClients: UwsSocket[] = []; + handshakes: Array = []; + + handleConnection(client: UwsSocket) { + this.connectionClients.push(client); + this.handshakes.push(client.handshake); + client.data = { label: 'set-in-connection-hook' }; + client.emit('welcome', { message: 'Hello!' }); + } + + handleDisconnect(client: UwsSocket) { + this.disconnectClients.push(client); + } + + handleWhoami(client: UwsSocket) { + this.handlerClients.push(client); + return { label: client.data?.label ?? null }; + } +} + +addMessageMetadata(LifecycleGateway.prototype.handleWhoami, 'whoami'); +addParamMetadata(LifecycleGateway.prototype, 'handleWhoami', [ + { index: 0, type: ParamType.CONNECTED_SOCKET }, +]); + +function connect(url: string): Promise { + return new Promise((resolve, reject) => { + const ws = new WebSocket(url); + ws.addEventListener('open', () => resolve(ws)); + ws.addEventListener('error', () => reject(new Error(`failed to connect to ${url}`))); + }); +} + +function nextMessage(ws: WebSocket): Promise<{ event: string; data?: unknown }> { + return new Promise((resolve, reject) => { + const timer = setTimeout(() => reject(new Error('timed out waiting for message')), 5000); + ws.addEventListener( + 'message', + (event) => { + clearTimeout(timer); + resolve(JSON.parse(String(event.data))); + }, + { once: true } + ); + }); +} + +function closed(ws: WebSocket): Promise { + return new Promise((resolve) => { + if (ws.readyState === WebSocket.CLOSED) { + resolve(); + return; + } + ws.addEventListener('close', () => resolve(), { once: true }); + }); +} + +async function waitFor(condition: () => boolean, timeoutMs = 5000): Promise { + const deadline = Date.now() + timeoutMs; + while (!condition()) { + if (Date.now() > deadline) { + throw new Error('timed out waiting for condition'); + } + await new Promise((resolve) => setTimeout(resolve, 20)); + } +} + +describe('Gateway lifecycle E2E', () => { + const port = 13390; + let adapter: UwsAdapter; + let gateway: LifecycleGateway; + + beforeEach(async () => { + gateway = new LifecycleGateway(); + adapter = new UwsAdapter(null, { port }); + await adapter.create(port); + adapter.registerGateway(gateway); + adapter.bindClientConnect(null as never, () => undefined); + await waitFor(() => adapter.getClientCount() === 0); + }); + + afterEach(() => { + adapter?.dispose(); + }); + + it('exposes the upgrade request on client.handshake', async () => { + const ws = await connect(`ws://localhost:${port}/chat?token=abc&room=42&room=43`); + await nextMessage(ws); // welcome + + expect(gateway.handshakes).toHaveLength(1); + const handshake = gateway.handshakes[0]; + expect(handshake).toBeDefined(); + expect(handshake!.url).toBe('/chat'); + expect(handshake!.query).toEqual({ token: 'abc', room: '42' }); + expect(handshake!.headers['host']).toBe(`localhost:${port}`); + expect(handshake!.headers['upgrade']).toBe('websocket'); + expect(handshake!.address.length).toBeGreaterThan(0); + + ws.close(); + await closed(ws); + }); + + it('passes the wrapped socket to handleConnection so the UwsSocket API works', async () => { + const ws = await connect(`ws://localhost:${port}/`); + + // handleConnection calls client.emit('welcome') - with the raw uWS socket + // this would throw because emit() does not exist there + const welcome = await nextMessage(ws); + expect(welcome).toEqual({ event: 'welcome', data: { message: 'Hello!' } }); + + ws.close(); + await closed(ws); + }); + + it('gives lifecycle hooks and message handlers the same socket instance', async () => { + const ws = await connect(`ws://localhost:${port}/`); + await nextMessage(ws); // welcome + + ws.send(JSON.stringify({ event: 'whoami' })); + const reply = await nextMessage(ws); + + // data attached in handleConnection is visible to the message handler + expect(reply).toEqual({ event: 'whoami', data: { label: 'set-in-connection-hook' } }); + expect(gateway.handlerClients[0]).toBe(gateway.connectionClients[0]); + + ws.close(); + await closed(ws); + await waitFor(() => gateway.disconnectClients.length === 1); + expect(gateway.disconnectClients[0]).toBe(gateway.connectionClients[0]); + }); +});