From c7e60f81e5d79ec404d0e6e386a5e8279bced25e Mon Sep 17 00:00:00 2001 From: Mayank Mehra Date: Mon, 24 Aug 2026 13:28:34 -0700 Subject: [PATCH 01/12] Explore application-typed handshake rejection codes --- protobuf/handshake.ts | 31 +++++-- router/client.ts | 19 ++-- router/handshake.ts | 55 ++++++++++-- router/server.ts | 31 +++++-- testUtil/fixtures/mockTransport.ts | 35 ++++++-- testUtil/fixtures/transports.ts | 57 ++++++++---- transport/client.ts | 33 +++++-- transport/events.ts | 38 +++++--- transport/impls/ws/client.ts | 4 +- transport/impls/ws/server.ts | 8 +- transport/index.ts | 1 + transport/message.test.ts | 46 ++++++++++ transport/message.ts | 54 ++++++++++- transport/server.ts | 45 +++++++--- transport/transport.test.ts | 140 +++++++++++++++++++++++++++++ transport/transport.ts | 27 +++--- 16 files changed, 530 insertions(+), 94 deletions(-) diff --git a/protobuf/handshake.ts b/protobuf/handshake.ts index f441ee4c..bd82f57d 100644 --- a/protobuf/handshake.ts +++ b/protobuf/handshake.ts @@ -27,23 +27,34 @@ type ConstructHandshake = () => | MessageInitShape | Promise>; -type ValidateHandshake = ( +type ValidateHandshake< + Schema extends DescMessage, + ParsedMetadata, + ApplicationErrorCode extends string = never, +> = ( metadata: MessageShape, previousParsedMetadata?: ParsedMetadata, from?: TransportClientId, ) => | ParsedMetadata | ProtobufHandshakeFailureCode - | Promise; + | ApplicationErrorCode + | Promise< + ParsedMetadata | ProtobufHandshakeFailureCode | ApplicationErrorCode + >; /** * Create client-side handshake options backed by a protobuf message type. */ -export function createClientHandshakeOptions( +export function createClientHandshakeOptions< + Schema extends DescMessage, + ApplicationErrorCode extends string = never, +>( schema: Schema, construct: ConstructHandshake, eager?: boolean, -): ClientHandshakeOptions { + rejectionCodes?: ReadonlyArray, +): ClientHandshakeOptions { return createTransportClientHandshakeOptions( HandshakeBytesSchema, async () => { @@ -52,6 +63,7 @@ export function createClientHandshakeOptions( return encodeMessageBytes(schema, metadata); }, eager, + rejectionCodes, ); } @@ -61,11 +73,17 @@ export function createClientHandshakeOptions( export function createServerHandshakeOptions< Schema extends DescMessage, ParsedMetadata extends object = object, + ApplicationErrorCode extends string = never, >( schema: Schema, - validate: ValidateHandshake, + validate: ValidateHandshake, expiry?: (parsedMetadata: ParsedMetadata) => Date | undefined, -): ServerHandshakeOptions { + rejectionCodes?: ReadonlyArray, +): ServerHandshakeOptions< + typeof HandshakeBytesSchema, + ParsedMetadata, + ApplicationErrorCode +> { return createTransportServerHandshakeOptions( HandshakeBytesSchema, async (metadata, previousParsedMetadata, from) => { @@ -79,5 +97,6 @@ export function createServerHandshakeOptions< return await validate(decoded, previousParsedMetadata, from); }, expiry, + rejectionCodes, ); } diff --git a/router/client.ts b/router/client.ts index 81d2bb7b..508738ef 100644 --- a/router/client.ts +++ b/router/client.ts @@ -18,7 +18,7 @@ import { closeStreamMessage, cancelMessage, } from '../transport/message'; -import type { Static } from 'typebox'; +import type { Static, TSchema } from 'typebox'; import { Err, Result, AnyResultSchema } from './result'; import { EventMap } from '../transport/events'; import { Connection } from '../transport/connection'; @@ -241,12 +241,16 @@ const defaultClientOptions: ClientOptions = { // We are using any here because the ServiceContext is a server-side implementation // detail that doesn't affect the client interface // eslint-disable-next-line @typescript-eslint/no-explicit-any -export function createClient>( - transport: ClientTransport, +export function createClient< + // eslint-disable-next-line @typescript-eslint/no-explicit-any + ServiceSchemaMap extends AnyServiceSchemaMap, + ApplicationErrorCode extends string = never, +>( + transport: ClientTransport, serverId: TransportClientId, providedClientOptions: Partial< ClientOptions & { - handshakeOptions: ClientHandshakeOptions; + handshakeOptions: ClientHandshakeOptions; } > = {}, ): Client { @@ -305,7 +309,8 @@ function mergeCallOptions( defaults: ClientOptions['defaultCallOptions'], caller: CallOptions | undefined, ): CallOptions { - const resolved = typeof defaults === 'function' ? defaults() : defaults ?? {}; + const resolved = + typeof defaults === 'function' ? defaults() : (defaults ?? {}); // Caller fields win: spread defaults first, caller second. return { ...resolved, ...caller }; @@ -317,9 +322,9 @@ type AnyProcReturn = | ReturnType> | ReturnType>; -function handleProc( +function handleProc( procType: ValidProcType, - transport: ClientTransport, + transport: ClientTransport, serverId: TransportClientId, init: Static, serviceName: string, diff --git a/router/handshake.ts b/router/handshake.ts index 427b86b6..d1b10bb8 100644 --- a/router/handshake.ts +++ b/router/handshake.ts @@ -8,20 +8,27 @@ type ConstructHandshake = () => | Static | Promise>; -type ValidateHandshake = ( +type ValidateHandshake< + T extends TSchema, + ParsedMetadata, + ApplicationErrorCode extends string = never, +> = ( metadata: Static, previousParsedMetadata?: ParsedMetadata, from?: TransportClientId, ) => | Static + | ApplicationErrorCode | ParsedMetadata | Promise< | Static + | ApplicationErrorCode | ParsedMetadata >; export interface ClientHandshakeOptions< MetadataSchema extends TSchema = TSchema, + ApplicationErrorCode extends string = never, > { /** * Schema for the metadata that the client sends to the server @@ -29,6 +36,14 @@ export interface ClientHandshakeOptions< */ schema: MetadataSchema; + /** + * Application-defined rejection codes the server may answer the handshake + * with, sent in the response's `code` field. Must match the server's + * {@link ServerHandshakeOptions.rejectionCodes}: an unconfigured code is + * rejected as a malformed handshake response. + */ + rejectionCodes?: ReadonlyArray; + /** * Gets the {@link HandshakeRequestMetadata} to send to the server. */ @@ -48,6 +63,7 @@ export interface ClientHandshakeOptions< export interface ServerHandshakeOptions< MetadataSchema extends TSchema = TSchema, ParsedMetadata extends object = object, + ApplicationErrorCode extends string = never, > { /** * Schema for the metadata that the server receives from the client @@ -55,6 +71,15 @@ export interface ServerHandshakeOptions< */ schema: MetadataSchema; + /** + * Application-defined rejection codes that {@link validate} may return. + * They travel in the handshake response's `code` field and are fatal like + * the built-in custom-handler codes. Clients must register the same codes + * in {@link ClientHandshakeOptions.rejectionCodes} or they reject the + * response as malformed. + */ + rejectionCodes?: ReadonlyArray; + /** * Parses the metadata sent by the client during the handshake into the * server-side {@link ParsedMetadata}, or returns a handshake failure code to @@ -67,7 +92,11 @@ export interface ServerHandshakeOptions< * confirm the presented id is the one the metadata authorizes before * returning parsed metadata. */ - validate: ValidateHandshake; + validate: ValidateHandshake< + MetadataSchema, + ParsedMetadata, + ApplicationErrorCode + >; /** * When the credential expires (or undefined if it never does). The server @@ -84,21 +113,33 @@ export interface ServerHandshakeOptions< export function createClientHandshakeOptions< MetadataSchema extends TSchema = TSchema, + ApplicationErrorCode extends string = never, >( schema: MetadataSchema, construct: ConstructHandshake, eager?: boolean, -): ClientHandshakeOptions { - return { schema, construct, eager }; + rejectionCodes?: ReadonlyArray, +): ClientHandshakeOptions { + return { schema, construct, eager, rejectionCodes }; } export function createServerHandshakeOptions< MetadataSchema extends TSchema = TSchema, ParsedMetadata extends object = object, + ApplicationErrorCode extends string = never, >( schema: MetadataSchema, - validate: ValidateHandshake, + validate: ValidateHandshake< + MetadataSchema, + ParsedMetadata, + ApplicationErrorCode + >, expiry?: (parsedMetadata: ParsedMetadata) => Date | undefined, -): ServerHandshakeOptions { - return { schema, validate, expiry }; + rejectionCodes?: ReadonlyArray, +): ServerHandshakeOptions< + MetadataSchema, + ParsedMetadata, + ApplicationErrorCode +> { + return { schema, validate, expiry, rejectionCodes }; } diff --git a/router/server.ts b/router/server.ts index e4548720..9f03e7db 100644 --- a/router/server.ts +++ b/router/server.ts @@ -112,12 +112,14 @@ class RiverServer< MetadataSchema extends TSchema, ParsedMetadata extends object, Services extends AnyServiceSchemaMap, + ApplicationErrorCode extends string = never, > implements Server { private transport: ServerTransport< Connection, MetadataSchema, - ParsedMetadata + ParsedMetadata, + ApplicationErrorCode >; private contextMap: Map; @@ -145,9 +147,18 @@ class RiverServer< private unregisterTransportListeners: () => void; constructor( - transport: ServerTransport, + transport: ServerTransport< + Connection, + MetadataSchema, + ParsedMetadata, + ApplicationErrorCode + >, services: Services, - handshakeOptions?: ServerHandshakeOptions, + handshakeOptions?: ServerHandshakeOptions< + MetadataSchema, + ParsedMetadata, + ApplicationErrorCode + >, extendedContext?: Context, maxCancelledStreamTombstonesPerSession = 200, middlewares: Array = [], @@ -1167,11 +1178,21 @@ export function createServer< // eslint-disable-next-line @typescript-eslint/no-explicit-any Services extends AnyServiceSchemaMap, Context extends MaybeDisposable = MaybeDisposable, + ApplicationErrorCode extends string = never, >( - transport: ServerTransport, + transport: ServerTransport< + Connection, + MetadataSchema, + ParsedMetadata, + ApplicationErrorCode + >, services: Services, providedServerOptions?: Partial<{ - handshakeOptions?: ServerHandshakeOptions; + handshakeOptions?: ServerHandshakeOptions< + MetadataSchema, + ParsedMetadata, + ApplicationErrorCode + >; extendedContext?: Context; /** * Maximum number of cancelled streams to keep track of to avoid diff --git a/testUtil/fixtures/mockTransport.ts b/testUtil/fixtures/mockTransport.ts index a4d25177..60bd1242 100644 --- a/testUtil/fixtures/mockTransport.ts +++ b/testUtil/fixtures/mockTransport.ts @@ -9,7 +9,10 @@ import { Duplex } from 'node:stream'; import { duplexPair } from '../duplex/duplexPair'; import { nanoid } from 'nanoid'; import type { TSchema } from 'typebox'; -import { ServerHandshakeOptions } from '../../router/handshake'; +import { + ClientHandshakeOptions, + ServerHandshakeOptions, +} from '../../router/handshake'; export class InMemoryConnection extends Connection { conn: Duplex; @@ -71,8 +74,10 @@ export function createMockTransportNetwork( // conn id -> [client->server, server->client] const connections = new Observable>({}); - const transports: Array> = []; - class MockClientTransport extends ClientTransport { + const transports: Array> = []; + class MockClientTransport< + ApplicationErrorCode extends string = never, + > extends ClientTransport { async createNewOutgoingConnection( to: TransportClientId, ): Promise { @@ -99,10 +104,12 @@ export function createMockTransportNetwork( class MockServerTransport< MetadataSchema extends TSchema = TSchema, ParsedMetadata extends object = object, + ApplicationErrorCode extends string = never, > extends ServerTransport< InMemoryConnection, MetadataSchema, - ParsedMetadata + ParsedMetadata, + ApplicationErrorCode > { subscribeCleanup: () => void; @@ -135,8 +142,14 @@ export function createMockTransportNetwork( } return { - getClientTransport: (id, handshakeOptions) => { - const clientTransport = new MockClientTransport(id, opts?.client); + getClientTransport: ( + id: TransportClientId, + handshakeOptions?: ClientHandshakeOptions, + ) => { + const clientTransport = new MockClientTransport( + id, + opts?.client, + ); if (handshakeOptions) { clientTransport.extendHandshake(handshakeOptions); } @@ -148,15 +161,21 @@ export function createMockTransportNetwork( getServerTransport: < MetadataSchema extends TSchema = TSchema, ParsedMetadata extends object = object, + ApplicationErrorCode extends string = never, >( id = 'SERVER', handshakeOptions: - | ServerHandshakeOptions + | ServerHandshakeOptions< + MetadataSchema, + ParsedMetadata, + ApplicationErrorCode + > | undefined, ) => { const serverTransport = new MockServerTransport< MetadataSchema, - ParsedMetadata + ParsedMetadata, + ApplicationErrorCode >(id, opts?.server); if (handshakeOptions) { serverTransport.extendHandshake(handshakeOptions); diff --git a/testUtil/fixtures/transports.ts b/testUtil/fixtures/transports.ts index 75822f68..0d339d67 100644 --- a/testUtil/fixtures/transports.ts +++ b/testUtil/fixtures/transports.ts @@ -30,17 +30,27 @@ export interface TestTransportOptions { } export interface TestSetupHelpers { - getClientTransport: ( + getClientTransport: ( id: TransportClientId, - handshakeOptions?: ClientHandshakeOptions, - ) => ClientTransport; + handshakeOptions?: ClientHandshakeOptions, + ) => ClientTransport; getServerTransport: < MetadataSchema extends TSchema = TSchema, ParsedMetadata extends object = object, + ApplicationErrorCode extends string = never, >( id?: TransportClientId, - handshakeOptions?: ServerHandshakeOptions, - ) => ServerTransport; + handshakeOptions?: ServerHandshakeOptions< + MetadataSchema, + ParsedMetadata, + ApplicationErrorCode + >, + ) => ServerTransport< + Connection, + MetadataSchema, + ParsedMetadata, + ApplicationErrorCode + >; simulatePhantomDisconnect: () => void; restartServer: () => Promise; cleanup: () => Promise | void; @@ -59,10 +69,11 @@ export const transports: Array = [ const port = await onWsServerReady(server); let wss = createWebSocketServer(server); + /* eslint-disable @typescript-eslint/no-explicit-any */ const transports: Array< - // eslint-disable-next-line @typescript-eslint/no-explicit-any - WebSocketClientTransport | WebSocketServerTransport + WebSocketClientTransport | WebSocketServerTransport > = []; + /* eslint-enable @typescript-eslint/no-explicit-any */ return { simulatePhantomDisconnect() { @@ -72,12 +83,19 @@ export const transports: Array = [ } } }, - getClientTransport: (id, handshakeOptions) => { - const clientTransport = new WebSocketClientTransport( - () => Promise.resolve(createLocalWebSocketClient(port)), - id, - opts?.client, - ); + getClientTransport: ( + id: TransportClientId, + handshakeOptions?: ClientHandshakeOptions< + TSchema, + ApplicationErrorCode + >, + ) => { + const clientTransport = + new WebSocketClientTransport( + () => Promise.resolve(createLocalWebSocketClient(port)), + id, + opts?.client, + ); if (handshakeOptions) { clientTransport.extendHandshake(handshakeOptions); @@ -99,15 +117,21 @@ export const transports: Array = [ getServerTransport: < MetadataSchema extends TSchema, ParsedMetadata extends object, + ApplicationErrorCode extends string = never, >( id = 'SERVER', handshakeOptions: - | ServerHandshakeOptions + | ServerHandshakeOptions< + MetadataSchema, + ParsedMetadata, + ApplicationErrorCode + > | undefined, ) => { const serverTransport = new WebSocketServerTransport< MetadataSchema, - ParsedMetadata + ParsedMetadata, + ApplicationErrorCode >(wss, id, opts?.server); serverTransport.bindLogger((msg, ctx, level) => { @@ -128,7 +152,8 @@ export const transports: Array = [ return serverTransport as ServerTransport< Connection, MetadataSchema, - ParsedMetadata + ParsedMetadata, + ApplicationErrorCode >; }, async restartServer() { diff --git a/transport/client.ts b/transport/client.ts index 4f9e440a..c6597fa0 100644 --- a/transport/client.ts +++ b/transport/client.ts @@ -3,7 +3,9 @@ import { ClientHandshakeOptions } from '../router/handshake'; import { validationErrorToRiverErrors } from '../router/errors'; import { ControlMessageHandshakeResponseSchema, + ControlMessageHandshakeResponseSchemaWithCodes, ControlMessageRehandshakeRequestSchema, + type HandshakeErrorCode, HandshakeErrorRetriableResponseCodes, OpaqueTransportMessage, TransportClientId, @@ -11,6 +13,8 @@ import { handshakeRequestMessage, rehandshakeResponseMessage, } from './message'; + +import type { TSchema } from 'typebox'; import { ClientTransportOptions, ProvidedClientTransportOptions, @@ -50,7 +54,8 @@ type ConstructedHandshakeMetadata = export abstract class ClientTransport< ConnType extends Connection, -> extends Transport { + ApplicationErrorCode extends string = never, +> extends Transport { /** * The options for this transport. */ @@ -69,7 +74,16 @@ export abstract class ClientTransport< /** * Optional handshake options for this client. */ - handshakeExtensions?: ClientHandshakeOptions; + handshakeExtensions?: ClientHandshakeOptions; + + /** + * Handshake response schema extended with the application-defined + * rejection codes, when any are registered. + */ + protected handshakeResponseSchema: + | typeof ControlMessageHandshakeResponseSchema + | ReturnType = + ControlMessageHandshakeResponseSchema; /** * Handshake-metadata constructions prefetched when a connection attempt begins @@ -98,8 +112,14 @@ export abstract class ClientTransport< this.retryBudget = new LeakyBucketRateLimit(this.options); } - extendHandshake(options: ClientHandshakeOptions) { + extendHandshake( + options: ClientHandshakeOptions, + ) { this.handshakeExtensions = options; + if (options.rejectionCodes?.length) { + this.handshakeResponseSchema = + ControlMessageHandshakeResponseSchemaWithCodes(options.rejectionCodes); + } } protected handleRehandshakeMessage(message: OpaqueTransportMessage): void { @@ -355,13 +375,13 @@ export abstract class ClientTransport< msg: OpaqueTransportMessage, ) { // invariant: msg is a handshake response - if (!Value.Check(ControlMessageHandshakeResponseSchema, msg.payload)) { + if (!Value.Check(this.handshakeResponseSchema, msg.payload)) { const reason = `received invalid handshake response`; this.rejectHandshakeResponse(session, reason, { ...session.loggingMetadata, transportMessage: msg, validationErrors: Value.Errors( - ControlMessageHandshakeResponseSchema, + this.handshakeResponseSchema, msg.payload, ).flatMap(validationErrorToRiverErrors), }); @@ -388,7 +408,8 @@ export abstract class ClientTransport< } else { this.protocolError({ type: ProtocolError.HandshakeFailed, - code: msg.payload.status.code, + code: msg.payload.status + .code as HandshakeErrorCode, message: reason, }); } diff --git a/transport/events.ts b/transport/events.ts index 5da1f81e..63a6e263 100644 --- a/transport/events.ts +++ b/transport/events.ts @@ -1,6 +1,5 @@ -import type { Static } from 'typebox'; import { Connection } from './connection'; -import { OpaqueTransportMessage, HandshakeErrorResponseCodes } from './message'; +import { OpaqueTransportMessage, HandshakeErrorCode } from './message'; import { Session, SessionState } from './sessionStateMachine'; import { SessionId } from './sessionStateMachine/common'; import { TransportStatus } from './transport'; @@ -16,7 +15,7 @@ export const ProtocolError = { export type ProtocolErrorType = (typeof ProtocolError)[keyof typeof ProtocolError]; -export interface EventMap { +export interface EventMap { message: OpaqueTransportMessage; sessionStatus: | { @@ -36,7 +35,7 @@ export interface EventMap { protocolError: | { type: (typeof ProtocolError)['HandshakeFailed']; - code: Static; + code: HandshakeErrorCode; message: string; } | { @@ -52,12 +51,18 @@ export interface EventMap { } export type EventTypes = keyof EventMap; -export type EventHandler = ( - event: EventMap[K], -) => unknown; +export type EventHandler< + K extends EventTypes, + ApplicationErrorCode extends string = never, +> = (event: EventMap[K]) => unknown; -export class EventDispatcher { - private eventListeners: { [K in T]?: Set> } = {}; +export class EventDispatcher< + T extends EventTypes, + ApplicationErrorCode extends string = never, +> { + private eventListeners: { + [K in T]?: Set>; + } = {}; removeAllListeners() { this.eventListeners = {}; @@ -67,7 +72,10 @@ export class EventDispatcher { return this.eventListeners[eventType]?.size ?? 0; } - addEventListener(eventType: K, handler: EventHandler) { + addEventListener( + eventType: K, + handler: EventHandler, + ) { if (!this.eventListeners[eventType]) { this.eventListeners[eventType] = new Set(); } @@ -75,14 +83,20 @@ export class EventDispatcher { this.eventListeners[eventType]?.add(handler); } - removeEventListener(eventType: K, handler: EventHandler) { + removeEventListener( + eventType: K, + handler: EventHandler, + ) { const handlers = this.eventListeners[eventType]; if (handlers) { this.eventListeners[eventType]?.delete(handler); } } - dispatchEvent(eventType: K, event: EventMap[K]) { + dispatchEvent( + eventType: K, + event: EventMap[K], + ) { const handlers = this.eventListeners[eventType]; if (handlers) { // copying ensures that adding more listeners in a handler doesn't diff --git a/transport/impls/ws/client.ts b/transport/impls/ws/client.ts index dea22b13..b6d716a7 100644 --- a/transport/impls/ws/client.ts +++ b/transport/impls/ws/client.ts @@ -9,7 +9,9 @@ import { WsLike } from './wslike'; * @class * @extends Transport */ -export class WebSocketClientTransport extends ClientTransport { +export class WebSocketClientTransport< + ApplicationErrorCode extends string = never, +> extends ClientTransport { /** * A function that returns a Promise that resolves to a websocket URL. */ diff --git a/transport/impls/ws/server.ts b/transport/impls/ws/server.ts index 9bcaddec..922bcddf 100644 --- a/transport/impls/ws/server.ts +++ b/transport/impls/ws/server.ts @@ -25,7 +25,13 @@ function cleanHeaders( export class WebSocketServerTransport< MetadataSchema extends TSchema = TSchema, ParsedMetadata extends object = object, -> extends ServerTransport { + ApplicationErrorCode extends string = never, +> extends ServerTransport< + WebSocketConnection, + MetadataSchema, + ParsedMetadata, + ApplicationErrorCode +> { wss: WebSocketServer; constructor( diff --git a/transport/index.ts b/transport/index.ts index 9bd911e1..c68bd4ed 100644 --- a/transport/index.ts +++ b/transport/index.ts @@ -28,6 +28,7 @@ export { isStreamClose, } from './message'; export type { + HandshakeErrorCode, TransportMessage, OpaqueTransportMessage, TransportClientId, diff --git a/transport/message.test.ts b/transport/message.test.ts index cc2426da..17747f51 100644 --- a/transport/message.test.ts +++ b/transport/message.test.ts @@ -1,6 +1,8 @@ import { TransportMessage } from '.'; import { ControlFlags, + ControlMessageHandshakeResponseSchema, + ControlMessageHandshakeResponseSchemaWithCodes, handshakeRequestMessage, handshakeResponseMessage, isAck, @@ -8,6 +10,7 @@ import { isStreamOpen, } from './message'; import { describe, test, expect } from 'vitest'; +import { Value } from 'typebox/value'; const msg = ( to: string, @@ -105,6 +108,49 @@ describe('message helpers', () => { expect(mFail.payload.status.ok).toBe(false); }); + test('handshake response schema with application codes', () => { + const extended = ControlMessageHandshakeResponseSchemaWithCodes([ + 'REPL_NOT_FOUND', + 'TOKEN_EXPIRED', + ] as const); + const rejection = { + type: 'HANDSHAKE_RESP', + status: { + ok: false, + reason: 'rejected by handshake handler', + code: 'REPL_NOT_FOUND', + }, + }; + const protocolFailure = { + type: 'HANDSHAKE_RESP', + status: { + ok: false, + reason: 'bad', + code: 'SESSION_STATE_MISMATCH', + }, + }; + const unknownCode = { + type: 'HANDSHAKE_RESP', + status: { + ok: false, + reason: 'bad', + code: 'NOT_A_REGISTERED_CODE', + }, + }; + + expect(Value.Check(extended, rejection)).toBe(true); + expect(Value.Check(extended, protocolFailure)).toBe(true); + expect(Value.Check(extended, unknownCode)).toBe(false); + + // the base schema only knows the protocol-level codes + expect(Value.Check(ControlMessageHandshakeResponseSchema, rejection)).toBe( + false, + ); + expect( + Value.Check(ControlMessageHandshakeResponseSchema, protocolFailure), + ).toBe(true); + }); + test('default message has no control flags set', () => { const m = msg('a', 'b', 'stream', { test: 1 }, 'svc', 'proc'); diff --git a/transport/message.ts b/transport/message.ts index 5dbd0cd2..67ceff8c 100644 --- a/transport/message.ts +++ b/transport/message.ts @@ -1,4 +1,4 @@ -import { Type, type TSchema, type Static } from 'typebox'; +import { Type, type TLiteral, type TSchema, type Static } from 'typebox'; import { PropagationContext } from '../tracing'; import { generateId } from './id'; // type-only: a value import closes a transport <-> router require cycle @@ -123,6 +123,14 @@ export const HandshakeErrorResponseCodes = Type.Union([ HandshakeErrorFatalResponseCodes, ]); +/** + * The protocol-level handshake error codes plus any application-defined + * rejection codes the application registered in its handshake options. + */ +export type HandshakeErrorCode = + | Static + | ApplicationErrorCode; + export const ControlMessageHandshakeResponseSchema = Type.Object({ type: Type.Literal('HANDSHAKE_RESP'), status: Type.Union([ @@ -138,6 +146,38 @@ export const ControlMessageHandshakeResponseSchema = Type.Object({ ]), }); +/** + * A handshake response schema that additionally accepts application-defined + * rejection codes. Both peers must be configured with the same codes: an + * unconfigured peer rejects an application code as a malformed response. + */ +export const ControlMessageHandshakeResponseSchemaWithCodes = < + const ApplicationErrorCodes extends readonly string[], +>( + applicationErrorCodes: ApplicationErrorCodes, +) => + Type.Object({ + type: Type.Literal('HANDSHAKE_RESP'), + status: Type.Union([ + Type.Object({ + ok: Type.Literal(true), + sessionId: Type.String(), + }), + Type.Object({ + ok: Type.Literal(false), + reason: Type.String(), + code: Type.Union([ + HandshakeErrorResponseCodes, + Type.Union( + applicationErrorCodes.map( + (code) => Type.Literal(code) as TLiteral, + ), + ), + ]), + }), + ]), + }); + /** * Reserved stream id for the follow-up handshake (re-handshake) control * messages, analogous to the reserved `heartbeat` stream id used for acks. @@ -264,7 +304,11 @@ export function handshakeResponseMessage({ }: { from: TransportClientId; to: TransportClientId; - status: Static['status']; + // the code may be an application-defined rejection code, which is only + // known to peers that registered it in their handshake options + status: + | { ok: true; sessionId: string } + | { ok: false; reason: string; code: string }; }): TransportMessage> { return { id: generateId(), @@ -276,8 +320,10 @@ export function handshakeResponseMessage({ controlFlags: 0, payload: { type: 'HANDSHAKE_RESP', - status, - } satisfies Static, + status: status as Static< + typeof ControlMessageHandshakeResponseSchema + >['status'], + }, }; } diff --git a/transport/server.ts b/transport/server.ts index c41e43ce..718dd4fb 100644 --- a/transport/server.ts +++ b/transport/server.ts @@ -5,7 +5,7 @@ import { ControlMessageHandshakeRequestSchema, ControlMessageRehandshakeResponseSchema, HandshakeErrorCustomHandlerFatalResponseCodes, - HandshakeErrorResponseCodes, + type HandshakeErrorCode, OpaqueTransportMessage, acceptedProtocolVersions, TransportClientId, @@ -20,7 +20,7 @@ import { } from './options'; import { DeleteSessionOptions, Transport } from './transport'; import { coerceErrorString } from './stringifyError'; -import type { Static, TSchema } from 'typebox'; +import type { TSchema } from 'typebox'; import { Value } from 'typebox/value'; import { ProtocolError } from './events'; import { Connection } from './connection'; @@ -36,7 +36,8 @@ export abstract class ServerTransport< ConnType extends Connection, MetadataSchema extends TSchema = TSchema, ParsedMetadata extends object = object, -> extends Transport { + ApplicationErrorCode extends string = never, +> extends Transport { /** * The options for this transport. */ @@ -45,7 +46,11 @@ export abstract class ServerTransport< /** * Optional handshake options for the server. */ - handshakeExtensions?: ServerHandshakeOptions; + handshakeExtensions?: ServerHandshakeOptions< + MetadataSchema, + ParsedMetadata, + ApplicationErrorCode + >; /** * A map of session handshake data for each session. @@ -72,11 +77,27 @@ export abstract class ServerTransport< } extendHandshake( - options: ServerHandshakeOptions, + options: ServerHandshakeOptions< + MetadataSchema, + ParsedMetadata, + ApplicationErrorCode + >, ) { this.handshakeExtensions = options; } + private isApplicationRejectionCode( + value: unknown, + ): value is ApplicationErrorCode { + return ( + typeof value === 'string' && + (this.handshakeExtensions?.rejectionCodes?.includes( + value as ApplicationErrorCode, + ) ?? + false) + ); + } + protected deletePendingSession( pendingSession: SessionWaitingForHandshake, ) { @@ -208,11 +229,13 @@ export abstract class ServerTransport< Value.Check( HandshakeErrorCustomHandlerFatalResponseCodes, parsedMetadataOrFailureCode, - ) + ) || + this.isApplicationRejectionCode(parsedMetadataOrFailureCode) ) { this.teardownForFailedRehandshake( session, 're-handshake metadata rejected by handshake handler', + parsedMetadataOrFailureCode as HandshakeErrorCode, ); return; @@ -246,6 +269,7 @@ export abstract class ServerTransport< private teardownForFailedRehandshake( session: ServerSession, reason: string, + code: HandshakeErrorCode = 'REJECTED_BY_CUSTOM_HANDLER', ) { if (session._isConsumed) { return; @@ -263,7 +287,7 @@ export abstract class ServerTransport< this.protocolError({ type: ProtocolError.HandshakeFailed, - code: 'REJECTED_BY_CUSTOM_HANDLER', + code, message: reason, }); this.deleteSession(session, { unhealthy: true }); @@ -354,7 +378,7 @@ export abstract class ServerTransport< session: SessionWaitingForHandshake, to: TransportClientId, reason: string, - code: Static, + code: HandshakeErrorCode, metadata: MessageMetadata, ) { session.conn.telemetry?.span.setStatus({ @@ -501,13 +525,14 @@ export abstract class ServerTransport< Value.Check( HandshakeErrorCustomHandlerFatalResponseCodes, parsedMetadataOrFailureCode, - ) + ) || + this.isApplicationRejectionCode(parsedMetadataOrFailureCode) ) { this.rejectHandshakeRequest( session, msg.from, 'rejected by handshake handler', - parsedMetadataOrFailureCode, + parsedMetadataOrFailureCode as HandshakeErrorCode, { ...session.loggingMetadata, connectedTo: msg.from, diff --git a/transport/transport.test.ts b/transport/transport.test.ts index ba1bc0d8..6026f238 100644 --- a/transport/transport.test.ts +++ b/transport/transport.test.ts @@ -25,6 +25,7 @@ import { Type } from 'typebox'; import { TestSetupHelpers } from '../testUtil/fixtures/transports'; import { createPostTestCleanups } from '../testUtil/fixtures/cleanup'; import { SessionState } from './sessionStateMachine'; +import { createServerHandshakeOptions } from '../router/handshake'; import { ProvidedClientTransportOptions, ProvidedTransportOptions, @@ -1945,5 +1946,144 @@ describe.each(testMatrix())( serverTransport, }); }); + + test('custom handler can reject with an application-defined code', async () => { + const schema = Type.Object({ + foo: Type.String(), + }); + + type ApplicationErrorCode = 'REPL_NOT_FOUND' | 'TOKEN_EXPIRED'; + const rejectionCodes: ReadonlyArray = [ + 'REPL_NOT_FOUND', + 'TOKEN_EXPIRED', + ]; + interface ParsedMetadata { + foo: string; + } + + // compile-time: undeclared codes cannot be returned by validate + expect( + createServerHandshakeOptions< + typeof schema, + ParsedMetadata, + ApplicationErrorCode + >( + schema, + // @ts-expect-error only declared rejection codes may be returned + async () => 'SOME_OTHER_CODE', + undefined, + rejectionCodes, + ), + ).toBeDefined(); + + const parse = vi.fn(async (): Promise => { + return 'REPL_NOT_FOUND'; + }); + const serverTransport = getServerTransport< + typeof schema, + ParsedMetadata, + ApplicationErrorCode + >('SERVER', { + schema, + validate: parse, + rejectionCodes, + }); + + const clientTransport = getClientTransport('client', { + schema, + construct: async () => ({ foo: 'foo' }), + rejectionCodes, + }); + + const clientHandshakeFailed = vi.fn(); + clientTransport.addEventListener('protocolError', clientHandshakeFailed); + const serverRejectedConnection = vi.fn(); + serverTransport.addEventListener( + 'protocolError', + serverRejectedConnection, + ); + clientTransport.connect(serverTransport.clientId); + + addPostTestCleanup(async () => { + clientTransport.removeEventListener( + 'protocolError', + clientHandshakeFailed, + ); + serverTransport.removeEventListener( + 'protocolError', + serverRejectedConnection, + ); + await cleanupTransports([clientTransport, serverTransport]); + }); + + await waitFor(() => { + expect(clientHandshakeFailed).toHaveBeenCalledTimes(1); + expect(clientHandshakeFailed).toHaveBeenCalledWith({ + type: ProtocolError.HandshakeFailed, + code: 'REPL_NOT_FOUND', + message: 'handshake failed: rejected by handshake handler', + }); + expect(serverRejectedConnection).toHaveBeenCalledWith({ + type: ProtocolError.HandshakeFailed, + code: 'REPL_NOT_FOUND', + message: 'rejected by handshake handler', + }); + }); + + await testFinishesCleanly({ + clientTransports: [clientTransport], + serverTransport, + }); + }); + + test('an application code is rejected by an unconfigured client', async () => { + const schema = Type.Object({ + foo: Type.String(), + }); + + type ApplicationErrorCode = 'REPL_NOT_FOUND'; + interface ParsedMetadata { + foo: string; + } + + const serverTransport = getServerTransport< + typeof schema, + ParsedMetadata, + ApplicationErrorCode + >('SERVER', { + schema, + validate: async () => 'REPL_NOT_FOUND', + rejectionCodes: ['REPL_NOT_FOUND'], + }); + + // the client did not register the application code: it must treat the + // response as malformed rather than accept an unknown code + const clientTransport = getClientTransport('client', { + schema, + construct: async () => ({ foo: 'foo' }), + }); + + const clientHandshakeFailed = vi.fn(); + clientTransport.addEventListener('protocolError', clientHandshakeFailed); + clientTransport.connect(serverTransport.clientId); + + addPostTestCleanup(async () => { + clientTransport.removeEventListener( + 'protocolError', + clientHandshakeFailed, + ); + await cleanupTransports([clientTransport, serverTransport]); + }); + + await waitFor(() => { + expect(clientTransport.sessions.size).toBe(0); + }); + expect(clientHandshakeFailed).not.toHaveBeenCalled(); + + await testFinishesCleanly({ + clientTransports: [clientTransport], + serverTransport, + }); + }); }, ); diff --git a/transport/transport.ts b/transport/transport.ts index 092ca8ff..eb41cf08 100644 --- a/transport/transport.ts +++ b/transport/transport.ts @@ -79,7 +79,10 @@ export interface SessionBackpressure { * ``` * @abstract */ -export abstract class Transport { +export abstract class Transport< + ConnType extends Connection, + ApplicationErrorCode extends string = never, +> { /** * The status of the transport. */ @@ -93,7 +96,7 @@ export abstract class Transport { /** * The event dispatcher for handling events of type EventTypes. */ - eventDispatcher: EventDispatcher; + eventDispatcher: EventDispatcher; /** * The options for this transport. @@ -149,10 +152,10 @@ export abstract class Transport { * @param the type of event to listen for * @param handler The message handler to add. */ - addEventListener>( - type: K, - handler: T, - ): void { + addEventListener< + K extends EventTypes, + T extends EventHandler, + >(type: K, handler: T): void { this.eventDispatcher.addEventListener(type, handler); } @@ -161,14 +164,16 @@ export abstract class Transport { * @param the type of event to un-listen on * @param handler The message handler to remove. */ - removeEventListener>( - type: K, - handler: T, - ): void { + removeEventListener< + K extends EventTypes, + T extends EventHandler, + >(type: K, handler: T): void { this.eventDispatcher.removeEventListener(type, handler); } - protected protocolError(message: EventMap['protocolError']) { + protected protocolError( + message: EventMap['protocolError'], + ) { this.eventDispatcher.dispatchEvent('protocolError', message); } From 90c1a153992c825684b48a71d131213475b03320 Mon Sep 17 00:00:00 2001 From: Mayank Mehra Date: Mon, 24 Aug 2026 14:55:49 -0700 Subject: [PATCH 02/12] Remove unused export and tidy imports --- transport/client.ts | 1 - transport/index.ts | 1 - 2 files changed, 2 deletions(-) diff --git a/transport/client.ts b/transport/client.ts index c6597fa0..df87174a 100644 --- a/transport/client.ts +++ b/transport/client.ts @@ -13,7 +13,6 @@ import { handshakeRequestMessage, rehandshakeResponseMessage, } from './message'; - import type { TSchema } from 'typebox'; import { ClientTransportOptions, diff --git a/transport/index.ts b/transport/index.ts index c68bd4ed..9bd911e1 100644 --- a/transport/index.ts +++ b/transport/index.ts @@ -28,7 +28,6 @@ export { isStreamClose, } from './message'; export type { - HandshakeErrorCode, TransportMessage, OpaqueTransportMessage, TransportClientId, From ba621ad7408935ba6df8144019186ba02b7e4f2d Mon Sep 17 00:00:00 2001 From: Mayank Mehra Date: Mon, 24 Aug 2026 15:05:45 -0700 Subject: [PATCH 03/12] Fix application error code inference --- protobuf/handshake.ts | 6 +++++- router/handshake.ts | 4 ++-- testUtil/fixtures/cleanup.ts | 25 +++++++++++++++---------- transport/transport.test.ts | 2 +- 4 files changed, 23 insertions(+), 14 deletions(-) diff --git a/protobuf/handshake.ts b/protobuf/handshake.ts index bd82f57d..8f9ee579 100644 --- a/protobuf/handshake.ts +++ b/protobuf/handshake.ts @@ -76,7 +76,11 @@ export function createServerHandshakeOptions< ApplicationErrorCode extends string = never, >( schema: Schema, - validate: ValidateHandshake, + validate: ValidateHandshake< + Schema, + ParsedMetadata, + NoInfer + >, expiry?: (parsedMetadata: ParsedMetadata) => Date | undefined, rejectionCodes?: ReadonlyArray, ): ServerHandshakeOptions< diff --git a/router/handshake.ts b/router/handshake.ts index d1b10bb8..7bafe395 100644 --- a/router/handshake.ts +++ b/router/handshake.ts @@ -95,7 +95,7 @@ export interface ServerHandshakeOptions< validate: ValidateHandshake< MetadataSchema, ParsedMetadata, - ApplicationErrorCode + NoInfer >; /** @@ -132,7 +132,7 @@ export function createServerHandshakeOptions< validate: ValidateHandshake< MetadataSchema, ParsedMetadata, - ApplicationErrorCode + NoInfer >, expiry?: (parsedMetadata: ParsedMetadata) => Date | undefined, rejectionCodes?: ReadonlyArray, diff --git a/testUtil/fixtures/cleanup.ts b/testUtil/fixtures/cleanup.ts index 423d5122..dc8b64cb 100644 --- a/testUtil/fixtures/cleanup.ts +++ b/testUtil/fixtures/cleanup.ts @@ -36,7 +36,9 @@ export async function advanceFakeTimersByConnectionBackoff() { await vi.advanceTimersByTimeAsync(500); } -export async function ensureTransportIsClean(t: Transport) { +export async function ensureTransportIsClean< + ApplicationErrorCode extends string = never, +>(t: Transport) { await advanceFakeTimersBySessionGrace(); await waitFor(() => expect( @@ -56,9 +58,9 @@ export function waitFor(cb: () => T | Promise) { return vi.waitFor(cb, waitUntilOptions); } -export async function ensureTransportBuffersAreEventuallyEmpty( - t: Transport, -) { +export async function ensureTransportBuffersAreEventuallyEmpty< + ApplicationErrorCode extends string = never, +>(t: Transport) { // wait for send buffers to be flushed // ignore heartbeat messages await waitFor(() => @@ -97,9 +99,10 @@ export async function ensureServerIsClean( ); } -export async function cleanupTransports( - transports: Array>, -) { +export async function cleanupTransports< + ConnType extends Connection, + ApplicationErrorCode extends string = never, +>(transports: Array>) { for (const t of transports) { if (t.getStatus() !== 'closed') { t.log?.info('*** end of test cleanup ***', { clientId: t.clientId }); @@ -108,16 +111,18 @@ export async function cleanupTransports( } } -export async function testFinishesCleanly({ +export async function testFinishesCleanly< + ApplicationErrorCode extends string = never, +>({ clientTransports, serverTransport, server, }: Partial<{ - clientTransports: Array>; + clientTransports: Array>; // MetadataSchema and ParsedMetadata are not used in this test, // so we can safely use any here // eslint-disable-next-line @typescript-eslint/no-explicit-any - serverTransport: ServerTransport; + serverTransport: ServerTransport; server: Server; }>) { // pre-close invariants diff --git a/transport/transport.test.ts b/transport/transport.test.ts index 6026f238..147ea589 100644 --- a/transport/transport.test.ts +++ b/transport/transport.test.ts @@ -2052,7 +2052,7 @@ describe.each(testMatrix())( ApplicationErrorCode >('SERVER', { schema, - validate: async () => 'REPL_NOT_FOUND', + validate: async (): Promise => 'REPL_NOT_FOUND', rejectionCodes: ['REPL_NOT_FOUND'], }); From f46e5c428fba5ea321773c3ade8b0fa13fe699c8 Mon Sep 17 00:00:00 2001 From: Mayank Mehra Date: Mon, 24 Aug 2026 17:23:11 -0700 Subject: [PATCH 04/12] Require literal application error codes --- protobuf/handshake.ts | 16 ++++++++++------ router/client.ts | 3 +-- router/handshake.ts | 24 ++++++++++++++---------- testUtil/fixtures/mockTransport.ts | 4 ++-- testUtil/fixtures/transports.ts | 8 ++++---- transport/message.ts | 10 +++++++++- transport/server.ts | 5 +++-- transport/transport.test.ts | 27 +++++++++++++++++---------- 8 files changed, 60 insertions(+), 37 deletions(-) diff --git a/protobuf/handshake.ts b/protobuf/handshake.ts index 8f9ee579..b12458ea 100644 --- a/protobuf/handshake.ts +++ b/protobuf/handshake.ts @@ -12,6 +12,8 @@ import { } from '../router/handshake'; import { HandshakeErrorCustomHandlerFatalResponseCodes, + type LiteralErrorCode, + type LiteralErrorCodes, type TransportClientId, } from '../transport/message'; import { decodeMessageBytes, encodeMessageBytes } from './shared'; @@ -38,9 +40,11 @@ type ValidateHandshake< ) => | ParsedMetadata | ProtobufHandshakeFailureCode - | ApplicationErrorCode + | LiteralErrorCode | Promise< - ParsedMetadata | ProtobufHandshakeFailureCode | ApplicationErrorCode + | ParsedMetadata + | ProtobufHandshakeFailureCode + | LiteralErrorCode >; /** @@ -48,12 +52,12 @@ type ValidateHandshake< */ export function createClientHandshakeOptions< Schema extends DescMessage, - ApplicationErrorCode extends string = never, + const ApplicationErrorCode extends string = never, >( schema: Schema, construct: ConstructHandshake, eager?: boolean, - rejectionCodes?: ReadonlyArray, + rejectionCodes?: LiteralErrorCodes, ): ClientHandshakeOptions { return createTransportClientHandshakeOptions( HandshakeBytesSchema, @@ -73,7 +77,7 @@ export function createClientHandshakeOptions< export function createServerHandshakeOptions< Schema extends DescMessage, ParsedMetadata extends object = object, - ApplicationErrorCode extends string = never, + const ApplicationErrorCode extends string = never, >( schema: Schema, validate: ValidateHandshake< @@ -82,7 +86,7 @@ export function createServerHandshakeOptions< NoInfer >, expiry?: (parsedMetadata: ParsedMetadata) => Date | undefined, - rejectionCodes?: ReadonlyArray, + rejectionCodes?: LiteralErrorCodes, ): ServerHandshakeOptions< typeof HandshakeBytesSchema, ParsedMetadata, diff --git a/router/client.ts b/router/client.ts index 508738ef..43b81300 100644 --- a/router/client.ts +++ b/router/client.ts @@ -309,8 +309,7 @@ function mergeCallOptions( defaults: ClientOptions['defaultCallOptions'], caller: CallOptions | undefined, ): CallOptions { - const resolved = - typeof defaults === 'function' ? defaults() : (defaults ?? {}); + const resolved = typeof defaults === 'function' ? defaults() : defaults ?? {}; // Caller fields win: spread defaults first, caller second. return { ...resolved, ...caller }; diff --git a/router/handshake.ts b/router/handshake.ts index 7bafe395..9c7604ac 100644 --- a/router/handshake.ts +++ b/router/handshake.ts @@ -1,6 +1,8 @@ import type { Static, TSchema } from 'typebox'; import { HandshakeErrorCustomHandlerFatalResponseCodes, + type LiteralErrorCode, + type LiteralErrorCodes, type TransportClientId, } from '../transport/message'; @@ -18,11 +20,11 @@ type ValidateHandshake< from?: TransportClientId, ) => | Static - | ApplicationErrorCode + | LiteralErrorCode | ParsedMetadata | Promise< | Static - | ApplicationErrorCode + | LiteralErrorCode | ParsedMetadata >; @@ -40,9 +42,10 @@ export interface ClientHandshakeOptions< * Application-defined rejection codes the server may answer the handshake * with, sent in the response's `code` field. Must match the server's * {@link ServerHandshakeOptions.rejectionCodes}: an unconfigured code is - * rejected as a malformed handshake response. + * rejected as a malformed handshake response. Pass a literal tuple (`as + * const`); broad `string[]` values are rejected by the type system. */ - rejectionCodes?: ReadonlyArray; + rejectionCodes?: LiteralErrorCodes; /** * Gets the {@link HandshakeRequestMetadata} to send to the server. @@ -76,9 +79,10 @@ export interface ServerHandshakeOptions< * They travel in the handshake response's `code` field and are fatal like * the built-in custom-handler codes. Clients must register the same codes * in {@link ClientHandshakeOptions.rejectionCodes} or they reject the - * response as malformed. + * response as malformed. Pass a literal tuple (`as const`); broad `string[]` + * values are rejected by the type system. */ - rejectionCodes?: ReadonlyArray; + rejectionCodes?: LiteralErrorCodes; /** * Parses the metadata sent by the client during the handshake into the @@ -113,12 +117,12 @@ export interface ServerHandshakeOptions< export function createClientHandshakeOptions< MetadataSchema extends TSchema = TSchema, - ApplicationErrorCode extends string = never, + const ApplicationErrorCode extends string = never, >( schema: MetadataSchema, construct: ConstructHandshake, eager?: boolean, - rejectionCodes?: ReadonlyArray, + rejectionCodes?: LiteralErrorCodes, ): ClientHandshakeOptions { return { schema, construct, eager, rejectionCodes }; } @@ -126,7 +130,7 @@ export function createClientHandshakeOptions< export function createServerHandshakeOptions< MetadataSchema extends TSchema = TSchema, ParsedMetadata extends object = object, - ApplicationErrorCode extends string = never, + const ApplicationErrorCode extends string = never, >( schema: MetadataSchema, validate: ValidateHandshake< @@ -135,7 +139,7 @@ export function createServerHandshakeOptions< NoInfer >, expiry?: (parsedMetadata: ParsedMetadata) => Date | undefined, - rejectionCodes?: ReadonlyArray, + rejectionCodes?: LiteralErrorCodes, ): ServerHandshakeOptions< MetadataSchema, ParsedMetadata, diff --git a/testUtil/fixtures/mockTransport.ts b/testUtil/fixtures/mockTransport.ts index 60bd1242..7792c5ad 100644 --- a/testUtil/fixtures/mockTransport.ts +++ b/testUtil/fixtures/mockTransport.ts @@ -142,7 +142,7 @@ export function createMockTransportNetwork( } return { - getClientTransport: ( + getClientTransport: ( id: TransportClientId, handshakeOptions?: ClientHandshakeOptions, ) => { @@ -161,7 +161,7 @@ export function createMockTransportNetwork( getServerTransport: < MetadataSchema extends TSchema = TSchema, ParsedMetadata extends object = object, - ApplicationErrorCode extends string = never, + const ApplicationErrorCode extends string = never, >( id = 'SERVER', handshakeOptions: diff --git a/testUtil/fixtures/transports.ts b/testUtil/fixtures/transports.ts index 0d339d67..a9d02831 100644 --- a/testUtil/fixtures/transports.ts +++ b/testUtil/fixtures/transports.ts @@ -30,14 +30,14 @@ export interface TestTransportOptions { } export interface TestSetupHelpers { - getClientTransport: ( + getClientTransport: ( id: TransportClientId, handshakeOptions?: ClientHandshakeOptions, ) => ClientTransport; getServerTransport: < MetadataSchema extends TSchema = TSchema, ParsedMetadata extends object = object, - ApplicationErrorCode extends string = never, + const ApplicationErrorCode extends string = never, >( id?: TransportClientId, handshakeOptions?: ServerHandshakeOptions< @@ -83,7 +83,7 @@ export const transports: Array = [ } } }, - getClientTransport: ( + getClientTransport: ( id: TransportClientId, handshakeOptions?: ClientHandshakeOptions< TSchema, @@ -117,7 +117,7 @@ export const transports: Array = [ getServerTransport: < MetadataSchema extends TSchema, ParsedMetadata extends object, - ApplicationErrorCode extends string = never, + const ApplicationErrorCode extends string = never, >( id = 'SERVER', handshakeOptions: diff --git a/transport/message.ts b/transport/message.ts index 67ceff8c..7e0b4a60 100644 --- a/transport/message.ts +++ b/transport/message.ts @@ -123,13 +123,21 @@ export const HandshakeErrorResponseCodes = Type.Union([ HandshakeErrorFatalResponseCodes, ]); +export type LiteralErrorCode = string extends Code + ? never + : Code; + +export type LiteralErrorCodes = ReadonlyArray< + LiteralErrorCode +>; + /** * The protocol-level handshake error codes plus any application-defined * rejection codes the application registered in its handshake options. */ export type HandshakeErrorCode = | Static - | ApplicationErrorCode; + | LiteralErrorCode; export const ControlMessageHandshakeResponseSchema = Type.Object({ type: Type.Literal('HANDSHAKE_RESP'), diff --git a/transport/server.ts b/transport/server.ts index 718dd4fb..b3cbd406 100644 --- a/transport/server.ts +++ b/transport/server.ts @@ -6,6 +6,7 @@ import { ControlMessageRehandshakeResponseSchema, HandshakeErrorCustomHandlerFatalResponseCodes, type HandshakeErrorCode, + type LiteralErrorCode, OpaqueTransportMessage, acceptedProtocolVersions, TransportClientId, @@ -88,11 +89,11 @@ export abstract class ServerTransport< private isApplicationRejectionCode( value: unknown, - ): value is ApplicationErrorCode { + ): value is LiteralErrorCode { return ( typeof value === 'string' && (this.handshakeExtensions?.rejectionCodes?.includes( - value as ApplicationErrorCode, + value as LiteralErrorCode, ) ?? false) ); diff --git a/transport/transport.test.ts b/transport/transport.test.ts index 147ea589..6cd158a1 100644 --- a/transport/transport.test.ts +++ b/transport/transport.test.ts @@ -1952,11 +1952,8 @@ describe.each(testMatrix())( foo: Type.String(), }); - type ApplicationErrorCode = 'REPL_NOT_FOUND' | 'TOKEN_EXPIRED'; - const rejectionCodes: ReadonlyArray = [ - 'REPL_NOT_FOUND', - 'TOKEN_EXPIRED', - ]; + const rejectionCodes = ['REPL_NOT_FOUND', 'TOKEN_EXPIRED'] as const; + type ApplicationErrorCode = (typeof rejectionCodes)[number]; interface ParsedMetadata { foo: string; } @@ -1976,6 +1973,17 @@ describe.each(testMatrix())( ), ).toBeDefined(); + const broadCodes: string[] = ['REPL_NOT_FOUND']; + expect( + createServerHandshakeOptions( + schema, + async () => ({ foo: 'foo' }), + undefined, + // @ts-expect-error application error codes must be string literals + broadCodes, + ), + ).toBeDefined(); + const parse = vi.fn(async (): Promise => { return 'REPL_NOT_FOUND'; }); @@ -2072,7 +2080,8 @@ describe.each(testMatrix())( 'protocolError', clientHandshakeFailed, ); - await cleanupTransports([clientTransport, serverTransport]); + await cleanupTransports([clientTransport]); + await cleanupTransports([serverTransport]); }); await waitFor(() => { @@ -2080,10 +2089,8 @@ describe.each(testMatrix())( }); expect(clientHandshakeFailed).not.toHaveBeenCalled(); - await testFinishesCleanly({ - clientTransports: [clientTransport], - serverTransport, - }); + await testFinishesCleanly({ clientTransports: [clientTransport] }); + await testFinishesCleanly({ serverTransport }); }); }, ); From 5dc9894b2ac7dd47b3983c5478be83d74aec3323 Mon Sep 17 00:00:00 2001 From: Mayank Mehra Date: Mon, 24 Aug 2026 17:30:07 -0700 Subject: [PATCH 05/12] Fix property test transport types --- __tests__/properties/session.property.test.ts | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/__tests__/properties/session.property.test.ts b/__tests__/properties/session.property.test.ts index b0a8df7a..5e3a8496 100644 --- a/__tests__/properties/session.property.test.ts +++ b/__tests__/properties/session.property.test.ts @@ -15,6 +15,9 @@ import { } from '../../router/handshake'; import { createClient } from '../../router/client'; import { createServer } from '../../router/server'; +import type { ClientTransport } from '../../transport/client'; +import type { Connection } from '../../transport/connection'; +import type { ServerTransport } from '../../transport/server'; import { closeAllConnections, numberOfConnections } from '../../testUtil'; import { createMockTransportNetwork } from '../../testUtil/fixtures/mockTransport'; import type { TestTransportOptions } from '../../testUtil/fixtures/transports'; @@ -152,12 +155,8 @@ const multiplexedSchedules: gs.Generator = gs.composite( function setup(opts?: TestTransportOptions): { network: ReturnType; - clientTransport: ReturnType< - ReturnType['getClientTransport'] - >; - serverTransport: ReturnType< - ReturnType['getServerTransport'] - >; + clientTransport: ClientTransport; + serverTransport: ServerTransport; client: ReturnType>; violations: Array; } { @@ -536,11 +535,8 @@ describe('re-handshake under faults', () => { throw new Error(`timed out waiting for ${what}`); } - const isConnected = ( - transport: ReturnType< - ReturnType['getClientTransport'] - >, - ) => numberOfConnections(transport) === 1; + const isConnected = (transport: ClientTransport) => + numberOfConnections(transport) === 1; const handshakeSchema = Type.Object({ token: Type.String() }); From 3b65b6611e664f46b20b807fb3a41a1dfb64e94e Mon Sep 17 00:00:00 2001 From: Mayank Mehra Date: Mon, 24 Aug 2026 17:39:33 -0700 Subject: [PATCH 06/12] Define application errors with TypeBox literals --- protobuf/handshake.ts | 10 +++++----- router/handshake.ts | 26 +++++++++++++------------- transport/client.ts | 6 ++++-- transport/message.test.ts | 5 +++-- transport/message.ts | 14 +++++--------- transport/server.ts | 8 +++----- transport/transport.test.ts | 27 +++++++++++++++++---------- 7 files changed, 50 insertions(+), 46 deletions(-) diff --git a/protobuf/handshake.ts b/protobuf/handshake.ts index b12458ea..7e96ac65 100644 --- a/protobuf/handshake.ts +++ b/protobuf/handshake.ts @@ -13,7 +13,7 @@ import { import { HandshakeErrorCustomHandlerFatalResponseCodes, type LiteralErrorCode, - type LiteralErrorCodes, + type LiteralErrorCodeSchemas, type TransportClientId, } from '../transport/message'; import { decodeMessageBytes, encodeMessageBytes } from './shared'; @@ -57,7 +57,7 @@ export function createClientHandshakeOptions< schema: Schema, construct: ConstructHandshake, eager?: boolean, - rejectionCodes?: LiteralErrorCodes, + rejectionCodeSchemas?: LiteralErrorCodeSchemas, ): ClientHandshakeOptions { return createTransportClientHandshakeOptions( HandshakeBytesSchema, @@ -67,7 +67,7 @@ export function createClientHandshakeOptions< return encodeMessageBytes(schema, metadata); }, eager, - rejectionCodes, + rejectionCodeSchemas, ); } @@ -86,7 +86,7 @@ export function createServerHandshakeOptions< NoInfer >, expiry?: (parsedMetadata: ParsedMetadata) => Date | undefined, - rejectionCodes?: LiteralErrorCodes, + rejectionCodeSchemas?: LiteralErrorCodeSchemas, ): ServerHandshakeOptions< typeof HandshakeBytesSchema, ParsedMetadata, @@ -105,6 +105,6 @@ export function createServerHandshakeOptions< return await validate(decoded, previousParsedMetadata, from); }, expiry, - rejectionCodes, + rejectionCodeSchemas, ); } diff --git a/router/handshake.ts b/router/handshake.ts index 9c7604ac..bc4ac966 100644 --- a/router/handshake.ts +++ b/router/handshake.ts @@ -2,7 +2,7 @@ import type { Static, TSchema } from 'typebox'; import { HandshakeErrorCustomHandlerFatalResponseCodes, type LiteralErrorCode, - type LiteralErrorCodes, + type LiteralErrorCodeSchemas, type TransportClientId, } from '../transport/message'; @@ -41,11 +41,11 @@ export interface ClientHandshakeOptions< /** * Application-defined rejection codes the server may answer the handshake * with, sent in the response's `code` field. Must match the server's - * {@link ServerHandshakeOptions.rejectionCodes}: an unconfigured code is - * rejected as a malformed handshake response. Pass a literal tuple (`as - * const`); broad `string[]` values are rejected by the type system. + * {@link ServerHandshakeOptions.rejectionCodeSchemas}: an unconfigured code + * is rejected as a malformed handshake response. Pass a tuple of TypeBox + * literals (`Type.Literal(...)`) with `as const`. */ - rejectionCodes?: LiteralErrorCodes; + rejectionCodeSchemas?: LiteralErrorCodeSchemas; /** * Gets the {@link HandshakeRequestMetadata} to send to the server. @@ -78,11 +78,11 @@ export interface ServerHandshakeOptions< * Application-defined rejection codes that {@link validate} may return. * They travel in the handshake response's `code` field and are fatal like * the built-in custom-handler codes. Clients must register the same codes - * in {@link ClientHandshakeOptions.rejectionCodes} or they reject the - * response as malformed. Pass a literal tuple (`as const`); broad `string[]` - * values are rejected by the type system. + * in {@link ClientHandshakeOptions.rejectionCodeSchemas} or they reject the + * response as malformed. Pass a tuple of TypeBox literals + * (`Type.Literal(...)`) with `as const`. */ - rejectionCodes?: LiteralErrorCodes; + rejectionCodeSchemas?: LiteralErrorCodeSchemas; /** * Parses the metadata sent by the client during the handshake into the @@ -122,9 +122,9 @@ export function createClientHandshakeOptions< schema: MetadataSchema, construct: ConstructHandshake, eager?: boolean, - rejectionCodes?: LiteralErrorCodes, + rejectionCodeSchemas?: LiteralErrorCodeSchemas, ): ClientHandshakeOptions { - return { schema, construct, eager, rejectionCodes }; + return { schema, construct, eager, rejectionCodeSchemas }; } export function createServerHandshakeOptions< @@ -139,11 +139,11 @@ export function createServerHandshakeOptions< NoInfer >, expiry?: (parsedMetadata: ParsedMetadata) => Date | undefined, - rejectionCodes?: LiteralErrorCodes, + rejectionCodeSchemas?: LiteralErrorCodeSchemas, ): ServerHandshakeOptions< MetadataSchema, ParsedMetadata, ApplicationErrorCode > { - return { schema, validate, expiry, rejectionCodes }; + return { schema, validate, expiry, rejectionCodeSchemas }; } diff --git a/transport/client.ts b/transport/client.ts index df87174a..ae9d008b 100644 --- a/transport/client.ts +++ b/transport/client.ts @@ -115,9 +115,11 @@ export abstract class ClientTransport< options: ClientHandshakeOptions, ) { this.handshakeExtensions = options; - if (options.rejectionCodes?.length) { + if (options.rejectionCodeSchemas?.length) { this.handshakeResponseSchema = - ControlMessageHandshakeResponseSchemaWithCodes(options.rejectionCodes); + ControlMessageHandshakeResponseSchemaWithCodes( + options.rejectionCodeSchemas, + ); } } diff --git a/transport/message.test.ts b/transport/message.test.ts index 17747f51..d8218390 100644 --- a/transport/message.test.ts +++ b/transport/message.test.ts @@ -10,6 +10,7 @@ import { isStreamOpen, } from './message'; import { describe, test, expect } from 'vitest'; +import { Type } from 'typebox'; import { Value } from 'typebox/value'; const msg = ( @@ -110,8 +111,8 @@ describe('message helpers', () => { test('handshake response schema with application codes', () => { const extended = ControlMessageHandshakeResponseSchemaWithCodes([ - 'REPL_NOT_FOUND', - 'TOKEN_EXPIRED', + Type.Literal('REPL_NOT_FOUND'), + Type.Literal('TOKEN_EXPIRED'), ] as const); const rejection = { type: 'HANDSHAKE_RESP', diff --git a/transport/message.ts b/transport/message.ts index 7e0b4a60..0b8a4de6 100644 --- a/transport/message.ts +++ b/transport/message.ts @@ -127,8 +127,8 @@ export type LiteralErrorCode = string extends Code ? never : Code; -export type LiteralErrorCodes = ReadonlyArray< - LiteralErrorCode +export type LiteralErrorCodeSchemas = ReadonlyArray< + TLiteral> >; /** @@ -160,9 +160,9 @@ export const ControlMessageHandshakeResponseSchema = Type.Object({ * unconfigured peer rejects an application code as a malformed response. */ export const ControlMessageHandshakeResponseSchemaWithCodes = < - const ApplicationErrorCodes extends readonly string[], + const ApplicationErrorCodeSchemas extends ReadonlyArray>, >( - applicationErrorCodes: ApplicationErrorCodes, + applicationErrorCodeSchemas: ApplicationErrorCodeSchemas, ) => Type.Object({ type: Type.Literal('HANDSHAKE_RESP'), @@ -176,11 +176,7 @@ export const ControlMessageHandshakeResponseSchemaWithCodes = < reason: Type.String(), code: Type.Union([ HandshakeErrorResponseCodes, - Type.Union( - applicationErrorCodes.map( - (code) => Type.Literal(code) as TLiteral, - ), - ), + Type.Union([...applicationErrorCodeSchemas]), ]), }), ]), diff --git a/transport/server.ts b/transport/server.ts index b3cbd406..8ec453b1 100644 --- a/transport/server.ts +++ b/transport/server.ts @@ -91,11 +91,9 @@ export abstract class ServerTransport< value: unknown, ): value is LiteralErrorCode { return ( - typeof value === 'string' && - (this.handshakeExtensions?.rejectionCodes?.includes( - value as LiteralErrorCode, - ) ?? - false) + this.handshakeExtensions?.rejectionCodeSchemas?.some((schema) => + Value.Check(schema, value), + ) ?? false ); } diff --git a/transport/transport.test.ts b/transport/transport.test.ts index 6cd158a1..c886e6b4 100644 --- a/transport/transport.test.ts +++ b/transport/transport.test.ts @@ -21,7 +21,7 @@ import { } from '../testUtil/fixtures/cleanup'; import { testMatrix } from '../testUtil/fixtures/matrix'; import { PartialTransportMessage } from './message'; -import { Type } from 'typebox'; +import { Type, type Static } from 'typebox'; import { TestSetupHelpers } from '../testUtil/fixtures/transports'; import { createPostTestCleanups } from '../testUtil/fixtures/cleanup'; import { SessionState } from './sessionStateMachine'; @@ -1952,8 +1952,13 @@ describe.each(testMatrix())( foo: Type.String(), }); - const rejectionCodes = ['REPL_NOT_FOUND', 'TOKEN_EXPIRED'] as const; - type ApplicationErrorCode = (typeof rejectionCodes)[number]; + const rejectionCodeSchemas = [ + Type.Literal('REPL_NOT_FOUND'), + Type.Literal('TOKEN_EXPIRED'), + ] as const; + + type ApplicationErrorCode = Static<(typeof rejectionCodeSchemas)[number]>; + interface ParsedMetadata { foo: string; } @@ -1969,18 +1974,20 @@ describe.each(testMatrix())( // @ts-expect-error only declared rejection codes may be returned async () => 'SOME_OTHER_CODE', undefined, - rejectionCodes, + rejectionCodeSchemas, ), ).toBeDefined(); - const broadCodes: string[] = ['REPL_NOT_FOUND']; + const broadCode: string = 'REPL_NOT_FOUND'; + const broadCodeSchemas = [Type.Literal(broadCode)]; + expect( createServerHandshakeOptions( schema, async () => ({ foo: 'foo' }), undefined, - // @ts-expect-error application error codes must be string literals - broadCodes, + // @ts-expect-error application error schemas must use literals + broadCodeSchemas, ), ).toBeDefined(); @@ -1994,13 +2001,13 @@ describe.each(testMatrix())( >('SERVER', { schema, validate: parse, - rejectionCodes, + rejectionCodeSchemas, }); const clientTransport = getClientTransport('client', { schema, construct: async () => ({ foo: 'foo' }), - rejectionCodes, + rejectionCodeSchemas, }); const clientHandshakeFailed = vi.fn(); @@ -2061,7 +2068,7 @@ describe.each(testMatrix())( >('SERVER', { schema, validate: async (): Promise => 'REPL_NOT_FOUND', - rejectionCodes: ['REPL_NOT_FOUND'], + rejectionCodeSchemas: [Type.Literal('REPL_NOT_FOUND')], }); // the client did not register the application code: it must treat the From ab631f103591f99d9ed2263f651264356c261c4f Mon Sep 17 00:00:00 2001 From: Mayank Mehra Date: Mon, 24 Aug 2026 17:47:14 -0700 Subject: [PATCH 07/12] Fix application error schema lint --- transport/transport.test.ts | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/transport/transport.test.ts b/transport/transport.test.ts index c886e6b4..f597c67b 100644 --- a/transport/transport.test.ts +++ b/transport/transport.test.ts @@ -1978,7 +1978,7 @@ describe.each(testMatrix())( ), ).toBeDefined(); - const broadCode: string = 'REPL_NOT_FOUND'; + const broadCode = String('REPL_NOT_FOUND'); const broadCodeSchemas = [Type.Literal(broadCode)]; expect( @@ -2056,7 +2056,10 @@ describe.each(testMatrix())( foo: Type.String(), }); - type ApplicationErrorCode = 'REPL_NOT_FOUND'; + const rejectionCodeSchema = Type.Literal('REPL_NOT_FOUND'); + + type ApplicationErrorCode = Static; + interface ParsedMetadata { foo: string; } @@ -2068,7 +2071,7 @@ describe.each(testMatrix())( >('SERVER', { schema, validate: async (): Promise => 'REPL_NOT_FOUND', - rejectionCodeSchemas: [Type.Literal('REPL_NOT_FOUND')], + rejectionCodeSchemas: [rejectionCodeSchema], }); // the client did not register the application code: it must treat the From 0eda02a82c31126a32f12e6d8ec0d7accd797056 Mon Sep 17 00:00:00 2001 From: Mayank Mehra Date: Mon, 24 Aug 2026 22:18:01 -0700 Subject: [PATCH 08/12] Derive rejection code types from TypeBox literal schemas --- protobuf/handshake.ts | 24 ++++++++-------- router/client.ts | 13 +++++---- router/handshake.ts | 34 +++++++++++------------ router/server.ts | 15 +++++----- testUtil/fixtures/cleanup.ts | 23 +++++++++------- testUtil/fixtures/mockTransport.ts | 27 ++++++++++-------- testUtil/fixtures/transports.ts | 35 ++++++++++++++---------- transport/client.ts | 11 ++++---- transport/events.ts | 26 +++++++++++------- transport/impls/ws/client.ts | 9 ++++-- transport/impls/ws/server.ts | 9 ++++-- transport/message.ts | 44 +++++++++++++++++++++--------- transport/server.ts | 21 +++++++------- transport/transport.test.ts | 26 +++++++++++------- transport/transport.ts | 11 ++++---- 15 files changed, 193 insertions(+), 135 deletions(-) diff --git a/protobuf/handshake.ts b/protobuf/handshake.ts index 7e96ac65..05c90fbf 100644 --- a/protobuf/handshake.ts +++ b/protobuf/handshake.ts @@ -11,9 +11,9 @@ import { type ServerHandshakeOptions, } from '../router/handshake'; import { + type ApplicationErrorCode, + type ApplicationErrorCodeSchemas, HandshakeErrorCustomHandlerFatalResponseCodes, - type LiteralErrorCode, - type LiteralErrorCodeSchemas, type TransportClientId, } from '../transport/message'; import { decodeMessageBytes, encodeMessageBytes } from './shared'; @@ -32,7 +32,7 @@ type ConstructHandshake = () => type ValidateHandshake< Schema extends DescMessage, ParsedMetadata, - ApplicationErrorCode extends string = never, + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], > = ( metadata: MessageShape, previousParsedMetadata?: ParsedMetadata, @@ -40,11 +40,11 @@ type ValidateHandshake< ) => | ParsedMetadata | ProtobufHandshakeFailureCode - | LiteralErrorCode + | ApplicationErrorCode | Promise< | ParsedMetadata | ProtobufHandshakeFailureCode - | LiteralErrorCode + | ApplicationErrorCode >; /** @@ -52,13 +52,13 @@ type ValidateHandshake< */ export function createClientHandshakeOptions< Schema extends DescMessage, - const ApplicationErrorCode extends string = never, + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], >( schema: Schema, construct: ConstructHandshake, eager?: boolean, - rejectionCodeSchemas?: LiteralErrorCodeSchemas, -): ClientHandshakeOptions { + rejectionCodeSchemas?: RejectionCodeSchemas, +): ClientHandshakeOptions { return createTransportClientHandshakeOptions( HandshakeBytesSchema, async () => { @@ -77,20 +77,20 @@ export function createClientHandshakeOptions< export function createServerHandshakeOptions< Schema extends DescMessage, ParsedMetadata extends object = object, - const ApplicationErrorCode extends string = never, + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], >( schema: Schema, validate: ValidateHandshake< Schema, ParsedMetadata, - NoInfer + NoInfer >, expiry?: (parsedMetadata: ParsedMetadata) => Date | undefined, - rejectionCodeSchemas?: LiteralErrorCodeSchemas, + rejectionCodeSchemas?: RejectionCodeSchemas, ): ServerHandshakeOptions< typeof HandshakeBytesSchema, ParsedMetadata, - ApplicationErrorCode + RejectionCodeSchemas > { return createTransportServerHandshakeOptions( HandshakeBytesSchema, diff --git a/router/client.ts b/router/client.ts index 43b81300..9f1c26ed 100644 --- a/router/client.ts +++ b/router/client.ts @@ -17,6 +17,7 @@ import { isStreamCancel, closeStreamMessage, cancelMessage, + type ApplicationErrorCodeSchemas, } from '../transport/message'; import type { Static, TSchema } from 'typebox'; import { Err, Result, AnyResultSchema } from './result'; @@ -244,13 +245,13 @@ const defaultClientOptions: ClientOptions = { export function createClient< // eslint-disable-next-line @typescript-eslint/no-explicit-any ServiceSchemaMap extends AnyServiceSchemaMap, - ApplicationErrorCode extends string = never, + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], >( - transport: ClientTransport, + transport: ClientTransport, serverId: TransportClientId, providedClientOptions: Partial< ClientOptions & { - handshakeOptions: ClientHandshakeOptions; + handshakeOptions: ClientHandshakeOptions; } > = {}, ): Client { @@ -321,9 +322,11 @@ type AnyProcReturn = | ReturnType> | ReturnType>; -function handleProc( +function handleProc< + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], +>( procType: ValidProcType, - transport: ClientTransport, + transport: ClientTransport, serverId: TransportClientId, init: Static, serviceName: string, diff --git a/router/handshake.ts b/router/handshake.ts index bc4ac966..3c603f12 100644 --- a/router/handshake.ts +++ b/router/handshake.ts @@ -1,8 +1,8 @@ import type { Static, TSchema } from 'typebox'; import { + type ApplicationErrorCode, + type ApplicationErrorCodeSchemas, HandshakeErrorCustomHandlerFatalResponseCodes, - type LiteralErrorCode, - type LiteralErrorCodeSchemas, type TransportClientId, } from '../transport/message'; @@ -13,24 +13,24 @@ type ConstructHandshake = () => type ValidateHandshake< T extends TSchema, ParsedMetadata, - ApplicationErrorCode extends string = never, + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], > = ( metadata: Static, previousParsedMetadata?: ParsedMetadata, from?: TransportClientId, ) => | Static - | LiteralErrorCode + | ApplicationErrorCode | ParsedMetadata | Promise< | Static - | LiteralErrorCode + | ApplicationErrorCode | ParsedMetadata >; export interface ClientHandshakeOptions< MetadataSchema extends TSchema = TSchema, - ApplicationErrorCode extends string = never, + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], > { /** * Schema for the metadata that the client sends to the server @@ -45,7 +45,7 @@ export interface ClientHandshakeOptions< * is rejected as a malformed handshake response. Pass a tuple of TypeBox * literals (`Type.Literal(...)`) with `as const`. */ - rejectionCodeSchemas?: LiteralErrorCodeSchemas; + rejectionCodeSchemas?: RejectionCodeSchemas; /** * Gets the {@link HandshakeRequestMetadata} to send to the server. @@ -66,7 +66,7 @@ export interface ClientHandshakeOptions< export interface ServerHandshakeOptions< MetadataSchema extends TSchema = TSchema, ParsedMetadata extends object = object, - ApplicationErrorCode extends string = never, + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], > { /** * Schema for the metadata that the server receives from the client @@ -82,7 +82,7 @@ export interface ServerHandshakeOptions< * response as malformed. Pass a tuple of TypeBox literals * (`Type.Literal(...)`) with `as const`. */ - rejectionCodeSchemas?: LiteralErrorCodeSchemas; + rejectionCodeSchemas?: RejectionCodeSchemas; /** * Parses the metadata sent by the client during the handshake into the @@ -99,7 +99,7 @@ export interface ServerHandshakeOptions< validate: ValidateHandshake< MetadataSchema, ParsedMetadata, - NoInfer + NoInfer >; /** @@ -117,33 +117,33 @@ export interface ServerHandshakeOptions< export function createClientHandshakeOptions< MetadataSchema extends TSchema = TSchema, - const ApplicationErrorCode extends string = never, + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], >( schema: MetadataSchema, construct: ConstructHandshake, eager?: boolean, - rejectionCodeSchemas?: LiteralErrorCodeSchemas, -): ClientHandshakeOptions { + rejectionCodeSchemas?: RejectionCodeSchemas, +): ClientHandshakeOptions { return { schema, construct, eager, rejectionCodeSchemas }; } export function createServerHandshakeOptions< MetadataSchema extends TSchema = TSchema, ParsedMetadata extends object = object, - const ApplicationErrorCode extends string = never, + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], >( schema: MetadataSchema, validate: ValidateHandshake< MetadataSchema, ParsedMetadata, - NoInfer + NoInfer >, expiry?: (parsedMetadata: ParsedMetadata) => Date | undefined, - rejectionCodeSchemas?: LiteralErrorCodeSchemas, + rejectionCodeSchemas?: RejectionCodeSchemas, ): ServerHandshakeOptions< MetadataSchema, ParsedMetadata, - ApplicationErrorCode + RejectionCodeSchemas > { return { schema, validate, expiry, rejectionCodeSchemas }; } diff --git a/router/server.ts b/router/server.ts index 9f03e7db..84ddd7ce 100644 --- a/router/server.ts +++ b/router/server.ts @@ -29,6 +29,7 @@ import { cancelMessage, ProtocolVersion, TransportClientId, + type ApplicationErrorCodeSchemas, } from '../transport/message'; import { ProcedureHandlerContext } from './context'; import { Logger } from '../logging/log'; @@ -112,14 +113,14 @@ class RiverServer< MetadataSchema extends TSchema, ParsedMetadata extends object, Services extends AnyServiceSchemaMap, - ApplicationErrorCode extends string = never, + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], > implements Server { private transport: ServerTransport< Connection, MetadataSchema, ParsedMetadata, - ApplicationErrorCode + RejectionCodeSchemas >; private contextMap: Map; @@ -151,13 +152,13 @@ class RiverServer< Connection, MetadataSchema, ParsedMetadata, - ApplicationErrorCode + RejectionCodeSchemas >, services: Services, handshakeOptions?: ServerHandshakeOptions< MetadataSchema, ParsedMetadata, - ApplicationErrorCode + RejectionCodeSchemas >, extendedContext?: Context, maxCancelledStreamTombstonesPerSession = 200, @@ -1178,20 +1179,20 @@ export function createServer< // eslint-disable-next-line @typescript-eslint/no-explicit-any Services extends AnyServiceSchemaMap, Context extends MaybeDisposable = MaybeDisposable, - ApplicationErrorCode extends string = never, + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], >( transport: ServerTransport< Connection, MetadataSchema, ParsedMetadata, - ApplicationErrorCode + RejectionCodeSchemas >, services: Services, providedServerOptions?: Partial<{ handshakeOptions?: ServerHandshakeOptions< MetadataSchema, ParsedMetadata, - ApplicationErrorCode + RejectionCodeSchemas >; extendedContext?: Context; /** diff --git a/testUtil/fixtures/cleanup.ts b/testUtil/fixtures/cleanup.ts index dc8b64cb..8ad5e3e7 100644 --- a/testUtil/fixtures/cleanup.ts +++ b/testUtil/fixtures/cleanup.ts @@ -9,7 +9,10 @@ import { Server } from '../../router'; import { AnyServiceSchemaMap, MaybeDisposable } from '../../router/services'; import { numberOfConnections, testingSessionOptions } from '..'; import { Value } from 'typebox/value'; -import { ControlMessageAckSchema } from '../../transport/message'; +import { + type ApplicationErrorCodeSchemas, + ControlMessageAckSchema, +} from '../../transport/message'; const waitUntilOptions = { timeout: 500, // account for possibility of conn backoff @@ -37,8 +40,8 @@ export async function advanceFakeTimersByConnectionBackoff() { } export async function ensureTransportIsClean< - ApplicationErrorCode extends string = never, ->(t: Transport) { + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], +>(t: Transport) { await advanceFakeTimersBySessionGrace(); await waitFor(() => expect( @@ -59,8 +62,8 @@ export function waitFor(cb: () => T | Promise) { } export async function ensureTransportBuffersAreEventuallyEmpty< - ApplicationErrorCode extends string = never, ->(t: Transport) { + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], +>(t: Transport) { // wait for send buffers to be flushed // ignore heartbeat messages await waitFor(() => @@ -101,8 +104,8 @@ export async function ensureServerIsClean( export async function cleanupTransports< ConnType extends Connection, - ApplicationErrorCode extends string = never, ->(transports: Array>) { + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], +>(transports: Array>) { for (const t of transports) { if (t.getStatus() !== 'closed') { t.log?.info('*** end of test cleanup ***', { clientId: t.clientId }); @@ -112,17 +115,17 @@ export async function cleanupTransports< } export async function testFinishesCleanly< - ApplicationErrorCode extends string = never, + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], >({ clientTransports, serverTransport, server, }: Partial<{ - clientTransports: Array>; + clientTransports: Array>; // MetadataSchema and ParsedMetadata are not used in this test, // so we can safely use any here // eslint-disable-next-line @typescript-eslint/no-explicit-any - serverTransport: ServerTransport; + serverTransport: ServerTransport; server: Server; }>) { // pre-close invariants diff --git a/testUtil/fixtures/mockTransport.ts b/testUtil/fixtures/mockTransport.ts index 7792c5ad..b6b8c952 100644 --- a/testUtil/fixtures/mockTransport.ts +++ b/testUtil/fixtures/mockTransport.ts @@ -1,4 +1,5 @@ import { Transport, TransportClientId } from '../../transport'; +import type { ApplicationErrorCodeSchemas } from '../../transport/message'; import { ClientTransport } from '../../transport/client'; import { Connection } from '../../transport/connection'; import { ServerTransport } from '../../transport/server'; @@ -74,10 +75,12 @@ export function createMockTransportNetwork( // conn id -> [client->server, server->client] const connections = new Observable>({}); - const transports: Array> = []; + const transports: Array< + Transport + > = []; class MockClientTransport< - ApplicationErrorCode extends string = never, - > extends ClientTransport { + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + > extends ClientTransport { async createNewOutgoingConnection( to: TransportClientId, ): Promise { @@ -104,12 +107,12 @@ export function createMockTransportNetwork( class MockServerTransport< MetadataSchema extends TSchema = TSchema, ParsedMetadata extends object = object, - ApplicationErrorCode extends string = never, + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], > extends ServerTransport< InMemoryConnection, MetadataSchema, ParsedMetadata, - ApplicationErrorCode + RejectionCodeSchemas > { subscribeCleanup: () => void; @@ -142,11 +145,13 @@ export function createMockTransportNetwork( } return { - getClientTransport: ( + getClientTransport: < + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + >( id: TransportClientId, - handshakeOptions?: ClientHandshakeOptions, + handshakeOptions?: ClientHandshakeOptions, ) => { - const clientTransport = new MockClientTransport( + const clientTransport = new MockClientTransport( id, opts?.client, ); @@ -161,21 +166,21 @@ export function createMockTransportNetwork( getServerTransport: < MetadataSchema extends TSchema = TSchema, ParsedMetadata extends object = object, - const ApplicationErrorCode extends string = never, + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], >( id = 'SERVER', handshakeOptions: | ServerHandshakeOptions< MetadataSchema, ParsedMetadata, - ApplicationErrorCode + RejectionCodeSchemas > | undefined, ) => { const serverTransport = new MockServerTransport< MetadataSchema, ParsedMetadata, - ApplicationErrorCode + RejectionCodeSchemas >(id, opts?.server); if (handshakeOptions) { serverTransport.extendHandshake(handshakeOptions); diff --git a/testUtil/fixtures/transports.ts b/testUtil/fixtures/transports.ts index a9d02831..abb5d40a 100644 --- a/testUtil/fixtures/transports.ts +++ b/testUtil/fixtures/transports.ts @@ -16,7 +16,10 @@ import { ProvidedClientTransportOptions, ProvidedServerTransportOptions, } from '../../transport/options'; -import { TransportClientId } from '../../transport/message'; +import { + type ApplicationErrorCodeSchemas, + TransportClientId, +} from '../../transport/message'; import { ClientTransport } from '../../transport/client'; import { Connection } from '../../transport/connection'; import { ServerTransport } from '../../transport/server'; @@ -30,26 +33,28 @@ export interface TestTransportOptions { } export interface TestSetupHelpers { - getClientTransport: ( + getClientTransport: < + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + >( id: TransportClientId, - handshakeOptions?: ClientHandshakeOptions, - ) => ClientTransport; + handshakeOptions?: ClientHandshakeOptions, + ) => ClientTransport; getServerTransport: < MetadataSchema extends TSchema = TSchema, ParsedMetadata extends object = object, - const ApplicationErrorCode extends string = never, + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], >( id?: TransportClientId, handshakeOptions?: ServerHandshakeOptions< MetadataSchema, ParsedMetadata, - ApplicationErrorCode + RejectionCodeSchemas >, ) => ServerTransport< Connection, MetadataSchema, ParsedMetadata, - ApplicationErrorCode + RejectionCodeSchemas >; simulatePhantomDisconnect: () => void; restartServer: () => Promise; @@ -83,15 +88,17 @@ export const transports: Array = [ } } }, - getClientTransport: ( + getClientTransport: < + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + >( id: TransportClientId, handshakeOptions?: ClientHandshakeOptions< TSchema, - ApplicationErrorCode + RejectionCodeSchemas >, ) => { const clientTransport = - new WebSocketClientTransport( + new WebSocketClientTransport( () => Promise.resolve(createLocalWebSocketClient(port)), id, opts?.client, @@ -117,21 +124,21 @@ export const transports: Array = [ getServerTransport: < MetadataSchema extends TSchema, ParsedMetadata extends object, - const ApplicationErrorCode extends string = never, + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], >( id = 'SERVER', handshakeOptions: | ServerHandshakeOptions< MetadataSchema, ParsedMetadata, - ApplicationErrorCode + RejectionCodeSchemas > | undefined, ) => { const serverTransport = new WebSocketServerTransport< MetadataSchema, ParsedMetadata, - ApplicationErrorCode + RejectionCodeSchemas >(wss, id, opts?.server); serverTransport.bindLogger((msg, ctx, level) => { @@ -153,7 +160,7 @@ export const transports: Array = [ Connection, MetadataSchema, ParsedMetadata, - ApplicationErrorCode + RejectionCodeSchemas >; }, async restartServer() { diff --git a/transport/client.ts b/transport/client.ts index ae9d008b..fde88058 100644 --- a/transport/client.ts +++ b/transport/client.ts @@ -5,6 +5,7 @@ import { ControlMessageHandshakeResponseSchema, ControlMessageHandshakeResponseSchemaWithCodes, ControlMessageRehandshakeRequestSchema, + type ApplicationErrorCodeSchemas, type HandshakeErrorCode, HandshakeErrorRetriableResponseCodes, OpaqueTransportMessage, @@ -53,8 +54,8 @@ type ConstructedHandshakeMetadata = export abstract class ClientTransport< ConnType extends Connection, - ApplicationErrorCode extends string = never, -> extends Transport { + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], +> extends Transport { /** * The options for this transport. */ @@ -73,7 +74,7 @@ export abstract class ClientTransport< /** * Optional handshake options for this client. */ - handshakeExtensions?: ClientHandshakeOptions; + handshakeExtensions?: ClientHandshakeOptions; /** * Handshake response schema extended with the application-defined @@ -112,7 +113,7 @@ export abstract class ClientTransport< } extendHandshake( - options: ClientHandshakeOptions, + options: ClientHandshakeOptions, ) { this.handshakeExtensions = options; if (options.rejectionCodeSchemas?.length) { @@ -410,7 +411,7 @@ export abstract class ClientTransport< this.protocolError({ type: ProtocolError.HandshakeFailed, code: msg.payload.status - .code as HandshakeErrorCode, + .code as HandshakeErrorCode, message: reason, }); } diff --git a/transport/events.ts b/transport/events.ts index 63a6e263..9e90ce5c 100644 --- a/transport/events.ts +++ b/transport/events.ts @@ -1,5 +1,9 @@ import { Connection } from './connection'; -import { OpaqueTransportMessage, HandshakeErrorCode } from './message'; +import type { + ApplicationErrorCodeSchemas, + HandshakeErrorCode, + OpaqueTransportMessage, +} from './message'; import { Session, SessionState } from './sessionStateMachine'; import { SessionId } from './sessionStateMachine/common'; import { TransportStatus } from './transport'; @@ -15,7 +19,9 @@ export const ProtocolError = { export type ProtocolErrorType = (typeof ProtocolError)[keyof typeof ProtocolError]; -export interface EventMap { +export interface EventMap< + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], +> { message: OpaqueTransportMessage; sessionStatus: | { @@ -35,7 +41,7 @@ export interface EventMap { protocolError: | { type: (typeof ProtocolError)['HandshakeFailed']; - code: HandshakeErrorCode; + code: HandshakeErrorCode; message: string; } | { @@ -53,15 +59,15 @@ export interface EventMap { export type EventTypes = keyof EventMap; export type EventHandler< K extends EventTypes, - ApplicationErrorCode extends string = never, -> = (event: EventMap[K]) => unknown; + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], +> = (event: EventMap[K]) => unknown; export class EventDispatcher< T extends EventTypes, - ApplicationErrorCode extends string = never, + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], > { private eventListeners: { - [K in T]?: Set>; + [K in T]?: Set>; } = {}; removeAllListeners() { @@ -74,7 +80,7 @@ export class EventDispatcher< addEventListener( eventType: K, - handler: EventHandler, + handler: EventHandler, ) { if (!this.eventListeners[eventType]) { this.eventListeners[eventType] = new Set(); @@ -85,7 +91,7 @@ export class EventDispatcher< removeEventListener( eventType: K, - handler: EventHandler, + handler: EventHandler, ) { const handlers = this.eventListeners[eventType]; if (handlers) { @@ -95,7 +101,7 @@ export class EventDispatcher< dispatchEvent( eventType: K, - event: EventMap[K], + event: EventMap[K], ) { const handlers = this.eventListeners[eventType]; if (handlers) { diff --git a/transport/impls/ws/client.ts b/transport/impls/ws/client.ts index b6d716a7..4e62efc7 100644 --- a/transport/impls/ws/client.ts +++ b/transport/impls/ws/client.ts @@ -1,5 +1,8 @@ import { ClientTransport } from '../../client'; -import { TransportClientId } from '../../message'; +import { + type ApplicationErrorCodeSchemas, + TransportClientId, +} from '../../message'; import { ProvidedClientTransportOptions } from '../../options'; import { WebSocketConnection } from './connection'; import { WsLike } from './wslike'; @@ -10,8 +13,8 @@ import { WsLike } from './wslike'; * @extends Transport */ export class WebSocketClientTransport< - ApplicationErrorCode extends string = never, -> extends ClientTransport { + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], +> extends ClientTransport { /** * A function that returns a Promise that resolves to a websocket URL. */ diff --git a/transport/impls/ws/server.ts b/transport/impls/ws/server.ts index 922bcddf..fb4bad9c 100644 --- a/transport/impls/ws/server.ts +++ b/transport/impls/ws/server.ts @@ -1,4 +1,7 @@ -import { TransportClientId } from '../../message'; +import { + type ApplicationErrorCodeSchemas, + TransportClientId, +} from '../../message'; import { WebSocketServer } from 'ws'; import { WebSocketConnection } from './connection'; import { WsLike } from './wslike'; @@ -25,12 +28,12 @@ function cleanHeaders( export class WebSocketServerTransport< MetadataSchema extends TSchema = TSchema, ParsedMetadata extends object = object, - ApplicationErrorCode extends string = never, + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], > extends ServerTransport< WebSocketConnection, MetadataSchema, ParsedMetadata, - ApplicationErrorCode + RejectionCodeSchemas > { wss: WebSocketServer; diff --git a/transport/message.ts b/transport/message.ts index 0b8a4de6..6988c8c6 100644 --- a/transport/message.ts +++ b/transport/message.ts @@ -123,21 +123,33 @@ export const HandshakeErrorResponseCodes = Type.Union([ HandshakeErrorFatalResponseCodes, ]); -export type LiteralErrorCode = string extends Code - ? never - : Code; +/** + * A tuple of TypeBox literal schemas declaring the application-defined + * handshake rejection codes, e.g. + * `[Type.Literal('REPL_NOT_FOUND'), Type.Literal('TOKEN_EXPIRED')] as const`. + */ +export type ApplicationErrorCodeSchemas = ReadonlyArray>; -export type LiteralErrorCodeSchemas = ReadonlyArray< - TLiteral> ->; +/** + * The union of codes declared by a tuple of rejection code schemas. + * Widened schemas (`TLiteral` rather than a specific literal) + * contribute no codes. + */ +export type ApplicationErrorCode< + RejectionCodeSchemas extends ApplicationErrorCodeSchemas, +> = string extends Static + ? never + : Static; /** * The protocol-level handshake error codes plus any application-defined * rejection codes the application registered in its handshake options. */ -export type HandshakeErrorCode = +export type HandshakeErrorCode< + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], +> = | Static - | LiteralErrorCode; + | ApplicationErrorCode; export const ControlMessageHandshakeResponseSchema = Type.Object({ type: Type.Literal('HANDSHAKE_RESP'), @@ -160,9 +172,9 @@ export const ControlMessageHandshakeResponseSchema = Type.Object({ * unconfigured peer rejects an application code as a malformed response. */ export const ControlMessageHandshakeResponseSchemaWithCodes = < - const ApplicationErrorCodeSchemas extends ReadonlyArray>, + RejectionCodeSchemas extends ApplicationErrorCodeSchemas, >( - applicationErrorCodeSchemas: ApplicationErrorCodeSchemas, + rejectionCodeSchemas: RejectionCodeSchemas, ) => Type.Object({ type: Type.Literal('HANDSHAKE_RESP'), @@ -176,7 +188,7 @@ export const ControlMessageHandshakeResponseSchemaWithCodes = < reason: Type.String(), code: Type.Union([ HandshakeErrorResponseCodes, - Type.Union([...applicationErrorCodeSchemas]), + ...rejectionCodeSchemas, ]), }), ]), @@ -301,7 +313,9 @@ export function handshakeRequestMessage({ */ export const SESSION_STATE_MISMATCH = 'session state mismatch'; -export function handshakeResponseMessage({ +export function handshakeResponseMessage< + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], +>({ from, to, status, @@ -312,7 +326,11 @@ export function handshakeResponseMessage({ // known to peers that registered it in their handshake options status: | { ok: true; sessionId: string } - | { ok: false; reason: string; code: string }; + | { + ok: false; + reason: string; + code: HandshakeErrorCode; + }; }): TransportMessage> { return { id: generateId(), diff --git a/transport/server.ts b/transport/server.ts index 8ec453b1..7b3f7365 100644 --- a/transport/server.ts +++ b/transport/server.ts @@ -5,8 +5,9 @@ import { ControlMessageHandshakeRequestSchema, ControlMessageRehandshakeResponseSchema, HandshakeErrorCustomHandlerFatalResponseCodes, + type ApplicationErrorCode, + type ApplicationErrorCodeSchemas, type HandshakeErrorCode, - type LiteralErrorCode, OpaqueTransportMessage, acceptedProtocolVersions, TransportClientId, @@ -37,8 +38,8 @@ export abstract class ServerTransport< ConnType extends Connection, MetadataSchema extends TSchema = TSchema, ParsedMetadata extends object = object, - ApplicationErrorCode extends string = never, -> extends Transport { + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], +> extends Transport { /** * The options for this transport. */ @@ -50,7 +51,7 @@ export abstract class ServerTransport< handshakeExtensions?: ServerHandshakeOptions< MetadataSchema, ParsedMetadata, - ApplicationErrorCode + RejectionCodeSchemas >; /** @@ -81,7 +82,7 @@ export abstract class ServerTransport< options: ServerHandshakeOptions< MetadataSchema, ParsedMetadata, - ApplicationErrorCode + RejectionCodeSchemas >, ) { this.handshakeExtensions = options; @@ -89,7 +90,7 @@ export abstract class ServerTransport< private isApplicationRejectionCode( value: unknown, - ): value is LiteralErrorCode { + ): value is ApplicationErrorCode { return ( this.handshakeExtensions?.rejectionCodeSchemas?.some((schema) => Value.Check(schema, value), @@ -234,7 +235,7 @@ export abstract class ServerTransport< this.teardownForFailedRehandshake( session, 're-handshake metadata rejected by handshake handler', - parsedMetadataOrFailureCode as HandshakeErrorCode, + parsedMetadataOrFailureCode as HandshakeErrorCode, ); return; @@ -268,7 +269,7 @@ export abstract class ServerTransport< private teardownForFailedRehandshake( session: ServerSession, reason: string, - code: HandshakeErrorCode = 'REJECTED_BY_CUSTOM_HANDLER', + code: HandshakeErrorCode = 'REJECTED_BY_CUSTOM_HANDLER', ) { if (session._isConsumed) { return; @@ -377,7 +378,7 @@ export abstract class ServerTransport< session: SessionWaitingForHandshake, to: TransportClientId, reason: string, - code: HandshakeErrorCode, + code: HandshakeErrorCode, metadata: MessageMetadata, ) { session.conn.telemetry?.span.setStatus({ @@ -531,7 +532,7 @@ export abstract class ServerTransport< session, msg.from, 'rejected by handshake handler', - parsedMetadataOrFailureCode as HandshakeErrorCode, + parsedMetadataOrFailureCode as HandshakeErrorCode, { ...session.loggingMetadata, connectedTo: msg.from, diff --git a/transport/transport.test.ts b/transport/transport.test.ts index f597c67b..7f13166f 100644 --- a/transport/transport.test.ts +++ b/transport/transport.test.ts @@ -1968,7 +1968,7 @@ describe.each(testMatrix())( createServerHandshakeOptions< typeof schema, ParsedMetadata, - ApplicationErrorCode + typeof rejectionCodeSchemas >( schema, // @ts-expect-error only declared rejection codes may be returned @@ -1982,11 +1982,15 @@ describe.each(testMatrix())( const broadCodeSchemas = [Type.Literal(broadCode)]; expect( - createServerHandshakeOptions( + createServerHandshakeOptions< + typeof schema, + ParsedMetadata, + typeof broadCodeSchemas + >( schema, - async () => ({ foo: 'foo' }), + // @ts-expect-error widened literal schemas contribute no codes + async () => broadCode, undefined, - // @ts-expect-error application error schemas must use literals broadCodeSchemas, ), ).toBeDefined(); @@ -1997,7 +2001,7 @@ describe.each(testMatrix())( const serverTransport = getServerTransport< typeof schema, ParsedMetadata, - ApplicationErrorCode + typeof rejectionCodeSchemas >('SERVER', { schema, validate: parse, @@ -2056,9 +2060,9 @@ describe.each(testMatrix())( foo: Type.String(), }); - const rejectionCodeSchema = Type.Literal('REPL_NOT_FOUND'); + const rejectionCodeSchemas = [Type.Literal('REPL_NOT_FOUND')] as const; - type ApplicationErrorCode = Static; + type ApplicationErrorCode = Static<(typeof rejectionCodeSchemas)[number]>; interface ParsedMetadata { foo: string; @@ -2067,11 +2071,11 @@ describe.each(testMatrix())( const serverTransport = getServerTransport< typeof schema, ParsedMetadata, - ApplicationErrorCode + typeof rejectionCodeSchemas >('SERVER', { schema, validate: async (): Promise => 'REPL_NOT_FOUND', - rejectionCodeSchemas: [rejectionCodeSchema], + rejectionCodeSchemas, }); // the client did not register the application code: it must treat the @@ -2100,7 +2104,9 @@ describe.each(testMatrix())( expect(clientHandshakeFailed).not.toHaveBeenCalled(); await testFinishesCleanly({ clientTransports: [clientTransport] }); - await testFinishesCleanly({ serverTransport }); + await testFinishesCleanly({ + serverTransport, + }); }); }, ); diff --git a/transport/transport.ts b/transport/transport.ts index eb41cf08..81f19960 100644 --- a/transport/transport.ts +++ b/transport/transport.ts @@ -1,4 +1,5 @@ import { + type ApplicationErrorCodeSchemas, OpaqueTransportMessage, PartialTransportMessage, TransportClientId, @@ -81,7 +82,7 @@ export interface SessionBackpressure { */ export abstract class Transport< ConnType extends Connection, - ApplicationErrorCode extends string = never, + RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], > { /** * The status of the transport. @@ -96,7 +97,7 @@ export abstract class Transport< /** * The event dispatcher for handling events of type EventTypes. */ - eventDispatcher: EventDispatcher; + eventDispatcher: EventDispatcher; /** * The options for this transport. @@ -154,7 +155,7 @@ export abstract class Transport< */ addEventListener< K extends EventTypes, - T extends EventHandler, + T extends EventHandler, >(type: K, handler: T): void { this.eventDispatcher.addEventListener(type, handler); } @@ -166,13 +167,13 @@ export abstract class Transport< */ removeEventListener< K extends EventTypes, - T extends EventHandler, + T extends EventHandler, >(type: K, handler: T): void { this.eventDispatcher.removeEventListener(type, handler); } protected protocolError( - message: EventMap['protocolError'], + message: EventMap['protocolError'], ) { this.eventDispatcher.dispatchEvent('protocolError', message); } From 72ee2125620c53b7c032fd9e80c84d26e476f55d Mon Sep 17 00:00:00 2001 From: Mayank Mehra Date: Mon, 24 Aug 2026 22:27:11 -0700 Subject: [PATCH 09/12] Simplify handshake response schema and rejection checks --- transport/message.ts | 49 ++++++++++++++++++-------------------------- transport/server.ts | 34 ++++++++++++------------------ 2 files changed, 33 insertions(+), 50 deletions(-) diff --git a/transport/message.ts b/transport/message.ts index 6988c8c6..dc2b1fcd 100644 --- a/transport/message.ts +++ b/transport/message.ts @@ -151,31 +151,7 @@ export type HandshakeErrorCode< | Static | ApplicationErrorCode; -export const ControlMessageHandshakeResponseSchema = Type.Object({ - type: Type.Literal('HANDSHAKE_RESP'), - status: Type.Union([ - Type.Object({ - ok: Type.Literal(true), - sessionId: Type.String(), - }), - Type.Object({ - ok: Type.Literal(false), - reason: Type.String(), - code: HandshakeErrorResponseCodes, - }), - ]), -}); - -/** - * A handshake response schema that additionally accepts application-defined - * rejection codes. Both peers must be configured with the same codes: an - * unconfigured peer rejects an application code as a malformed response. - */ -export const ControlMessageHandshakeResponseSchemaWithCodes = < - RejectionCodeSchemas extends ApplicationErrorCodeSchemas, ->( - rejectionCodeSchemas: RejectionCodeSchemas, -) => +const handshakeResponseSchema = (code: Code) => Type.Object({ type: Type.Literal('HANDSHAKE_RESP'), status: Type.Union([ @@ -186,14 +162,29 @@ export const ControlMessageHandshakeResponseSchemaWithCodes = < Type.Object({ ok: Type.Literal(false), reason: Type.String(), - code: Type.Union([ - HandshakeErrorResponseCodes, - ...rejectionCodeSchemas, - ]), + code, }), ]), }); +export const ControlMessageHandshakeResponseSchema = handshakeResponseSchema( + HandshakeErrorResponseCodes, +); + +/** + * A handshake response schema that additionally accepts application-defined + * rejection codes. Both peers must be configured with the same codes: an + * unconfigured peer rejects an application code as a malformed response. + */ +export const ControlMessageHandshakeResponseSchemaWithCodes = < + RejectionCodeSchemas extends ApplicationErrorCodeSchemas, +>( + rejectionCodeSchemas: RejectionCodeSchemas, +) => + handshakeResponseSchema( + Type.Union([HandshakeErrorResponseCodes, ...rejectionCodeSchemas]), + ); + /** * Reserved stream id for the follow-up handshake (re-handshake) control * messages, analogous to the reserved `heartbeat` stream id used for acks. diff --git a/transport/server.ts b/transport/server.ts index 7b3f7365..5c1a6808 100644 --- a/transport/server.ts +++ b/transport/server.ts @@ -22,7 +22,7 @@ import { } from './options'; import { DeleteSessionOptions, Transport } from './transport'; import { coerceErrorString } from './stringifyError'; -import type { TSchema } from 'typebox'; +import type { Static, TSchema } from 'typebox'; import { Value } from 'typebox/value'; import { ProtocolError } from './events'; import { Connection } from './connection'; @@ -88,13 +88,17 @@ export abstract class ServerTransport< this.handshakeExtensions = options; } - private isApplicationRejectionCode( + private isRejectionCode( value: unknown, - ): value is ApplicationErrorCode { + ): value is + | Static + | ApplicationErrorCode { return ( - this.handshakeExtensions?.rejectionCodeSchemas?.some((schema) => + Value.Check(HandshakeErrorCustomHandlerFatalResponseCodes, value) || + (this.handshakeExtensions?.rejectionCodeSchemas?.some((schema) => Value.Check(schema, value), - ) ?? false + ) ?? + false) ); } @@ -225,17 +229,11 @@ export abstract class ServerTransport< return; } - if ( - Value.Check( - HandshakeErrorCustomHandlerFatalResponseCodes, - parsedMetadataOrFailureCode, - ) || - this.isApplicationRejectionCode(parsedMetadataOrFailureCode) - ) { + if (this.isRejectionCode(parsedMetadataOrFailureCode)) { this.teardownForFailedRehandshake( session, 're-handshake metadata rejected by handshake handler', - parsedMetadataOrFailureCode as HandshakeErrorCode, + parsedMetadataOrFailureCode, ); return; @@ -521,18 +519,12 @@ export abstract class ServerTransport< } // handler rejected the connection - if ( - Value.Check( - HandshakeErrorCustomHandlerFatalResponseCodes, - parsedMetadataOrFailureCode, - ) || - this.isApplicationRejectionCode(parsedMetadataOrFailureCode) - ) { + if (this.isRejectionCode(parsedMetadataOrFailureCode)) { this.rejectHandshakeRequest( session, msg.from, 'rejected by handshake handler', - parsedMetadataOrFailureCode as HandshakeErrorCode, + parsedMetadataOrFailureCode, { ...session.loggingMetadata, connectedTo: msg.from, From c858e4abd15e6048dcc2a01aeb2c2099fb1eae42 Mon Sep 17 00:00:00 2001 From: Will Ernst Date: Tue, 25 Aug 2026 10:17:42 -0700 Subject: [PATCH 10/12] Refactor naming, use default generics only for client apis, rejection codes fully in typebox --- __tests__/protobuf.test.ts | 22 +++++++---- protobuf/client.ts | 31 ++++++++++----- protobuf/handshake.ts | 28 +++++++------- protobuf/server.ts | 32 ++++++++++++--- router/client.ts | 14 +++---- router/handshake.ts | 57 +++++++++++++-------------- router/server.ts | 16 ++++---- testUtil/fixtures/cleanup.ts | 31 ++++++++------- testUtil/fixtures/mockTransport.ts | 34 ++++++++-------- testUtil/fixtures/transports.ts | 28 +++++++------- testUtil/index.ts | 37 ++++++++++-------- transport/client.ts | 24 ++++++------ transport/events.ts | 27 +++++++------ transport/impls/ws/client.ts | 6 +-- transport/impls/ws/server.ts | 6 +-- transport/impls/ws/ws.test.ts | 29 ++++++++++++++ transport/message.test.ts | 12 +++--- transport/message.ts | 62 +++++++++++++++++------------- transport/server.ts | 31 ++++++++------- transport/transport.test.ts | 45 +++++++++++----------- transport/transport.ts | 17 ++++---- 21 files changed, 337 insertions(+), 252 deletions(-) diff --git a/__tests__/protobuf.test.ts b/__tests__/protobuf.test.ts index 6ca580a6..12e31c38 100644 --- a/__tests__/protobuf.test.ts +++ b/__tests__/protobuf.test.ts @@ -1,5 +1,6 @@ /* eslint-disable @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-unsafe-call, @typescript-eslint/no-unsafe-assignment, @typescript-eslint/no-unsafe-return, @typescript-eslint/no-unsafe-argument */ import { beforeEach, describe, expect, test, vi } from 'vitest'; +import { Type } from 'typebox'; import { BinaryCodec, NaiveJsonCodec } from '../codec'; import { type ClientError, @@ -298,25 +299,29 @@ describe.each(protobufRouterMatrix)( }); test('protobuf handshake metadata is decoded through the router helpers', async () => { + const rejectionCodeSchema = Type.Union([Type.Literal('TOKEN_EXPIRED')]); const clientHandshakeOptions = createClientHandshakeOptions( AuthHandshakeSchema, () => ({ token: 'let-me-in' }), + undefined, + rejectionCodeSchema, ); const serverHandshakeOptions = createServerHandshakeOptions( AuthHandshakeSchema, (metadata) => ({ token: metadata.token, }), + undefined, + rejectionCodeSchema, ); - const clientTransport = getClientTransport( - 'client', - clientHandshakeOptions, - ); - const serverTransport = getServerTransport( - 'SERVER', - serverHandshakeOptions, - ); + const clientTransport = + getClientTransport('client'); + const serverTransport = getServerTransport< + (typeof serverHandshakeOptions)['schema'], + { token: string }, + typeof rejectionCodeSchema + >('SERVER', undefined); const TypedProtoService = createProtoService(); const testSvc = TypedProtoService.define(TestService, { echo: (request, ctx) => @@ -335,6 +340,7 @@ describe.each(protobufRouterMatrix)( TestService, clientTransport, serverTransport.clientId, + { handshakeOptions: clientHandshakeOptions }, ); await expect(client.echo({ text: 'hello' })).resolves.toMatchObject({ diff --git a/protobuf/client.ts b/protobuf/client.ts index 9984759e..00947531 100644 --- a/protobuf/client.ts +++ b/protobuf/client.ts @@ -6,6 +6,7 @@ import type { MessageInitShape, MessageShape, } from '@bufbuild/protobuf'; +import type { TSchema } from 'typebox'; import { Value } from 'typebox/value'; import { ClientTransport } from '../transport/client'; import { Connection } from '../transport/connection'; @@ -13,6 +14,7 @@ import { EventMap } from '../transport/events'; import { ControlFlags, ControlMessageCloseSchema, + type CustomHandshakeErrorCodeSchema, OpaqueTransportMessage, TransportClientId, cancelMessage, @@ -75,13 +77,16 @@ interface StartedMethodCall< /** * Creates a protobuf client for a single protobuf service descriptor. */ -export function createClient( +export function createClient< + Service extends DescService, + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema = never, +>( service: Service, - transport: ClientTransport, + transport: ClientTransport, serverId: TransportClientId, providedClientOptions: Partial< ClientOptions & { - handshakeOptions: ClientHandshakeOptions; + handshakeOptions: ClientHandshakeOptions; } > = {}, ): ProtobufClient { @@ -111,10 +116,13 @@ export function createClient( return client as ProtobufClient; } -function createMethodCaller( +function createMethodCaller< + Method extends DescMethod, + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema, +>( service: DescService, method: Method, - transport: ClientTransport, + transport: ClientTransport, serverId: TransportClientId, clientOptions: ClientOptions, ): ClientMethod { @@ -235,9 +243,11 @@ function createMethodCaller( } } -function connectOnInvokeIfNeeded( +function connectOnInvokeIfNeeded< + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema, +>( clientOptions: ClientOptions, - transport: ClientTransport, + transport: ClientTransport, serverId: TransportClientId, ) { if (clientOptions.connectOnInvoke && !transport.sessions.has(serverId)) { @@ -245,10 +255,13 @@ function connectOnInvokeIfNeeded( } } -function startMethodCall( +function startMethodCall< + Method extends DescMethod, + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema, +>( service: DescService, method: Method, - transport: ClientTransport, + transport: ClientTransport, serverId: TransportClientId, initialPayload: Uint8Array, procClosesWithInit: boolean, diff --git a/protobuf/handshake.ts b/protobuf/handshake.ts index 05c90fbf..c87d2d39 100644 --- a/protobuf/handshake.ts +++ b/protobuf/handshake.ts @@ -11,8 +11,8 @@ import { type ServerHandshakeOptions, } from '../router/handshake'; import { - type ApplicationErrorCode, - type ApplicationErrorCodeSchemas, + type CustomHandshakeErrorCode, + type CustomHandshakeErrorCodeSchema, HandshakeErrorCustomHandlerFatalResponseCodes, type TransportClientId, } from '../transport/message'; @@ -32,7 +32,7 @@ type ConstructHandshake = () => type ValidateHandshake< Schema extends DescMessage, ParsedMetadata, - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema, > = ( metadata: MessageShape, previousParsedMetadata?: ParsedMetadata, @@ -40,11 +40,11 @@ type ValidateHandshake< ) => | ParsedMetadata | ProtobufHandshakeFailureCode - | ApplicationErrorCode + | CustomHandshakeErrorCode | Promise< | ParsedMetadata | ProtobufHandshakeFailureCode - | ApplicationErrorCode + | CustomHandshakeErrorCode >; /** @@ -52,13 +52,13 @@ type ValidateHandshake< */ export function createClientHandshakeOptions< Schema extends DescMessage, - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema = never, >( schema: Schema, construct: ConstructHandshake, eager?: boolean, - rejectionCodeSchemas?: RejectionCodeSchemas, -): ClientHandshakeOptions { + rejectionCodeSchema?: RejectionCodeSchema, +): ClientHandshakeOptions { return createTransportClientHandshakeOptions( HandshakeBytesSchema, async () => { @@ -67,7 +67,7 @@ export function createClientHandshakeOptions< return encodeMessageBytes(schema, metadata); }, eager, - rejectionCodeSchemas, + rejectionCodeSchema, ); } @@ -77,20 +77,20 @@ export function createClientHandshakeOptions< export function createServerHandshakeOptions< Schema extends DescMessage, ParsedMetadata extends object = object, - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema = never, >( schema: Schema, validate: ValidateHandshake< Schema, ParsedMetadata, - NoInfer + NoInfer >, expiry?: (parsedMetadata: ParsedMetadata) => Date | undefined, - rejectionCodeSchemas?: RejectionCodeSchemas, + rejectionCodeSchema?: RejectionCodeSchema, ): ServerHandshakeOptions< typeof HandshakeBytesSchema, ParsedMetadata, - RejectionCodeSchemas + RejectionCodeSchema > { return createTransportServerHandshakeOptions( HandshakeBytesSchema, @@ -105,6 +105,6 @@ export function createServerHandshakeOptions< return await validate(decoded, previousParsedMetadata, from); }, expiry, - rejectionCodeSchemas, + rejectionCodeSchema, ); } diff --git a/protobuf/server.ts b/protobuf/server.ts index 164e65f1..f848aa13 100644 --- a/protobuf/server.ts +++ b/protobuf/server.ts @@ -16,6 +16,7 @@ import { EventMap } from '../transport/events'; import { ControlFlags, ControlMessageCloseSchema, + type CustomHandshakeErrorCodeSchema, OpaqueTransportMessage, TransportClientId, cancelMessage, @@ -133,11 +134,13 @@ export type Middleware = ( export interface ServerOptions< MetadataSchema extends TSchema, ParsedMetadata extends object, + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema = never, > { readonly extendedContext?: object; readonly handshakeOptions?: ServerHandshakeOptions< MetadataSchema, - ParsedMetadata + ParsedMetadata, + RejectionCodeSchema >; readonly middlewares?: Array>; readonly maxCancelledStreamTombstonesPerSession?: number; @@ -146,6 +149,7 @@ export interface ServerOptions< class ProtobufServer< MetadataSchema extends TSchema, ParsedMetadata extends object, + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema, > implements Server { readonly streams: Map; @@ -153,7 +157,8 @@ class ProtobufServer< private readonly transport: ServerTransport< Connection, MetadataSchema, - ParsedMetadata + ParsedMetadata, + RejectionCodeSchema >; private readonly methods: Map; @@ -171,9 +176,18 @@ class ProtobufServer< private unregisterTransportListeners: () => void; constructor( - transport: ServerTransport, + transport: ServerTransport< + Connection, + MetadataSchema, + ParsedMetadata, + RejectionCodeSchema + >, services: ReadonlyArray, - options: ServerOptions = {}, + options: ServerOptions< + MetadataSchema, + ParsedMetadata, + RejectionCodeSchema + > = {}, ) { this.transport = transport; this.log = transport.log; @@ -979,10 +993,16 @@ class LRUSet { export function createServer< MetadataSchema extends TSchema, ParsedMetadata extends object, + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema = never, >( - transport: ServerTransport, + transport: ServerTransport< + Connection, + MetadataSchema, + ParsedMetadata, + RejectionCodeSchema + >, services: ReadonlyArray, - options?: ServerOptions, + options?: ServerOptions, ): Server { return new ProtobufServer(transport, services, options); } diff --git a/router/client.ts b/router/client.ts index 9f1c26ed..a14d4454 100644 --- a/router/client.ts +++ b/router/client.ts @@ -17,7 +17,7 @@ import { isStreamCancel, closeStreamMessage, cancelMessage, - type ApplicationErrorCodeSchemas, + type CustomHandshakeErrorCodeSchema, } from '../transport/message'; import type { Static, TSchema } from 'typebox'; import { Err, Result, AnyResultSchema } from './result'; @@ -245,13 +245,13 @@ const defaultClientOptions: ClientOptions = { export function createClient< // eslint-disable-next-line @typescript-eslint/no-explicit-any ServiceSchemaMap extends AnyServiceSchemaMap, - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema = never, >( - transport: ClientTransport, + transport: ClientTransport, serverId: TransportClientId, providedClientOptions: Partial< ClientOptions & { - handshakeOptions: ClientHandshakeOptions; + handshakeOptions: ClientHandshakeOptions; } > = {}, ): Client { @@ -322,11 +322,9 @@ type AnyProcReturn = | ReturnType> | ReturnType>; -function handleProc< - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], ->( +function handleProc( procType: ValidProcType, - transport: ClientTransport, + transport: ClientTransport, serverId: TransportClientId, init: Static, serviceName: string, diff --git a/router/handshake.ts b/router/handshake.ts index 3c603f12..1b41a0e0 100644 --- a/router/handshake.ts +++ b/router/handshake.ts @@ -1,7 +1,7 @@ import type { Static, TSchema } from 'typebox'; import { - type ApplicationErrorCode, - type ApplicationErrorCodeSchemas, + type CustomHandshakeErrorCode, + type CustomHandshakeErrorCodeSchema, HandshakeErrorCustomHandlerFatalResponseCodes, type TransportClientId, } from '../transport/message'; @@ -13,24 +13,24 @@ type ConstructHandshake = () => type ValidateHandshake< T extends TSchema, ParsedMetadata, - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema, > = ( metadata: Static, previousParsedMetadata?: ParsedMetadata, from?: TransportClientId, ) => | Static - | ApplicationErrorCode + | CustomHandshakeErrorCode | ParsedMetadata | Promise< | Static - | ApplicationErrorCode + | CustomHandshakeErrorCode | ParsedMetadata >; export interface ClientHandshakeOptions< MetadataSchema extends TSchema = TSchema, - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema = never, > { /** * Schema for the metadata that the client sends to the server @@ -39,13 +39,13 @@ export interface ClientHandshakeOptions< schema: MetadataSchema; /** - * Application-defined rejection codes the server may answer the handshake + * Custom rejection codes the server may answer the handshake * with, sent in the response's `code` field. Must match the server's - * {@link ServerHandshakeOptions.rejectionCodeSchemas}: an unconfigured code - * is rejected as a malformed handshake response. Pass a tuple of TypeBox - * literals (`Type.Literal(...)`) with `as const`. + * {@link ServerHandshakeOptions.rejectionCodeSchema}: an unconfigured code + * is rejected as a malformed handshake response. Pass a TypeBox union of + * literals. */ - rejectionCodeSchemas?: RejectionCodeSchemas; + rejectionCodeSchema?: RejectionCodeSchema; /** * Gets the {@link HandshakeRequestMetadata} to send to the server. @@ -66,7 +66,7 @@ export interface ClientHandshakeOptions< export interface ServerHandshakeOptions< MetadataSchema extends TSchema = TSchema, ParsedMetadata extends object = object, - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema = never, > { /** * Schema for the metadata that the server receives from the client @@ -75,14 +75,13 @@ export interface ServerHandshakeOptions< schema: MetadataSchema; /** - * Application-defined rejection codes that {@link validate} may return. + * Custom rejection codes that {@link validate} may return. * They travel in the handshake response's `code` field and are fatal like * the built-in custom-handler codes. Clients must register the same codes - * in {@link ClientHandshakeOptions.rejectionCodeSchemas} or they reject the - * response as malformed. Pass a tuple of TypeBox literals - * (`Type.Literal(...)`) with `as const`. + * in {@link ClientHandshakeOptions.rejectionCodeSchema} or they reject the + * response as malformed. Pass a TypeBox union of literals. */ - rejectionCodeSchemas?: RejectionCodeSchemas; + rejectionCodeSchema?: RejectionCodeSchema; /** * Parses the metadata sent by the client during the handshake into the @@ -99,7 +98,7 @@ export interface ServerHandshakeOptions< validate: ValidateHandshake< MetadataSchema, ParsedMetadata, - NoInfer + NoInfer >; /** @@ -117,33 +116,29 @@ export interface ServerHandshakeOptions< export function createClientHandshakeOptions< MetadataSchema extends TSchema = TSchema, - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema = never, >( schema: MetadataSchema, construct: ConstructHandshake, eager?: boolean, - rejectionCodeSchemas?: RejectionCodeSchemas, -): ClientHandshakeOptions { - return { schema, construct, eager, rejectionCodeSchemas }; + rejectionCodeSchema?: RejectionCodeSchema, +): ClientHandshakeOptions { + return { schema, construct, eager, rejectionCodeSchema }; } export function createServerHandshakeOptions< MetadataSchema extends TSchema = TSchema, ParsedMetadata extends object = object, - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema = never, >( schema: MetadataSchema, validate: ValidateHandshake< MetadataSchema, ParsedMetadata, - NoInfer + NoInfer >, expiry?: (parsedMetadata: ParsedMetadata) => Date | undefined, - rejectionCodeSchemas?: RejectionCodeSchemas, -): ServerHandshakeOptions< - MetadataSchema, - ParsedMetadata, - RejectionCodeSchemas -> { - return { schema, validate, expiry, rejectionCodeSchemas }; + rejectionCodeSchema?: RejectionCodeSchema, +): ServerHandshakeOptions { + return { schema, validate, expiry, rejectionCodeSchema }; } diff --git a/router/server.ts b/router/server.ts index 84ddd7ce..b8ba4616 100644 --- a/router/server.ts +++ b/router/server.ts @@ -29,7 +29,7 @@ import { cancelMessage, ProtocolVersion, TransportClientId, - type ApplicationErrorCodeSchemas, + type CustomHandshakeErrorCodeSchema, } from '../transport/message'; import { ProcedureHandlerContext } from './context'; import { Logger } from '../logging/log'; @@ -113,14 +113,14 @@ class RiverServer< MetadataSchema extends TSchema, ParsedMetadata extends object, Services extends AnyServiceSchemaMap, - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema, > implements Server { private transport: ServerTransport< Connection, MetadataSchema, ParsedMetadata, - RejectionCodeSchemas + RejectionCodeSchema >; private contextMap: Map; @@ -152,13 +152,13 @@ class RiverServer< Connection, MetadataSchema, ParsedMetadata, - RejectionCodeSchemas + RejectionCodeSchema >, services: Services, handshakeOptions?: ServerHandshakeOptions< MetadataSchema, ParsedMetadata, - RejectionCodeSchemas + RejectionCodeSchema >, extendedContext?: Context, maxCancelledStreamTombstonesPerSession = 200, @@ -1179,20 +1179,20 @@ export function createServer< // eslint-disable-next-line @typescript-eslint/no-explicit-any Services extends AnyServiceSchemaMap, Context extends MaybeDisposable = MaybeDisposable, - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema = never, >( transport: ServerTransport< Connection, MetadataSchema, ParsedMetadata, - RejectionCodeSchemas + RejectionCodeSchema >, services: Services, providedServerOptions?: Partial<{ handshakeOptions?: ServerHandshakeOptions< MetadataSchema, ParsedMetadata, - RejectionCodeSchemas + RejectionCodeSchema >; extendedContext?: Context; /** diff --git a/testUtil/fixtures/cleanup.ts b/testUtil/fixtures/cleanup.ts index 8ad5e3e7..32d8abb5 100644 --- a/testUtil/fixtures/cleanup.ts +++ b/testUtil/fixtures/cleanup.ts @@ -8,9 +8,10 @@ import { import { Server } from '../../router'; import { AnyServiceSchemaMap, MaybeDisposable } from '../../router/services'; import { numberOfConnections, testingSessionOptions } from '..'; +import type { TSchema } from 'typebox'; import { Value } from 'typebox/value'; import { - type ApplicationErrorCodeSchemas, + type CustomHandshakeErrorCodeSchema, ControlMessageAckSchema, } from '../../transport/message'; @@ -40,8 +41,8 @@ export async function advanceFakeTimersByConnectionBackoff() { } export async function ensureTransportIsClean< - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], ->(t: Transport) { + HandshakeFailureCode extends string, +>(t: Transport) { await advanceFakeTimersBySessionGrace(); await waitFor(() => expect( @@ -62,8 +63,8 @@ export function waitFor(cb: () => T | Promise) { } export async function ensureTransportBuffersAreEventuallyEmpty< - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], ->(t: Transport) { + HandshakeFailureCode extends string, +>(t: Transport) { // wait for send buffers to be flushed // ignore heartbeat messages await waitFor(() => @@ -104,8 +105,8 @@ export async function ensureServerIsClean( export async function cleanupTransports< ConnType extends Connection, - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], ->(transports: Array>) { + HandshakeFailureCode extends string, +>(transports: Array>) { for (const t of transports) { if (t.getStatus() !== 'closed') { t.log?.info('*** end of test cleanup ***', { clientId: t.clientId }); @@ -115,17 +116,21 @@ export async function cleanupTransports< } export async function testFinishesCleanly< - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + MetadataSchema extends TSchema, + ParsedMetadata extends object, + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema, >({ clientTransports, serverTransport, server, }: Partial<{ - clientTransports: Array>; - // MetadataSchema and ParsedMetadata are not used in this test, - // so we can safely use any here - // eslint-disable-next-line @typescript-eslint/no-explicit-any - serverTransport: ServerTransport; + clientTransports: Array>; + serverTransport: ServerTransport< + Connection, + MetadataSchema, + ParsedMetadata, + RejectionCodeSchema + >; server: Server; }>) { // pre-close invariants diff --git a/testUtil/fixtures/mockTransport.ts b/testUtil/fixtures/mockTransport.ts index b6b8c952..fd7369f9 100644 --- a/testUtil/fixtures/mockTransport.ts +++ b/testUtil/fixtures/mockTransport.ts @@ -1,5 +1,5 @@ import { Transport, TransportClientId } from '../../transport'; -import type { ApplicationErrorCodeSchemas } from '../../transport/message'; +import type { CustomHandshakeErrorCodeSchema } from '../../transport/message'; import { ClientTransport } from '../../transport/client'; import { Connection } from '../../transport/connection'; import { ServerTransport } from '../../transport/server'; @@ -75,12 +75,10 @@ export function createMockTransportNetwork( // conn id -> [client->server, server->client] const connections = new Observable>({}); - const transports: Array< - Transport - > = []; + const transports: Array> = []; class MockClientTransport< - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], - > extends ClientTransport { + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema, + > extends ClientTransport { async createNewOutgoingConnection( to: TransportClientId, ): Promise { @@ -105,14 +103,14 @@ export function createMockTransportNetwork( } class MockServerTransport< - MetadataSchema extends TSchema = TSchema, - ParsedMetadata extends object = object, - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + MetadataSchema extends TSchema, + ParsedMetadata extends object, + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema, > extends ServerTransport< InMemoryConnection, MetadataSchema, ParsedMetadata, - RejectionCodeSchemas + RejectionCodeSchema > { subscribeCleanup: () => void; @@ -146,12 +144,12 @@ export function createMockTransportNetwork( return { getClientTransport: < - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema, >( id: TransportClientId, - handshakeOptions?: ClientHandshakeOptions, + handshakeOptions?: ClientHandshakeOptions, ) => { - const clientTransport = new MockClientTransport( + const clientTransport = new MockClientTransport( id, opts?.client, ); @@ -164,23 +162,23 @@ export function createMockTransportNetwork( return clientTransport; }, getServerTransport: < - MetadataSchema extends TSchema = TSchema, - ParsedMetadata extends object = object, - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + MetadataSchema extends TSchema, + ParsedMetadata extends object, + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema, >( id = 'SERVER', handshakeOptions: | ServerHandshakeOptions< MetadataSchema, ParsedMetadata, - RejectionCodeSchemas + RejectionCodeSchema > | undefined, ) => { const serverTransport = new MockServerTransport< MetadataSchema, ParsedMetadata, - RejectionCodeSchemas + RejectionCodeSchema >(id, opts?.server); if (handshakeOptions) { serverTransport.extendHandshake(handshakeOptions); diff --git a/testUtil/fixtures/transports.ts b/testUtil/fixtures/transports.ts index abb5d40a..e4a4b3a5 100644 --- a/testUtil/fixtures/transports.ts +++ b/testUtil/fixtures/transports.ts @@ -17,7 +17,7 @@ import { ProvidedServerTransportOptions, } from '../../transport/options'; import { - type ApplicationErrorCodeSchemas, + type CustomHandshakeErrorCodeSchema, TransportClientId, } from '../../transport/message'; import { ClientTransport } from '../../transport/client'; @@ -34,27 +34,27 @@ export interface TestTransportOptions { export interface TestSetupHelpers { getClientTransport: < - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema = never, >( id: TransportClientId, - handshakeOptions?: ClientHandshakeOptions, - ) => ClientTransport; + handshakeOptions?: ClientHandshakeOptions, + ) => ClientTransport; getServerTransport: < MetadataSchema extends TSchema = TSchema, ParsedMetadata extends object = object, - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema = never, >( id?: TransportClientId, handshakeOptions?: ServerHandshakeOptions< MetadataSchema, ParsedMetadata, - RejectionCodeSchemas + RejectionCodeSchema >, ) => ServerTransport< Connection, MetadataSchema, ParsedMetadata, - RejectionCodeSchemas + RejectionCodeSchema >; simulatePhantomDisconnect: () => void; restartServer: () => Promise; @@ -89,16 +89,16 @@ export const transports: Array = [ } }, getClientTransport: < - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema, >( id: TransportClientId, handshakeOptions?: ClientHandshakeOptions< TSchema, - RejectionCodeSchemas + RejectionCodeSchema >, ) => { const clientTransport = - new WebSocketClientTransport( + new WebSocketClientTransport( () => Promise.resolve(createLocalWebSocketClient(port)), id, opts?.client, @@ -124,21 +124,21 @@ export const transports: Array = [ getServerTransport: < MetadataSchema extends TSchema, ParsedMetadata extends object, - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema, >( id = 'SERVER', handshakeOptions: | ServerHandshakeOptions< MetadataSchema, ParsedMetadata, - RejectionCodeSchemas + RejectionCodeSchema > | undefined, ) => { const serverTransport = new WebSocketServerTransport< MetadataSchema, ParsedMetadata, - RejectionCodeSchemas + RejectionCodeSchema >(wss, id, opts?.server); serverTransport.bindLogger((msg, ctx, level) => { @@ -160,7 +160,7 @@ export const transports: Array = [ Connection, MetadataSchema, ParsedMetadata, - RejectionCodeSchemas + RejectionCodeSchema >; }, async restartServer() { diff --git a/testUtil/index.ts b/testUtil/index.ts index da0c883c..a113fd7b 100644 --- a/testUtil/index.ts +++ b/testUtil/index.ts @@ -2,6 +2,7 @@ import NodeWs, { WebSocketServer } from 'ws'; import http from 'node:http'; import type { Static } from 'typebox'; import { + type CustomHandshakeErrorCodeSchema, OpaqueTransportMessage, PartialTransportMessage, currentProtocolVersion, @@ -18,7 +19,6 @@ import { SessionState } from '../transport/sessionStateMachine/common'; import { SessionStateGraph } from '../transport/sessionStateMachine/transitions'; import { BaseErrorSchemaType } from '../router/errors'; import { ClientTransport } from '../transport/client'; -import { ServerTransport } from '../transport/server'; import { getTracer } from '../tracing'; export { @@ -194,9 +194,11 @@ export function dummySession() { ); } -export function getClientSendFn( - clientTransport: ClientTransport, - serverTransport: ServerTransport, +export function getClientSendFn< + ClientRejectionCodeSchema extends CustomHandshakeErrorCodeSchema, +>( + clientTransport: ClientTransport, + serverTransport: { clientId: string }, ) { const session = clientTransport.sessions.get(serverTransport.clientId) ?? @@ -208,9 +210,9 @@ export function getClientSendFn( ); } -export function getServerSendFn( - serverTransport: ServerTransport, - clientTransport: ClientTransport, +export function getServerSendFn( + serverTransport: Transport, + clientTransport: { clientId: string }, ) { const session = serverTransport.sessions.get(clientTransport.clientId); if (!session) { @@ -223,9 +225,10 @@ export function getServerSendFn( ); } -export function getTransportConnections( - transport: Transport, -): Array { +export function getTransportConnections< + ConnType extends Connection, + HandshakeFailureCode extends string, +>(transport: Transport): Array { const connections = []; for (const session of transport.sessions.values()) { if (session.state === SessionState.Connected) { @@ -236,15 +239,17 @@ export function getTransportConnections( return connections; } -export function numberOfConnections( - transport: Transport, -): number { +export function numberOfConnections< + ConnType extends Connection, + HandshakeFailureCode extends string, +>(transport: Transport): number { return getTransportConnections(transport).length; } -export function closeAllConnections( - transport: Transport, -) { +export function closeAllConnections< + ConnType extends Connection, + HandshakeFailureCode extends string, +>(transport: Transport) { for (const conn of getTransportConnections(transport)) { conn.close(); } diff --git a/transport/client.ts b/transport/client.ts index fde88058..83e1a9b4 100644 --- a/transport/client.ts +++ b/transport/client.ts @@ -5,7 +5,7 @@ import { ControlMessageHandshakeResponseSchema, ControlMessageHandshakeResponseSchemaWithCodes, ControlMessageRehandshakeRequestSchema, - type ApplicationErrorCodeSchemas, + type CustomHandshakeErrorCodeSchema, type HandshakeErrorCode, HandshakeErrorRetriableResponseCodes, OpaqueTransportMessage, @@ -54,8 +54,8 @@ type ConstructedHandshakeMetadata = export abstract class ClientTransport< ConnType extends Connection, - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], -> extends Transport { + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema = never, +> extends Transport> { /** * The options for this transport. */ @@ -74,10 +74,10 @@ export abstract class ClientTransport< /** * Optional handshake options for this client. */ - handshakeExtensions?: ClientHandshakeOptions; + handshakeExtensions?: ClientHandshakeOptions; /** - * Handshake response schema extended with the application-defined + * Handshake response schema extended with the custom * rejection codes, when any are registered. */ protected handshakeResponseSchema: @@ -112,17 +112,17 @@ export abstract class ClientTransport< this.retryBudget = new LeakyBucketRateLimit(this.options); } - extendHandshake( - options: ClientHandshakeOptions, - ) { + extendHandshake = ( + options: ClientHandshakeOptions, + ) => { this.handshakeExtensions = options; - if (options.rejectionCodeSchemas?.length) { + if (options.rejectionCodeSchema) { this.handshakeResponseSchema = ControlMessageHandshakeResponseSchemaWithCodes( - options.rejectionCodeSchemas, + options.rejectionCodeSchema, ); } - } + }; protected handleRehandshakeMessage(message: OpaqueTransportMessage): void { if (!Value.Check(ControlMessageRehandshakeRequestSchema, message.payload)) { @@ -411,7 +411,7 @@ export abstract class ClientTransport< this.protocolError({ type: ProtocolError.HandshakeFailed, code: msg.payload.status - .code as HandshakeErrorCode, + .code as HandshakeErrorCode, message: reason, }); } diff --git a/transport/events.ts b/transport/events.ts index 9e90ce5c..df708516 100644 --- a/transport/events.ts +++ b/transport/events.ts @@ -1,7 +1,6 @@ import { Connection } from './connection'; import type { - ApplicationErrorCodeSchemas, - HandshakeErrorCode, + BuiltInHandshakeErrorCode, OpaqueTransportMessage, } from './message'; import { Session, SessionState } from './sessionStateMachine'; @@ -19,8 +18,12 @@ export const ProtocolError = { export type ProtocolErrorType = (typeof ProtocolError)[keyof typeof ProtocolError]; +/** + * Transport events. `HandshakeFailureCode` is the full set of codes observable + * on handshake-failed protocol errors, including built-in and custom codes. + */ export interface EventMap< - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + HandshakeFailureCode extends string = BuiltInHandshakeErrorCode, > { message: OpaqueTransportMessage; sessionStatus: @@ -41,7 +44,7 @@ export interface EventMap< protocolError: | { type: (typeof ProtocolError)['HandshakeFailed']; - code: HandshakeErrorCode; + code: HandshakeFailureCode; message: string; } | { @@ -59,15 +62,15 @@ export interface EventMap< export type EventTypes = keyof EventMap; export type EventHandler< K extends EventTypes, - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], -> = (event: EventMap[K]) => unknown; + HandshakeFailureCode extends string = BuiltInHandshakeErrorCode, +> = (event: EventMap[K]) => unknown; export class EventDispatcher< T extends EventTypes, - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + HandshakeFailureCode extends string, > { private eventListeners: { - [K in T]?: Set>; + [K in T]?: Set>; } = {}; removeAllListeners() { @@ -80,18 +83,18 @@ export class EventDispatcher< addEventListener( eventType: K, - handler: EventHandler, + handler: EventHandler, ) { if (!this.eventListeners[eventType]) { this.eventListeners[eventType] = new Set(); } - this.eventListeners[eventType]?.add(handler); + this.eventListeners[eventType].add(handler); } removeEventListener( eventType: K, - handler: EventHandler, + handler: EventHandler, ) { const handlers = this.eventListeners[eventType]; if (handlers) { @@ -101,7 +104,7 @@ export class EventDispatcher< dispatchEvent( eventType: K, - event: EventMap[K], + event: EventMap[K], ) { const handlers = this.eventListeners[eventType]; if (handlers) { diff --git a/transport/impls/ws/client.ts b/transport/impls/ws/client.ts index 4e62efc7..27dc6bc0 100644 --- a/transport/impls/ws/client.ts +++ b/transport/impls/ws/client.ts @@ -1,6 +1,6 @@ import { ClientTransport } from '../../client'; import { - type ApplicationErrorCodeSchemas, + type CustomHandshakeErrorCodeSchema, TransportClientId, } from '../../message'; import { ProvidedClientTransportOptions } from '../../options'; @@ -13,8 +13,8 @@ import { WsLike } from './wslike'; * @extends Transport */ export class WebSocketClientTransport< - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], -> extends ClientTransport { + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema = never, +> extends ClientTransport { /** * A function that returns a Promise that resolves to a websocket URL. */ diff --git a/transport/impls/ws/server.ts b/transport/impls/ws/server.ts index fb4bad9c..e1df56eb 100644 --- a/transport/impls/ws/server.ts +++ b/transport/impls/ws/server.ts @@ -1,5 +1,5 @@ import { - type ApplicationErrorCodeSchemas, + type CustomHandshakeErrorCodeSchema, TransportClientId, } from '../../message'; import { WebSocketServer } from 'ws'; @@ -28,12 +28,12 @@ function cleanHeaders( export class WebSocketServerTransport< MetadataSchema extends TSchema = TSchema, ParsedMetadata extends object = object, - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema = never, > extends ServerTransport< WebSocketConnection, MetadataSchema, ParsedMetadata, - RejectionCodeSchemas + RejectionCodeSchema > { wss: WebSocketServer; diff --git a/transport/impls/ws/ws.test.ts b/transport/impls/ws/ws.test.ts index 285e66e7..f1f4c3f2 100644 --- a/transport/impls/ws/ws.test.ts +++ b/transport/impls/ws/ws.test.ts @@ -1,5 +1,6 @@ import http from 'node:http'; import { describe, test, expect, beforeEach } from 'vitest'; +import { Type } from 'typebox'; import { createWebSocketServer, onWsServerReady, @@ -23,6 +24,7 @@ import { import { PartialTransportMessage } from '../../message'; import type NodeWs from 'ws'; import { createPostTestCleanups } from '../../../testUtil/fixtures/cleanup'; +import { createClientHandshakeOptions } from '../../../router/handshake'; describe('sending and receiving across websockets works', async () => { let server: http.Server; @@ -42,6 +44,33 @@ describe('sending and receiving across websockets works', async () => { }; }); + test('custom handshake codes require an explicit transport type', () => { + const rejectionCodeSchema = Type.Union([Type.Literal('TOKEN_EXPIRED')]); + const handshakeOptions = createClientHandshakeOptions( + Type.Object({}), + () => ({}), + undefined, + rejectionCodeSchema, + ); + const getWs = () => { + throw new Error('not called'); + }; + const defaultTransport = new WebSocketClientTransport(getWs, 'client'); + + // @ts-expect-error a default transport cannot be widened to custom codes + const widenedTransport: WebSocketClientTransport< + typeof rejectionCodeSchema + > = defaultTransport; + + const typedTransport = new WebSocketClientTransport< + typeof rejectionCodeSchema + >(getWs, 'client'); + typedTransport.extendHandshake(handshakeOptions); + + expect(typedTransport.handshakeExtensions).toBe(handshakeOptions); + expect(widenedTransport).toBe(defaultTransport); + }); + test('basic send/receive', async () => { const clientTransport = new WebSocketClientTransport( () => Promise.resolve(createLocalWebSocketClient(port)), diff --git a/transport/message.test.ts b/transport/message.test.ts index d8218390..61f20add 100644 --- a/transport/message.test.ts +++ b/transport/message.test.ts @@ -109,11 +109,13 @@ describe('message helpers', () => { expect(mFail.payload.status.ok).toBe(false); }); - test('handshake response schema with application codes', () => { - const extended = ControlMessageHandshakeResponseSchemaWithCodes([ - Type.Literal('REPL_NOT_FOUND'), - Type.Literal('TOKEN_EXPIRED'), - ] as const); + test('handshake response schema with custom error codes', () => { + const extended = ControlMessageHandshakeResponseSchemaWithCodes( + Type.Union([ + Type.Literal('REPL_NOT_FOUND'), + Type.Literal('TOKEN_EXPIRED'), + ]), + ); const rejection = { type: 'HANDSHAKE_RESP', status: { diff --git a/transport/message.ts b/transport/message.ts index dc2b1fcd..e06e4aac 100644 --- a/transport/message.ts +++ b/transport/message.ts @@ -1,4 +1,10 @@ -import { Type, type TLiteral, type TSchema, type Static } from 'typebox'; +import { + Type, + type TLiteral, + type TSchema, + type TUnion, + type Static, +} from 'typebox'; import { PropagationContext } from '../tracing'; import { generateId } from './id'; // type-only: a value import closes a transport <-> router require cycle @@ -124,32 +130,36 @@ export const HandshakeErrorResponseCodes = Type.Union([ ]); /** - * A tuple of TypeBox literal schemas declaring the application-defined - * handshake rejection codes, e.g. - * `[Type.Literal('REPL_NOT_FOUND'), Type.Literal('TOKEN_EXPIRED')] as const`. + * A TypeBox union of literals declaring the custom handshake rejection codes. */ -export type ApplicationErrorCodeSchemas = ReadonlyArray>; +export type CustomHandshakeErrorCodeSchema = TUnion>>; /** - * The union of codes declared by a tuple of rejection code schemas. - * Widened schemas (`TLiteral` rather than a specific literal) - * contribute no codes. + * The union of codes declared by a rejection code schema. Widened literal + * schemas (`TLiteral` rather than a specific literal) contribute no + * codes. */ -export type ApplicationErrorCode< - RejectionCodeSchemas extends ApplicationErrorCodeSchemas, -> = string extends Static +export type CustomHandshakeErrorCode< + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema, +> = string extends Static ? never - : Static; + : Static; /** - * The protocol-level handshake error codes plus any application-defined - * rejection codes the application registered in its handshake options. + * The protocol-level handshake error codes River can emit without any custom + * rejection codes. + */ +export type BuiltInHandshakeErrorCode = Static< + typeof HandshakeErrorResponseCodes +>; + +/** + * The protocol-level handshake error codes plus any custom rejection codes + * registered in the handshake options. */ export type HandshakeErrorCode< - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], -> = - | Static - | ApplicationErrorCode; + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema, +> = BuiltInHandshakeErrorCode | CustomHandshakeErrorCode; const handshakeResponseSchema = (code: Code) => Type.Object({ @@ -172,17 +182,17 @@ export const ControlMessageHandshakeResponseSchema = handshakeResponseSchema( ); /** - * A handshake response schema that additionally accepts application-defined + * A handshake response schema that additionally accepts custom * rejection codes. Both peers must be configured with the same codes: an - * unconfigured peer rejects an application code as a malformed response. + * unconfigured peer rejects a custom code as a malformed response. */ export const ControlMessageHandshakeResponseSchemaWithCodes = < - RejectionCodeSchemas extends ApplicationErrorCodeSchemas, + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema, >( - rejectionCodeSchemas: RejectionCodeSchemas, + rejectionCodeSchema: RejectionCodeSchema, ) => handshakeResponseSchema( - Type.Union([HandshakeErrorResponseCodes, ...rejectionCodeSchemas]), + Type.Union([HandshakeErrorResponseCodes, rejectionCodeSchema]), ); /** @@ -305,7 +315,7 @@ export function handshakeRequestMessage({ export const SESSION_STATE_MISMATCH = 'session state mismatch'; export function handshakeResponseMessage< - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema, >({ from, to, @@ -313,14 +323,14 @@ export function handshakeResponseMessage< }: { from: TransportClientId; to: TransportClientId; - // the code may be an application-defined rejection code, which is only + // the code may be a custom rejection code, which is only // known to peers that registered it in their handshake options status: | { ok: true; sessionId: string } | { ok: false; reason: string; - code: HandshakeErrorCode; + code: HandshakeErrorCode; }; }): TransportMessage> { return { diff --git a/transport/server.ts b/transport/server.ts index 5c1a6808..6f69fdc2 100644 --- a/transport/server.ts +++ b/transport/server.ts @@ -5,8 +5,8 @@ import { ControlMessageHandshakeRequestSchema, ControlMessageRehandshakeResponseSchema, HandshakeErrorCustomHandlerFatalResponseCodes, - type ApplicationErrorCode, - type ApplicationErrorCodeSchemas, + type CustomHandshakeErrorCode, + type CustomHandshakeErrorCodeSchema, type HandshakeErrorCode, OpaqueTransportMessage, acceptedProtocolVersions, @@ -38,8 +38,8 @@ export abstract class ServerTransport< ConnType extends Connection, MetadataSchema extends TSchema = TSchema, ParsedMetadata extends object = object, - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], -> extends Transport { + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema = never, +> extends Transport> { /** * The options for this transport. */ @@ -51,7 +51,7 @@ export abstract class ServerTransport< handshakeExtensions?: ServerHandshakeOptions< MetadataSchema, ParsedMetadata, - RejectionCodeSchemas + RejectionCodeSchema >; /** @@ -78,27 +78,26 @@ export abstract class ServerTransport< }); } - extendHandshake( + extendHandshake = ( options: ServerHandshakeOptions< MetadataSchema, ParsedMetadata, - RejectionCodeSchemas + RejectionCodeSchema >, - ) { + ) => { this.handshakeExtensions = options; - } + }; private isRejectionCode( value: unknown, ): value is | Static - | ApplicationErrorCode { + | CustomHandshakeErrorCode { return ( Value.Check(HandshakeErrorCustomHandlerFatalResponseCodes, value) || - (this.handshakeExtensions?.rejectionCodeSchemas?.some((schema) => - Value.Check(schema, value), - ) ?? - false) + (this.handshakeExtensions?.rejectionCodeSchema + ? Value.Check(this.handshakeExtensions.rejectionCodeSchema, value) + : false) ); } @@ -267,7 +266,7 @@ export abstract class ServerTransport< private teardownForFailedRehandshake( session: ServerSession, reason: string, - code: HandshakeErrorCode = 'REJECTED_BY_CUSTOM_HANDLER', + code: HandshakeErrorCode = 'REJECTED_BY_CUSTOM_HANDLER', ) { if (session._isConsumed) { return; @@ -376,7 +375,7 @@ export abstract class ServerTransport< session: SessionWaitingForHandshake, to: TransportClientId, reason: string, - code: HandshakeErrorCode, + code: HandshakeErrorCode, metadata: MessageMetadata, ) { session.conn.telemetry?.span.setStatus({ diff --git a/transport/transport.test.ts b/transport/transport.test.ts index 7f13166f..dcecc26e 100644 --- a/transport/transport.test.ts +++ b/transport/transport.test.ts @@ -1947,17 +1947,17 @@ describe.each(testMatrix())( }); }); - test('custom handler can reject with an application-defined code', async () => { + test('custom handler can reject with a registered handshake error code', async () => { const schema = Type.Object({ foo: Type.String(), }); - const rejectionCodeSchemas = [ + const rejectionCodeSchema = Type.Union([ Type.Literal('REPL_NOT_FOUND'), Type.Literal('TOKEN_EXPIRED'), - ] as const; + ]); - type ApplicationErrorCode = Static<(typeof rejectionCodeSchemas)[number]>; + type CustomHandshakeErrorCode = Static; interface ParsedMetadata { foo: string; @@ -1968,50 +1968,50 @@ describe.each(testMatrix())( createServerHandshakeOptions< typeof schema, ParsedMetadata, - typeof rejectionCodeSchemas + typeof rejectionCodeSchema >( schema, // @ts-expect-error only declared rejection codes may be returned async () => 'SOME_OTHER_CODE', undefined, - rejectionCodeSchemas, + rejectionCodeSchema, ), ).toBeDefined(); const broadCode = String('REPL_NOT_FOUND'); - const broadCodeSchemas = [Type.Literal(broadCode)]; + const broadCodeSchema = Type.Union([Type.Literal(broadCode)]); expect( createServerHandshakeOptions< typeof schema, ParsedMetadata, - typeof broadCodeSchemas + typeof broadCodeSchema >( schema, // @ts-expect-error widened literal schemas contribute no codes async () => broadCode, undefined, - broadCodeSchemas, + broadCodeSchema, ), ).toBeDefined(); - const parse = vi.fn(async (): Promise => { + const parse = vi.fn(async (): Promise => { return 'REPL_NOT_FOUND'; }); const serverTransport = getServerTransport< typeof schema, ParsedMetadata, - typeof rejectionCodeSchemas + typeof rejectionCodeSchema >('SERVER', { schema, validate: parse, - rejectionCodeSchemas, + rejectionCodeSchema, }); const clientTransport = getClientTransport('client', { schema, construct: async () => ({ foo: 'foo' }), - rejectionCodeSchemas, + rejectionCodeSchema, }); const clientHandshakeFailed = vi.fn(); @@ -2055,14 +2055,14 @@ describe.each(testMatrix())( }); }); - test('an application code is rejected by an unconfigured client', async () => { + test('an unregistered handshake error code is rejected by an unconfigured client', async () => { const schema = Type.Object({ foo: Type.String(), }); - const rejectionCodeSchemas = [Type.Literal('REPL_NOT_FOUND')] as const; + const rejectionCodeSchema = Type.Union([Type.Literal('REPL_NOT_FOUND')]); - type ApplicationErrorCode = Static<(typeof rejectionCodeSchemas)[number]>; + type CustomHandshakeErrorCode = Static; interface ParsedMetadata { foo: string; @@ -2071,14 +2071,15 @@ describe.each(testMatrix())( const serverTransport = getServerTransport< typeof schema, ParsedMetadata, - typeof rejectionCodeSchemas + typeof rejectionCodeSchema >('SERVER', { schema, - validate: async (): Promise => 'REPL_NOT_FOUND', - rejectionCodeSchemas, + validate: async (): Promise => + 'REPL_NOT_FOUND', + rejectionCodeSchema, }); - // the client did not register the application code: it must treat the + // the client did not register the custom code: it must treat the // response as malformed rather than accept an unknown code const clientTransport = getClientTransport('client', { schema, @@ -2104,9 +2105,7 @@ describe.each(testMatrix())( expect(clientHandshakeFailed).not.toHaveBeenCalled(); await testFinishesCleanly({ clientTransports: [clientTransport] }); - await testFinishesCleanly({ - serverTransport, - }); + await testFinishesCleanly({ serverTransport }); }); }, ); diff --git a/transport/transport.ts b/transport/transport.ts index 81f19960..390f5d47 100644 --- a/transport/transport.ts +++ b/transport/transport.ts @@ -1,5 +1,5 @@ import { - type ApplicationErrorCodeSchemas, + type BuiltInHandshakeErrorCode, OpaqueTransportMessage, PartialTransportMessage, TransportClientId, @@ -82,7 +82,7 @@ export interface SessionBackpressure { */ export abstract class Transport< ConnType extends Connection, - RejectionCodeSchemas extends ApplicationErrorCodeSchemas = [], + HandshakeFailureCode extends string = BuiltInHandshakeErrorCode, > { /** * The status of the transport. @@ -97,7 +97,7 @@ export abstract class Transport< /** * The event dispatcher for handling events of type EventTypes. */ - eventDispatcher: EventDispatcher; + eventDispatcher: EventDispatcher; /** * The options for this transport. @@ -118,7 +118,10 @@ export abstract class Transport< providedOptions?: ProvidedTransportOptions, ) { this.options = { ...defaultTransportOptions, ...providedOptions }; - this.eventDispatcher = new EventDispatcher(); + this.eventDispatcher = new EventDispatcher< + EventTypes, + HandshakeFailureCode + >(); this.clientId = clientId; this.status = 'open'; this.sessions = new Map(); @@ -155,7 +158,7 @@ export abstract class Transport< */ addEventListener< K extends EventTypes, - T extends EventHandler, + T extends EventHandler, >(type: K, handler: T): void { this.eventDispatcher.addEventListener(type, handler); } @@ -167,13 +170,13 @@ export abstract class Transport< */ removeEventListener< K extends EventTypes, - T extends EventHandler, + T extends EventHandler, >(type: K, handler: T): void { this.eventDispatcher.removeEventListener(type, handler); } protected protocolError( - message: EventMap['protocolError'], + message: EventMap['protocolError'], ) { this.eventDispatcher.dispatchEvent('protocolError', message); } From 1d9c1d47bcc2b24a72f99da17a2e0e3d1ab947f9 Mon Sep 17 00:00:00 2001 From: Mayank Mehra Date: Tue, 25 Aug 2026 14:22:42 -0700 Subject: [PATCH 11/12] Fix typecheck for generic-keyed listener map access --- transport/events.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transport/events.ts b/transport/events.ts index df708516..f33fcfa4 100644 --- a/transport/events.ts +++ b/transport/events.ts @@ -89,7 +89,7 @@ export class EventDispatcher< this.eventListeners[eventType] = new Set(); } - this.eventListeners[eventType].add(handler); + this.eventListeners[eventType]?.add(handler); } removeEventListener( From 0838b78ea974b6445fb87b87af49e062248cdaff Mon Sep 17 00:00:00 2001 From: Mayank Mehra Date: Tue, 25 Aug 2026 14:29:40 -0700 Subject: [PATCH 12/12] Simplify event listener initialization --- transport/events.ts | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/transport/events.ts b/transport/events.ts index f33fcfa4..b8b70fa7 100644 --- a/transport/events.ts +++ b/transport/events.ts @@ -85,11 +85,13 @@ export class EventDispatcher< eventType: K, handler: EventHandler, ) { - if (!this.eventListeners[eventType]) { - this.eventListeners[eventType] = new Set(); + let listeners = this.eventListeners[eventType]; + if (!listeners) { + listeners = new Set(); + this.eventListeners[eventType] = listeners; } - this.eventListeners[eventType]?.add(handler); + listeners.add(handler); } removeEventListener(