From c604e5f7d3a307fbde73ac7eaa7a9dc8d0d48f6a Mon Sep 17 00:00:00 2001 From: Mathieu Colmon Date: Fri, 4 Sep 2026 22:15:28 +0200 Subject: [PATCH] fix: serialize browser relay replacements --- README.md | 15 ++- src/browser-client.ts | 64 ++++++++-- src/internal/browser-socket.ts | 4 +- src/internal/runtime.ts | 8 +- src/internal/socket.ts | 12 +- src/internal/user-session.ts | 22 +++- test/browser-lifecycle.test.ts | 209 +++++++++++++++++++++++++++++++++ test/browser-socket.test.ts | 85 +++++++++++++- test/fakes/relay.ts | 36 +++++- 9 files changed, 426 insertions(+), 29 deletions(-) diff --git a/README.md b/README.md index 6450169..514c73a 100644 --- a/README.md +++ b/README.md @@ -217,12 +217,15 @@ browser storage, log, error, `HELLO`, or `REAUTH` frame. The control plane returns an up-to-five-minute Miakapp access token atomically with its authoritative relay URL. MiakAPI sends only that audience-bound token -to the returned relay. If a renewal selects a different relay, the client closes -the old session and opens the replacement with the already-issued credential; it -does not expose the new token to the old relay or repeat the exchange. Stop and -discard the client immediately when the Firebase user signs out or the selected -home changes. Relay routing changes arrive through credentials and do not require -mutating the client options. +to the returned relay. If a renewal selects a different relay, the client marks +the old session stale and closes it. Before any automatic replacement connection, +including recovery from a transport or protocol failure, the client waits for the +native transport close event. If closure is not confirmed within ten seconds, the +client stops fail-closed instead of opening overlapping relay sockets. A routing +handoff uses the already-issued credential; it does not expose the new token to +the old relay or repeat the exchange. Stop and discard the client immediately +when the Firebase user signs out or the selected home changes. Relay routing +changes arrive through credentials and do not require mutating the client options. Audience binding limits credential replay; it does not encrypt home traffic from the selected relay. Users should still choose an operator they trust with the diff --git a/src/browser-client.ts b/src/browser-client.ts index ae15eb6..7018b6d 100644 --- a/src/browser-client.ts +++ b/src/browser-client.ts @@ -306,6 +306,10 @@ class BrowserClientImpl implements BrowserClient, UserStateHost, UserCallHost { connectionEnd?.resolve({ failure }); }, }, + (createdSession) => { + this.#session = createdSession; + this.#sessionRelayUrl = credential.relayUrl; + }, ); if (this.#loopController.signal.aborted) { session.terminate(); @@ -314,12 +318,9 @@ class BrowserClientImpl implements BrowserClient, UserStateHost, UserCallHost { } if (credential.expiresAtMs <= this.#runtime.now()) { session.terminate(); - session.detach(); throw browserUnavailable('Browser relay credential expired during authentication'); } this.#reconnectAttempt = 0; - this.#session = session; - this.#sessionRelayUrl = credential.relayUrl; this.#setHomeStatus(Object.freeze({ enrolled: session.welcome.readySession.enrolled, coordinators: session.welcome.readySession.coordinators, @@ -366,12 +367,26 @@ class BrowserClientImpl implements BrowserClient, UserStateHost, UserCallHost { ? error : browserUnavailable('Browser relay connection attempt failed') }; } - this.#disconnectSession(); + const endedSession = this.#session; + this.#deactivateSession(); if (this.#loopController.signal.aborted) break; if (end.failure !== undefined) this.#emitFailure(end.failure); if (this.#loopController.signal.aborted) break; this.#transition('reconnecting', undefined, end.failure); if (this.#loopController.signal.aborted) break; + if (endedSession !== undefined + && !endedSession.transportClosed + && !await this.#waitForTransportClose(endedSession)) { + if (!this.#loopController.signal.aborted) { + const failure = browserUnavailable('Browser relay transport close timed out'); + this.#emitFailure(failure); + this.#startDeferred?.reject(failure); + void this.stop(); + } + break; + } + this.#disconnectSession(); + if (this.#loopController.signal.aborted) break; if (end.handoffCredential !== undefined) { pendingCredential = end.handoffCredential; reason = 'reconnect'; @@ -478,8 +493,8 @@ class BrowserClientImpl implements BrowserClient, UserStateHost, UserCallHost { if (this.#sessionEnd === undefined || this.#sessionEnd.settled) { throw browserUnavailable('Browser relay handoff state is unavailable'); } - this.#sessionEnd.resolve({ handoffCredential: credential }); session.terminate(); + this.#sessionEnd.resolve({ handoffCredential: credential }); return; } const remaining = Math.min(currentExpiresAtMs, credential.expiresAtMs) @@ -513,6 +528,27 @@ class BrowserClientImpl implements BrowserClient, UserStateHost, UserCallHost { && this.#status !== 'draining'; } + async #waitForTransportClose(session: UserRelaySession): Promise { + const interrupted = createDeferred(); + const abort = () => interrupted.resolve(false); + if (this.#loopController.signal.aborted) return false; + this.#loopController.signal.addEventListener('abort', abort, { once: true }); + if (this.#loopController.signal.aborted) abort(); + const timeout = this.#runtime.setTimer( + () => interrupted.resolve(false), + SESSION_PHASE_TIMEOUT_MS, + ); + try { + return await Promise.race([ + session.waitForTransportClose().then(() => true), + interrupted.promise, + ]); + } finally { + timeout.cancel(); + this.#loopController.signal.removeEventListener('abort', abort); + } + } + #clearReauthentication(): void { this.#reauthTimer?.cancel(); this.#reauthTimer = undefined; @@ -625,13 +661,7 @@ class BrowserClientImpl implements BrowserClient, UserStateHost, UserCallHost { } #disconnectSession(): void { - this.#bootstrapTimer?.cancel(); - this.#bootstrapTimer = undefined; - this.#clearReauthentication(); - this.#abortCredentialRequest(); - this.state.disconnected(); - this.calls.disconnected(); - this.#markHomeStale(); + this.#deactivateSession(); this.#session?.detach(); this.#session = undefined; this.#sessionRelayUrl = undefined; @@ -643,6 +673,16 @@ class BrowserClientImpl implements BrowserClient, UserStateHost, UserCallHost { this.#topicReady = false; } + #deactivateSession(): void { + this.#bootstrapTimer?.cancel(); + this.#bootstrapTimer = undefined; + this.#clearReauthentication(); + this.#abortCredentialRequest(); + this.state.disconnected(); + this.calls.disconnected(); + this.#markHomeStale(); + } + #emitFailure(failure: BrowserClientFailure): void { this.#errorListeners.emit(failure, () => { safeBrowserLog(this.#options.logger, { level: 'error', event: 'error_listener_failed' }); diff --git a/src/internal/browser-socket.ts b/src/internal/browser-socket.ts index 709a6da..657f02c 100644 --- a/src/internal/browser-socket.ts +++ b/src/internal/browser-socket.ts @@ -213,6 +213,7 @@ export class BrowserSocketFactory implements SocketFactory { url: string, handlers: SocketHandlers, signal: AbortSignal, + onSocket?: (socket: ManagedSocket) => void, ): Promise { if (signal.aborted) throw signal.reason; const ManagedWebSocket = nativeConstructor(); @@ -222,12 +223,13 @@ export class BrowserSocketFactory implements SocketFactory { signal, this.#now, ); + onSocket?.(socket); try { await socket.ready(); return socket; } catch (error) { - socket.detach(); socket.terminate(); + if (onSocket === undefined) socket.detach(); throw error; } } diff --git a/src/internal/runtime.ts b/src/internal/runtime.ts index 5074bc5..755c171 100644 --- a/src/internal/runtime.ts +++ b/src/internal/runtime.ts @@ -15,7 +15,13 @@ export interface ManagedSocket { } export interface SocketFactory { - connect(url: string, handlers: SocketHandlers, signal: AbortSignal): Promise; + /** Transfers a created transport to its owner before asynchronous readiness work begins. */ + connect( + url: string, + handlers: SocketHandlers, + signal: AbortSignal, + onSocket?: (socket: ManagedSocket) => void, + ): Promise; } export interface RuntimeTimer { diff --git a/src/internal/socket.ts b/src/internal/socket.ts index ea40763..dd5be58 100644 --- a/src/internal/socket.ts +++ b/src/internal/socket.ts @@ -150,6 +150,7 @@ export class WsSocketFactory implements SocketFactory { url: string, handlers: SocketHandlers, signal: AbortSignal, + onSocket?: (socket: ManagedSocket) => void, ): Promise { if (signal.aborted) throw signal.reason; const options: BoundedClientOptions = { @@ -165,8 +166,15 @@ export class WsSocketFactory implements SocketFactory { handlers, signal, ); - await managed.ready(); - return managed; + onSocket?.(managed); + try { + await managed.ready(); + return managed; + } catch (error) { + managed.terminate(); + if (onSocket === undefined) managed.detach(); + throw error; + } } } diff --git a/src/internal/user-session.ts b/src/internal/user-session.ts index afed1f0..1403cc7 100644 --- a/src/internal/user-session.ts +++ b/src/internal/user-session.ts @@ -130,6 +130,7 @@ export class UserRelaySession { readonly #callbacks: UserRelaySessionCallbacks; readonly #now: () => number; readonly #protocol = new UserProtocolSession(); + readonly #transportClosed = createDeferred(); readonly #welcome = createDeferred(); readonly #queuedFrames: Frame[] = []; #queuedFrameBytes = 0; @@ -152,9 +153,13 @@ export class UserRelaySession { token: string, signal: AbortSignal, callbacks: UserRelaySessionCallbacks, + onCreated: (session: UserRelaySession) => void, ): Promise { const session = new UserRelaySession(callbacks, () => runtime.now()); + onCreated(session); const handshake = childAbortController(signal); + const abortWelcome = () => session.#welcome.reject(handshake.controller.signal.reason); + handshake.controller.signal.addEventListener('abort', abortWelcome, { once: true }); const timeout = runtime.setTimer(() => { handshake.controller.abort(new Error('Browser relay handshake timed out')); }, HANDSHAKE_TIMEOUT_MS); @@ -168,6 +173,10 @@ export class UserRelaySession { relayUrl, handlers, handshake.controller.signal, + (socket) => { + session.#socket = socket; + if (session.#closed) socket.terminate(); + }, ); session.#connectedAtMs = runtime.now(); await session.#socket.write(session.#protocol.encode({ @@ -178,10 +187,10 @@ export class UserRelaySession { return session; } catch (error) { session.terminate(); - session.detach(); throw error; } finally { timeout.cancel(); + handshake.controller.signal.removeEventListener('abort', abortWelcome); handshake.dispose(); } } @@ -232,6 +241,16 @@ export class UserRelaySession { this.#socket?.terminate(); } + waitForTransportClose(): Promise { + return this.#socket === undefined + ? Promise.resolve() + : this.#transportClosed.promise; + } + + get transportClosed(): boolean { + return this.#socket === undefined || this.#transportClosed.settled; + } + detach(): void { this.#socket?.detach(); } @@ -270,6 +289,7 @@ export class UserRelaySession { const wasClosed = this.#closed; this.#closed = true; this.#protocol.close(); + this.#transportClosed.resolve(undefined); this.#welcome.reject(browserProtocolFailure('Relay closed before WELCOME')); if (!wasClosed) this.#callbacks.closed(code, reason); } diff --git a/test/browser-lifecycle.test.ts b/test/browser-lifecycle.test.ts index 949b354..e7760c7 100644 --- a/test/browser-lifecycle.test.ts +++ b/test/browser-lifecycle.test.ts @@ -72,6 +72,63 @@ describe('browser lifecycle', () => { await harness.client.stop(); }); + test('waits for a failed WELCOME transport to close before reconnecting', async () => { + const harness = createBrowserTestHarness(); + harness.runtime.queueRandom(0); + const started = harness.client.start(); + void started.catch(() => undefined); + const connection = await harness.relay.connectionAt(0); + await connection.nextClientFrame(Opcode.Hello); + const oldTransport = connection.deferClientTermination(); + connection.send({ + opcode: Opcode.Welcome, + payload: [2, 0, 41, connection.epoch, true, [], [262_144, 128, 256, 1_048_576], 2_000_000], + }); + await flushMicrotasks(); + expect(harness.client.status).toBe('reconnecting'); + expect(harness.relay.connectCount).toBe(1); + expect(harness.relay.openConnectionCount).toBe(1); + + oldTransport.release(); + await flushMicrotasks(); + await harness.runtime.advanceBy(0); + const replacement = await harness.relay.connectionAt(1); + await replacement.nextClientFrame(Opcode.Hello); + expect(harness.relay.socketHighWater).toBe(1); + await harness.client.stop(); + }); + + test('fails closed when a missing WELCOME transport remains closing', async () => { + const harness = createBrowserTestHarness(); + const failures: string[] = []; + harness.client.errors.subscribe(({ kind }) => failures.push(kind)); + const started = harness.client.start(); + void started.catch(() => undefined); + const connection = await harness.relay.connectionAt(0); + await connection.nextClientFrame(Opcode.Hello); + const oldTransport = connection.deferClientTermination(); + + await harness.runtime.advanceBy(10_000); + expect(harness.client.status).toBe('reconnecting'); + expect(harness.relay.connectCount).toBe(1); + expect(harness.relay.openConnectionCount).toBe(1); + expect(failures).toEqual(['unavailable']); + + await harness.runtime.advanceBy(9_999); + expect(harness.client.status).toBe('reconnecting'); + expect(harness.relay.connectCount).toBe(1); + expect(failures).toEqual(['unavailable']); + + await harness.runtime.advanceBy(1); + expect(harness.client.status).toBe('stopped'); + expect(harness.relay.connectCount).toBe(1); + expect(failures).toEqual(['unavailable', 'unavailable']); + + oldTransport.release(); + await flushMicrotasks(); + expect(harness.relay.openConnectionCount).toBe(0); + }); + test('reauthenticates on the same socket from the verified lease', async () => { const harness = createBrowserTestHarness({ expiresAtMs: 1_010_000 }); const { connection } = await startBrowserReady(harness); @@ -117,7 +174,15 @@ describe('browser lifecycle', () => { sendUserBootstrap(oldConnection); await started; + const oldTransport = oldConnection.deferClientTermination(); await runtime.advanceBy(5_000); + expect(relay.connectCount).toBe(1); + expect(relay.openConnectionCount).toBe(1); + expect(client.status).toBe('reconnecting'); + expect(client.state.snapshot()?.stale).toBe(true); + + oldTransport.release(); + await flushMicrotasks(); const replacement = await relay.connectionAt(1); const replacementHello = await replacement.nextClientFrame(Opcode.Hello); expect(replacementHello.payload[4]).toBe('user.handoff.signature'); @@ -137,6 +202,150 @@ describe('browser lifecycle', () => { await client.stop(); }); + test('fails closed when a previous relay transport cannot close during handoff', async () => { + const relay = new FakeRelay({ autoWelcome: false, expiresAtMs: 1_010_000 }); + const runtime = new FakeRuntime(relay); + const failures: string[] = []; + const client = createBrowserClientWithRuntime({ + homeId: 'test-home', + credentialProvider: { + async getCredential({ reason }) { + return { + relayUrl: reason === 'initial' + ? 'wss://old-relay.test/miakapp/ws' + : 'wss://new-relay.test/miakapp/ws', + accessToken: `user.${reason}.signature`, + expiresAtMs: 1_100_000, + }; + }, + }, + }, runtime); + client.errors.subscribe(({ kind }) => failures.push(kind)); + + const started = client.start(); + const oldConnection = await relay.connectionAt(0); + await oldConnection.nextClientFrame(Opcode.Hello); + sendUserBootstrap(oldConnection); + await started; + + const oldTransport = oldConnection.deferClientTermination(); + await runtime.advanceBy(5_000); + expect(client.status).toBe('reconnecting'); + expect(failures).toEqual([]); + await runtime.advanceBy(9_999); + expect(relay.connectCount).toBe(1); + expect(relay.openConnectionCount).toBe(1); + expect(client.status).toBe('reconnecting'); + expect(failures).toEqual([]); + + await runtime.advanceBy(1); + expect(relay.connectCount).toBe(1); + expect(client.status).toBe('stopped'); + expect(failures).toEqual(['unavailable']); + + oldTransport.release(); + await flushMicrotasks(); + expect(relay.openConnectionCount).toBe(0); + }); + + test('waits for a failed relay transport before reconnecting to a changed URL', async () => { + const relay = new FakeRelay({ autoWelcome: false }); + const runtime = new FakeRuntime(relay); + runtime.queueRandom(0); + const client = createBrowserClientWithRuntime({ + homeId: 'test-home', + credentialProvider: { + async getCredential({ reason }) { + return { + relayUrl: reason === 'initial' + ? 'wss://old-relay.test/miakapp/ws' + : 'wss://new-relay.test/miakapp/ws', + accessToken: `user.${reason}.signature`, + expiresAtMs: 2_000_000, + }; + }, + }, + }, runtime); + + const started = client.start(); + const oldConnection = await relay.connectionAt(0); + await oldConnection.nextClientFrame(Opcode.Hello); + sendUserBootstrap(oldConnection); + await started; + + const oldTransport = oldConnection.deferClientTermination(); + oldConnection.send({ opcode: Opcode.ReauthOk, payload: [999, 2_000_000] }); + await flushMicrotasks(); + expect(client.status).toBe('reconnecting'); + expect(client.state.snapshot()?.stale).toBe(true); + expect(relay.connectCount).toBe(1); + expect(relay.openConnectionCount).toBe(1); + + oldTransport.release(); + await flushMicrotasks(); + await runtime.advanceBy(0); + const replacement = await relay.connectionAt(1); + const replacementHello = await replacement.nextClientFrame(Opcode.Hello); + expect(replacementHello.payload[4]).toBe('user.reconnect.signature'); + expect(relay.connectUrls).toEqual([ + 'wss://old-relay.test/miakapp/ws', + 'wss://new-relay.test/miakapp/ws', + ]); + expect(relay.socketHighWater).toBe(1); + + sendUserBootstrap(replacement, { revision: 2 }); + await flushMicrotasks(); + await client.stop(); + }); + + test('fails closed when a failed relay transport cannot close before reconnect', async () => { + const relay = new FakeRelay({ autoWelcome: false }); + const runtime = new FakeRuntime(relay); + const failures: string[] = []; + const client = createBrowserClientWithRuntime({ + homeId: 'test-home', + credentialProvider: { + async getCredential({ reason }) { + return { + relayUrl: reason === 'initial' + ? 'wss://old-relay.test/miakapp/ws' + : 'wss://new-relay.test/miakapp/ws', + accessToken: `user.${reason}.signature`, + expiresAtMs: 2_000_000, + }; + }, + }, + }, runtime); + client.errors.subscribe(({ kind }) => failures.push(kind)); + + const started = client.start(); + const oldConnection = await relay.connectionAt(0); + await oldConnection.nextClientFrame(Opcode.Hello); + sendUserBootstrap(oldConnection); + await started; + + const oldTransport = oldConnection.deferClientTermination(); + oldConnection.send({ opcode: Opcode.ReauthOk, payload: [999, 2_000_000] }); + await flushMicrotasks(); + expect(client.status).toBe('reconnecting'); + expect(failures).toEqual(['protocol']); + + await runtime.advanceBy(9_999); + expect(relay.connectCount).toBe(1); + expect(relay.openConnectionCount).toBe(1); + expect(client.status).toBe('reconnecting'); + expect(failures).toEqual(['protocol']); + + await runtime.advanceBy(1); + expect(relay.connectCount).toBe(1); + expect(client.status).toBe('stopped'); + expect(failures).toEqual(['protocol', 'unavailable']); + + oldTransport.release(); + await flushMicrotasks(); + expect(relay.openConnectionCount).toBe(0); + }); + test('rejects a credential that expires while the relay handshake is pending', async () => { const relay = new FakeRelay({ autoWelcome: false }); const runtime = new FakeRuntime(relay); diff --git a/test/browser-socket.test.ts b/test/browser-socket.test.ts index d9b1674..a17a196 100644 --- a/test/browser-socket.test.ts +++ b/test/browser-socket.test.ts @@ -1,6 +1,8 @@ import { afterEach, beforeEach, describe, expect, test } from 'bun:test'; +import { createBrowserClientWithRuntime } from '../src/browser-client.js'; import { BrowserSocketFactory } from '../src/internal/browser-socket.js'; import type { SocketHandlers } from '../src/internal/runtime.js'; +import { FakeRuntime, flushMicrotasks } from './fakes/runtime.js'; type NativeListener = (event: never) => void; @@ -38,10 +40,11 @@ class MockNativeWebSocket { close(): void { this.closeCount += 1; - this.readyState = 3; + this.readyState = 2; } emit(type: string, event: unknown = {}): void { + if (type === 'close') this.readyState = 3; for (const listener of [...(this.#listeners.get(type) ?? [])]) { listener(event as never); } @@ -101,6 +104,84 @@ describe('browser socket', () => { expect(socket?.listenerCount()).toBe(0); }); + test('waits for a tracked pre-open transport to close before replacement', async () => { + const runtime = new FakeRuntime(new BrowserSocketFactory(), 1_000_000); + runtime.queueRandom(0); + const client = createBrowserClientWithRuntime({ + homeId: 'test-home', + credentialProvider: { + async getCredential({ reason }) { + return { + relayUrl: 'wss://relay.test/ws', + accessToken: `user.${reason}.signature`, + expiresAtMs: 2_000_000, + }; + }, + }, + }, runtime); + const started = client.start(); + void started.catch(() => undefined); + await flushMicrotasks(); + const first = MockNativeWebSocket.instances[0]; + if (first === undefined) throw new Error('Expected the first native WebSocket'); + + first.emit('error'); + await flushMicrotasks(); + expect(client.status).toBe('reconnecting'); + expect(first.readyState).toBe(2); + expect(MockNativeWebSocket.instances).toHaveLength(1); + + first.emit('close', { code: 1006, reason: 'failed' }); + await flushMicrotasks(); + await runtime.advanceBy(0); + expect(first.readyState).toBe(3); + expect(MockNativeWebSocket.instances).toHaveLength(2); + await client.stop(); + }); + + test('stops after ten seconds when a tracked pre-open transport stays closing', async () => { + const runtime = new FakeRuntime(new BrowserSocketFactory(), 1_000_000); + const failures: string[] = []; + const client = createBrowserClientWithRuntime({ + homeId: 'test-home', + credentialProvider: { + async getCredential() { + return { + relayUrl: 'wss://relay.test/ws', + accessToken: 'user.initial.signature', + expiresAtMs: 2_000_000, + }; + }, + }, + }, runtime); + client.errors.subscribe(({ message }) => failures.push(message)); + const started = client.start(); + void started.catch(() => undefined); + await flushMicrotasks(); + const socket = MockNativeWebSocket.instances[0]; + if (socket === undefined) throw new Error('Expected a native WebSocket'); + + socket.emit('error'); + await flushMicrotasks(); + expect(client.status).toBe('reconnecting'); + expect(failures).toEqual(['Browser relay connection failed']); + await runtime.advanceBy(9_999); + expect(client.status).toBe('reconnecting'); + expect(MockNativeWebSocket.instances).toHaveLength(1); + expect(socket.listenerCount()).toBe(4); + + await runtime.advanceBy(1); + await expect(started).rejects.toThrow('Browser relay transport close timed out'); + expect(client.status).toBe('stopped'); + expect(MockNativeWebSocket.instances).toHaveLength(1); + expect(socket.readyState).toBe(2); + expect(socket.listenerCount()).toBe(0); + expect(failures).toEqual([ + 'Browser relay connection failed', + 'Browser relay transport close timed out', + ]); + }); + test('terminates a relay that exceeds the rolling inbound byte budget', async () => { let now = 1_000; let messages = 0; @@ -132,7 +213,7 @@ describe('browser socket', () => { expect(messages).toBe(8); expect(failures).toEqual(['Relay exceeded the browser inbound budget']); expect(socket.closeCount).toBe(1); - expect(socket.readyState).toBe(3); + expect(socket.readyState).toBe(2); }); }); diff --git a/test/fakes/relay.ts b/test/fakes/relay.ts index 1da896b..734c5a9 100644 --- a/test/fakes/relay.ts +++ b/test/fakes/relay.ts @@ -69,10 +69,12 @@ function dictionary(names: readonly string[], firstId: number): ProtocolValue[] class FakeManagedSocket implements ManagedSocket { readonly #connection: FakeRelayConnection; #closed = false; + #closing = false; #detached = false; #bufferedBytes = 0; #nextWriteError: Error | undefined; #nextWriteCompletion: Promise | undefined; + #terminateCompletion: Promise | undefined; constructor(connection: FakeRelayConnection) { this.#connection = connection; @@ -97,6 +99,13 @@ class FakeManagedSocket implements ManagedSocket { this.#nextWriteCompletion = completion; } + deferTermination(completion: Promise): void { + if (this.#terminateCompletion !== undefined) { + throw new Error('Synthetic termination is already deferred'); + } + this.#terminateCompletion = completion; + } + async write(bytes: Uint8Array): Promise { if (this.#closed) throw new Error('Synthetic socket is closed'); const failure = this.#nextWriteError; @@ -111,11 +120,20 @@ class FakeManagedSocket implements ManagedSocket { close(code = 1000, reason = ''): void { if (this.#closed) return; this.#closed = true; - if (!this.#detached) this.#connection.notifyClientClose(code, reason); + this.#connection.notifyClientClose(code, reason, !this.#detached); } terminate(): void { - this.close(1006, 'terminated'); + if (this.#closed || this.#closing) return; + this.#closing = true; + const completion = this.#terminateCompletion; + this.#terminateCompletion = undefined; + if (completion === undefined) { + this.close(1006, 'terminated'); + return; + } + const close = () => this.close(1006, 'terminated'); + void completion.then(close, close); } detach(): void { @@ -160,6 +178,12 @@ export class FakeRelayConnection { }; } + deferClientTermination(): { release(): void } { + const completion = createDeferred(); + this.socket.deferTermination(completion.promise); + return { release: () => completion.resolve(undefined) }; + } + receiveClientBytes(bytes: Uint8Array): void { const frame = decodeFrame(bytes); const waiter = this.#waiters.shift(); @@ -170,8 +194,10 @@ export class FakeRelayConnection { } } - notifyClientClose(_code: number, _reason: string): void { + notifyClientClose(code: number, reason: string, notifyHandler: boolean): void { + if (this.#serverClosed) return; this.#markClosed(); + if (notifyHandler) this.#handlers.close(code, reason); } async nextClientFrame(expectedOpcode?: number): Promise { @@ -316,6 +342,7 @@ export class FakeRelay implements SocketFactory { url: string, handlers: SocketHandlers, signal: AbortSignal, + onSocket?: (socket: ManagedSocket) => void, ): Promise { if (signal.aborted) throw signal.reason; this.#connectUrls.push(url); @@ -335,7 +362,8 @@ export class FakeRelay implements SocketFactory { this.#connections.push(connection); this.#openConnections += 1; this.#socketHighWater = Math.max(this.#socketHighWater, this.#openConnections); - signal.addEventListener('abort', () => connection.close(1006, 'aborted'), { once: true }); + onSocket?.(connection.socket); + signal.addEventListener('abort', () => connection.socket.terminate(), { once: true }); const waiter = this.#connectionWaiters.shift(); waiter?.(connection); return connection.socket;