fix(websocket): capture upgrade request and pass wrapped socket to lifecycle hooks

src/websocket/adapter/uws.adapter.ts

@@ -9,6 +9,7 @@
   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 @@
  */
 interface ExtendedWebSocket extends uWS.WebSocket<WebSocketClient> {
   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<string, string>;
 }
 
 /**
@@ -113,74 +126,74 @@
   private gatewaySet = new WeakSet<object>(); // Track registered instances
   private bindMessageHandlersCalled = false;
 
   constructor(appInstance: unknown, options?: UwsAdapterOptions) {
     // Note: appInstance is required by NestJS WebSocketAdapter interface but not used
     // in this implementation. NestJS passes the app instance for adapters that need
     // access to the application context, but UwsAdapter uses explicit gateway registration
     // via registerGateway() instead.
 
     // Apply default options
     // Note: port will be set in create() using NestJS-provided port as fallback
     this.options = {
       maxPayloadLength: options?.maxPayloadLength ?? 16 * 1024,
       idleTimeout: options?.idleTimeout ?? 60,
       maxLifetime: options?.maxLifetime ?? 0, // 0 = disabled (no lifetime limit)
       compression: options?.compression ?? uWS.SHARED_COMPRESSOR,
       port: options?.port ?? 8099, // Default to 8099 if not provided
       path: options?.path ?? '/*',
       maxBackpressure: options?.maxBackpressure ?? 1024 * 1024, // 1MB default
       closeOnBackpressureLimit: options?.closeOnBackpressureLimit ?? false,
       sendPingsAutomatically: options?.sendPingsAutomatically ?? true,
       cors: options?.cors,
       // SSL/TLS fields - carried through as-is (undefined when not supplied)
       cert_file_name: options?.cert_file_name,
       key_file_name: options?.key_file_name,
       passphrase: options?.passphrase,
       dh_params_file_name: options?.dh_params_file_name,
       ssl_prefer_low_memory_usage: options?.ssl_prefer_low_memory_usage,
     };
 
     // Use provided uWS App if available (v2.0.0+ for HTTP + WebSocket integration)
     if (options?.uwsApp) {
       this.app = options.uwsApp;
       this.isSharedApp = true;
       this.logger.log('UwsAdapter using provided uWS App instance (shared with HTTP)');
 
       // Warn if port is also specified (it will be ignored in shared mode)
       if (options?.port) {
         this.logger.warn(
           'Both uwsApp and port options provided. Port option is ignored when using shared app - ' +
             'the HTTP server manages the listening port.'
         );
       }
 
       // Warn if SSL options are also specified (TLS is controlled by the shared app)
       if (
         options.cert_file_name ||
         options.key_file_name ||
         options.passphrase ||
         options.dh_params_file_name ||
         options.ssl_prefer_low_memory_usage !== undefined
       ) {
         this.logger.warn(
           'SSL/TLS options are ignored when uwsApp is provided - the shared uWS instance controls TLS configuration.'
         );
       }
     }
 
     // Initialize handler executor with optional ModuleRef for DI support
     // Accept both our ModuleRef interface and NestJS ModuleRef, auto-wrap if needed
     // Normalize null to undefined for consistent "no DI" representation
     const rawModuleRef = options?.moduleRef ?? undefined;
     let moduleRef: ModuleRef | undefined = rawModuleRef;
     if (rawModuleRef && !(rawModuleRef instanceof DefaultModuleRef)) {
       // External Nestjs ModuleRef provided, wrap it with our adapter.
       moduleRef = NestJsModuleRef.create(rawModuleRef as unknown as NestModuleRef);
     }
     this.handlerExecutor = new HandlerExecutor({ moduleRef });
 
     this.logger.log('UwsAdapter initialized');
   }
 
   /**
    * Create the uWebSockets.js server
@@ -252,6 +265,33 @@
         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<string, string> = {};
+          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<WebSocketClient>) => {
           const extWs = ws as ExtendedWebSocket;
           const id = this.generateId();
@@ -265,13 +305,17 @@
             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 @@
         close: (ws: uWS.WebSocket<WebSocketClient>, _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 @@
           }
 
           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 @@
     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<string, string> = {};
+
+    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

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<TData = unknown, TEmitData = unknown> 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<WebSocketClient>;
   private roomManager: RoomManager;
   private broadcastFn: (

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<string, string>;
+
+  /**
+   * Headers sent with the upgrade request (header names are lowercase)
+   */
+  headers: Record<string, string>;
+
+  /**
+   * 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<TData = unknown, TEmitData = unknown> {
    */
   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.

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<TestSocketData>[] = [];
+  disconnectClients: UwsSocket<TestSocketData>[] = [];
+  handlerClients: UwsSocket<TestSocketData>[] = [];
+  handshakes: Array<SocketHandshake | undefined> = [];
+
+  handleConnection(client: UwsSocket<TestSocketData>) {
+    this.connectionClients.push(client);
+    this.handshakes.push(client.handshake);
+    client.data = { label: 'set-in-connection-hook' };
+    client.emit('welcome', { message: 'Hello!' });
+  }
+
+  handleDisconnect(client: UwsSocket<TestSocketData>) {
+    this.disconnectClients.push(client);
+  }
+
+  handleWhoami(client: UwsSocket<TestSocketData>) {
+    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<WebSocket> {
+  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<void> {
+  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<void> {
+  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]);
+  });
+});