diff --git a/packages/ui/src/lib/relay/handshake.ts b/packages/ui/src/lib/relay/handshake.ts index 35497eac..a2fef028 100644 --- a/packages/ui/src/lib/relay/handshake.ts +++ b/packages/ui/src/lib/relay/handshake.ts @@ -44,7 +44,7 @@ export type HandshakeAction = | { type: 'send-text'; text: string } // `replyText`, when present, must be sent to the peer before any encrypted frame. // `batch` is the negotiated frame-batching capability for the session. - | { type: 'established'; channel: EstablishedChannelCrypto; batch: boolean; replyText?: string } + | { type: 'established'; channel: EstablishedChannelCrypto; batch: boolean; flowControl: boolean; replyText?: string } | { type: 'ignore' } | { type: 'fail'; closeCode: number; reason: string }; @@ -60,8 +60,9 @@ const parseHandshakeMessage = (raw: string): E2eeHelloMessage | E2eeReadyMessage if (message.v !== RELAY_PROTOCOL_VERSION) return null; // Unknown/missing capability flag = false = legacy behavior. const batch = message.batch === true; + const flowControl = message.flowControl === true; if (message.t === 'ready') { - return { t: 'ready', v: RELAY_PROTOCOL_VERSION, batch }; + return { t: 'ready', v: RELAY_PROTOCOL_VERSION, batch, flowControl }; } if ( message.t === 'hello' && @@ -75,6 +76,7 @@ const parseHandshakeMessage = (raw: string): E2eeHelloMessage | E2eeReadyMessage clientPubJwk: message.clientPubJwk as JsonWebKey, nonce: message.nonce, batch, + flowControl, }; } return null; @@ -97,6 +99,7 @@ export interface ClientHandshake { export interface ClientHandshakeOptions { /** Advertise frame batching. Default true; set false to force legacy behavior. */ batch?: boolean; + flowControl?: boolean; } // hostEncPubJwk comes from the pairing offer (QR / deep link) and is the trust @@ -114,8 +117,9 @@ export const createClientHandshake = async ( v: RELAY_PROTOCOL_VERSION, clientPubJwk: await exportPublicKeyJwk(ephemeralKeyPair.publicKey), nonce: bytesToBase64Url(nonce), - ...(localBatch ? { batch: true } : {}), }; + if (localBatch) hello.batch = true; + if (options.flowControl !== false) hello.flowControl = true; let established = false; return { helloText: JSON.stringify(hello), @@ -143,6 +147,7 @@ export const createClientHandshake = async ( type: 'established', // Batching runs only if both peers advertised it. batch: localBatch && message.batch === true, + flowControl: options.flowControl !== false && message.flowControl === true, channel: { encryptor: createFrameEncryptor(keys.clientToHost), decryptor: createFrameDecryptor(keys.hostToClient), @@ -161,6 +166,7 @@ export interface HostHandshake { export interface HostHandshakeOptions { /** Support frame batching. Default true; set false to force legacy behavior. */ batch?: boolean; + flowControl?: boolean; } export const createHostHandshake = ( @@ -216,13 +222,15 @@ export const createHostHandshake = ( const ready: E2eeReadyMessage = { t: 'ready', v: RELAY_PROTOCOL_VERSION, - ...(negotiatedBatch ? { batch: true } : {}), }; + if (negotiatedBatch) ready.batch = true; + if (options.flowControl !== false && message.flowControl === true) ready.flowControl = true; readyText = JSON.stringify(ready); established = true; return { type: 'established', batch: negotiatedBatch, + flowControl: ready.flowControl === true, replyText: readyText, channel: { encryptor: createFrameEncryptor(keys.hostToClient), diff --git a/packages/ui/src/lib/relay/protocol.ts b/packages/ui/src/lib/relay/protocol.ts index b14d34df..b6ff89b8 100644 --- a/packages/ui/src/lib/relay/protocol.ts +++ b/packages/ui/src/lib/relay/protocol.ts @@ -47,6 +47,7 @@ export const TunnelFrameType = { WsClose: 10, Ping: 11, Pong: 12, + DeliveryAck: 13, } as const; export type TunnelFrameTypeValue = (typeof TunnelFrameType)[keyof typeof TunnelFrameType]; @@ -97,6 +98,8 @@ export interface E2eeHelloMessage { // Capability advertisement: the client can pack multiple tunnel frames into // one encrypted WS message. Missing/false = legacy (one frame per message). batch?: boolean; + /** Client supports cumulative downstream delivery acknowledgements. */ + flowControl?: boolean; } export interface E2eeReadyMessage { @@ -105,6 +108,8 @@ export interface E2eeReadyMessage { // Host echoes `batch: true` only when it also supports batching AND the client // advertised it. Batching is enabled for the session only if both agree. batch?: boolean; + /** Enabled only when both peers support downstream flow control. */ + flowControl?: boolean; } // Relay-assigned WebSocket close codes. @@ -119,4 +124,3 @@ export const RelayCloseCode = { RekeyMismatch: 1008, ChannelFailure: 1011, } as const; - diff --git a/packages/ui/src/lib/relay/tunnel-client.ts b/packages/ui/src/lib/relay/tunnel-client.ts index 24a9d985..8f55a743 100644 --- a/packages/ui/src/lib/relay/tunnel-client.ts +++ b/packages/ui/src/lib/relay/tunnel-client.ts @@ -25,6 +25,7 @@ import { encodeFragmentedMessage, encodeJsonPayload, encodeTunnelFrame, + encodeDeliveryAck, type OutboundFrameBatcher, type TunnelFrame, } from './tunnel-codec'; @@ -162,6 +163,8 @@ export interface RelayTunnelClientOptions { batchWindowMs?: number; /** Advertise frame batching in the handshake. Default true. */ batch?: boolean; + /** Advertise downstream delivery acknowledgements. Default true. */ + flowControl?: boolean; reconnectBaseDelayMs?: number; reconnectMaxDelayMs?: number; hiddenOrOfflineMaxDelayMs?: number; @@ -350,7 +353,7 @@ export const createRelayTunnelClient = (options: RelayTunnelClientOptions): Rela let handshake; try { - handshake = await createClientHandshake(options.hostEncPubJwk, { batch: advertiseBatch }); + handshake = await createClientHandshake(options.hostEncPubJwk, { batch: advertiseBatch, flowControl: options.flowControl }); } catch (error) { if (generation !== attemptGeneration || closed) return; failAttempt(generation, toError(error), true); @@ -380,12 +383,20 @@ export const createRelayTunnelClient = (options: RelayTunnelClientOptions): Rela let channel: ActiveChannel | null = null; let cryptoChannel: EstablishedChannelCrypto | null = null; let batchNegotiated = false; + let flowControlNegotiated = false; + let receivedBytes = 0; + let acknowledgedBytes = 0; + let ackTimer: ReturnType | null = null; let batcher: OutboundFrameBatcher | null = null; // Idle tracking: updated on any non-Ping/Pong frame in EITHER direction. // Ping/Pong are excluded so the keepalive can't sustain itself. let lastActivityAt = Date.now(); const cleanupTimers = (): void => { + if (ackTimer !== null) { + clearTimeout(ackTimer); + ackTimer = null; + } if (helloInterval !== null) { clearInterval(helloInterval); helloInterval = null; @@ -444,9 +455,10 @@ export const createRelayTunnelClient = (options: RelayTunnelClientOptions): Rela } }; - const establish = (crypto: EstablishedChannelCrypto, batch: boolean): void => { + const establish = (crypto: EstablishedChannelCrypto, batch: boolean, flowControl: boolean): void => { cryptoChannel = crypto; batchNegotiated = batch; + flowControlNegotiated = flowControl; if (helloInterval !== null) { clearInterval(helloInterval); helloInterval = null; @@ -467,10 +479,11 @@ export const createRelayTunnelClient = (options: RelayTunnelClientOptions): Rela .then(async () => { if (channelObj.dead) return; const encrypted = await crypto.encryptor.encrypt(plaintext); + if (channelObj.dead) return; wire.send(encrypted); }) .catch(() => { - // Send failures surface via wire close; do not break the chain. + failAttemptLocal(new Error('relay encrypt/send failed')); }); }; const localBatcher = batch @@ -485,7 +498,7 @@ export const createRelayTunnelClient = (options: RelayTunnelClientOptions): Rela send(frame: Uint8Array): void { if (channelObj.dead) return; const frameType = frame[0] & ~TUNNEL_FRAGMENT_FLAG; - if (frameType !== TunnelFrameType.Ping && frameType !== TunnelFrameType.Pong) { + if (frameType !== TunnelFrameType.Ping && frameType !== TunnelFrameType.Pong && frameType !== TunnelFrameType.DeliveryAck) { lastActivityAt = Date.now(); } if (localBatcher) localBatcher.enqueue(frame); @@ -514,6 +527,14 @@ export const createRelayTunnelClient = (options: RelayTunnelClientOptions): Rela }, pingIntervalMs); }; + const acknowledgeDelivery = (): void => { + if (ackTimer !== null) clearTimeout(ackTimer); + ackTimer = null; + if (!channel || channel.dead || receivedBytes === acknowledgedBytes) return; + acknowledgedBytes = receivedBytes; + channel.send(encodeTunnelFrame(TunnelFrameType.DeliveryAck, 0, encodeDeliveryAck(receivedBytes))); + }; + const handleTunnelFrame = (channelObj: ActiveChannel, plaintext: Uint8Array): void => { let frame: TunnelFrame; try { @@ -522,6 +543,15 @@ export const createRelayTunnelClient = (options: RelayTunnelClientOptions): Rela failAttemptLocal(toError(error)); return; } + if (frame.frameType === TunnelFrameType.DeliveryAck) { + failAttemptLocal(new Error('unexpected downstream delivery acknowledgement')); + return; + } + if (flowControlNegotiated && frame.streamId !== 0) { + // Count even late/cancelled streams: they still consumed sender credit. + receivedBytes += plaintext.length; + if (ackTimer === null) ackTimer = setTimeout(acknowledgeDelivery, 10); + } // Any received frame proves the tunnel is alive — clear the pong deadline. if (pongDeadline !== null) { clearTimeout(pongDeadline); @@ -576,7 +606,7 @@ export const createRelayTunnelClient = (options: RelayTunnelClientOptions): Rela const action = await handshake.handleText(data); if (action.type === 'established') { if (cryptoChannel) return; - establish(action.channel, action.batch); + establish(action.channel, action.batch, action.flowControl); } else if (action.type === 'fail') { failAttemptLocal(new Error(`relay handshake failed: ${action.reason}`)); } @@ -605,23 +635,20 @@ export const createRelayTunnelClient = (options: RelayTunnelClientOptions): Rela failAttemptLocal(toError(error)); return; } - if (batchNegotiated) { - // One encrypted message may carry several tunnel frames; dispatch - // each in order through the same per-frame handling as legacy. - let frames: Uint8Array[]; - try { - frames = decodeFrameBatch(plaintext); - } catch (error) { - failAttemptLocal(toError(error)); - return; - } - for (const frame of frames) { - if (settled || generation !== attemptGeneration || currentChannel.dead) return; - handleTunnelFrame(currentChannel, frame); - } + let frames: Uint8Array[]; + try { + frames = batchNegotiated ? decodeFrameBatch(plaintext) : [plaintext]; + } catch (error) { + failAttemptLocal(toError(error)); return; } - handleTunnelFrame(currentChannel, plaintext); + for (const frame of frames) { + if (settled || generation !== attemptGeneration || currentChannel.dead) return; + handleTunnelFrame(currentChannel, frame); + } + // One ACK per received batch, rather than one per fragment. Small + // tails use the timer so a final frame cannot strand sender credit. + if (flowControlNegotiated && receivedBytes - acknowledgedBytes >= 16 * 1024) acknowledgeDelivery(); }) .catch((error: unknown) => { failAttemptLocal(toError(error)); diff --git a/packages/ui/src/lib/relay/tunnel-codec.ts b/packages/ui/src/lib/relay/tunnel-codec.ts index 2c390abe..97109e4b 100644 --- a/packages/ui/src/lib/relay/tunnel-codec.ts +++ b/packages/ui/src/lib/relay/tunnel-codec.ts @@ -18,6 +18,14 @@ import { const MAX_STREAM_ID = 0xffffffff; +/** Cumulative raw tunnel-frame bytes, excluding stream zero and batch/crypto overhead. */ +export const encodeDeliveryAck = (receivedBytes: number): Uint8Array => { + if (!Number.isSafeInteger(receivedBytes) || receivedBytes < 0) throw new Error('invalid delivery acknowledgement'); + const payload = new Uint8Array(8); + new DataView(payload.buffer).setBigUint64(0, BigInt(receivedBytes)); + return payload; +}; + export class TunnelCodecError extends Error { constructor(message: string) { super(message); diff --git a/packages/web/server/lib/relay/DOCUMENTATION.md b/packages/web/server/lib/relay/DOCUMENTATION.md index 2ce5f53a..0049079e 100644 --- a/packages/web/server/lib/relay/DOCUMENTATION.md +++ b/packages/web/server/lib/relay/DOCUMENTATION.md @@ -25,6 +25,7 @@ Host side (`packages/web/server/lib/relay/`): - `host-client.js` — the long-lived connection manager: one outbound control connection to the relay, a per-client data connection for each connected device, reconnect/backoff, and the E2EE responder handshake per connection. - `host-lock.js` — the per-machine host claim. Every local instance sharing the data dir shares the relay identity (same serverId), so concurrent relay hosts evict each other at the relay worker (`4001: Control replaced`) and paired devices land on whichever local process won last. The claim file (`/relay-host.lock`, `{ pid }`) makes this deterministic: `service.js` only starts the host when no LIVE process holds the claim (stale claims from dead pids are ignored), goes to `standby` otherwise, and a 30s watcher both takes over when the claimant dies and stands down when another process claims. A standby watcher waits a 2-minute grace after the claim frees before taking over, so a cleanly restarting host (app update/relaunch) — which reclaims at boot with no wait — always wins the restart window over a bystander instance. Explicit user intent — creating a pairing link or hitting `/relay/enable` — force-claims; the previous holder's watcher sees the takeover and backs off instead of fighting. Instances created with `allowPassiveHost: false` (dev servers via `OPENCHAMBER_RELAY_HOST=off`, the Electron dev shell via `OPENCHAMBER_ELECTRON_DEV`; `OPENCHAMBER_RELAY_HOST=on` overrides) never start the host passively at all — boot, demand reconcile, and watcher takeover leave them in `standby`; only explicit enable/pairing hosts there. The claim is cooperative (the relay worker still enforces the single host slot); it only decides which process keeps retrying. - `tunnel-host.js` — the per-connection dispatcher: decrypts tunnel frames and forwards HTTP/SSE/WS to the local server over loopback, then streams responses back. Enforces a path allowlist and never injects credentials. +- `downstream-scheduler.js` owns end-client delivery credit and round-robin transmission across response streams. It selects plaintext frames before encryption. - `e2ee.js`, `tunnel-codec.js` — host-side (JS) mirrors of the shared crypto and framing (see "Two implementations" below). Client side (`packages/ui/src/lib/relay/`): @@ -47,6 +48,53 @@ The host dispatcher restricts tunneled traffic to explicit path allowlists (one Request bodies crossing the tunnel are buffered on the host and forwarded to loopback only once the client's `StreamEnd` frame arrives (bodies above ~512 KB stream live instead). A body whose frames were lost in transit therefore never reaches the loopback server as an empty/truncated chunked body — the host aborts the stream and the client sees an ambiguous transport failure it can retry, instead of the loopback server's bare `400` ("Failed to send message (400)" from the mobile app). +## Downstream flow control + +`hello` and `ready` negotiate `flowControl: true` independently of `batch`. +When either peer omits the flag, the host keeps the legacy transmission path. +This capability controls host-to-client traffic only. Uploads retain their +existing request-body behavior. The Cloudflare broker needs no changes. + +The client sends encrypted `DeliveryAck` frames on stream zero. Their payload is +an eight-byte unsigned big-endian cumulative count of received raw tunnel-frame +bytes on nonzero streams, including the five-byte frame header. Batch envelopes, +encryption overhead and stream-zero keepalives are excluded. The count resets +with each handshake. ACKs cover complete frames and include late frames for +cancelled streams, since those frames still consumed credit. The host rejects +duplicate, regressing, partial-frame or beyond-sent acknowledgements. + +The client acknowledges after decoding and dispatching transport data, without +waiting for React rendering. It coalesces acknowledgements by byte count with a +short timer for the tail. Keepalives and ACK processing never await downstream +credit, which would deadlock the channel. + +The scheduler starts with a small byte window. A backlogged sender can grow its +window when ACK latency stays near the best observed round trip, and reduces it +when delivery delay rises. Sparse traffic cannot inflate the window for a later +bulk burst. The window has a hard upper bound; a sudden bandwidth drop can still +delay bytes already sent, but cannot create an unlimited network backlog. + +HTTP bodies and WS messages are sliced into small frames. The scheduler rotates +streams, preserves each stream's order, and batches selected frames before the +serialized encryption step. Producers offer a bounded group of slices and await +transmission before reading more. The legacy timed batcher remains in use only +without negotiated flow control. + +HTTP/SSE backpressure reaches the loopback response reader. Node's `ws` pauses +loopback socket reads while output is waiting. Bun's `ws` shim cannot pause, so +the dispatcher instead enforces a connection-wide pending WS byte/message cap +and explicitly aborts the offending substream on overflow. It never silently +discards deltas. WS output messages are serialized through their last fragment, +and the close frame follows pending output. Cancellation releases queued work; +channel teardown releases all producers and acknowledgement state. + +Regression coverage lives in `downstream-scheduler.test.js`, `flow-control.test.js` +and `cross-compat.test.js`. The end-to-end fixture uses the real client, host, +crypto and loopback requests with a byte-paced relay leg. It compares delivery +ordering, queue growth and a concurrent small response, including both legacy +fallback directions, unbatched operation, cancellation and fragmented WS output. +It does not replace a physical iOS/LTE check. + ## Authentication model - The tunnel is **transport only**. The OpenChamber server still authenticates every tunneled request exactly as it authenticates a direct remote client. The relay path grants reachability, not authorization. diff --git a/packages/web/server/lib/relay/cross-compat.test.js b/packages/web/server/lib/relay/cross-compat.test.js index 67918f29..1721f471 100644 --- a/packages/web/server/lib/relay/cross-compat.test.js +++ b/packages/web/server/lib/relay/cross-compat.test.js @@ -12,16 +12,42 @@ import { decodeTunnelFrame as jsDecode, encodeFrameBatch as jsEncodeBatch, encodeTunnelFrame as jsEncode, + decodeDeliveryAck, } from './tunnel-codec.js'; import { decodeFrameBatch as tsDecodeBatch, decodeTunnelFrame as tsDecode, encodeFrameBatch as tsEncodeBatch, encodeTunnelFrame as tsEncode, + encodeDeliveryAck, } from '../../../../ui/src/lib/relay/tunnel-codec.ts'; import { TunnelFrameType as TsFrameType } from '../../../../ui/src/lib/relay/protocol.ts'; describe('relay JS-host <-> TS-client cross compatibility', () => { + it('negotiates flow control independently of batching, with legacy fallback on either side', async () => { + const keys = await generateEcdhKeyPair(); + const pub = await exportPublicKeyJwk(keys.publicKey); + for (const batch of [false, true]) { + for (const clientFlow of [false, true]) { + for (const hostFlow of [false, true]) { + const client = await createClientHandshake(pub, { batch, flowControl: clientFlow }); + const host = createHostHandshake(keys.privateKey, { batch, flowControl: hostFlow }); + const ready = await host.handleText(client.helloText); + const established = await client.handleText(ready.replyText); + expect(ready.flowControl).toBe(clientFlow && hostFlow); + expect(established.flowControl).toBe(clientFlow && hostFlow); + expect(established.batch).toBe(batch); + expect((await host.handleText(client.helloText)).text).toBe(ready.replyText); + } + } + } + expect(JsFrameType).toEqual(TsFrameType); + for (const bytes of [0, 8197, 2 ** 32 + 17, Number.MAX_SAFE_INTEGER]) { + expect(decodeDeliveryAck(encodeDeliveryAck(bytes))).toBe(bytes); + } + expect(() => decodeDeliveryAck(new Uint8Array(7))).toThrow(); + expect(() => decodeDeliveryAck(new Uint8Array(8).fill(255))).toThrow(); + }); it('completes a handshake and exchanges frames both ways', async () => { const hostKeys = await generateEcdhKeyPair(); const hostPubJwk = await exportPublicKeyJwk(hostKeys.publicKey); diff --git a/packages/web/server/lib/relay/downstream-scheduler.js b/packages/web/server/lib/relay/downstream-scheduler.js new file mode 100644 index 00000000..0ca4663d --- /dev/null +++ b/packages/web/server/lib/relay/downstream-scheduler.js @@ -0,0 +1,154 @@ +// Downstream credit belongs to the end client, not the relay's TCP connection. +// Producers await send(), so HTTP readers stop before building a plaintext queue. +import { decodeTunnelFrame } from './tunnel-codec.js'; +import { MAX_PLAINTEXT_FRAME_BYTES } from './e2ee.js'; + +export const DOWNSTREAM_CHUNK_BYTES = 8 * 1024; +const MIN_WINDOW_BYTES = 64 * 1024; +const MAX_WINDOW_BYTES = 1024 * 1024; +const MAX_QUEUED_BYTES = 4 * 1024 * 1024; +const MAX_QUEUED_FRAMES = 4096; +const MAX_OUTSTANDING_FRAMES = 4096; +const QUEUE_DELAY_TARGET_MS = 100; + +export const createDownstreamScheduler = ({ sendBatch, onError, maxBatchFrames = 32, now = () => performance.now() }) => { + const queues = new Map(); + const outstanding = []; + let queuedBytes = 0; + let queuedFrames = 0; + let sentBytes = 0; + let acknowledgedBytes = 0; + let windowBytes = MIN_WINDOW_BYTES; + let minRtt = Infinity; + let adjustmentBoundary = MIN_WINDOW_BYTES; + let windowLimited = false; + let draining = false; + let closed = false; + let sending = []; + + const close = () => { + closed = true; + for (const queue of queues.values()) for (const entry of queue) entry.resolve(); + for (const entry of sending) entry.resolve(); + sending = []; + queues.clear(); + outstanding.length = 0; + queuedBytes = 0; + queuedFrames = 0; + }; + + const fail = error => { + close(); + onError(error); + }; + + const drain = async () => { + if (draining || closed) return; + draining = true; + try { + while (!closed) { + const selected = []; + let batchBytes = 1; + while (selected.length < maxBatchFrames) { + // Stream zero carries only tiny keepalives; it must not wait for credit. + const next = queues.has(0) ? [0, queues.get(0)] : queues.entries().next().value; + if (!next) break; + const [streamId, queue] = next; + const entry = queue[0]; + // Do not skip a larger head indefinitely in favour of smaller frames. + if (streamId !== 0 && sentBytes - acknowledgedBytes + entry.frame.length > windowBytes) { + windowLimited = true; + break; + } + if (streamId !== 0 && outstanding.length >= MAX_OUTSTANDING_FRAMES) break; + if (batchBytes + 4 + entry.frame.length > MAX_PLAINTEXT_FRAME_BYTES) break; + queue.shift(); + queues.delete(streamId); + if (queue.length) queues.set(streamId, queue); + queuedBytes -= entry.frame.length; + queuedFrames -= 1; + if (streamId !== 0) { + sentBytes += entry.frame.length; + outstanding.push({ end: sentBytes, at: now() }); + } + selected.push(entry); + batchBytes += 4 + entry.frame.length; + } + if (!selected.length) break; + sending = selected; + try { + // Selection precedes encryption. Never reorder encrypted IV counters. + await sendBatch(selected.map(entry => entry.frame)); + } finally { + for (const entry of selected) entry.resolve(); + sending = []; + } + } + } catch (error) { + fail(error); + } finally { + draining = false; + } + }; + + return { + send(frame) { + if (closed) return Promise.resolve(); + const { streamId } = decodeTunnelFrame(frame); + if (queuedBytes + frame.length > MAX_QUEUED_BYTES || queuedFrames >= MAX_QUEUED_FRAMES) { + fail(new Error('relay downstream queue limit exceeded')); + return Promise.resolve(); + } + return new Promise(resolve => { + const queue = queues.get(streamId) ?? []; + queue.push({ frame, resolve, streamId }); + queues.set(streamId, queue); + queuedBytes += frame.length; + queuedFrames += 1; + // Let concurrent streams join this round before selecting the next frame. + queueMicrotask(() => { void drain(); }); + }); + }, + acknowledge(bytes) { + if (closed) return; + if (!Number.isSafeInteger(bytes) || bytes <= acknowledgedBytes || bytes > sentBytes) { + throw new Error('invalid relay delivery acknowledgement'); + } + // An ACK may cover several complete frames, never a partial frame. + const endIndex = outstanding.findIndex(entry => entry.end === bytes); + if (endIndex < 0) throw new Error('relay acknowledgement is not a frame boundary'); + const rtt = Math.max(1, now() - outstanding[endIndex].at); + outstanding.splice(0, endIndex + 1); + acknowledgedBytes = bytes; + minRtt = Math.min(minRtt, rtt); + // Grow on a clear path, reduce when the client's delivery delay grows. + // Adjust once per window, not once per frame in a burst of ACKs. + if (bytes >= adjustmentBoundary) { + if (rtt > minRtt + QUEUE_DELAY_TARGET_MS) { + windowBytes = Math.max(MIN_WINDOW_BYTES, Math.floor(windowBytes / 2)); + } else if (windowLimited) { + // Sparse token traffic must not inflate a future bash burst's window. + windowBytes = Math.min(MAX_WINDOW_BYTES, windowBytes * 2); + } + windowLimited = false; + adjustmentBoundary = bytes + windowBytes; + } + void drain(); + }, + cancel(streamId) { + // Already-selected bytes may still arrive and must still be ACKed, but + // cancelled producers need not wait for an in-progress encryption call. + for (const entry of sending) if (entry.streamId === streamId) entry.resolve(); + const queue = queues.get(streamId); + if (!queue) return; + queues.delete(streamId); + for (const entry of queue) { + queuedBytes -= entry.frame.length; + queuedFrames -= 1; + entry.resolve(); + } + void drain(); + }, + close, + }; +}; diff --git a/packages/web/server/lib/relay/downstream-scheduler.test.js b/packages/web/server/lib/relay/downstream-scheduler.test.js new file mode 100644 index 00000000..90eb4bfa --- /dev/null +++ b/packages/web/server/lib/relay/downstream-scheduler.test.js @@ -0,0 +1,158 @@ +import { describe, expect, test } from 'bun:test'; +import { createDownstreamScheduler, DOWNSTREAM_CHUNK_BYTES } from './downstream-scheduler.js'; +import { encodeTunnelFrame, decodeTunnelFrame, TunnelFrameType } from './tunnel-codec.js'; + +const tick = () => new Promise(resolve => setTimeout(resolve, 0)); +const body = (streamId, value = 0, size = DOWNSTREAM_CHUNK_BYTES) => + encodeTunnelFrame(TunnelFrameType.HttpBody, streamId, new Uint8Array(size).fill(value)); + +function fixture() { + const sent = []; + const errors = []; + let time = 0; + const scheduler = createDownstreamScheduler({ + sendBatch: async frames => { sent.push(...frames); }, + onError: error => errors.push(error.message), + now: () => time, + }); + return { scheduler, sent, errors, advance: ms => { time += ms; } }; +} + +describe('relay downstream scheduling', () => { + test('stops without end-client ACKs even when the relay socket accepts everything', async () => { + const { scheduler, sent } = fixture(); + const pending = Array.from({ length: 100 }, () => scheduler.send(body(1))); + await tick(); + const bytes = sent.reduce((sum, frame) => sum + frame.length, 0); + expect(bytes).toBeLessThanOrEqual(64 * 1024); + expect(bytes).toBeGreaterThan(48 * 1024); + await tick(); + expect(sent.reduce((sum, frame) => sum + frame.length, 0)).toBe(bytes); + scheduler.acknowledge(bytes); + await tick(); + expect(sent.reduce((sum, frame) => sum + frame.length, 0)).toBeGreaterThan(bytes); + scheduler.close(); + await Promise.all(pending); + }); + + test('rotates streams while preserving every delta and each stream end', async () => { + const { scheduler, sent } = fixture(); + const pending = []; + for (let value = 0; value < 3; value++) pending.push(scheduler.send(body(1, value))); + pending.push(scheduler.send(encodeTunnelFrame(TunnelFrameType.StreamEnd, 1, new Uint8Array()))); + pending.push(scheduler.send(body(3, 99, 10))); + await Promise.all(pending); + expect(sent.map(frame => decodeTunnelFrame(frame).streamId)).toEqual([1, 3, 1, 1, 1]); + expect(sent.filter(frame => decodeTunnelFrame(frame).streamId === 1).map(frame => decodeTunnelFrame(frame).payload[0])).toEqual([0, 1, 2, undefined]); + scheduler.close(); + }); + + test('keepalives bypass credit; cancellation and close release blocked producers', async () => { + const { scheduler, sent } = fixture(); + const pending = Array.from({ length: 20 }, () => scheduler.send(body(1))); + await tick(); + const before = sent.length; + await scheduler.send(encodeTunnelFrame(TunnelFrameType.Pong, 0, new Uint8Array())); + expect(sent.length).toBe(before + 1); + scheduler.cancel(1); + await Promise.all(pending); + const blocked = scheduler.send(body(3)); + scheduler.close(); + await blocked; + // Teardown makes late ACKs harmless. + scheduler.acknowledge(123); + }); + + test('rejects forged, replayed, regressing and partial-frame ACKs', async () => { + const { scheduler, sent } = fixture(); + await scheduler.send(body(1)); + const bytes = sent[0].length; + for (const value of [0, -1, 1, bytes + 1, NaN, Infinity]) { + expect(() => scheduler.acknowledge(value)).toThrow(); + } + scheduler.acknowledge(bytes); + expect(() => scheduler.acknowledge(bytes)).toThrow(); + scheduler.close(); + }); + + test('bounds pending memory and reports overflow instead of silently dropping deltas', async () => { + const { scheduler, errors } = fixture(); + const pending = Array.from({ length: 600 }, () => scheduler.send(body(1))); + await Promise.all(pending); + expect(errors).toEqual(['relay downstream queue limit exceeded']); + }); + + test('send failure settles selected and queued producers', async () => { + const errors = []; + const scheduler = createDownstreamScheduler({ + sendBatch: async () => { throw new Error('wire failed'); }, + onError: error => errors.push(error.message), + }); + await Promise.all([scheduler.send(body(1)), scheduler.send(body(3))]); + expect(errors).toEqual(['wire failed']); + }); + + test('teardown releases a producer even while encryption is still pending', async () => { + let release; + let started = false; + const scheduler = createDownstreamScheduler({ + sendBatch: () => new Promise(resolve => { started = true; release = resolve; }), + onError: error => { throw error; }, + }); + const pending = scheduler.send(body(1)); + await tick(); + expect(started).toBe(true); + scheduler.close(); + await pending; + release(); + await tick(); + }); + + test('sparse traffic cannot inflate the window for a later bulk burst', async () => { + const { scheduler, sent, advance } = fixture(); + let acknowledged = 0; + for (let i = 0; i < 100; i++) { + const frame = body(1); + await scheduler.send(frame); + advance(50); + acknowledged += frame.length; + scheduler.acknowledge(acknowledged); + } + const before = sent.length; + const pending = Array.from({ length: 100 }, () => scheduler.send(body(1))); + await tick(); + expect(sent.slice(before).reduce((sum, frame) => sum + frame.length, 0)).toBeLessThanOrEqual(64 * 1024); + scheduler.close(); + await Promise.all(pending); + }); + + test('adapts to a clear high-RTT path, then contracts when delivery backs up', async () => { + const { scheduler, sent, advance, errors } = fixture(); + let acknowledged = 0; + let maxFlight = 0; + const pending = Array.from({ length: 400 }, () => scheduler.send(body(1))); + for (let round = 0; round < 8; round++) { + await tick(); + const delivered = sent.reduce((sum, frame) => sum + frame.length, 0); + maxFlight = Math.max(maxFlight, delivered - acknowledged); + advance(200); + scheduler.acknowledge(delivered); + acknowledged = delivered; + } + expect(maxFlight).toBeGreaterThan(256 * 1024); + // Keep the sender backlogged while ACK latency jumps to a second. + for (let round = 0; round < 12; round++) { + for (let i = 0; i < 30; i++) pending.push(scheduler.send(body(1))); + await tick(); + const delivered = sent.reduce((sum, frame) => sum + frame.length, 0); + advance(1000); + scheduler.acknowledge(delivered); + acknowledged = delivered; + } + await tick(); + expect(sent.reduce((sum, frame) => sum + frame.length, 0) - acknowledged).toBeLessThanOrEqual(64 * 1024); + expect(errors).toEqual([]); + scheduler.close(); + await Promise.all(pending); + }); +}); diff --git a/packages/web/server/lib/relay/e2ee.js b/packages/web/server/lib/relay/e2ee.js index 00853af0..cd632c6b 100644 --- a/packages/web/server/lib/relay/e2ee.js +++ b/packages/web/server/lib/relay/e2ee.js @@ -252,11 +252,12 @@ const parseHandshakeMessage = (raw) => { if (parsed.v !== RELAY_PROTOCOL_VERSION) return null; // Unknown/missing capability flag = false = legacy behavior. const batch = parsed.batch === true; + const flowControl = parsed.flowControl === true; if (parsed.t === 'ready') { - return { t: 'ready', v: RELAY_PROTOCOL_VERSION, batch }; + return { t: 'ready', v: RELAY_PROTOCOL_VERSION, batch, flowControl }; } if (parsed.t === 'hello' && typeof parsed.nonce === 'string' && typeof parsed.clientPubJwk === 'object' && parsed.clientPubJwk !== null) { - return { t: 'hello', v: RELAY_PROTOCOL_VERSION, clientPubJwk: parsed.clientPubJwk, nonce: parsed.nonce, batch }; + return { t: 'hello', v: RELAY_PROTOCOL_VERSION, clientPubJwk: parsed.clientPubJwk, nonce: parsed.nonce, batch, flowControl }; } return null; }; @@ -275,7 +276,7 @@ const failClosed = (reason) => ({ * { type: 'ignore' } — drop the frame * { type: 'fail', closeCode, reason } — close the socket with closeCode * @param {CryptoKey} hostEncPrivateKey long-lived ECDH private key - * @param {{ batch?: boolean }} [options] `batch` defaults true; set false to force legacy behavior + * @param {{ batch?: boolean, flowControl?: boolean }} [options] Capabilities default true. */ export const createHostHandshake = (hostEncPrivateKey, options = {}) => { const localBatch = options.batch !== false; @@ -321,15 +322,16 @@ export const createHostHandshake = (hostEncPrivateKey, options = {}) => { acceptedClientKeyFingerprint = fingerprint; // Batching runs only if both peers advertised it. negotiatedBatch = localBatch && message.batch === true; - readyText = JSON.stringify( - negotiatedBatch - ? { t: 'ready', v: RELAY_PROTOCOL_VERSION, batch: true } - : { t: 'ready', v: RELAY_PROTOCOL_VERSION }, - ); + const flowControl = options.flowControl !== false && message.flowControl === true; + const ready = { t: 'ready', v: RELAY_PROTOCOL_VERSION }; + if (negotiatedBatch) ready.batch = true; + if (flowControl) ready.flowControl = true; + readyText = JSON.stringify(ready); established = true; return { type: 'established', batch: negotiatedBatch, + flowControl, replyText: readyText, channel: { encryptor: createFrameEncryptor(keys.hostToClient), diff --git a/packages/web/server/lib/relay/flow-control.test.js b/packages/web/server/lib/relay/flow-control.test.js new file mode 100644 index 00000000..4601b5f2 --- /dev/null +++ b/packages/web/server/lib/relay/flow-control.test.js @@ -0,0 +1,181 @@ +import { describe, expect, test } from 'bun:test'; +import http from 'node:http'; +import { WebSocketServer } from 'ws'; +import { startRelayHost } from './host-client.js'; +import { generateEcdhKeyPair, exportPublicKeyJwk } from './e2ee.js'; +import { createRelayTunnelClient } from '../../../../ui/src/lib/relay/tunnel-client.ts'; + +const waitFor = async predicate => { + for (let attempt = 0; attempt < 1000; attempt++) { + if (predicate()) return; + await new Promise(resolve => setTimeout(resolve, 5)); + } + throw new Error('test condition timed out'); +}; + +async function exercise({ clientFlow = true, hostFlow = true, batch = true, cancel = false, websocket = false } = {}) { + const total = 512 * 1024; + const content = Buffer.alloc(total, 'a'); + content.fill('b', total / 2); + const upstream = http.createServer((req, res) => { + // HTTP bearer ownership is unchanged by flow control. + if (req.headers.authorization !== 'Bearer fixture-token') { res.writeHead(401); res.end(); return; } + res.end(req.url === '/health' ? 'ok' : content); + }); + const upstreamWs = new WebSocketServer({ noServer: true }); + upstream.on('upgrade', (req, socket, head) => { + // Mirror the real URL-token and loopback-origin gates for tunneled WS. + const query = new URL(req.url, 'http://localhost'); + if (query.searchParams.get('oc_url_token') !== 'fixture-url-token' || req.headers.origin !== `http://127.0.0.1:${upstream.address().port}`) { + socket.end('HTTP/1.1 403 Forbidden\r\n\r\n'); + return; + } + upstreamWs.handleUpgrade(req, socket, head, ws => { + ws.send(content.subarray(0, total / 2)); + ws.send(content.subarray(total / 2)); + // Let Bun finish the local write turn, then close while the throttled + // tunnel still has most of the two messages queued. + setTimeout(() => ws.close(1000, 'complete'), 20); + }); + }); + await new Promise(resolve => upstream.listen(0, '127.0.0.1', resolve)); + const relay = new WebSocketServer({ host: '127.0.0.1', port: 0 }); + await new Promise(resolve => relay.on('listening', resolve)); + let control; + let clientSocket; + let hostData; + const waiting = []; + const down = []; + let pendingBytes = 0; + let peakBytes = 0; + // This fixture brokers an isolated room only; it does not replace production + // relay authentication. Neither WebSocket leg fails or drops a frame. + relay.on('connection', (socket, req) => { + const role = new URL(req.url, 'http://localhost').searchParams.get('role'); + if (role === 'host-control') { + control = socket; + socket.send(JSON.stringify({ type: 'sync', connectionIds: clientSocket ? ['fixture'] : [] })); + } else if (role === 'client') { + clientSocket = socket; + control?.send(JSON.stringify({ type: 'connected', connectionId: 'fixture' })); + socket.on('message', (data, binary) => { + if (hostData) hostData.send(data, { binary }); + else waiting.push({ data, binary }); + }); + } else { + hostData = socket; + for (const frame of waiting.splice(0)) socket.send(frame.data, { binary: frame.binary }); + socket.on('message', (data, binary) => { + down.push({ data, binary }); + pendingBytes += data.length; + peakBytes = Math.max(peakBytes, pendingBytes); + }); + } + }); + // Drain at a byte budget, not a frame count: frame-size changes cannot make + // the test's effective bandwidth change. Preserve whole WS messages in order. + let credit = 0; + const drain = setInterval(() => { + credit = Math.min(128 * 1024, credit + 4096); + while (down.length && down[0].data.length <= credit) { + const frame = down.shift(); + credit -= frame.data.length; + pendingBytes -= frame.data.length; + clientSocket.send(frame.data, { binary: frame.binary }); + } + }, 5); + const keys = await generateEcdhKeyPair(); + const relayUrl = `ws://127.0.0.1:${relay.address().port}/ws`; + const host = startRelayHost({ + relayUrl, localPort: upstream.address().port, batch, flowControl: hostFlow, + identity: { serverId: 'fixture', hostEncPrivateKey: keys.privateKey, signRelayAuth: () => ({ ts: 0, sig: '', pk: '' }) }, + }); + const client = createRelayTunnelClient({ relayUrl, serverId: 'fixture', hostEncPubJwk: await exportPublicKeyJwk(keys.publicKey), batch, flowControl: clientFlow }); + const states = []; + client.subscribeStatus(status => states.push(status.state)); + const headers = { authorization: 'Bearer fixture-token' }; + let received = 0; + let finished = false; + let wsMessages = 0; + let streamFailure; + const chunks = []; + const abort = new AbortController(); + try { + await waitFor(() => client.getStatus().state === 'connected'); + let stream; + if (websocket) { + const socket = client.openWebSocket('/api/terminal/ws?oc_url_token=fixture-url-token'); + stream = new Promise((resolve, reject) => { + socket.onmessage = event => { + const bytes = new Uint8Array(event.data); + chunks.push(bytes); + received += bytes.length; + wsMessages++; + }; + socket.onclose = event => { + finished = true; + if (event.code !== 1000) reject(new Error(event.reason)); + else resolve(); + }; + }); + } else { + const response = await client.fetch('/api/global/event', { headers, signal: abort.signal }); + stream = (async () => { + try { + for await (const chunk of response.body) { chunks.push(chunk); received += chunk.length; } + } catch (error) { + if (!cancel || error.name !== 'AbortError') throw error; + } finally { finished = true; } + })(); + } + stream = stream.catch(error => { streamFailure = error; }); + await waitFor(() => received > 0 || streamFailure); + if (streamFailure) throw streamFailure; + if (cancel) abort.abort(); + const response = await client.fetch('/health', { headers, signal: AbortSignal.timeout(5000) }); + expect(await response.text()).toBe('ok'); + const bytesAtProbe = received; + await stream; + expect(streamFailure).toBeUndefined(); + expect(finished).toBe(true); + if (!cancel) expect(Buffer.concat(chunks).equals(content)).toBe(true); + if (websocket) expect(wsMessages).toBe(2); + expect(states).toEqual(['connected']); + return { bytesAtProbe, peakBytes, total }; + } finally { + client.close(); + host.stop(); + clearInterval(drain); + for (const ws of relay.clients) ws.terminate(); + relay.close(); + for (const ws of upstreamWs.clients) ws.terminate(); + upstreamWs.close(); + upstream.closeAllConnections(); + upstream.close(); + } +} + +describe('end-to-end downstream credit', () => { + test('small HTTP response overtakes bulk output without losing bytes or reconnecting', async () => { + const result = await exercise(); + expect(result.bytesAtProbe).toBeLessThan(result.total / 2); + expect(result.peakBytes).toBeLessThan(192 * 1024); + }); + test('works without batching', async () => { + const result = await exercise({ batch: false }); + expect(result.bytesAtProbe).toBeLessThan(result.total / 2); + }); + test('cancelling a blocked stream leaves unrelated requests usable', async () => { + const result = await exercise({ cancel: true }); + expect(result.bytesAtProbe).toBeLessThan(result.total); + }); + test('fragmented WS messages remain complete and precede the close', async () => { + await exercise({ websocket: true }); + }); + for (const legacy of [{ clientFlow: false }, { hostFlow: false }]) { + test(`legacy fallback ${JSON.stringify(legacy)} preserves delivery`, async () => { + const result = await exercise(legacy); + expect(result.bytesAtProbe).toBe(result.total); + }); + } +}); diff --git a/packages/web/server/lib/relay/host-client.js b/packages/web/server/lib/relay/host-client.js index 58f29091..ed5f1bde 100644 --- a/packages/web/server/lib/relay/host-client.js +++ b/packages/web/server/lib/relay/host-client.js @@ -6,8 +6,9 @@ import { WebSocket } from 'ws'; import { RELAY_PROTOCOL_VERSION, RelayCloseCode, createHostHandshake } from './e2ee.js'; -import { createOutboundFrameBatcher, decodeFrameBatch } from './tunnel-codec.js'; +import { createOutboundFrameBatcher, decodeFrameBatch, decodeTunnelFrame, decodeDeliveryAck, encodeFrameBatch, TunnelFrameType } from './tunnel-codec.js'; import { createTunnelHost } from './tunnel-host.js'; +import { createDownstreamScheduler, DOWNSTREAM_CHUNK_BYTES } from './downstream-scheduler.js'; const BACKOFF_BASE_MS = 1000; const BACKOFF_CAP_MS = 30000; @@ -48,7 +49,7 @@ const resolveBatchWindowMs = (option) => { * logger?: Pick, * }} options */ -export const startRelayHost = ({ relayUrl, identity, localPort, getLocalPort, onStatus, logger = console, batchWindowMs, batch }) => { +export const startRelayHost = ({ relayUrl, identity, localPort, getLocalPort, onStatus, logger = console, batchWindowMs, batch, flowControl }) => { const resolveLocalPort = typeof getLocalPort === 'function' ? getLocalPort : () => localPort; const localBatch = batch !== false; const resolvedBatchWindowMs = resolveBatchWindowMs(batchWindowMs); @@ -95,6 +96,7 @@ export const startRelayHost = ({ relayUrl, identity, localPort, getLocalPort, on dataSockets.delete(connectionId); if (entry.openTimer) clearTimeout(entry.openTimer); entry.batcher?.dispose(); + entry.scheduler?.close(); entry.tunnel?.close(); try { if (entry.socket.readyState === WebSocket.OPEN || entry.socket.readyState === WebSocket.CONNECTING) { @@ -118,14 +120,14 @@ export const startRelayHost = ({ relayUrl, identity, localPort, getLocalPort, on return; } - const entry = { socket, tunnel: null, openTimer: null, batcher: null, lastActivityAt: Date.now() }; + const entry = { socket, tunnel: null, openTimer: null, batcher: null, scheduler: null, lastActivityAt: Date.now() }; dataSockets.set(connectionId, entry); entry.openTimer = setTimeout(() => { logger.warn('[Relay] host-data socket open timeout'); teardownDataSocket(connectionId); }, DATA_SOCKET_OPEN_TIMEOUT_MS); - const handshake = createHostHandshake(identity.hostEncPrivateKey, { batch: localBatch }); + const handshake = createHostHandshake(identity.hostEncPrivateKey, { batch: localBatch, flowControl }); let channel = null; let batchNegotiated = false; // Serialize async message handling so encrypted frame order (and the @@ -140,11 +142,14 @@ export const startRelayHost = ({ relayUrl, identity, localPort, getLocalPort, on .then(async () => { if (dataSockets.get(connectionId) !== entry || socket.readyState !== WebSocket.OPEN || !channel) return; const encrypted = await channel.encryptor.encrypt(plaintext); + if (dataSockets.get(connectionId) !== entry || socket.readyState !== WebSocket.OPEN) return; socket.send(encrypted, { binary: true }); }) .catch((error) => { logger.warn(`[Relay] host-data send failed: ${error?.message ?? error}`); + failChannel(RelayCloseCode.ChannelFailure, 'send failed'); }); + return sendChain; }; const failChannel = (closeCode, reason) => { @@ -167,17 +172,25 @@ export const startRelayHost = ({ relayUrl, identity, localPort, getLocalPort, on } else if (action.type === 'established') { channel = action.channel; batchNegotiated = action.batch === true; - entry.batcher = batchNegotiated + entry.scheduler = action.flowControl ? createDownstreamScheduler({ + sendBatch: frames => sendEncryptedPlaintext(batchNegotiated ? encodeFrameBatch(frames) : frames[0]), + maxBatchFrames: batchNegotiated ? 32 : 1, + onError: () => failChannel(RelayCloseCode.ChannelFailure, 'downstream queue failed'), + }) : null; + entry.batcher = batchNegotiated && !entry.scheduler ? createOutboundFrameBatcher({ windowMs: resolvedBatchWindowMs, sendBatch: sendEncryptedPlaintext }) : null; entry.tunnel = createTunnelHost({ connectionId, getLocalPort: resolveLocalPort, getBufferedAmount: () => socket.bufferedAmount, + responseChunkBytes: entry.scheduler ? DOWNSTREAM_CHUNK_BYTES : undefined, + cancelPendingFrames: streamId => entry.scheduler?.cancel(streamId), sendFrame: (plaintextFrame) => { if (dataSockets.get(connectionId) !== entry || socket.readyState !== WebSocket.OPEN) return; + if (entry.scheduler) return entry.scheduler.send(plaintextFrame); if (entry.batcher) entry.batcher.enqueue(plaintextFrame); - else sendEncryptedPlaintext(plaintextFrame); + else return sendEncryptedPlaintext(plaintextFrame); }, }); if (action.replyText) socket.send(action.replyText); @@ -205,16 +218,32 @@ export const startRelayHost = ({ relayUrl, identity, localPort, getLocalPort, on // in order through the same per-frame handling as legacy. for (const frame of decodeFrameBatch(plaintext)) { if (dataSockets.get(connectionId) !== entry) return; - await entry.tunnel.handleFrame(frame); + await dispatchFrame(frame); } } else { - await entry.tunnel.handleFrame(plaintext); + await dispatchFrame(plaintext); } } catch (error) { logger.warn(`[Relay] tunnel frame handling failed: ${error?.message ?? error}`); + failChannel(RelayCloseCode.ChannelFailure, 'invalid tunnel frame'); } }; + const dispatchFrame = (plaintext) => { + const frame = decodeTunnelFrame(plaintext); + if (frame.frameType === TunnelFrameType.DeliveryAck) { + if (!entry.scheduler || frame.streamId !== 0 || frame.hasMoreFragments) { + throw new Error('unexpected delivery acknowledgement'); + } + entry.scheduler.acknowledge(decodeDeliveryAck(frame.payload)); + return; + } + // Never await outbound credit on the receive chain: ACKs use this chain too. + void entry.tunnel.handleFrame(plaintext).catch(() => { + failChannel(RelayCloseCode.ChannelFailure, 'invalid tunnel frame'); + }); + }; + socket.on('open', () => { if (entry.openTimer) { clearTimeout(entry.openTimer); diff --git a/packages/web/server/lib/relay/tunnel-codec.js b/packages/web/server/lib/relay/tunnel-codec.js index f402bc30..9a290883 100644 --- a/packages/web/server/lib/relay/tunnel-codec.js +++ b/packages/web/server/lib/relay/tunnel-codec.js @@ -34,6 +34,7 @@ export const TunnelFrameType = { WsClose: 10, Ping: 11, Pong: 12, + DeliveryAck: 13, }; const TUNNEL_FRAME_TYPE_VALUES = new Set(Object.values(TunnelFrameType)); @@ -43,6 +44,14 @@ export const isTunnelFrameType = (value) => TUNNEL_FRAME_TYPE_VALUES.has(value); const MAX_STREAM_ID = 0xffffffff; +/** Decode the client's cumulative raw tunnel-frame byte count. */ +export const decodeDeliveryAck = (payload) => { + if (payload.length !== 8) throw new Error('invalid delivery acknowledgement'); + const bytes = Number(new DataView(payload.buffer, payload.byteOffset, payload.byteLength).getBigUint64(0)); + if (!Number.isSafeInteger(bytes)) throw new Error('invalid delivery acknowledgement'); + return bytes; +}; + export class TunnelCodecError extends Error { constructor(message) { super(message); diff --git a/packages/web/server/lib/relay/tunnel-host.js b/packages/web/server/lib/relay/tunnel-host.js index 91c31b68..f9297309 100644 --- a/packages/web/server/lib/relay/tunnel-host.js +++ b/packages/web/server/lib/relay/tunnel-host.js @@ -15,7 +15,6 @@ import { createFragmentAssembler, decodeJsonPayload, decodeTunnelFrame, - encodeFragmentedMessage, encodeJsonPayload, encodeTunnelFrame, } from './tunnel-codec.js'; @@ -74,6 +73,11 @@ const BODY_BUFFER_MAX_BYTES = 512 * 1024; // completes, so a stalled tunnel converts into an ambiguous transport failure // (which the client already retries) instead of a hung loopback request. const BODY_DELIVERY_TIMEOUT_MS = 15_000; +const MAX_PENDING_WS_BYTES = 16 * 1024 * 1024; +const MAX_PENDING_WS_MESSAGES = 1024; +// Node ws supports read backpressure. Bun's ws shim does not implement pause; +// there the bounded queue fails the affected substream explicitly on overflow. +const CAN_PAUSE_WS_READS = !process.versions.bun; const sleep = (ms) => new Promise((resolve) => setTimeout(resolve, ms)); @@ -99,13 +103,17 @@ const isWsClosePayload = (parsed) => Boolean(parsed && typeof parsed === 'object * sendFrame: (plaintextFrame: Uint8Array) => void | Promise, * getBufferedAmount: () => number, * bodyDeliveryTimeoutMs?: number, + * responseChunkBytes?: number, + * cancelPendingFrames?: (streamId: number) => void, * }} deps */ -export const createTunnelHost = ({ connectionId, getLocalPort, sendFrame, getBufferedAmount, bodyDeliveryTimeoutMs = BODY_DELIVERY_TIMEOUT_MS }) => { +export const createTunnelHost = ({ connectionId, getLocalPort, sendFrame, getBufferedAmount, bodyDeliveryTimeoutMs = BODY_DELIVERY_TIMEOUT_MS, responseChunkBytes = MAX_TUNNEL_PAYLOAD_BYTES, cancelPendingFrames = () => {} }) => { /** @type {Map} */ const streams = new Map(); const assembler = createFragmentAssembler(); let closed = false; + let pendingWsBytes = 0; + let pendingWsMessages = 0; const send = async (frame) => { if (closed) return; @@ -128,6 +136,7 @@ export const createTunnelHost = ({ connectionId, getLocalPort, sendFrame, getBuf const stream = streams.get(streamId); if (!stream) return; dropStream(streamId); + cancelPendingFrames(streamId); if (stream.kind === 'http') { try { stream.body?.error(new Error(String(reason ?? 'aborted'))); @@ -205,17 +214,27 @@ export const createTunnelHost = ({ connectionId, getLocalPort, sendFrame, getBuf if (STRIPPED_RESPONSE_HEADERS.has(name)) continue; responseHeaders[name] = value; } - await sendJson(TunnelFrameType.HttpResponse, streamId, { status: response.status, headers: responseHeaders }); + if (closed || stream.abort.signal.aborted) { + await response.body?.cancel(); + return; + } try { + await sendJson(TunnelFrameType.HttpResponse, streamId, { status: response.status, headers: responseHeaders }); if (response.body) { for await (const chunk of response.body) { if (closed || stream.abort.signal.aborted) return; const bytes = chunk instanceof Uint8Array ? chunk : new Uint8Array(chunk); - for (const piece of chunkPayload(bytes, MAX_TUNNEL_PAYLOAD_BYTES)) { + const pieces = chunkPayload(bytes, responseChunkBytes); + // Offer at most one plaintext-frame budget at a time. The scheduler + // can batch small slices and interleave streams without buffering the + // entire source or encrypting hundreds of tiny messages separately. + const groupSize = Math.max(1, Math.floor(MAX_TUNNEL_PAYLOAD_BYTES / responseChunkBytes)); + for (let offset = 0; offset < pieces.length; offset += groupSize) { await waitForBackpressure(stream.abort.signal); if (closed || stream.abort.signal.aborted) return; - await send(encodeTunnelFrame(TunnelFrameType.HttpBody, streamId, piece)); + await Promise.all(pieces.slice(offset, offset + groupSize).map(piece => + send(encodeTunnelFrame(TunnelFrameType.HttpBody, streamId, piece)))); } } } @@ -448,6 +467,7 @@ export const createTunnelHost = ({ connectionId, getLocalPort, sendFrame, getBuf try { socket = new WebSocket(url, open.protocols, { headers: dialHeaders, + maxPayload: MAX_PENDING_WS_BYTES, }); } catch (error) { void sendAbort(streamId, error?.message ?? 'ws dial failed'); @@ -455,6 +475,8 @@ export const createTunnelHost = ({ connectionId, getLocalPort, sendFrame, getBuf } const stream = { kind: 'ws', socket, opened: false }; streams.set(streamId, stream); + let outputChain = Promise.resolve(); + let socketPendingMessages = 0; socket.on('open', () => { if (streams.get(streamId) !== stream) return; @@ -465,23 +487,47 @@ export const createTunnelHost = ({ connectionId, getLocalPort, sendFrame, getBuf if (streams.get(streamId) !== stream || closed) return; const bytes = Buffer.isBuffer(data) ? new Uint8Array(data) : new Uint8Array(Buffer.concat(data)); const frameType = isBinary ? TunnelFrameType.WsBinary : TunnelFrameType.WsText; - void (async () => { - for (const frame of encodeFragmentedMessage(frameType, streamId, bytes)) { + if (pendingWsBytes + bytes.length > MAX_PENDING_WS_BYTES || pendingWsMessages >= MAX_PENDING_WS_MESSAGES) { + abortLocalStream(streamId, 'upstream WebSocket exceeded downstream queue limit'); + void sendAbort(streamId, 'upstream WebSocket exceeded downstream queue limit'); + return; + } + pendingWsBytes += bytes.length; + pendingWsMessages += 1; + socketPendingMessages += 1; + if (CAN_PAUSE_WS_READS) socket.pause(); + outputChain = outputChain.then(async () => { + // Serialize entire messages, including their fragments. A later close + // must also wait here or it can overtake the final output. + const chunks = chunkPayload(bytes, responseChunkBytes); + const groupSize = Math.max(1, Math.floor(MAX_TUNNEL_PAYLOAD_BYTES / responseChunkBytes)); + for (let offset = 0; offset < chunks.length; offset += groupSize) { await waitForBackpressure(null); if (streams.get(streamId) !== stream || closed) return; - await send(frame); + await Promise.all(chunks.slice(offset, offset + groupSize).map((chunk, index) => + send(encodeTunnelFrame(frameType, streamId, chunk, offset + index < chunks.length - 1)))); } - })(); + }).catch(() => { + abortLocalStream(streamId, 'upstream WebSocket forwarding failed'); + void sendAbort(streamId, 'upstream WebSocket forwarding failed'); + }).finally(() => { + pendingWsBytes -= bytes.length; + pendingWsMessages -= 1; + socketPendingMessages -= 1; + if (CAN_PAUSE_WS_READS && socketPendingMessages === 0 && streams.get(streamId) === stream && socket.readyState === WebSocket.OPEN) socket.resume(); + }); }); socket.on('close', (code, reasonBuffer) => { - if (streams.get(streamId) !== stream) return; - dropStream(streamId); - const reason = reasonBuffer ? reasonBuffer.toString('utf8') : ''; - if (stream.opened) { - void sendJson(TunnelFrameType.WsClose, streamId, { code: code || 1000, reason }); - } else { - void sendAbort(streamId, reason || `upstream ws closed (${code || 'no code'})`); - } + void outputChain.then(async () => { + if (streams.get(streamId) !== stream) return; + dropStream(streamId); + const reason = reasonBuffer ? reasonBuffer.toString('utf8') : ''; + if (stream.opened) { + await sendJson(TunnelFrameType.WsClose, streamId, { code: code || 1000, reason }); + } else { + await sendAbort(streamId, reason || `upstream ws closed (${code || 'no code'})`); + } + }); }); socket.on('error', (error) => { if (streams.get(streamId) !== stream) return; @@ -512,6 +558,7 @@ export const createTunnelHost = ({ connectionId, getLocalPort, sendFrame, getBuf const stream = streams.get(streamId); if (!stream || stream.kind !== 'ws') return; dropStream(streamId); + cancelPendingFrames(streamId); let close = { code: 1000, reason: '' }; try { close = decodeJsonPayload(payload, isWsClosePayload); @@ -520,6 +567,8 @@ export const createTunnelHost = ({ connectionId, getLocalPort, sendFrame, getBuf } const code = Number.isInteger(close.code) && close.code >= 1000 && close.code <= 4999 ? close.code : 1000; try { + // A paused receiver still needs to read the peer's close handshake. + if (CAN_PAUSE_WS_READS) stream.socket.resume(); stream.socket.close(code, typeof close.reason === 'string' ? close.reason : ''); } catch { stream.socket.terminate();