diff --git a/docs/architecture.md b/docs/architecture.md index 5b492eb01..62f80b63b 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -252,6 +252,10 @@ Terminal I/O is sent as binary WebSocket frames decoded by `decodeTerminalStream Terminal PTY size is last-interacting-client-wins. A client claims the PTY size only when its terminal viewport genuinely changes size or the user focuses/taps the terminal. Passive rendering work — attaching, restoring visibility, font settling, renderer refits, or just looking at a visible terminal — must not send a resize frame. The server does not broadcast resize ownership; the resized PTY redraws through normal output, and every attached client renders that output in its own local viewport. There is also a separate file-transfer binary frame format in the same directory, used for download/upload streams. +File downloads keep the existing `FileBegin`/`FileChunk`/`FileEnd` framing and stream 256 KiB chunks +from one stable file handle. Each transfer awaits completion of its own physical WebSocket send before +reading the next chunk; it is scoped to the requesting physical socket and does not queue unrelated +messages or transfers. ### Compatibility rules diff --git a/packages/relay/src/encrypted-channel.test.ts b/packages/relay/src/encrypted-channel.test.ts index 90ef5e0e3..b2f79c494 100644 --- a/packages/relay/src/encrypted-channel.test.ts +++ b/packages/relay/src/encrypted-channel.test.ts @@ -1,5 +1,10 @@ import { describe, it, expect, vi } from "vitest"; -import { createClientChannel, createDaemonChannel, Transport } from "./encrypted-channel.js"; +import { + createClientChannel, + createDaemonChannel, + EncryptedChannel, + Transport, +} from "./encrypted-channel.js"; import { deriveSharedKey, encrypt, @@ -47,6 +52,64 @@ async function waitForAsyncDelivery(): Promise { } describe("EncryptedChannel", () => { + it("rejects the daemon handshake when the ready frame fails to send", async () => { + const daemonKeyPair = generateKeyPair(); + const clientKeyPair = generateKeyPair(); + const transport: Transport = { + send: () => Promise.reject(new Error("ready send failed")), + close: () => undefined, + onmessage: null, + onclose: null, + onerror: null, + }; + const channel = createDaemonChannel(transport, daemonKeyPair); + + transport.onmessage?.({ + data: JSON.stringify({ + type: "e2ee_hello", + key: exportPublicKey(clientKeyPair.publicKey), + }), + isBinary: false, + }); + + await expect(channel).rejects.toThrow("ready send failed"); + }); + + it("waits for transport send completion", async () => { + let completeSend: (() => void) | undefined; + const transport: Transport = { + send: () => + new Promise((resolve) => { + completeSend = resolve; + }), + close: () => undefined, + onmessage: null, + onclose: null, + onerror: null, + }; + const first = generateKeyPair(); + const second = generateKeyPair(); + const channel = new EncryptedChannel( + transport, + deriveSharedKey(first.secretKey, second.publicKey), + {}, + { binaryCiphertext: true }, + ); + channel.setState("open"); + let completed = false; + + const sending = channel.send(new Uint8Array([1, 2, 3]).buffer).then(() => { + completed = true; + return undefined; + }); + await Promise.resolve(); + expect(completed).toBe(false); + + completeSend?.(); + await sending; + expect(completed).toBe(true); + }); + it("establishes encrypted channel between daemon and client", async () => { const [daemonTransport, clientTransport] = createMockTransportPair(); @@ -191,6 +254,58 @@ describe("EncryptedChannel", () => { } }); + it("reports rejected handshake hello sends", async () => { + const daemonKeyPair = generateKeyPair(); + const daemonPubKeyB64 = exportPublicKey(daemonKeyPair.publicKey); + const transport: Transport = { + send: () => Promise.reject(new Error("hello send failed")), + close: vi.fn(), + onmessage: null, + onclose: null, + onerror: null, + }; + const onerror = vi.fn(); + + await createClientChannel(transport, daemonPubKeyB64, { onerror }); + await Promise.resolve(); + + expect(onerror).toHaveBeenCalledTimes(1); + expect((onerror.mock.calls[0][0] as Error).message).toBe("hello send failed"); + transport.onclose?.(1000, "closed"); + }); + + it("reports rejected sends while flushing the handshake backlog", async () => { + const daemonKeyPair = generateKeyPair(); + const daemonPubKeyB64 = exportPublicKey(daemonKeyPair.publicKey); + let sendAttempts = 0; + const transport: Transport = { + send: () => { + sendAttempts += 1; + return sendAttempts === 1 ? undefined : Promise.reject(new Error("backlog send failed")); + }, + close: vi.fn(), + onmessage: null, + onclose: null, + onerror: null, + }; + const onerror = vi.fn(); + const channel = await createClientChannel(transport, daemonPubKeyB64, { onerror }); + await channel.send(new ArrayBuffer(8)); + + transport.onmessage?.({ + data: JSON.stringify({ + type: "e2ee_ready", + capabilities: { binaryCiphertext: true }, + }), + isBinary: false, + }); + await waitForAsyncDelivery(); + + expect(onerror).toHaveBeenCalledTimes(1); + expect((onerror.mock.calls[0][0] as Error).message).toBe("backlog send failed"); + expect(transport.close).toHaveBeenCalledWith(1011, "backlog send failed"); + }); + it("fails handshake on invalid hello", async () => { const [daemonTransport] = createMockTransportPair(); diff --git a/packages/relay/src/encrypted-channel.ts b/packages/relay/src/encrypted-channel.ts index f25dc3d05..de953616a 100644 --- a/packages/relay/src/encrypted-channel.ts +++ b/packages/relay/src/encrypted-channel.ts @@ -19,7 +19,7 @@ import { import { arrayBufferToBase64, base64ToArrayBuffer } from "./base64.js"; export interface Transport { - send(data: string | ArrayBuffer): void; + send(data: string | ArrayBuffer): void | Promise; close(code?: number, reason?: string): void; onmessage: ((message: TransportMessage) => void) | null; onclose: ((code: number, reason: string) => void) | null; @@ -169,7 +169,10 @@ export async function createClientChannel( }; const sendHello = () => { try { - transport.send(helloText); + const result = transport.send(helloText); + if (result) { + void result.catch(emitSendError); + } return true; } catch (error) { // This can happen during daemon restarts while the socket transitions @@ -263,11 +266,7 @@ export async function createDaemonChannel( const sharedKey = deriveSharedKey(daemonKeyPair.secretKey, clientPublicKey); const binaryCiphertext = supportsBinaryCiphertext(msg); - const channel = new EncryptedChannel(transport, sharedKey, events, { - daemonKeyPair, - binaryCiphertext, - }); - transport.send( + await transport.send( JSON.stringify({ type: "e2ee_ready", ...(binaryCiphertext @@ -276,6 +275,10 @@ export async function createDaemonChannel( } satisfies E2EEReadyMessage), ); + const channel = new EncryptedChannel(transport, sharedKey, events, { + daemonKeyPair, + binaryCiphertext, + }); channel.setState("open"); events.onopen?.(); @@ -354,7 +357,14 @@ export class EncryptedChannel { this.state = "open"; this.events.onopen?.(); for (const cb of this.onOpenCallbacks) cb(); - await this.flushPendingSends(); + try { + await this.flushPendingSends(); + } catch (error) { + const err = error instanceof Error ? error : new Error(String(error)); + this.events.onerror?.(err); + this.state = "closed"; + this.transport.close(1011, err.message); + } } } catch { // ignore non-ready handshake traffic @@ -455,12 +465,12 @@ export class EncryptedChannel { const ciphertext = encrypt(this.sharedKey, data); if (this.options.binaryCiphertext && data instanceof ArrayBuffer) { - this.transport.send(ciphertext); + await this.transport.send(ciphertext); return; } // COMPAT(binaryCiphertext): added in v0.2.3, remove base64 binary sends // after 2027-01-27 once the supported peer floor includes negotiation. - this.transport.send(arrayBufferToBase64(ciphertext)); + await this.transport.send(arrayBufferToBase64(ciphertext)); } outboundWireByteLength(data: string | ArrayBuffer): number { @@ -489,7 +499,7 @@ export class EncryptedChannel { // "ready" but do not re-key. Re-keying here would desync // the channel and cause decrypt failures. if (keysEqual(nextSharedKey, this.sharedKey)) { - this.transport.send( + await this.transport.send( JSON.stringify({ type: "e2ee_ready", ...(this.options.binaryCiphertext diff --git a/packages/server/src/server/file-explorer/service.test.ts b/packages/server/src/server/file-explorer/service.test.ts index 1025f395e..07f451e83 100644 --- a/packages/server/src/server/file-explorer/service.test.ts +++ b/packages/server/src/server/file-explorer/service.test.ts @@ -1,8 +1,13 @@ -import { chmod, mkdtemp, rm, stat, writeFile } from "node:fs/promises"; +import { appendFile, chmod, mkdtemp, rm, stat, truncate, writeFile } from "node:fs/promises"; import os from "node:os"; import path from "node:path"; import { describe, expect, it } from "vitest"; -import { getExplorerFileVersion, readExplorerFile, writeExplorerFile } from "./service.js"; +import { + getExplorerFileVersion, + readExplorerFile, + streamExplorerFile, + writeExplorerFile, +} from "./service.js"; async function createHomeTempDir(prefix: string): Promise { return mkdtemp(path.join(os.homedir(), prefix)); @@ -192,6 +197,127 @@ describe("file explorer service", () => { } }); + it("fails a stream when the file grows after its revision is advertised", async () => { + const root = await createTempDir("paseo-file-stream-growth-"); + + try { + const filePath = path.join(root, "growing.log"); + const initial = Buffer.alloc(300 * 1024, 0x61); + await writeFile(filePath, initial); + await expect( + streamExplorerFile({ root, relativePath: "growing.log" }, async (file) => { + await appendFile(filePath, Buffer.alloc(300 * 1024, 0x62)); + for await (const _chunk of file.chunks) { + // Consume through the advertised prefix before validating the revision. + } + }), + ).rejects.toThrow("File changed during transfer"); + } finally { + await rm(root, { recursive: true, force: true }); + } + }); + + it("fails a stream when the file shrinks below its advertised size", async () => { + const root = await createTempDir("paseo-file-stream-truncate-"); + + try { + const filePath = path.join(root, "shrinking.log"); + await writeFile(filePath, Buffer.alloc(300 * 1024, 0x61)); + + await expect( + streamExplorerFile({ root, relativePath: "shrinking.log" }, async (file) => { + await truncate(filePath, 100 * 1024); + for await (const _chunk of file.chunks) { + // Consume until the stream detects the premature EOF. + } + }), + ).rejects.toThrow("File changed during transfer"); + } finally { + await rm(root, { recursive: true, force: true }); + } + }); + + it("fails a stream when the file is overwritten in place", async () => { + const root = await createTempDir("paseo-file-stream-overwrite-"); + + try { + const filePath = path.join(root, "changing.log"); + const initial = Buffer.alloc(600 * 1024, 0x61); + await writeFile(filePath, initial); + + await expect( + streamExplorerFile({ root, relativePath: "changing.log" }, async (file) => { + let chunkIndex = 0; + for await (const _chunk of file.chunks) { + chunkIndex += 1; + if (chunkIndex === 1) { + const replacement = Buffer.alloc(initial.byteLength, 0x62); + await writeFile(filePath, replacement); + } + } + }), + ).rejects.toThrow("File changed during transfer"); + } finally { + await rm(root, { recursive: true, force: true }); + } + }); + + it("classifies sampled text when UTF-8 crosses the sample boundary", async () => { + const root = await createTempDir("paseo-file-stream-utf8-"); + + try { + const content = Buffer.concat([Buffer.alloc(8191, 0x61), Buffer.from("€"), Buffer.from("z")]); + await writeFile(path.join(root, "sample.txt"), content); + let kind: string | undefined; + let encoding: string | undefined; + + await streamExplorerFile({ root, relativePath: "sample.txt" }, async (file) => { + kind = file.kind; + encoding = file.encoding; + }); + + expect(kind).toBe("text"); + expect(encoding).toBe("utf-8"); + } finally { + await rm(root, { recursive: true, force: true }); + } + }); + + it("rejects incomplete UTF-8 when the whole file was sampled", async () => { + const root = await createTempDir("paseo-file-stream-invalid-utf8-"); + + try { + await writeFile(path.join(root, "invalid.txt"), Buffer.from([0x61, 0xe2, 0x82])); + let kind: string | undefined; + + await streamExplorerFile({ root, relativePath: "invalid.txt" }, async (file) => { + kind = file.kind; + }); + + expect(kind).toBe("binary"); + } finally { + await rm(root, { recursive: true, force: true }); + } + }); + + it("detects binary bytes beyond the initial classification block", async () => { + const root = await createTempDir("paseo-file-stream-late-binary-"); + + try { + const content = Buffer.concat([Buffer.alloc(8192, 0x61), Buffer.from([0xff])]); + await writeFile(path.join(root, "late-binary.unknown"), content); + let kind: string | undefined; + + await streamExplorerFile({ root, relativePath: "late-binary.unknown" }, async (file) => { + kind = file.kind; + }); + + expect(kind).toBe("binary"); + } finally { + await rm(root, { recursive: true, force: true }); + } + }); + it("expands a ~ prefix in relative paths against the user home directory", async () => { const root = await createHomeTempDir(".paseo-file-explorer-home-"); diff --git a/packages/server/src/server/file-explorer/service.ts b/packages/server/src/server/file-explorer/service.ts index 45a8d74ca..a33007b59 100644 --- a/packages/server/src/server/file-explorer/service.ts +++ b/packages/server/src/server/file-explorer/service.ts @@ -76,12 +76,24 @@ export interface FileExplorerFileBytes { revision: string; } +export interface FileExplorerFileStream { + path: string; + kind: ExplorerFileKind; + encoding: "utf-8" | "binary"; + mimeType: string; + size: number; + modifiedAt: string; + revision: string; + chunks: AsyncIterable; +} + const TEXT_MIME_TYPES: Record = { ".json": "application/json", }; const DEFAULT_TEXT_MIME_TYPE = "text/plain"; const FILE_TYPE_SAMPLE_BYTES = 8192; +export const FILE_EXPLORER_STREAM_CHUNK_BYTES = 256 * 1024; export const MAX_EDITABLE_FILE_BYTES = 1024 * 1024; const READ_FILE_OPEN_FLAGS = process.platform === "win32" ? constants.O_RDONLY : constants.O_RDONLY | constants.O_NOFOLLOW; @@ -275,6 +287,109 @@ export async function readExplorerFileBytes({ } } +export async function streamExplorerFile( + { root, relativePath }: ReadFileParams, + consume: (file: FileExplorerFileStream) => Promise, +): Promise { + const filePath = await resolveScopedPath({ root, relativePath }); + const handle = await openFileForRead(filePath.resolvedPath); + + try { + const stats = await handle.stat({ bigint: true }); + if (!stats.isFile()) { + throw new Error("Requested path is not a file"); + } + + const advertisedSize = Number(stats.size); + const advertisedRevision = fileRevision(stats); + const ext = path.extname(filePath.resolvedPath).toLowerCase(); + const isImage = ext in IMAGE_MIME_TYPES; + const isBinary = isImage || (await isFileHandleBinary(handle, advertisedSize)); + let kind: ExplorerFileKind = "text"; + let mimeType = textMimeTypeForExtension(ext); + if (isImage) { + kind = "image"; + mimeType = IMAGE_MIME_TYPES[ext]; + } else if (isBinary) { + kind = "binary"; + mimeType = "application/octet-stream"; + } + + await consume({ + path: normalizeRelativePath({ root, targetPath: filePath.requestedPath }), + kind, + encoding: isBinary ? "binary" : "utf-8", + mimeType, + size: advertisedSize, + modifiedAt: stats.mtime.toISOString(), + revision: advertisedRevision, + chunks: readFileHandleChunks(handle, advertisedSize, advertisedRevision), + }); + } finally { + await handle.close(); + } +} + +async function isFileHandleBinary(handle: FileHandle, advertisedSize: number): Promise { + if (advertisedSize === 0) return false; + + const decoder = new TextDecoder("utf-8", { fatal: true }); + let position = 0; + let suspiciousBytes = 0; + while (position < advertisedSize) { + const block = Buffer.allocUnsafe( + Math.min(FILE_EXPLORER_STREAM_CHUNK_BYTES, advertisedSize - position), + ); + const { bytesRead } = await handle.read(block, 0, block.byteLength, position); + if (bytesRead === 0) { + throw new Error("File changed during transfer"); + } + const bytes = block.subarray(0, bytesRead); + for (const byte of bytes) { + if (byte === 0) return true; + const isControl = byte < 32 && byte !== 9 && byte !== 10 && byte !== 13; + if (isControl || byte === 127) suspiciousBytes += 1; + } + try { + decoder.decode(bytes, { stream: true }); + } catch { + return true; + } + position += bytesRead; + } + + try { + decoder.decode(); + } catch { + return true; + } + return suspiciousBytes / advertisedSize > 0.3; +} + +async function* readFileHandleChunks( + handle: FileHandle, + advertisedSize: number, + advertisedRevision: string, +): AsyncIterable { + let position = 0; + while (position < advertisedSize) { + const chunk = Buffer.allocUnsafe( + Math.min(FILE_EXPLORER_STREAM_CHUNK_BYTES, advertisedSize - position), + ); + const { bytesRead } = await handle.read(chunk, 0, chunk.byteLength, position); + if (bytesRead === 0) { + throw new Error("File changed during transfer"); + } + position += bytesRead; + yield chunk.subarray(0, bytesRead); + } + + const finalStats = await handle.stat({ bigint: true }); + if (fileRevision(finalStats) !== advertisedRevision) { + throw new Error("File changed during transfer"); + } +} + export async function getExplorerFileVersion({ root, relativePath, diff --git a/packages/server/src/server/relay-transport.test.ts b/packages/server/src/server/relay-transport.test.ts index fb1ddfe4b..06775e5c5 100644 --- a/packages/server/src/server/relay-transport.test.ts +++ b/packages/server/src/server/relay-transport.test.ts @@ -1,5 +1,7 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import type pino from "pino"; +import { createClientChannel, type Transport } from "@getpaseo/relay/e2ee"; +import { exportPublicKey, generateKeyPair } from "@getpaseo/relay"; import { startRelayTransport } from "./relay-transport"; function createMockLogger() { @@ -32,7 +34,10 @@ class FakeRelayWebSocket { sent: Array = []; terminateCalls = 0; pingCalls = 0; + deferSendCompletion = false; + onSend: ((data: string | Uint8Array | ArrayBuffer) => void) | null = null; private readonly listeners = new Map void>>(); + private readonly pendingSendCallbacks: Array<(error?: Error) => void> = []; constructor(readonly url: string) {} @@ -61,11 +66,22 @@ class FakeRelayWebSocket { this.emit("close", 1006, ""); } - send(data: string | Uint8Array | ArrayBuffer) { + send(data: string | Uint8Array | ArrayBuffer, callback?: (error?: Error) => void) { if (this.readyState !== FakeRelayWebSocket.OPEN) { throw new Error(`WebSocket not open (readyState=${this.readyState})`); } this.sent.push(data); + this.onSend?.(data); + if (!callback) return; + if (this.deferSendCompletion) { + this.pendingSendCallbacks.push(callback); + return; + } + callback(); + } + + completeNextSend() { + this.pendingSendCallbacks.shift()?.(); } ping() { @@ -80,8 +96,8 @@ class FakeRelayWebSocket { this.emit("open"); } - message(data: unknown) { - this.emit("message", data); + message(data: unknown, isBinary = data instanceof ArrayBuffer || ArrayBuffer.isView(data)) { + this.emit("message", data, isBinary); } pong() { @@ -239,6 +255,80 @@ describe("relay-transport control lifecycle", () => { ]); }); + test("encrypted sends wait for the physical data socket callback", async () => { + const logger = createMockLogger(); + const daemonKeyPair = generateKeyPair(); + let resolveAttached: ((socket: unknown) => void) | undefined; + const attached = new Promise((resolve) => { + resolveAttached = resolve; + }); + const controller = startRelayTransport({ + logger: logger as unknown as pino.Logger, + attachSocket: async (socket) => resolveAttached?.(socket), + relayEndpoint: "relay.paseo.sh:443", + relayUseTls: true, + serverId: "srv_test", + daemonKeyPair, + createWebSocket: relay.createWebSocket, + }); + controllers.push(controller); + + const control = relay.sockets[0]; + control.open(); + control.message(JSON.stringify({ type: "sync", connectionIds: [] }), false); + control.message(JSON.stringify({ type: "connected", connectionId: "clt_test" }), false); + + const dataSocket = relay.sockets[1]; + dataSocket.deferSendCompletion = true; + dataSocket.open(); + let clientTransport: Transport; + clientTransport = { + send: (data) => dataSocket.message(data, data instanceof ArrayBuffer), + close: () => undefined, + onmessage: null, + onclose: null, + onerror: null, + }; + dataSocket.onSend = (data) => { + clientTransport.onmessage?.({ + data: data instanceof Uint8Array ? data.slice().buffer : data, + isBinary: data instanceof ArrayBuffer || data instanceof Uint8Array, + }); + }; + let resolveClientOpen: (() => void) | undefined; + const clientOpen = new Promise((resolve) => { + resolveClientOpen = resolve; + }); + await createClientChannel(clientTransport, exportPublicKey(daemonKeyPair.publicKey), { + onopen: () => resolveClientOpen?.(), + }); + + let attachedCompleted = false; + void attached.then(() => { + attachedCompleted = true; + return undefined; + }); + await clientOpen; + await Promise.resolve(); + expect(attachedCompleted).toBe(false); + dataSocket.completeNextSend(); + const encryptedSocket = (await attached) as { + send: (data: Uint8Array) => void | Promise; + }; + let completed = false; + + const sending = Promise.resolve(encryptedSocket.send(new Uint8Array([1, 2, 3]))).then(() => { + completed = true; + return undefined; + }); + await Promise.resolve(); + expect(completed).toBe(false); + + dataSocket.completeNextSend(); + await sending; + expect(completed).toBe(true); + }); + test("uses relayUseTls for control and data socket URLs", () => { const logger = createMockLogger(); const controller = startRelayTransport({ diff --git a/packages/server/src/server/relay-transport.ts b/packages/server/src/server/relay-transport.ts index b39fef1a2..b60b01bd0 100644 --- a/packages/server/src/server/relay-transport.ts +++ b/packages/server/src/server/relay-transport.ts @@ -28,7 +28,7 @@ export interface RelayTransportController { interface RelaySocketLike { readyState: number; bufferedAmount?: number; - send: (data: string | Uint8Array | ArrayBuffer) => void; + send: (data: string | Uint8Array | ArrayBuffer, callback?: (error?: Error) => void) => void; close: (code?: number, reason?: string) => void; terminate?: () => void; on: (event: "message" | "close" | "error", listener: (...args: unknown[]) => void) => void; @@ -467,16 +467,23 @@ function createRelayTransportAdapter( logger: pino.Logger, ): RelayTransport { const relayTransport: RelayTransport = { - send: (data) => { - try { - socket.send(data); - } catch (err) { - // Socket likely transitioned to closed between checks; let onclose/onerror - // drive cleanup. Without this guard the synchronous throw would propagate - // up as an uncaughtException and take down the daemon. - logger.warn({ err }, "relay_socket_send_failed"); - } - }, + send: (data) => + new Promise((resolve, reject) => { + try { + socket.send(data, (error) => { + if (!error) { + resolve(); + return; + } + logger.warn({ err: error }, "relay_socket_send_failed"); + reject(error); + }); + } catch (error) { + const err = error instanceof Error ? error : new Error(String(error)); + logger.warn({ err }, "relay_socket_send_failed"); + reject(err); + } + }), close: (code?: number, reason?: string) => socket.close(code, reason), onmessage: null, onclose: null, diff --git a/packages/server/src/server/session.ts b/packages/server/src/server/session.ts index 04f04fe6f..7cf2f24c3 100644 --- a/packages/server/src/server/session.ts +++ b/packages/server/src/server/session.ts @@ -405,6 +405,7 @@ export interface SessionOptions { onMessage: (msg: SessionOutboundMessage) => void; onMessageToSource?: (source: object, msg: SessionOutboundMessage) => void; onBinaryMessage?: (frame: Uint8Array) => void; + onBinaryMessageToSource?: (source: object, frame: Uint8Array) => Promise; getTransportBufferedAmount?: () => number | null; onLifecycleIntent?: (intent: SessionLifecycleIntent) => void; onWorkspaceRecovered?: (workspace: PersistedWorkspaceRecord) => Promise; @@ -568,6 +569,9 @@ export class Session { | ((source: object, msg: SessionOutboundMessage) => void) | null; private readonly onBinaryMessage: ((frame: Uint8Array) => void) | null; + private readonly onBinaryMessageToSource: + | ((source: object, frame: Uint8Array) => Promise) + | null; private readonly getTransportBufferedAmount: () => number | null; private readonly onLifecycleIntent: ((intent: SessionLifecycleIntent) => void) | null; private readonly onWorkspaceRecovered: @@ -646,6 +650,7 @@ export class Session { onMessage, onMessageToSource, onBinaryMessage, + onBinaryMessageToSource, getTransportBufferedAmount, onLifecycleIntent, onWorkspaceRecovered, @@ -698,6 +703,7 @@ export class Session { this.onMessage = onMessage; this.onMessageToSource = onMessageToSource ?? null; this.onBinaryMessage = onBinaryMessage ?? null; + this.onBinaryMessageToSource = onBinaryMessageToSource ?? null; this.getTransportBufferedAmount = getTransportBufferedAmount ?? (() => 0); this.onLifecycleIntent = onLifecycleIntent ?? null; this.onWorkspaceRecovered = onWorkspaceRecovered ?? null; @@ -711,8 +717,8 @@ export class Session { }); this.workspaceFilesSession = new WorkspaceFilesSession({ host: { - emit: (msg) => this.emit(msg), - emitBinary: (frame) => this.emitBinary(frame), + emit: (msg, source) => this.emitForSource(msg, source), + emitBinary: (frame, source) => this.emitBinaryForFileTransfer(frame, source), hasBinaryChannel: () => this.onBinaryMessage !== null, }, downloadTokenStore, @@ -1786,7 +1792,7 @@ export class Session { this.dispatchCheckoutMessage(msg) ?? this.dispatchWorkspaceRecoveryMessage(msg) ?? this.dispatchWorkspaceAndProjectMessage(msg) ?? - this.dispatchWorkspaceFileMessage(msg) ?? + this.dispatchWorkspaceFileMessage(msg, source) ?? this.dispatchProviderMessage(msg) ?? this.dispatchTerminalMessage(msg) ?? this.dispatchChatScheduleLoopMessage(msg) ?? @@ -2106,10 +2112,13 @@ export class Session { } } - private dispatchWorkspaceFileMessage(msg: SessionInboundMessage): Promise | undefined { + private dispatchWorkspaceFileMessage( + msg: SessionInboundMessage, + source?: object, + ): Promise | undefined { switch (msg.type) { case "file_explorer_request": - return this.workspaceFilesSession.handleFileExplorerRequest(msg); + return this.workspaceFilesSession.handleFileExplorerRequest(msg, source); case "fs.file.subscribe.request": return this.workspaceFilesSession.handleFileSubscribeRequest(msg); case "fs.file.unsubscribe.request": @@ -6532,6 +6541,22 @@ export class Session { } } + private async emitBinaryForFileTransfer(frame: Uint8Array, source?: object): Promise { + if (source && this.onBinaryMessageToSource) { + await this.onBinaryMessageToSource(source, frame); + return; + } + this.emitBinary(frame); + } + + private emitForSource(msg: SessionOutboundMessage, source?: object): void { + if (source && this.onMessageToSource) { + this.onMessageToSource(source, msg); + return; + } + this.emit(msg); + } + /** * Clean up session resources */ diff --git a/packages/server/src/server/session/files/workspace-files-session.test.ts b/packages/server/src/server/session/files/workspace-files-session.test.ts index cfc45c5c7..0b09c8774 100644 --- a/packages/server/src/server/session/files/workspace-files-session.test.ts +++ b/packages/server/src/server/session/files/workspace-files-session.test.ts @@ -30,13 +30,21 @@ function makeDir(prefix: string): string { return dir; } -function makeSubsystem(options: { hasBinaryChannel?: boolean } = {}) { +function makeSubsystem( + options: { + hasBinaryChannel?: boolean; + emitBinary?: (frame: Uint8Array) => Promise | void; + } = {}, +) { const emitted: SessionOutboundMessage[] = []; const binary: Uint8Array[] = []; let hasBinary = options.hasBinaryChannel ?? false; const host: WorkspaceFilesSessionHost = { emit: (msg) => emitted.push(msg), - emitBinary: (frame) => binary.push(frame), + emitBinary: async (frame) => { + binary.push(frame); + await options.emitBinary?.(frame); + }, hasBinaryChannel: () => hasBinary, }; const paseoHome = makeDir("workspace-files-home-"); @@ -136,6 +144,74 @@ describe("WorkspaceFilesSession", () => { ]); }); + test("streams a real file larger than the socket limit as paced ordered chunks", async () => { + const cwd = makeDir("workspace-files-large-binary-"); + const fileBytes = Buffer.alloc(8 * 1024 * 1024 + 123); + for (let index = 0; index < fileBytes.length; index += 1) { + fileBytes[index] = index % 251; + } + writeFileSync(join(cwd, "large.bin"), fileBytes); + + let releaseFirstChunk: (() => void) | undefined; + const firstChunkSent = new Promise((resolve) => { + releaseFirstChunk = resolve; + }); + let chunkSends = 0; + const { subsystem, emitted, binary } = makeSubsystem({ + hasBinaryChannel: true, + emitBinary: async (frame) => { + if (decodeFileTransferFrame(frame)?.opcode !== FileTransferOpcode.FileChunk) return; + chunkSends += 1; + if (chunkSends === 1) await firstChunkSent; + }, + }); + + const transfer = subsystem.handleFileExplorerRequest({ + type: "file_explorer_request", + cwd, + path: "large.bin", + mode: "file", + requestId: "req-large-binary", + acceptBinary: true, + }); + + await expect.poll(() => chunkSends).toBe(1); + expect(binary.map((frame) => decodeFileTransferFrame(frame)?.opcode)).toEqual([ + FileTransferOpcode.FileBegin, + FileTransferOpcode.FileChunk, + ]); + + await subsystem.handleFileExplorerRequest({ + type: "file_explorer_request", + cwd, + path: ".", + mode: "list", + requestId: "req-unrelated-list", + }); + expect(emitted).toEqual([ + expect.objectContaining({ + type: "file_explorer_response", + payload: expect.objectContaining({ requestId: "req-unrelated-list", error: null }), + }), + ]); + + releaseFirstChunk?.(); + await transfer; + + const frames = binary.map((frame) => decodeFileTransferFrame(frame)); + const chunks = frames.flatMap((frame) => + frame?.opcode === FileTransferOpcode.FileChunk ? [frame.payload] : [], + ); + expect(chunks.length).toBeGreaterThan(1); + expect(chunks.every((chunk) => chunk.byteLength <= 256 * 1024)).toBe(true); + expect( + Buffer.compare(Buffer.concat(chunks.map((chunk) => Buffer.from(chunk))), fileBytes), + ).toBe(0); + expect(frames.at(0)?.opcode).toBe(FileTransferOpcode.FileBegin); + expect(frames.at(-1)?.opcode).toBe(FileTransferOpcode.FileEnd); + expect(emitted).toHaveLength(1); + }, 30_000); + test("rejects an empty file-explorer cwd with an error envelope", async () => { const { subsystem, emitted } = makeSubsystem(); diff --git a/packages/server/src/server/session/files/workspace-files-session.ts b/packages/server/src/server/session/files/workspace-files-session.ts index 740f74a61..dd7d7d04d 100644 --- a/packages/server/src/server/session/files/workspace-files-session.ts +++ b/packages/server/src/server/session/files/workspace-files-session.ts @@ -21,7 +21,7 @@ import { getDownloadableFileInfo, listDirectoryEntries, readExplorerFile, - readExplorerFileBytes, + streamExplorerFile, writeExplorerFile, } from "../../file-explorer/service.js"; import { workspaceFileObserver, type FileObserver } from "../../file-explorer/observer.js"; @@ -34,8 +34,8 @@ import { getProjectIcon } from "../../../utils/project-icon.js"; * — old clients without a binary channel fall back to inline JSON file content. */ export interface WorkspaceFilesSessionHost { - emit(msg: SessionOutboundMessage): void; - emitBinary(frame: Uint8Array): void; + emit(msg: SessionOutboundMessage, source?: object): void; + emitBinary(frame: Uint8Array, source?: object): Promise; hasBinaryChannel(): boolean; } @@ -136,22 +136,25 @@ export class WorkspaceFilesSession { this.fileSubscriptions.clear(); } - async handleFileExplorerRequest(request: FileExplorerRequest): Promise { + async handleFileExplorerRequest(request: FileExplorerRequest, source?: object): Promise { const { cwd: workspaceCwd, path: requestedPath = ".", mode, requestId } = request; const cwd = workspaceCwd.trim(); if (!cwd) { - this.host.emit({ - type: "file_explorer_response", - payload: { - cwd: workspaceCwd, - path: requestedPath, - mode, - directory: null, - file: null, - error: "cwd is required", - requestId, + this.host.emit( + { + type: "file_explorer_response", + payload: { + cwd: workspaceCwd, + path: requestedPath, + mode, + directory: null, + file: null, + error: "cwd is required", + requestId, + }, }, - }); + source, + ); return; } @@ -162,69 +165,77 @@ export class WorkspaceFilesSession { relativePath: requestedPath, }); - this.host.emit({ - type: "file_explorer_response", - payload: { - cwd, - path: directory.path, - mode, - directory, - file: null, - error: null, - requestId, + this.host.emit( + { + type: "file_explorer_response", + payload: { + cwd, + path: directory.path, + mode, + directory, + file: null, + error: null, + requestId, + }, }, - }); + source, + ); } else { if (request.acceptBinary && this.host.hasBinaryChannel()) { - const file = await readExplorerFileBytes({ - root: cwd, - relativePath: requestedPath, + await streamExplorerFile({ root: cwd, relativePath: requestedPath }, async (file) => { + await this.host.emitBinary( + encodeFileTransferFrame({ + opcode: FileTransferOpcode.FileBegin, + requestId, + metadata: { + mime: file.mimeType, + size: file.size, + encoding: file.encoding, + modifiedAt: file.modifiedAt, + revision: file.revision, + }, + }), + source, + ); + for await (const chunk of file.chunks) { + await this.host.emitBinary( + encodeFileTransferFrame({ + opcode: FileTransferOpcode.FileChunk, + requestId, + payload: chunk, + }), + source, + ); + } + await this.host.emitBinary( + encodeFileTransferFrame({ + opcode: FileTransferOpcode.FileEnd, + requestId, + }), + source, + ); }); - - this.host.emitBinary( - encodeFileTransferFrame({ - opcode: FileTransferOpcode.FileBegin, - requestId, - metadata: { - mime: file.mimeType, - size: file.size, - encoding: file.encoding, - modifiedAt: file.modifiedAt, - revision: file.revision, - }, - }), - ); - this.host.emitBinary( - encodeFileTransferFrame({ - opcode: FileTransferOpcode.FileChunk, - requestId, - payload: file.bytes, - }), - ); - this.host.emitBinary( - encodeFileTransferFrame({ - opcode: FileTransferOpcode.FileEnd, - requestId, - }), - ); } else { const file = await readExplorerFile({ root: cwd, relativePath: requestedPath, }); - this.host.emit({ - type: "file_explorer_response", - payload: { - cwd, - path: file.path, - mode, - directory: null, - file, - error: null, - requestId, + this.host.emit( + { + type: "file_explorer_response", + payload: { + cwd, + path: file.path, + mode, + directory: null, + file, + error: null, + requestId, + }, }, - }); + source, + ); } } } catch (error) { @@ -232,18 +243,21 @@ export class WorkspaceFilesSession { { err: error, cwd, path: requestedPath }, `Failed to fulfill file explorer request for workspace ${cwd}`, ); - this.host.emit({ - type: "file_explorer_response", - payload: { - cwd, - path: requestedPath, - mode, - directory: null, - file: null, - error: getErrorMessage(error), - requestId, + this.host.emit( + { + type: "file_explorer_response", + payload: { + cwd, + path: requestedPath, + mode, + directory: null, + file: null, + error: getErrorMessage(error), + requestId, + }, }, - }); + source, + ); } } diff --git a/packages/server/src/server/websocket-server.file-transfer.e2e.test.ts b/packages/server/src/server/websocket-server.file-transfer.e2e.test.ts new file mode 100644 index 000000000..5131456ee --- /dev/null +++ b/packages/server/src/server/websocket-server.file-transfer.e2e.test.ts @@ -0,0 +1,168 @@ +import { mkdtempSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, expect, test } from "vitest"; +import { WebSocket, type RawData } from "ws"; +import { + decodeFileTransferFrame, + FileTransferOpcode, +} from "@getpaseo/protocol/binary-frames/index"; +import { createTestPaseoDaemon, type TestPaseoDaemon } from "./test-utils/index.js"; +import { WSOutboundMessageSchema, type WSOutboundMessage } from "./messages.js"; + +const TEST_TIMEOUT_MS = 30_000; +const FILE_SIZE = 8 * 1024 * 1024 + 123; + +let daemon: TestPaseoDaemon | undefined; +const temporaryDirectories: string[] = []; +const sockets: WebSocket[] = []; + +afterEach(async () => { + for (const socket of sockets.splice(0)) socket.terminate(); + await daemon?.close(); + daemon = undefined; + for (const directory of temporaryDirectories.splice(0)) { + rmSync(directory, { recursive: true, force: true }); + } +}); + +test( + "a large file stays ordered, source-scoped, and does not block another socket", + async () => { + const cwd = mkdtempSync(join(tmpdir(), "paseo-large-file-transfer-")); + temporaryDirectories.push(cwd); + const expected = Buffer.alloc(FILE_SIZE); + for (let index = 0; index < expected.length; index += 1) expected[index] = index % 251; + writeFileSync(join(cwd, "large.bin"), expected); + + daemon = await createTestPaseoDaemon(); + const source = await connectSocket(daemon.port, "shared-file-client"); + const unrelated = await connectSocket(daemon.port, "shared-file-client"); + sockets.push(source, unrelated); + + let unrelatedBinaryFrames = 0; + unrelated.on("message", (_data, isBinary) => { + if (isBinary) unrelatedBinaryFrames += 1; + }); + + const transfer = receiveFileTransfer(source, "req-large-file"); + source.send( + JSON.stringify({ + type: "session", + message: { + type: "file_explorer_request", + cwd, + path: "large.bin", + mode: "file", + acceptBinary: true, + requestId: "req-large-file", + }, + }), + ); + + await sendAndWait( + unrelated, + { + type: "session", + message: { type: "ping", requestId: "req-unrelated", clientSentAt: 1 }, + }, + (message) => + message.type === "session" && + message.message.type === "pong" && + message.message.payload.requestId === "req-unrelated", + ); + + const frames = await transfer; + const chunks = frames.flatMap((frame) => + frame.opcode === FileTransferOpcode.FileChunk ? [frame.payload] : [], + ); + expect(frames[0]?.opcode).toBe(FileTransferOpcode.FileBegin); + expect(frames.at(-1)?.opcode).toBe(FileTransferOpcode.FileEnd); + expect(chunks.length).toBeGreaterThan(1); + expect(chunks.every((chunk) => chunk.byteLength <= 256 * 1024)).toBe(true); + expect(Buffer.compare(Buffer.concat(chunks.map((chunk) => Buffer.from(chunk))), expected)).toBe( + 0, + ); + expect(source.readyState).toBe(WebSocket.OPEN); + expect(unrelated.readyState).toBe(WebSocket.OPEN); + expect(unrelatedBinaryFrames).toBe(0); + }, + TEST_TIMEOUT_MS, +); + +async function connectSocket(port: number, clientId: string): Promise { + const socket = new WebSocket(`ws://127.0.0.1:${port}/ws`); + await new Promise((resolve, reject) => { + socket.once("open", resolve); + socket.once("error", reject); + }); + await sendAndWait( + socket, + { type: "hello", clientId, clientType: "browser", protocolVersion: 1 }, + (message) => + message.type === "session" && + message.message.type === "status" && + message.message.payload.status === "server_info", + ); + return socket; +} + +function receiveFileTransfer(source: WebSocket, requestId: string) { + return new Promise>[]>( + (resolve, reject) => { + const frames: NonNullable>[] = []; + const timeout = setTimeout(() => { + cleanup(); + reject(new Error("Timed out waiting for file transfer")); + }, TEST_TIMEOUT_MS); + const onMessage = (data: RawData, isBinary: boolean) => { + if (!isBinary) return; + const frame = decodeFileTransferFrame(new Uint8Array(data as Buffer)); + if (!frame || frame.requestId !== requestId) return; + frames.push(frame); + if (frame.opcode === FileTransferOpcode.FileEnd) { + cleanup(); + resolve(frames); + } + }; + const onClose = () => { + cleanup(); + reject(new Error("Socket closed during file transfer")); + }; + const cleanup = () => { + clearTimeout(timeout); + source.off("message", onMessage); + source.off("close", onClose); + }; + source.on("message", onMessage); + source.on("close", onClose); + }, + ); +} + +function sendAndWait( + socket: WebSocket, + message: unknown, + matches: (message: WSOutboundMessage) => boolean, +): Promise { + const response = new Promise((resolve, reject) => { + const timeout = setTimeout(() => { + cleanup(); + reject(new Error("Timed out waiting for WebSocket message")); + }, TEST_TIMEOUT_MS); + const onMessage = (data: RawData, isBinary: boolean) => { + if (isBinary) return; + const parsed = WSOutboundMessageSchema.safeParse(JSON.parse(data.toString())); + if (!parsed.success || !matches(parsed.data)) return; + cleanup(); + resolve(parsed.data); + }; + const cleanup = () => { + clearTimeout(timeout); + socket.off("message", onMessage); + }; + socket.on("message", onMessage); + }); + socket.send(JSON.stringify(message)); + return response; +} diff --git a/packages/server/src/server/websocket-server.ts b/packages/server/src/server/websocket-server.ts index 104eb08e9..c54c4c5fc 100644 --- a/packages/server/src/server/websocket-server.ts +++ b/packages/server/src/server/websocket-server.ts @@ -96,6 +96,7 @@ import { outboundFrameByteLength, physicalSocketHasCapacity, sendBoundedPhysicalFrame, + sendBoundedPhysicalFrameAndWait, } from "./websocket/physical-socket.js"; const WS_CLOSE_DAEMON_AUTH_FAILED = 4401; @@ -371,7 +372,10 @@ function getBrowserHostCapability( export interface WebSocketLike { readyState: number; bufferedAmount?: number; - send: (data: string | Uint8Array | ArrayBuffer) => void; + send: ( + data: string | Uint8Array | ArrayBuffer, + callback?: (error?: Error) => void, + ) => void | Promise; close: (code?: number, reason?: string) => void; terminate?: () => void; on: (event: "message" | "close" | "error", listener: (...args: unknown[]) => void) => void; @@ -423,6 +427,7 @@ interface SocketSessionOptions { onMessage: (message: SessionOutboundMessage) => void; onMessageToSource?: (source: object, message: SessionOutboundMessage) => void; onBinaryMessage?: (frame: Uint8Array) => void; + onBinaryMessageToSource?: (source: object, frame: Uint8Array) => Promise; getTransportBufferedAmount?: () => number | null; onLifecycleIntent?: (intent: SessionLifecycleIntent) => void; hubExecutionAgents?: HubExecutionAgents; @@ -1054,6 +1059,23 @@ export class VoiceAssistantWebSocketServer { }); } + private async sendBinaryToClientAndWait(ws: WebSocketLike, frame: Uint8Array): Promise { + try { + const sent = await sendBoundedPhysicalFrameAndWait({ + socket: ws, + frame, + onHighWater: () => this.closeAtOutboundHighWater(ws), + }); + if (!sent) { + throw new Error("Physical WebSocket is not open"); + } + this.runtimeMetrics.recordOutboundBinaryFrame(ws.bufferedAmount); + } catch (err) { + this.logger.warn({ err }, "ws_send_failed"); + throw err; + } + } + private sendFrameToClient( ws: WebSocketLike, frame: string | Uint8Array, @@ -1226,6 +1248,12 @@ export class VoiceAssistantWebSocketServer { } this.sendBinaryToConnection(connection, frame); }, + onBinaryMessageToSource: async (source, frame) => { + if (!connection || !connection.sockets.has(source as WebSocketLike)) { + throw new Error("File transfer source socket is no longer attached"); + } + await this.sendBinaryToClientAndWait(source as WebSocketLike, frame); + }, getTransportBufferedAmount: () => { if (!connection) { return null; @@ -1271,6 +1299,7 @@ export class VoiceAssistantWebSocketServer { onMessage: options.onMessage, onMessageToSource: options.onMessageToSource, onBinaryMessage: options.onBinaryMessage, + onBinaryMessageToSource: options.onBinaryMessageToSource, getTransportBufferedAmount: options.getTransportBufferedAmount, onLifecycleIntent: options.onLifecycleIntent, logger: options.connectionLogger.child({ module: "session" }), diff --git a/packages/server/src/server/websocket/encrypted-relay-socket.test.ts b/packages/server/src/server/websocket/encrypted-relay-socket.test.ts index 53b230c56..2e50d7530 100644 --- a/packages/server/src/server/websocket/encrypted-relay-socket.test.ts +++ b/packages/server/src/server/websocket/encrypted-relay-socket.test.ts @@ -40,10 +40,11 @@ class BlockingChannel implements EncryptedRelayChannel { test("negotiated binary ciphertext accepts the exact hard bound and rejects one byte over", async () => { const channel = new BlockingChannel(); let terminations = 0; + let transportBufferedAmount = 0; const socket = createEncryptedRelaySocket({ channel, emitter: new EventEmitter(), - getTransportBufferedAmount: () => 0, + getTransportBufferedAmount: () => transportBufferedAmount, terminateTransport: () => { terminations += 1; }, @@ -51,9 +52,11 @@ test("negotiated binary ciphertext accepts the exact hard bound and rejects one socket.send(new Uint8Array(MAX_PHYSICAL_SOCKET_BUFFERED_BYTES - 40)); expect(channel.sent).toHaveLength(1); + transportBufferedAmount = MAX_PHYSICAL_SOCKET_BUFFERED_BYTES; expect(socket.bufferedAmount).toBe(MAX_PHYSICAL_SOCKET_BUFFERED_BYTES); - socket.send(new Uint8Array(1)); + const rejected = socket.send(new Uint8Array(1)); + await expect(rejected).rejects.toThrow("outbound high-water mark"); expect(channel.sent).toHaveLength(1); expect(terminations).toBe(1); @@ -64,7 +67,7 @@ test("negotiated binary ciphertext accepts the exact hard bound and rejects one await Promise.resolve(); }); -test("underlying relay backpressure rejects binary before encryption and terminates physically", () => { +test("underlying relay backpressure rejects binary before encryption and terminates physically", async () => { const channel = new BlockingChannel(); let terminations = 0; const socket = createEncryptedRelaySocket({ @@ -76,36 +79,14 @@ test("underlying relay backpressure rejects binary before encryption and termina }, }); - socket.send(new Uint8Array(1)); + const rejected = socket.send(new Uint8Array(1)); + await expect(rejected).rejects.toThrow("outbound high-water mark"); expect(channel.sent).toEqual([]); expect(channel.closes).toEqual([]); expect(terminations).toBe(1); }); -test("pending encryption and underlying relay backpressure share one hard bound", () => { - const channel = new BlockingChannel(); - let transportBufferedAmount = 3 * 1024 * 1024; - let terminations = 0; - const socket = createEncryptedRelaySocket({ - channel, - emitter: new EventEmitter(), - getTransportBufferedAmount: () => transportBufferedAmount, - terminateTransport: () => { - terminations += 1; - }, - }); - - socket.send(new Uint8Array(3 * 1024 * 1024)); - expect(channel.sent).toHaveLength(1); - - transportBufferedAmount = 6 * 1024 * 1024; - socket.send(new Uint8Array(1)); - - expect(channel.sent).toHaveLength(1); - expect(terminations).toBe(1); -}); - test("explicit encrypted-socket termination forcibly terminates the relay transport", () => { const channel = new BlockingChannel(); let terminations = 0; @@ -124,3 +105,42 @@ test("explicit encrypted-socket termination forcibly terminates the relay transp expect(channel.closes).toEqual([]); expect(socket.readyState).toBe(3); }); + +test("encrypted sends report physical completion through the returned promise", async () => { + const channel = new BlockingChannel(); + const socket = createEncryptedRelaySocket({ + channel, + emitter: new EventEmitter(), + getTransportBufferedAmount: () => 0, + terminateTransport: () => undefined, + }); + let completed = false; + + const sending = socket.send(new Uint8Array([1])); + if (!sending) throw new Error("Expected an awaitable encrypted send"); + void sending.then(() => (completed = true)); + await Promise.resolve(); + expect(completed).toBe(false); + + channel.drain(); + await sending; + expect(completed).toBe(true); + expect(socket.bufferedAmount).toBe(0); +}); + +test("encrypted sockets do not double-count bytes already buffered by the transport", () => { + const channel = new BlockingChannel(); + let transportBufferedAmount = 0; + const socket = createEncryptedRelaySocket({ + channel, + emitter: new EventEmitter(), + getTransportBufferedAmount: () => transportBufferedAmount, + terminateTransport: () => undefined, + }); + const payload = new Uint8Array(3 * 1024 * 1024); + + void socket.send(payload); + transportBufferedAmount = payload.byteLength + 40; + + expect(socket.bufferedAmount).toBe(payload.byteLength + 40); +}); diff --git a/packages/server/src/server/websocket/encrypted-relay-socket.ts b/packages/server/src/server/websocket/encrypted-relay-socket.ts index c225ed3ec..3b76a4971 100644 --- a/packages/server/src/server/websocket/encrypted-relay-socket.ts +++ b/packages/server/src/server/websocket/encrypted-relay-socket.ts @@ -11,7 +11,7 @@ export interface EncryptedRelayChannel { export interface EncryptedRelaySocket { readonly readyState: number; readonly bufferedAmount: number; - send: (data: string | Uint8Array | ArrayBuffer) => void; + send: (data: string | Uint8Array | ArrayBuffer) => void | Promise; close: (code?: number, reason?: string) => void; terminate: () => void; on: (event: "message" | "close" | "error", listener: (...args: unknown[]) => void) => void; @@ -26,7 +26,6 @@ export function createEncryptedRelaySocket(params: { }): EncryptedRelaySocket { const { channel, emitter, getTransportBufferedAmount, terminateTransport } = params; let readyState = 1; - let pendingEncryptedBytes = 0; channel.setState("open"); @@ -51,26 +50,25 @@ export function createEncryptedRelaySocket(params: { return readyState; }, get bufferedAmount() { - return pendingEncryptedBytes + (getTransportBufferedAmount() ?? 0); + return getTransportBufferedAmount() ?? 0; }, send: (data) => { - if (readyState !== 1) return; + if (readyState !== 1) { + return Promise.reject(new Error("Encrypted relay socket is not open")); + } const outbound = normalizeRelaySendPayload(data); const outboundBytes = channel.outboundWireByteLength(outbound); - const queuedBytes = pendingEncryptedBytes + (getTransportBufferedAmount() ?? 0); + const queuedBytes = getTransportBufferedAmount() ?? 0; if (queuedBytes + outboundBytes > MAX_PHYSICAL_SOCKET_BUFFERED_BYTES) { terminate(); - return; + return Promise.reject( + new Error("Encrypted relay socket exceeded its outbound high-water mark"), + ); } - pendingEncryptedBytes += outboundBytes; - void channel - .send(outbound) - .catch((error) => { - emitter.emit("error", error); - }) - .finally(() => { - pendingEncryptedBytes -= outboundBytes; - }); + return channel.send(outbound).catch((error) => { + emitter.emit("error", error); + throw error; + }); }, close, terminate, diff --git a/packages/server/src/server/websocket/physical-socket.test.ts b/packages/server/src/server/websocket/physical-socket.test.ts index 9b334cc11..c81b1410d 100644 --- a/packages/server/src/server/websocket/physical-socket.test.ts +++ b/packages/server/src/server/websocket/physical-socket.test.ts @@ -4,6 +4,7 @@ import { ApplicationSocketLease, MAX_PHYSICAL_SOCKET_BUFFERED_BYTES, sendBoundedPhysicalFrame, + sendBoundedPhysicalFrameAndWait, } from "./physical-socket.js"; test("sockets remain exempt until they send an application ping", () => { @@ -66,3 +67,56 @@ test("the shared physical send boundary rejects binary above the hard bound", () expect(sent).toEqual([]); expect(terminated).toBe(true); }); + +test("the awaitable physical send resolves only when that frame send completes", async () => { + const sent: Array = []; + let completeSend: (() => void) | undefined; + const socket = { + readyState: 1, + bufferedAmount: 0, + send: (_data: string | Uint8Array | ArrayBuffer, callback?: (error?: Error) => void) => { + sent.push(_data); + if (callback) completeSend = () => callback(); + }, + }; + let completed = false; + + const sending = sendBoundedPhysicalFrameAndWait({ + socket, + frame: new Uint8Array([1, 2, 3]), + onHighWater: () => undefined, + }).then(() => { + return (completed = true); + }); + + await Promise.resolve(); + expect(completed).toBe(false); + expect( + sendBoundedPhysicalFrame({ + socket, + frame: "unrelated", + onHighWater: () => undefined, + }), + ).toBe(true); + expect(sent).toEqual([new Uint8Array([1, 2, 3]), "unrelated"]); + completeSend?.(); + await sending; + expect(completed).toBe(true); +}); + +test("the awaitable physical send rejects callback errors", async () => { + const socket = { + readyState: 1, + bufferedAmount: 0, + send: (_data: string | Uint8Array | ArrayBuffer, callback?: (error?: Error) => void) => + callback?.(new Error("send failed")), + }; + + await expect( + sendBoundedPhysicalFrameAndWait({ + socket, + frame: new Uint8Array([1]), + onHighWater: () => undefined, + }), + ).rejects.toThrow("send failed"); +}); diff --git a/packages/server/src/server/websocket/physical-socket.ts b/packages/server/src/server/websocket/physical-socket.ts index eed29ae35..491a7c612 100644 --- a/packages/server/src/server/websocket/physical-socket.ts +++ b/packages/server/src/server/websocket/physical-socket.ts @@ -50,7 +50,39 @@ export function outboundFrameByteLength(data: string | Uint8Array | ArrayBuffer) interface BoundedPhysicalSocket { readyState: number; bufferedAmount?: number; - send: (data: string | Uint8Array | ArrayBuffer) => void; + send: ( + data: string | Uint8Array | ArrayBuffer, + callback?: (error?: Error) => void, + ) => void | Promise; +} + +export async function sendBoundedPhysicalFrameAndWait(params: { + socket: BoundedPhysicalSocket; + frame: string | Uint8Array | ArrayBuffer; + frameBytes?: number; + onHighWater: () => void; +}): Promise { + const { socket, frame, frameBytes = outboundFrameByteLength(frame), onHighWater } = params; + if (socket.readyState !== 1) return false; + if (!physicalSocketHasCapacity(socket, frameBytes)) { + onHighWater(); + return false; + } + + await new Promise((resolve, reject) => { + let callbackUsed = false; + const result = socket.send(frame, (error) => { + callbackUsed = true; + if (error) reject(error); + else resolve(); + }); + if (result && typeof result.then === "function") { + result.then(resolve, reject); + } else if (socket.send.length < 2 && !callbackUsed) { + resolve(); + } + }); + return true; } export function physicalSocketHasCapacity( @@ -73,6 +105,9 @@ export function sendBoundedPhysicalFrame(params: { onHighWater(); return false; } - socket.send(frame); + const result = socket.send(frame); + if (result && typeof result.then === "function") { + void result.catch(() => undefined); + } return true; }