From aa27295a74077897ef59116699f51d2b5c7195be Mon Sep 17 00:00:00 2001 From: Armaan Aggarwal Date: Wed, 2 Sep 2026 00:36:15 -0700 Subject: [PATCH 1/3] Expose WebSocket upgrade context to handshake validation --- protobuf/handshake.ts | 11 ++++-- router/handshake.ts | 4 +++ transport/connection.ts | 6 +++- transport/impls/ws/connection.ts | 12 ++----- transport/impls/ws/server.ts | 22 ++++++++++-- transport/impls/ws/ws.test.ts | 62 +++++++++++++++++++++++++++++++- transport/server.ts | 2 ++ transport/transport.test.ts | 22 +++++++----- 8 files changed, 117 insertions(+), 24 deletions(-) diff --git a/protobuf/handshake.ts b/protobuf/handshake.ts index c87d2d39..2e8e7bba 100644 --- a/protobuf/handshake.ts +++ b/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, previousParsedMetadata?: ParsedMetadata, from?: TransportClientId, + connectionExtras?: ConnectionExtras, ) => | ParsedMetadata | ProtobufHandshakeFailureCode @@ -94,7 +96,7 @@ export function createServerHandshakeOptions< > { return createTransportServerHandshakeOptions( HandshakeBytesSchema, - async (metadata, previousParsedMetadata, from) => { + async (metadata, previousParsedMetadata, from, connectionExtras) => { let decoded; try { decoded = decodeMessageBytes(schema, metadata); @@ -102,7 +104,12 @@ export function createServerHandshakeOptions< return 'REJECTED_BY_CUSTOM_HANDLER' as ProtobufHandshakeFailureCode; } - return await validate(decoded, previousParsedMetadata, from); + return await validate( + decoded, + previousParsedMetadata, + from, + connectionExtras, + ); }, expiry, rejectionCodeSchema, diff --git a/router/handshake.ts b/router/handshake.ts index 1b41a0e0..1e1d6b0c 100644 --- a/router/handshake.ts +++ b/router/handshake.ts @@ -5,6 +5,7 @@ import { HandshakeErrorCustomHandlerFatalResponseCodes, type TransportClientId, } from '../transport/message'; +import type { ConnectionExtras } from '../transport/connection'; type ConstructHandshake = () => | Static @@ -18,6 +19,7 @@ type ValidateHandshake< metadata: Static, previousParsedMetadata?: ParsedMetadata, from?: TransportClientId, + connectionExtras?: ConnectionExtras, ) => | Static | CustomHandshakeErrorCode @@ -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, diff --git a/transport/connection.ts b/transport/connection.ts index 14a9f567..007faa19 100644 --- a/transport/connection.ts +++ b/transport/connection.ts @@ -2,6 +2,8 @@ import { TelemetryInfo } from '../tracing'; import { MessageMetadata } from '../logging'; import { generateId } from './id'; +export type ConnectionExtras = Record; + /** * 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 { diff --git a/transport/impls/ws/connection.ts b/transport/impls/ws/connection.ts index bee2ce39..43b4ac62 100644 --- a/transport/impls/ws/connection.ts +++ b/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 { - headers: Record; -} - 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 diff --git a/transport/impls/ws/server.ts b/transport/impls/ws/server.ts index e1df56eb..2dc7cceb 100644 --- a/transport/impls/ws/server.ts +++ b/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,26 @@ 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); }; diff --git a/transport/impls/ws/ws.test.ts b/transport/impls/ws/ws.test.ts index f1f4c3f2..e56eb043 100644 --- a/transport/impls/ws/ws.test.ts +++ b/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,64 @@ 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 | undefined, + ) => ({ token: metadata.token }), + ); + const createConnectionExtras = vi + .fn() + .mockReturnValueOnce({ verifiedUserId: 'user-1' }) + .mockReturnValueOnce({ verifiedUserId: 'user-2' }); + const serverTransport = new WebSocketServerTransport( + 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 }); diff --git a/transport/server.ts b/transport/server.ts index 6f69fdc2..7c12bf8b 100644 --- a/transport/server.ts +++ b/transport/server.ts @@ -214,6 +214,7 @@ export abstract class ServerTransport< metadata, previousParsedMetadata, from, + session.conn.extras, ); } catch (err) { // teardownForFailedRehandshake no-ops if this session was already replaced @@ -494,6 +495,7 @@ export abstract class ServerTransport< msg.payload.metadata, previousParsedMetadata, msg.from, + session.conn.extras, ); } catch (err) { this.rejectHandshakeRequest( diff --git a/transport/transport.test.ts b/transport/transport.test.ts index dcecc26e..1329f9de 100644 --- a/transport/transport.test.ts +++ b/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], From 0cf9e3ca6a97e69fa9dc8db690b645e72b9f2d81 Mon Sep 17 00:00:00 2001 From: Armaan Aggarwal Date: Wed, 2 Sep 2026 00:45:13 -0700 Subject: [PATCH 2/3] Fix rehandshake extras typing --- transport/impls/ws/server.ts | 3 +-- transport/impls/ws/ws.test.ts | 10 ++++------ transport/server.ts | 6 +++++- 3 files changed, 10 insertions(+), 9 deletions(-) diff --git a/transport/impls/ws/server.ts b/transport/impls/ws/server.ts index 2dc7cceb..9cfd2a59 100644 --- a/transport/impls/ws/server.ts +++ b/transport/impls/ws/server.ts @@ -53,8 +53,7 @@ export class WebSocketServerTransport< wss: WebSocketServer, clientId: TransportClientId, providedOptions?: ProvidedServerTransportOptions, - createConnectionExtras: WebSocketConnectionExtrasFactory = - defaultConnectionExtrasFactory, + createConnectionExtras: WebSocketConnectionExtrasFactory = defaultConnectionExtrasFactory, ) { super(clientId, providedOptions); this.wss = wss; diff --git a/transport/impls/ws/ws.test.ts b/transport/impls/ws/ws.test.ts index e56eb043..0377bac6 100644 --- a/transport/impls/ws/ws.test.ts +++ b/transport/impls/ws/ws.test.ts @@ -112,12 +112,10 @@ describe('sending and receiving across websockets works', async () => { .fn() .mockReturnValueOnce({ verifiedUserId: 'user-1' }) .mockReturnValueOnce({ verifiedUserId: 'user-2' }); - const serverTransport = new WebSocketServerTransport( - wss, - 'SERVER', - undefined, - createConnectionExtras, - ); + const serverTransport = new WebSocketServerTransport< + typeof schema, + { token: string } + >(wss, 'SERVER', undefined, createConnectionExtras); serverTransport.extendHandshake({ schema, validate }); const clientTransport = new WebSocketClientTransport( () => Promise.resolve(createLocalWebSocketClient(port)), diff --git a/transport/server.ts b/transport/server.ts index 7c12bf8b..f25d8d2d 100644 --- a/transport/server.ts +++ b/transport/server.ts @@ -207,6 +207,10 @@ export abstract class ServerTransport< } const previousParsedMetadata = this.sessionHandshakeMetadata.get(from); + const connectionExtras = + session.state === SessionState.Connected + ? session.conn.extras + : undefined; let parsedMetadataOrFailureCode; try { @@ -214,7 +218,7 @@ export abstract class ServerTransport< metadata, previousParsedMetadata, from, - session.conn.extras, + connectionExtras, ); } catch (err) { // teardownForFailedRehandshake no-ops if this session was already replaced From 7c7322c454ad527b48e750a964ac87d651e3978e Mon Sep 17 00:00:00 2001 From: Armaan Aggarwal Date: Wed, 2 Sep 2026 01:45:09 -0700 Subject: [PATCH 3/3] Document WebSocket handshake extras --- README.md | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index e185fd34..2d84c5b5 100644 --- a/README.md +++ b/README.md @@ -799,12 +799,13 @@ createClient(clientTransport, 'SERVER', { createServer(serverTransport, services, { handshakeOptions: createServerHandshakeOptions( handshakeSchema, - (metadata, previousMetadata, from) => { + (metadata, previousMetadata, from, connectionExtras) => { // the type of this function is // ( // metadata: Static, // previousMetadata?: ParsedMetadata, // from?: TransportClientId, + // connectionExtras?: Record, // ) => // | '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