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]); + }); +});