Expose WebSocket upgrade context to handshake validation

README.md

@@ -799,12 +799,13 @@ createClient<typeof services>(clientTransport, 'SERVER', {
 createServer(serverTransport, services, {
   handshakeOptions: createServerHandshakeOptions(
     handshakeSchema,
-    (metadata, previousMetadata, from) => {
+    (metadata, previousMetadata, from, connectionExtras) => {
       // the type of this function is
       // (
       //   metadata: Static<typeof handshakeSchema>,
       //   previousMetadata?: ParsedMetadata,
       //   from?: TransportClientId,
+      //   connectionExtras?: Record<string, unknown>,
       // ) =>
       //   | 'REJECTED_BY_CUSTOM_HANDLER' | 'REJECTED_UNSUPPORTED_CLIENT' (if you reject it)
       //   | ParsedMetadata (if you allow it)
@@ -813,13 +814,28 @@ createServer(serverTransport, services, {
       // next time a connection happens on the same session, previousMetadata will
       // be populated with the last returned value. `from` is the client id the peer
       // presented in its handshake — check it against what the metadata authorizes
-      // before returning parsed metadata.
+      // before returning parsed metadata. `connectionExtras` contains context from
+      // the transport connection that carried this handshake.
       return { parsedToken: metadata.token };
     },
   ),
 });
 ```
 
+For WebSocket servers, `connectionExtras` is `{ headers: cleanedUpgradeHeaders }` by default.
+Pass an extras factory as the fourth transport constructor argument to replace this value with verified upgrade data:
+
+```ts
+const transport = new WebSocketServerTransport(
+  wss,
+  'SERVER',
+  undefined,
+  (ws, req) => ({ identity: getVerifiedUpgradeIdentity(ws, req) }),
+);
+```
+
+The extras factory is synchronous. Complete asynchronous authentication before the WebSocket connection event, then read its result in the factory.
+
 `createClientHandshakeOptions` also takes an optional third `eager` argument. When set, the
 client constructs handshake metadata as soon as it starts dialing, so a slow `construct`
 (e.g. fetching a fresh token) overlaps establishing the connection instead of running after

protobuf/handshake.ts

@@ -18,6 +18,7 @@ import {
 } from '../transport/message';
 import { decodeMessageBytes, encodeMessageBytes } from './shared';
 import { Uint8ArrayType } from '../customSchemas';
+import type { ConnectionExtras } from '../transport/connection';
 
 const HandshakeBytesSchema = Uint8ArrayType();
 
@@ -37,6 +38,7 @@ type ValidateHandshake<
   metadata: MessageShape<Schema>,
   previousParsedMetadata?: ParsedMetadata,
   from?: TransportClientId,
+  connectionExtras?: ConnectionExtras,
 ) =>
   | ParsedMetadata
   | ProtobufHandshakeFailureCode
@@ -94,15 +96,20 @@ export function createServerHandshakeOptions<
 > {
   return createTransportServerHandshakeOptions(
     HandshakeBytesSchema,
-    async (metadata, previousParsedMetadata, from) => {
+    async (metadata, previousParsedMetadata, from, connectionExtras) => {
       let decoded;
       try {
         decoded = decodeMessageBytes(schema, metadata);
       } catch {
         return 'REJECTED_BY_CUSTOM_HANDLER' as ProtobufHandshakeFailureCode;
       }
 
-      return await validate(decoded, previousParsedMetadata, from);
+      return await validate(
+        decoded,
+        previousParsedMetadata,
+        from,
+        connectionExtras,
+      );
     },
     expiry,
     rejectionCodeSchema,

router/handshake.ts

@@ -5,6 +5,7 @@ import {
   HandshakeErrorCustomHandlerFatalResponseCodes,
   type TransportClientId,
 } from '../transport/message';
+import type { ConnectionExtras } from '../transport/connection';
 
 type ConstructHandshake<T extends TSchema> = () =>
   | Static<T>
@@ -18,6 +19,7 @@ type ValidateHandshake<
   metadata: Static<T>,
   previousParsedMetadata?: ParsedMetadata,
   from?: TransportClientId,
+  connectionExtras?: ConnectionExtras,
 ) =>
   | Static<typeof HandshakeErrorCustomHandlerFatalResponseCodes>
   | CustomHandshakeErrorCode<RejectionCodeSchema>
@@ -94,6 +96,8 @@ export interface ServerHandshakeOptions<
    * @param from - The client id the peer presented in its handshake. Use it to
    *   confirm the presented id is the one the metadata authorizes before
    *   returning parsed metadata.
+   * @param connectionExtras - Context attached to the current transport
+   *   connection, if any.
    */
   validate: ValidateHandshake<
     MetadataSchema,

transport/connection.ts

@@ -2,6 +2,8 @@ import { TelemetryInfo } from '../tracing';
 import { MessageMetadata } from '../logging';
 import { generateId } from './id';
 
+export type ConnectionExtras = Record<string, unknown>;
+
 /**
  * A connection is the actual raw underlying transport connection.
  * It's responsible for dispatching to/from the actual connection itself
@@ -11,9 +13,11 @@ import { generateId } from './id';
 export abstract class Connection {
   id: string;
   telemetry?: TelemetryInfo;
+  extras?: ConnectionExtras;
 
-  constructor() {
+  constructor(extras?: ConnectionExtras) {
     this.id = `conn-${generateId()}`; // for debugging, no collision safety needed
+    this.extras = extras;
   }
 
   get loggingMetadata(): MessageMetadata {

transport/impls/ws/connection.ts

@@ -1,10 +1,6 @@
-import { Connection } from '../../connection';
+import { Connection, type ConnectionExtras } from '../../connection';
 import { WsLike } from './wslike';
 
-interface ConnectionInfoExtras extends Record<string, unknown> {
-  headers: Record<string, string>;
-}
-
 const WS_HEALTHY_CLOSE_CODE = 1000;
 
 export class WebSocketCloseError extends Error {
@@ -20,7 +16,6 @@ export class WebSocketCloseError extends Error {
 
 export class WebSocketConnection extends Connection {
   ws: WsLike;
-  extras?: ConnectionInfoExtras;
 
   get loggingMetadata() {
     const metadata = super.loggingMetadata;
@@ -31,10 +26,9 @@ export class WebSocketConnection extends Connection {
     return metadata;
   }
 
-  constructor(ws: WsLike, extras?: ConnectionInfoExtras) {
-    super();
+  constructor(ws: WsLike, extras?: ConnectionExtras) {
+    super(extras);
     this.ws = ws;
-    this.extras = extras;
     this.ws.binaryType = 'arraybuffer';
 
     // Websockets are kinda shitty, they emit error events with no

transport/impls/ws/server.ts

@@ -9,6 +9,7 @@ import { ServerTransport } from '../../server';
 import { ProvidedServerTransportOptions } from '../../options';
 import { type IncomingMessage } from 'http';
 import type { TSchema } from 'typebox';
+import type { ConnectionExtras } from '../../connection';
 
 function cleanHeaders(
   headers: IncomingMessage['headers'],
@@ -25,6 +26,16 @@ function cleanHeaders(
   return cleanedHeaders;
 }
 
+export type WebSocketConnectionExtrasFactory = (
+  ws: WsLike,
+  req: IncomingMessage,
+) => ConnectionExtras;
+
+const defaultConnectionExtrasFactory: WebSocketConnectionExtrasFactory = (
+  _ws,
+  req,
+) => ({ headers: cleanHeaders(req.headersDistinct) });
+
 export class WebSocketServerTransport<
   MetadataSchema extends TSchema = TSchema,
   ParsedMetadata extends object = object,
@@ -36,21 +47,25 @@ export class WebSocketServerTransport<
   RejectionCodeSchema
 > {
   wss: WebSocketServer;
+  private readonly createConnectionExtras: WebSocketConnectionExtrasFactory;
 
   constructor(
     wss: WebSocketServer,
     clientId: TransportClientId,
     providedOptions?: ProvidedServerTransportOptions,
+    createConnectionExtras: WebSocketConnectionExtrasFactory = defaultConnectionExtrasFactory,
   ) {
     super(clientId, providedOptions);
     this.wss = wss;
+    this.createConnectionExtras = createConnectionExtras;
     this.wss.on('connection', this.connectionHandler);
   }
 
   connectionHandler = (ws: WsLike, req: IncomingMessage) => {
-    const conn = new WebSocketConnection(ws, {
-      headers: cleanHeaders(req.headersDistinct),
-    });
+    const conn = new WebSocketConnection(
+      ws,
+      this.createConnectionExtras(ws, req),
+    );
 
     this.handleConnection(conn);
   };

transport/impls/ws/ws.test.ts

@@ -1,5 +1,5 @@
 import http from 'node:http';
-import { describe, test, expect, beforeEach } from 'vitest';
+import { describe, test, expect, beforeEach, vi } from 'vitest';
 import { Type } from 'typebox';
 import {
   createWebSocketServer,
@@ -8,6 +8,7 @@ import {
   createDummyTransportMessage,
   payloadToTransportMessage,
   createLocalWebSocketClient,
+  closeAllConnections,
   numberOfConnections,
   getTransportConnections,
   getClientSendFn,
@@ -16,6 +17,7 @@ import {
 import { WebSocketServerTransport } from './server';
 import { WebSocketClientTransport } from './client';
 import {
+  advanceFakeTimersByConnectionBackoff,
   advanceFakeTimersBySessionGrace,
   cleanupTransports,
   testFinishesCleanly,
@@ -96,6 +98,62 @@ describe('sending and receiving across websockets works', async () => {
     });
   });
 
+  test('handshake validation receives connection extras', async () => {
+    const schema = Type.Object({ token: Type.String() });
+    const validate = vi.fn(
+      async (
+        metadata: { token: string },
+        _previousMetadata: { token: string } | undefined,
+        _from: string | undefined,
+        _connectionExtras: Record<string, unknown> | undefined,
+      ) => ({ token: metadata.token }),
+    );
+    const createConnectionExtras = vi
+      .fn()
+      .mockReturnValueOnce({ verifiedUserId: 'user-1' })
+      .mockReturnValueOnce({ verifiedUserId: 'user-2' });
+    const serverTransport = new WebSocketServerTransport<
+      typeof schema,
+      { token: string }
+    >(wss, 'SERVER', undefined, createConnectionExtras);
+    serverTransport.extendHandshake({ schema, validate });
+    const clientTransport = new WebSocketClientTransport(
+      () => Promise.resolve(createLocalWebSocketClient(port)),
+      'client',
+    );
+    clientTransport.extendHandshake({
+      schema,
+      construct: async () => ({ token: 'token' }),
+    });
+
+    clientTransport.connect(serverTransport.clientId);
+    addPostTestCleanup(async () => {
+      await cleanupTransports([clientTransport, serverTransport]);
+    });
+
+    await waitFor(() => expect(validate).toHaveBeenCalledTimes(1));
+    expect(serverTransport.requestRehandshake(clientTransport.clientId)).toBe(
+      true,
+    );
+    await waitFor(() => expect(validate).toHaveBeenCalledTimes(2));
+
+    closeAllConnections(clientTransport);
+    await waitFor(() => expect(numberOfConnections(clientTransport)).toBe(0));
+    await waitFor(() => expect(numberOfConnections(serverTransport)).toBe(0));
+    await advanceFakeTimersByConnectionBackoff();
+    await waitFor(() => expect(validate).toHaveBeenCalledTimes(3));
+    expect(validate.mock.calls.map((call) => call[3])).toEqual([
+      { verifiedUserId: 'user-1' },
+      { verifiedUserId: 'user-1' },
+      { verifiedUserId: 'user-2' },
+    ]);
+
+    await testFinishesCleanly({
+      clientTransports: [clientTransport],
+      serverTransport,
+    });
+  });
+
   test('sending respects to/from fields', async () => {
     const makeDummyMessage = (message: string): PartialTransportMessage => {
       return payloadToTransportMessage({ message });

transport/server.ts

@@ -207,13 +207,18 @@ export abstract class ServerTransport<
     }
 
     const previousParsedMetadata = this.sessionHandshakeMetadata.get(from);
+    const connectionExtras =
+      session.state === SessionState.Connected
+        ? session.conn.extras
+        : undefined;
 
     let parsedMetadataOrFailureCode;
     try {
       parsedMetadataOrFailureCode = await handshakeExtensions.validate(
         metadata,
         previousParsedMetadata,
         from,
+        connectionExtras,
       );
     } catch (err) {
       // teardownForFailedRehandshake no-ops if this session was already replaced
@@ -494,6 +499,7 @@ export abstract class ServerTransport<
           msg.payload.metadata,
           previousParsedMetadata,
           msg.from,
+          session.conn.extras,
         );
       } catch (err) {
         this.rejectHandshakeRequest(

transport/transport.test.ts

@@ -1487,7 +1487,7 @@ describe.each(testMatrix())(
       };
     });
 
-    test('handshakes and stores parsed metadata in session', async () => {
+    test('handshakes with a three-argument validator and stores metadata', async () => {
       const schema = Type.Object({
         kept: Type.String(),
         discarded: Type.String(),
@@ -1498,9 +1498,15 @@ describe.each(testMatrix())(
       }
 
       const get = vi.fn(async () => ({ kept: 'kept', discarded: 'discarded' }));
-      const parse = vi.fn(async (metadata: Metadata) => ({
-        kept: metadata.kept,
-      }));
+      const parse = vi.fn(
+        async (
+          metadata: Metadata,
+          _previousMetadata?: Metadata,
+          _from?: string,
+        ) => ({
+          kept: metadata.kept,
+        }),
+      );
 
       const serverTransport = getServerTransport('SERVER', {
         schema,
@@ -1842,14 +1848,14 @@ describe.each(testMatrix())(
       await waitFor(() => expect(serverTransport.sessions.size).toBe(1));
       expect(construct).toHaveBeenCalledTimes(1);
       expect(validate).toHaveBeenCalledTimes(1);
-      expect(validate).toHaveBeenCalledWith(
+      expect(validate.mock.calls[0]?.slice(0, 3)).toEqual([
         {
           kept: 'kept',
           discarded: 'discarded',
         },
         undefined,
         clientTransport.clientId,
-      );
+      ]);
 
       const session = serverTransport.sessions.get(clientTransport.clientId);
       assert(session);
@@ -1870,7 +1876,7 @@ describe.each(testMatrix())(
       await waitFor(() => expect(numberOfConnections(serverTransport)).toBe(1));
 
       expect(validate).toHaveBeenCalledTimes(2);
-      expect(validate).toHaveBeenCalledWith(
+      expect(validate.mock.calls[1]?.slice(0, 3)).toEqual([
         {
           kept: 'kept',
           discarded: 'discarded',
@@ -1879,7 +1885,7 @@ describe.each(testMatrix())(
           kept: 'kept',
         },
         clientTransport.clientId,
-      );
+      ]);
 
       await testFinishesCleanly({
         clientTransports: [clientTransport],