From f881f9ae324df9e6a91888a72aa78fbf029be9c4 Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Sun, 10 May 2026 15:53:54 +0800 Subject: [PATCH] Refactor relay transport socket tests (#882) --- .../server/src/server/relay-transport.test.ts | 253 +++++++++--------- packages/server/src/server/relay-transport.ts | 37 ++- 2 files changed, 157 insertions(+), 133 deletions(-) diff --git a/packages/server/src/server/relay-transport.test.ts b/packages/server/src/server/relay-transport.test.ts index 517781451..f6abbb12c 100644 --- a/packages/server/src/server/relay-transport.test.ts +++ b/packages/server/src/server/relay-transport.test.ts @@ -1,124 +1,115 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; - -const wsMock = vi.hoisted(() => { - class MockWebSocket { - static readonly CONNECTING = 0; - static readonly OPEN = 1; - static readonly CLOSING = 2; - static readonly CLOSED = 3; - static instances: MockWebSocket[] = []; - - readonly url: string; - readonly options: unknown; - readyState = MockWebSocket.CONNECTING; - sent: string[] = []; - terminateCalls = 0; - private listeners = new Map void>>(); - - constructor(url: string, options?: unknown) { - this.url = url; - this.options = options; - MockWebSocket.instances.push(this); - } - - static reset() { - MockWebSocket.instances = []; - } - - on(event: string, listener: (...args: unknown[]) => void) { - const handlers = this.listeners.get(event) ?? []; - handlers.push(listener); - this.listeners.set(event, handlers); - return this; - } - - once(event: string, listener: (...args: unknown[]) => void) { - const wrapped = (...args: unknown[]) => { - this.off(event, wrapped); - listener(...args); - }; - return this.on(event, wrapped); - } - - close(code?: number, reason?: string) { - this.readyState = MockWebSocket.CLOSED; - this.emit("close", code ?? 1000, reason ?? ""); - } - - terminate() { - this.terminateCalls += 1; - this.readyState = MockWebSocket.CLOSED; - this.emit("close", 1006, ""); - } - - send(data: string) { - if (this.readyState !== MockWebSocket.OPEN) { - throw new Error(`WebSocket not open (readyState=${this.readyState})`); - } - this.sent.push(data); - } - - open() { - this.readyState = MockWebSocket.OPEN; - this.emit("open"); - } - - message(data: unknown) { - this.emit("message", data); - } - - error(err: unknown) { - this.emit("error", err); - } - - private off(event: string, listener: (...args: unknown[]) => void) { - const handlers = this.listeners.get(event) ?? []; - this.listeners.set( - event, - handlers.filter((handler) => handler !== listener), - ); - } - - private emit(event: string, ...args: unknown[]) { - const handlers = this.listeners.get(event) ?? []; - for (const handler of handlers.slice()) { - handler(...args); - } - } - } - - return { MockWebSocket }; -}); - -vi.mock("ws", () => ({ - default: wsMock.MockWebSocket, - WebSocket: wsMock.MockWebSocket, -})); - import type pino from "pino"; import { startRelayTransport } from "./relay-transport"; function createMockLogger() { + const messages: { level: "debug" | "info" | "warn" | "error"; args: unknown[] }[] = []; const logger = { - child: vi.fn(() => logger), - debug: vi.fn(), - info: vi.fn(), - warn: vi.fn(), - error: vi.fn(), + messages, + child: () => logger, + debug: (...args: unknown[]) => messages.push({ level: "debug", args }), + info: (...args: unknown[]) => messages.push({ level: "info", args }), + warn: (...args: unknown[]) => messages.push({ level: "warn", args }), + error: (...args: unknown[]) => messages.push({ level: "error", args }), }; return logger; } -function hasLogMessage(mockFn: ReturnType, message: string): boolean { - return mockFn.mock.calls.some((call) => call.some((arg) => arg === message)); +type TestLogger = ReturnType; + +function hasLogMessage(logger: TestLogger, level: "info" | "warn", message: string): boolean { + return logger.messages.some((entry) => { + return entry.level === level && entry.args.some((arg) => arg === message); + }); +} + +class FakeRelayWebSocket { + static readonly CONNECTING = 0; + static readonly OPEN = 1; + static readonly CLOSED = 3; + + readyState = FakeRelayWebSocket.CONNECTING; + sent: Array = []; + terminateCalls = 0; + private readonly listeners = new Map void>>(); + + constructor(readonly url: string) {} + + on(event: string, listener: (...args: unknown[]) => void) { + const handlers = this.listeners.get(event) ?? []; + handlers.push(listener); + this.listeners.set(event, handlers); + } + + once(event: string, listener: (...args: unknown[]) => void) { + const wrapped = (...args: unknown[]) => { + this.off(event, wrapped); + listener(...args); + }; + this.on(event, wrapped); + } + + close(code?: number, reason?: string) { + this.readyState = FakeRelayWebSocket.CLOSED; + this.emit("close", code ?? 1000, reason ?? ""); + } + + terminate() { + this.terminateCalls += 1; + this.readyState = FakeRelayWebSocket.CLOSED; + this.emit("close", 1006, ""); + } + + send(data: string | Uint8Array | ArrayBuffer) { + if (this.readyState !== FakeRelayWebSocket.OPEN) { + throw new Error(`WebSocket not open (readyState=${this.readyState})`); + } + this.sent.push(data); + } + + open() { + this.readyState = FakeRelayWebSocket.OPEN; + this.emit("open"); + } + + message(data: unknown) { + this.emit("message", data); + } + + private off(event: string, listener: (...args: unknown[]) => void) { + const handlers = this.listeners.get(event) ?? []; + this.listeners.set( + event, + handlers.filter((handler) => handler !== listener), + ); + } + + private emit(event: string, ...args: unknown[]) { + const handlers = this.listeners.get(event) ?? []; + for (const handler of handlers.slice()) { + handler(...args); + } + } +} + +function createFakeWebSockets() { + const sockets: FakeRelayWebSocket[] = []; + return { + sockets, + createWebSocket(url: string) { + const socket = new FakeRelayWebSocket(url); + sockets.push(socket); + return socket; + }, + }; } describe("relay-transport control lifecycle", () => { const controllers: Array<{ stop: () => Promise }> = []; - const MockWebSocket = wsMock.MockWebSocket; + let relay: ReturnType; beforeEach(() => { - MockWebSocket.reset(); + relay = createFakeWebSockets(); }); afterEach(async () => { @@ -135,18 +126,19 @@ describe("relay-transport control lifecycle", () => { relayEndpoint: "relay.paseo.sh:443", relayUseTls: true, serverId: "srv_test", + createWebSocket: relay.createWebSocket, }); controllers.push(controller); - const control = MockWebSocket.instances[0]; + const control = relay.sockets[0]; expect(control).toBeDefined(); control.open(); - expect(hasLogMessage(logger.info, "relay_control_connected")).toBe(false); + expect(hasLogMessage(logger, "info", "relay_control_connected")).toBe(false); expect(control.sent.length).toBeGreaterThan(0); control.message(JSON.stringify({ type: "pong", ts: Date.now() })); - expect(hasLogMessage(logger.info, "relay_control_connected")).toBe(true); + expect(hasLogMessage(logger, "info", "relay_control_connected")).toBe(true); }); test("terminates and reconnects when control socket opens but never becomes ready", () => { @@ -158,18 +150,19 @@ describe("relay-transport control lifecycle", () => { relayEndpoint: "relay.paseo.sh:443", relayUseTls: true, serverId: "srv_test", + createWebSocket: relay.createWebSocket, }); controllers.push(controller); - const firstControl = MockWebSocket.instances[0]; + const firstControl = relay.sockets[0]; firstControl.open(); vi.advanceTimersByTime(8_000); - expect(hasLogMessage(logger.warn, "relay_control_ready_timeout_terminating")).toBe(true); + expect(hasLogMessage(logger, "warn", "relay_control_ready_timeout_terminating")).toBe(true); expect(firstControl.terminateCalls).toBe(1); vi.advanceTimersByTime(1_000); - expect(MockWebSocket.instances.length).toBeGreaterThanOrEqual(2); + expect(relay.sockets.length).toBeGreaterThanOrEqual(2); }); test("terminates stale control sockets in under one minute", () => { @@ -181,47 +174,56 @@ describe("relay-transport control lifecycle", () => { relayEndpoint: "relay.paseo.sh:443", relayUseTls: true, serverId: "srv_test", + createWebSocket: relay.createWebSocket, }); controllers.push(controller); - const control = MockWebSocket.instances[0]; + const control = relay.sockets[0]; control.open(); control.message(JSON.stringify({ type: "pong", ts: Date.now() })); - logger.warn.mockClear(); + logger.messages.length = 0; vi.advanceTimersByTime(40_000); - expect(hasLogMessage(logger.warn, "relay_control_stale_terminating")).toBe(true); + expect(hasLogMessage(logger, "warn", "relay_control_stale_terminating")).toBe(true); expect(control.terminateCalls).toBe(1); }); test("passes stable relay external session metadata when attaching data socket", async () => { const logger = createMockLogger(); - const attachSocket = vi.fn(async () => {}); + const attachedSockets: unknown[] = []; + const attachedMetadata: unknown[] = []; + const attachSocket = async (socket: unknown, metadata: unknown) => { + attachedSockets.push(socket); + attachedMetadata.push(metadata); + }; const controller = startRelayTransport({ logger: logger as unknown as pino.Logger, attachSocket, relayEndpoint: "relay.paseo.sh:443", relayUseTls: true, serverId: "srv_test", + createWebSocket: relay.createWebSocket, }); controllers.push(controller); - const control = MockWebSocket.instances[0]; + const control = relay.sockets[0]; control.open(); control.message(JSON.stringify({ type: "pong", ts: Date.now() })); control.message(JSON.stringify({ type: "connected", connectionId: "clt_test" })); - const dataSocket = MockWebSocket.instances[1]; + const dataSocket = relay.sockets[1]; expect(dataSocket).toBeDefined(); dataSocket.open(); await Promise.resolve(); - expect(attachSocket).toHaveBeenCalledTimes(1); - expect(attachSocket).toHaveBeenCalledWith(dataSocket, { - transport: "relay", - externalSessionKey: "session:clt_test", - }); + expect(attachedSockets).toEqual([dataSocket]); + expect(attachedMetadata).toEqual([ + { + transport: "relay", + externalSessionKey: "session:clt_test", + }, + ]); }); test("uses relayUseTls for control and data socket URLs", () => { @@ -232,15 +234,16 @@ describe("relay-transport control lifecycle", () => { relayEndpoint: "[::1]:443", relayUseTls: true, serverId: "srv_test", + createWebSocket: relay.createWebSocket, }); controllers.push(controller); - const control = MockWebSocket.instances[0]; + const control = relay.sockets[0]; control.open(); control.message(JSON.stringify({ type: "pong", ts: Date.now() })); control.message(JSON.stringify({ type: "connected", connectionId: "clt_test" })); - expect(MockWebSocket.instances[0]?.url).toMatch(/^wss:\/\/\[::1\]\/ws\?/); - expect(MockWebSocket.instances[1]?.url).toMatch(/^wss:\/\/\[::1\]\/ws\?/); + expect(relay.sockets[0]?.url).toMatch(/^wss:\/\/\[::1\]\/ws\?/); + expect(relay.sockets[1]?.url).toMatch(/^wss:\/\/\[::1\]\/ws\?/); }); }); diff --git a/packages/server/src/server/relay-transport.ts b/packages/server/src/server/relay-transport.ts index fe9cf198c..60401182b 100644 --- a/packages/server/src/server/relay-transport.ts +++ b/packages/server/src/server/relay-transport.ts @@ -18,6 +18,7 @@ interface RelayTransportOptions { relayUseTls: boolean; serverId: string; daemonKeyPair?: KeyPair; + createWebSocket?: RelayWebSocketFactory; } export interface RelayTransportController { @@ -32,6 +33,16 @@ interface RelaySocketLike { once: (event: "close" | "error", listener: (...args: unknown[]) => void) => void; } +interface RelayWebSocketLike extends RelaySocketLike { + terminate: () => void; + on: ( + event: "open" | "message" | "close" | "error", + listener: (...args: unknown[]) => void, + ) => void; +} + +type RelayWebSocketFactory = (url: string) => RelayWebSocketLike; + type ControlMessage = | { type: "sync"; connectionIds: string[] } | { type: "connected"; connectionId: string } @@ -42,6 +53,11 @@ type ControlMessage = const CONTROL_PING_INTERVAL_MS = 10_000; const CONTROL_STALE_TIMEOUT_MS = 30_000; const CONTROL_READY_TIMEOUT_MS = 8_000; +const RELAY_WEBSOCKET_OPTIONS = { handshakeTimeout: 10_000, perMessageDeflate: false } as const; + +function createDefaultRelayWebSocket(url: string): RelayWebSocketLike { + return new WebSocket(url, RELAY_WEBSOCKET_OPTIONS); +} function normalizeRelaySendPayload(data: string | Uint8Array | ArrayBuffer): string | ArrayBuffer { if (typeof data === "string") return data; @@ -106,14 +122,15 @@ export function startRelayTransport({ relayUseTls, serverId, daemonKeyPair, + createWebSocket = createDefaultRelayWebSocket, }: RelayTransportOptions): RelayTransportController { const relayLogger = logger.child({ module: "relay-transport" }); let stopped = false; - let controlWs: WebSocket | null = null; + let controlWs: RelayWebSocketLike | null = null; let reconnectTimeout: ReturnType | null = null; let reconnectAttempt = 0; - const dataSockets = new Map(); // connectionId -> ws + const dataSockets = new Map(); // connectionId -> ws let controlKeepaliveInterval: ReturnType | null = null; let controlReadyTimeout: ReturnType | null = null; let controlLastSeenAt = 0; @@ -161,7 +178,7 @@ export function startRelayTransport({ serverId, role: "server", }); - const socket = new WebSocket(url, { handshakeTimeout: 10_000, perMessageDeflate: false }); + const socket = createWebSocket(url); controlWs = socket; let controlConnected = false; @@ -338,7 +355,7 @@ export function startRelayTransport({ role: "server", connectionId, }); - const socket = new WebSocket(url, { handshakeTimeout: 10_000, perMessageDeflate: false }); + const socket = createWebSocket(url); dataSockets.set(connectionId, socket); let attached = false; @@ -397,7 +414,7 @@ export function startRelayTransport({ } async function attachEncryptedSocket( - socket: WebSocket, + socket: RelayWebSocketLike, daemonKeyPair: KeyPair, logger: pino.Logger, attachSocket: (ws: RelaySocketLike, metadata?: ExternalSocketMetadata) => Promise, @@ -426,7 +443,10 @@ async function attachEncryptedSocket( } } -function createRelayTransportAdapter(socket: WebSocket, logger: pino.Logger): RelayTransport { +function createRelayTransportAdapter( + socket: RelayWebSocketLike, + logger: pino.Logger, +): RelayTransport { const relayTransport: RelayTransport = { send: (data) => { try { @@ -445,10 +465,11 @@ function createRelayTransportAdapter(socket: WebSocket, logger: pino.Logger): Re }; socket.on("message", (data, isBinary) => { - relayTransport.onmessage?.(normalizeMessageData(data, isBinary)); + relayTransport.onmessage?.(normalizeMessageData(data, isBinary === true)); }); socket.on("close", (code, reason) => { - relayTransport.onclose?.(code, reason.toString()); + const closeCode = typeof code === "number" ? code : 1006; + relayTransport.onclose?.(closeCode, String(reason ?? "")); }); socket.on("error", (err) => { relayTransport.onerror?.(err instanceof Error ? err : new Error(String(err)));