fix(relay): keep bulk streaming from blocking client requests
Negotiate downstream delivery credit and schedule response streams fairly before encryption. Bound queued output, preserve deltas and WebSocket close ordering, and retain legacy peer compatibility. Validated with 81 relay tests, workspace type-check and lint, web build and mobile assets, and slow-link tests through the production relay. Confirmed on LTE by the maintainer.
This commit is contained in:
@@ -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),
|
||||
|
||||
@@ -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;
|
||||
|
||||
|
||||
@@ -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<typeof setTimeout> | 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));
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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 (`<data-dir>/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.
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
};
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
@@ -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),
|
||||
|
||||
@@ -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);
|
||||
});
|
||||
}
|
||||
});
|
||||
@@ -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<Console, 'warn'>,
|
||||
* }} 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);
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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<void>,
|
||||
* 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<number, { kind: 'http', abort: AbortController, body: { enqueue(payload: Uint8Array): void, close(): void, error(error: Error): void } | null, noBody: boolean } | { kind: 'ws', socket: WebSocket, opened: boolean }>} */
|
||||
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();
|
||||
|
||||
Reference in New Issue
Block a user