diff --git a/packages/relay/src/cloudflare-adapter.ts b/packages/relay/src/cloudflare-adapter.ts index 5c670bd79..21b36ac7b 100644 --- a/packages/relay/src/cloudflare-adapter.ts +++ b/packages/relay/src/cloudflare-adapter.ts @@ -66,16 +66,70 @@ interface CFResponseInit extends ResponseInit { export class RelayDurableObject { private state: DurableObjectState; + private pendingClientFrames = new Map>(); constructor(state: DurableObjectState) { this.state = state; } + private bufferClientFrame(clientId: string, message: string | ArrayBuffer): void { + const existing = this.pendingClientFrames.get(clientId) ?? []; + existing.push(message); + // Prevent unbounded memory growth if a daemon never connects. + if (existing.length > 200) { + existing.splice(0, existing.length - 200); + } + this.pendingClientFrames.set(clientId, existing); + } + + private flushClientFrames(clientId: string, serverWs: WebSocket): void { + const frames = this.pendingClientFrames.get(clientId); + if (!frames || frames.length === 0) return; + this.pendingClientFrames.delete(clientId); + for (const frame of frames) { + try { + serverWs.send(frame); + } catch { + // If we can't flush, re-buffer and let the daemon re-establish. + this.bufferClientFrame(clientId, frame); + break; + } + } + } + + private listConnectedClientIds(): string[] { + const out = new Set(); + for (const ws of this.state.getWebSockets("client")) { + try { + const attachment = (ws as WebSocketWithAttachment).deserializeAttachment() as RelaySessionAttachment | null; + if (attachment?.role === "client" && typeof attachment.clientId === "string" && attachment.clientId) { + out.add(attachment.clientId); + } + } catch { + // ignore + } + } + return Array.from(out); + } + + private notifyControls(message: unknown): void { + const text = JSON.stringify(message); + for (const ws of this.state.getWebSockets("server-control")) { + try { + ws.send(text); + } catch { + // ignore + } + } + } + async fetch(request: Request): Promise { const url = new URL(request.url); const role = url.searchParams.get("role") as ConnectionRole | null; const serverId = url.searchParams.get("serverId"); + const clientIdRaw = url.searchParams.get("clientId"); + const clientId = typeof clientIdRaw === "string" ? clientIdRaw.trim() : ""; if (!role || (role !== "server" && role !== "client")) { return new Response("Missing or invalid role parameter", { status: 400 }); @@ -85,33 +139,83 @@ export class RelayDurableObject { return new Response("Missing serverId parameter", { status: 400 }); } + // v2 relay protocol: clients must provide a clientId so the daemon can create + // an independent E2EE channel per client connection. + if (role === "client" && !clientId) { + return new Response("Missing clientId parameter", { status: 400 }); + } + const upgradeHeader = request.headers.get("Upgrade"); if (!upgradeHeader || upgradeHeader.toLowerCase() !== "websocket") { return new Response("Expected WebSocket upgrade", { status: 426 }); } - // Close any existing connection with the same role - const existingConnections = this.state.getWebSockets(role); - for (const ws of existingConnections) { - ws.close(1008, "Replaced by new connection"); + const isServerControl = role === "server" && !clientId; + const isServerData = role === "server" && !!clientId; + + // Close any existing connection with the same identity. + // - server-control: single per serverId + // - server-data: single per clientId + // - client: single per clientId + if (isServerControl) { + for (const ws of this.state.getWebSockets("server-control")) { + ws.close(1008, "Replaced by new connection"); + } + } else if (isServerData) { + for (const ws of this.state.getWebSockets(`server:${clientId}`)) { + ws.close(1008, "Replaced by new connection"); + } + } else { + for (const ws of this.state.getWebSockets(`client:${clientId}`)) { + ws.close(1008, "Replaced by new connection"); + } } // Create WebSocket pair const pair = new (globalThis as unknown as { WebSocketPair: new () => WebSocketPair }).WebSocketPair(); const [client, server] = [pair[0], pair[1]]; - // Accept with hibernation support, tagged by role - this.state.acceptWebSocket(server, [role]); + const tags: string[] = []; + if (role === "client") { + tags.push("client", `client:${clientId}`); + } else if (isServerControl) { + tags.push("server-control"); + } else { + tags.push("server", `server:${clientId}`); + } + + // Accept with hibernation support, tagged for lookup. + this.state.acceptWebSocket(server, tags); // Store attachment for hibernation recovery const attachment: RelaySessionAttachment = { serverId, role, + clientId: clientId || null, createdAt: Date.now(), }; (server as WebSocketWithAttachment).serializeAttachment(attachment); - console.log(`[Relay DO] ${role} connected to session ${serverId}`); + console.log( + `[Relay DO] ${role}${isServerControl ? "(control)" : ""}${isServerData ? `(data:${clientId})` : role === "client" ? `(${clientId})` : ""} connected to session ${serverId}` + ); + + if (role === "client") { + this.notifyControls({ type: "client_connected", clientId }); + } + + if (isServerControl) { + // Send current client list so the daemon can attach existing clients. + try { + server.send(JSON.stringify({ type: "sync", clientIds: this.listConnectedClientIds() })); + } catch { + // ignore + } + } + + if (isServerData && clientId) { + this.flushClientFrames(clientId, server); + } return new Response(null, { status: 101, @@ -129,15 +233,35 @@ export class RelayDurableObject { return; } - const { role } = attachment; - const targetRole = role === "server" ? "client" : "server"; - const targets = this.state.getWebSockets(targetRole); + const { role, clientId } = attachment; + if (!clientId) { + // Control channel: ignore payloads (daemon can use it for pings later). + return; + } + if (role === "client") { + const servers = this.state.getWebSockets(`server:${clientId}`); + if (servers.length === 0) { + this.bufferClientFrame(clientId, message); + return; + } + for (const target of servers) { + try { + target.send(message); + } catch (error) { + console.error(`[Relay DO] Failed to forward client->server(${clientId}):`, error); + } + } + return; + } + + // server data socket -> client + const targets = this.state.getWebSockets(`client:${clientId}`); for (const target of targets) { try { target.send(message); } catch (error) { - console.error(`[Relay DO] Failed to forward to ${targetRole}:`, error); + console.error(`[Relay DO] Failed to forward server->client(${clientId}):`, error); } } } @@ -155,11 +279,33 @@ export class RelayDurableObject { if (!attachment) return; console.log( - `[Relay DO] ${attachment.role} disconnected from session ${attachment.serverId} (${code}: ${reason})` + `[Relay DO] ${attachment.role}${attachment.clientId ? `(${attachment.clientId})` : ""} disconnected from session ${attachment.serverId} (${code}: ${reason})` ); - // Note: The other connection remains open. - // The client/server will detect the peer is gone when they try to send. + if (attachment.role === "client" && attachment.clientId) { + this.pendingClientFrames.delete(attachment.clientId); + // Close the matching server-data socket so the daemon can clean up quickly. + for (const serverWs of this.state.getWebSockets(`server:${attachment.clientId}`)) { + try { + serverWs.close(1001, "Client disconnected"); + } catch { + // ignore + } + } + this.notifyControls({ type: "client_disconnected", clientId: attachment.clientId }); + return; + } + + if (attachment.role === "server" && attachment.clientId) { + // Force the client to reconnect and re-handshake when the daemon side drops. + for (const clientWs of this.state.getWebSockets(`client:${attachment.clientId}`)) { + try { + clientWs.close(1012, "Server disconnected"); + } catch { + // ignore + } + } + } } /** diff --git a/packages/relay/src/e2e.test.ts b/packages/relay/src/e2e.test.ts index 672eeada7..601468929 100644 --- a/packages/relay/src/e2e.test.ts +++ b/packages/relay/src/e2e.test.ts @@ -1,6 +1,6 @@ import { describe, it, expect, beforeAll, afterAll } from "vitest"; import { WebSocket } from "ws"; -import { createRelayServer, RelayServer } from "./node-adapter.js"; +import { createRelayServer, type RelayServer } from "./node-adapter.js"; import { generateKeyPair, exportPublicKey, @@ -26,6 +26,7 @@ describe("E2E Relay with E2EE", () => { it("full flow: daemon and client exchange encrypted messages through relay", async () => { const serverId = "test-session-" + Date.now(); + const clientId = "clt_test_" + Date.now() + "_" + Math.random().toString(36).slice(2); // === DAEMON SIDE === // Generate keypair (public key goes in QR) @@ -34,14 +35,14 @@ describe("E2E Relay with E2EE", () => { // QR would contain: { serverId, daemonPubKeyB64, relay: { endpoint } } - // Daemon connects to relay as "server" role - const daemonWs = new WebSocket( + // Daemon connects to relay as "server" control role + const daemonControlWs = new WebSocket( `ws://127.0.0.1:${TEST_PORT}/ws?serverId=${serverId}&role=server` ); await new Promise((resolve, reject) => { - daemonWs.on("open", resolve); - daemonWs.on("error", reject); + daemonControlWs.on("open", resolve); + daemonControlWs.on("error", reject); }); // === CLIENT SIDE === @@ -57,9 +58,9 @@ describe("E2E Relay with E2EE", () => { daemonPubKeyOnClient ); - // Client connects to relay as "client" role + // Client connects to relay as "client" role (must include clientId in v2) const clientWs = new WebSocket( - `ws://127.0.0.1:${TEST_PORT}/ws?serverId=${serverId}&role=client` + `ws://127.0.0.1:${TEST_PORT}/ws?serverId=${serverId}&role=client&clientId=${clientId}` ); await new Promise((resolve, reject) => { @@ -67,6 +68,30 @@ describe("E2E Relay with E2EE", () => { clientWs.on("error", reject); }); + // Wait for relay to notify daemon control socket, then open the per-client server-data socket. + await new Promise((resolve, reject) => { + const timeout = setTimeout(() => reject(new Error("timed out waiting for client_connected")), 5000); + daemonControlWs.on("message", (raw) => { + try { + const msg = JSON.parse(raw.toString()); + if (msg && msg.type === "client_connected" && msg.clientId === clientId) { + clearTimeout(timeout); + resolve(); + } + } catch { + // ignore + } + }); + }); + + const daemonWs = new WebSocket( + `ws://127.0.0.1:${TEST_PORT}/ws?serverId=${serverId}&role=server&clientId=${clientId}` + ); + await new Promise((resolve, reject) => { + daemonWs.on("open", resolve); + daemonWs.on("error", reject); + }); + // Client sends hello with its public key (this message is NOT encrypted - it's the handshake) const helloMsg = JSON.stringify({ type: "hello", key: clientPubKeyB64 }); clientWs.send(helloMsg); @@ -150,6 +175,7 @@ describe("E2E Relay with E2EE", () => { it("relay only sees opaque bytes after handshake", async () => { const serverId = "opaque-test-" + Date.now(); + const clientId = "clt_opaque_" + Date.now() + "_" + Math.random().toString(36).slice(2); // Setup keys const daemonKeyPair = await generateKeyPair(); @@ -170,18 +196,59 @@ describe("E2E Relay with E2EE", () => { daemonPubKey ); - // Connect both - const daemonWs = new WebSocket( + const daemonControlWs = new WebSocket( `ws://127.0.0.1:${TEST_PORT}/ws?serverId=${serverId}&role=server` ); - const clientWs = new WebSocket( - `ws://127.0.0.1:${TEST_PORT}/ws?serverId=${serverId}&role=client` - ); + await new Promise((r) => daemonControlWs.on("open", r)); - await Promise.all([ - new Promise((r) => daemonWs.on("open", r)), - new Promise((r) => clientWs.on("open", r)), - ]); + const waitForClientSeen = new Promise((resolve, reject) => { + const timeout = setTimeout( + () => reject(new Error("timed out waiting for client_connected")), + 5000 + ); + const onMessage = (raw: unknown) => { + try { + const text = + typeof raw === "string" + ? raw + : raw && typeof (raw as any).toString === "function" + ? (raw as any).toString() + : ""; + const msg = JSON.parse(text); + if (msg?.type === "client_connected" && msg.clientId === clientId) { + clearTimeout(timeout); + daemonControlWs.off("message", onMessage); + resolve(); + return; + } + if (msg?.type === "sync" && Array.isArray(msg.clientIds) && msg.clientIds.includes(clientId)) { + clearTimeout(timeout); + daemonControlWs.off("message", onMessage); + resolve(); + } + } catch { + // ignore + } + }; + daemonControlWs.on("message", onMessage); + }); + + const clientWs = new WebSocket( + `ws://127.0.0.1:${TEST_PORT}/ws?serverId=${serverId}&role=client&clientId=${clientId}` + ); + await new Promise((r) => clientWs.on("open", r)); + await waitForClientSeen; + + const daemonWs = new WebSocket( + `ws://127.0.0.1:${TEST_PORT}/ws?serverId=${serverId}&role=server&clientId=${clientId}` + ); + await new Promise((r) => daemonWs.on("open", r)); + + // Handshake (not encrypted) + clientWs.send(JSON.stringify({ type: "hello", key: clientPubKeyB64 })); + await new Promise((resolve) => { + daemonWs.once("message", () => resolve()); + }); // Send encrypted secret const secret = "This is a secret that relay cannot read"; @@ -207,6 +274,7 @@ describe("E2E Relay with E2EE", () => { ); expect(decrypted).toBe(secret); + daemonControlWs.close(); daemonWs.close(); clientWs.close(); }); diff --git a/packages/relay/src/live-relay.e2e.test.ts b/packages/relay/src/live-relay.e2e.test.ts index 62a57be1d..fd83e4b29 100644 --- a/packages/relay/src/live-relay.e2e.test.ts +++ b/packages/relay/src/live-relay.e2e.test.ts @@ -36,8 +36,14 @@ describe("Live relay (relay.paseo.sh) E2E", () => { await withRetry( async () => { const serverId = `live-${Date.now()}-${Math.random().toString(16).slice(2)}`; - const serverUrl = `${RELAY_BASE_URL}/ws?serverId=${encodeURIComponent(serverId)}&role=server`; - const clientUrl = `${RELAY_BASE_URL}/ws?serverId=${encodeURIComponent(serverId)}&role=client`; + const clientId = `clt_live_${Date.now()}_${Math.random().toString(16).slice(2)}`; + const serverControlUrl = `${RELAY_BASE_URL}/ws?serverId=${encodeURIComponent(serverId)}&role=server`; + const serverDataUrl = `${RELAY_BASE_URL}/ws?serverId=${encodeURIComponent( + serverId + )}&role=server&clientId=${encodeURIComponent(clientId)}`; + const clientUrl = `${RELAY_BASE_URL}/ws?serverId=${encodeURIComponent( + serverId + )}&role=client&clientId=${encodeURIComponent(clientId)}`; // === Key setup === const daemonKeyPair = await generateKeyPair(); @@ -53,8 +59,9 @@ describe("Live relay (relay.paseo.sh) E2E", () => { ); // === Connect === - const daemonWs = new WebSocket(serverUrl); + const daemonControlWs = new WebSocket(serverControlUrl); const clientWs = new WebSocket(clientUrl); + let daemonWs: WebSocket | null = null; const waitOpen = (ws: WebSocket, label: string) => new Promise((resolve, reject) => { @@ -73,7 +80,31 @@ describe("Live relay (relay.paseo.sh) E2E", () => { }); try { - await Promise.all([waitOpen(daemonWs, "server"), waitOpen(clientWs, "client")]); + await Promise.all([ + waitOpen(daemonControlWs, "server-control"), + waitOpen(clientWs, "client"), + ]); + + await new Promise((resolve, reject) => { + const timeout = setTimeout( + () => reject(new Error("Timed out waiting for client_connected")), + 10_000 + ); + daemonControlWs.on("message", (raw) => { + try { + const msg = JSON.parse(raw.toString()); + if (msg && msg.type === "client_connected" && msg.clientId === clientId) { + clearTimeout(timeout); + resolve(); + } + } catch { + // ignore + } + }); + }); + + daemonWs = new WebSocket(serverDataUrl); + await waitOpen(daemonWs, "server-data"); // === Handshake === // Client sends hello with its public key (not encrypted). @@ -84,7 +115,7 @@ describe("Live relay (relay.paseo.sh) E2E", () => { () => reject(new Error("Timed out waiting for hello")), 10_000 ); - daemonWs.once("message", (data) => { + daemonWs!.once("message", (data) => { clearTimeout(timeout); resolve(data.toString()); }); @@ -117,7 +148,7 @@ describe("Live relay (relay.paseo.sh) E2E", () => { () => reject(new Error("Timed out waiting for encrypted message")), 10_000 ); - daemonWs.once("message", (data) => { + daemonWs!.once("message", (data) => { clearTimeout(timeout); resolve(data as Buffer); }); @@ -139,7 +170,7 @@ describe("Live relay (relay.paseo.sh) E2E", () => { daemonSharedKey, plaintextFromDaemon ); - daemonWs.send(Buffer.from(ciphertextFromDaemon)); + daemonWs!.send(Buffer.from(ciphertextFromDaemon)); const clientReceivedCiphertext = await new Promise( (resolve, reject) => { @@ -164,7 +195,8 @@ describe("Live relay (relay.paseo.sh) E2E", () => { ); expect(decryptedOnClient).toBe(plaintextFromDaemon); } finally { - daemonWs.close(); + daemonControlWs.close(); + daemonWs?.close(); clientWs.close(); } }, diff --git a/packages/relay/src/node-adapter.ts b/packages/relay/src/node-adapter.ts index 430d609d7..cf97cd922 100644 --- a/packages/relay/src/node-adapter.ts +++ b/packages/relay/src/node-adapter.ts @@ -1,8 +1,7 @@ import http from "http"; import type { Socket } from "net"; import { WebSocketServer, WebSocket as NodeWebSocket } from "ws"; -import { Relay } from "./relay.js"; -import type { ConnectionRole, RelayConnection } from "./types.js"; +import type { ConnectionRole } from "./types.js"; export interface NodeRelayServerConfig { port: number; @@ -12,23 +11,62 @@ export interface NodeRelayServerConfig { /** * Standalone Node.js relay server for self-hosting. * - * This is a separate process that bridges daemon↔client connections. - * Use this when you want to self-host a relay instead of using Cloudflare. + * v2 protocol: + * - Daemon connects a single control socket: + * ws://host:port/ws?serverId=abc&role=server + * - Each app client connects with a clientId: + * ws://host:port/ws?serverId=abc&role=client&clientId=clt_... + * - For every connected clientId, the daemon opens a dedicated server-data socket: + * ws://host:port/ws?serverId=abc&role=server&clientId=clt_... * - * Usage: - * ```ts - * const server = createRelayServer({ port: 8080 }); - * await server.start(); - * ``` - * - * Clients connect via: - * - ws://host:port/ws?serverId=abc&role=server - * - ws://host:port/ws?serverId=abc&role=client + * This allows multiple independent clients, each with its own E2EE handshake. */ export interface RelayServer { start(): Promise; stop(): Promise; - getRelay(): Relay; +} + +type Session = { + serverId: string; + control: NodeWebSocket | null; + clients: Map>; // clientId -> sockets + servers: Map; // clientId -> server-data socket + pending: Map>; // clientId -> buffered frames +}; + +function bufferFromWsData(data: unknown): Buffer { + if (Buffer.isBuffer(data)) return data; + + if (Array.isArray(data)) { + return Buffer.concat(data.map(bufferFromWsData)); + } + + if (data instanceof ArrayBuffer) { + return Buffer.from(data); + } + + if (ArrayBuffer.isView(data)) { + return Buffer.from(data.buffer, data.byteOffset, data.byteLength); + } + + if (typeof data === "string") { + return Buffer.from(data, "utf8"); + } + + return Buffer.from(String(data), "utf8"); +} + +function normalizeWsMessage(data: unknown, isBinary: boolean): string | ArrayBuffer { + if (!isBinary) { + if (typeof data === "string") return data; + return bufferFromWsData(data).toString("utf8"); + } + + if (data instanceof ArrayBuffer) return data; + const buffer = bufferFromWsData(data); + const out = new Uint8Array(buffer.byteLength); + out.set(buffer); + return out.buffer; } export function createRelayServer(config: NodeRelayServerConfig): RelayServer { @@ -36,12 +74,57 @@ export function createRelayServer(config: NodeRelayServerConfig): RelayServer { const logFrames = process.env.PASEO_RELAY_LOG_FRAMES === "1"; const logUpgrades = process.env.PASEO_RELAY_LOG_UPGRADES === "1"; - const relay = new Relay({ - onSessionCreated: (id) => console.log(`[Relay] Session created: ${id}`), - onSessionBridged: (id) => console.log(`[Relay] Session bridged: ${id}`), - onSessionClosed: (id) => console.log(`[Relay] Session closed: ${id}`), - onError: (id, err) => console.error(`[Relay] Session ${id} error:`, err), - }); + const sessions = new Map(); + + const getSession = (serverId: string): Session => { + let session = sessions.get(serverId); + if (!session) { + session = { + serverId, + control: null, + clients: new Map(), + servers: new Map(), + pending: new Map(), + }; + sessions.set(serverId, session); + console.log(`[Relay] Session created: ${serverId}`); + } + return session; + }; + + const notifyControl = (session: Session, msg: unknown) => { + if (!session.control || session.control.readyState !== NodeWebSocket.OPEN) return; + try { + session.control.send(JSON.stringify(msg)); + } catch { + // ignore + } + }; + + const listClientIds = (session: Session): string[] => Array.from(session.clients.keys()); + + const bufferFrame = (session: Session, clientId: string, frame: string | ArrayBuffer) => { + const existing = session.pending.get(clientId) ?? []; + existing.push(frame); + if (existing.length > 200) existing.splice(0, existing.length - 200); + session.pending.set(clientId, existing); + }; + + const flushFrames = (session: Session, clientId: string) => { + const server = session.servers.get(clientId); + if (!server || server.readyState !== NodeWebSocket.OPEN) return; + const frames = session.pending.get(clientId); + if (!frames || frames.length === 0) return; + session.pending.delete(clientId); + for (const frame of frames) { + try { + server.send(frame); + } catch { + bufferFrame(session, clientId, frame); + break; + } + } + }; const httpServer = http.createServer((req, res) => { if (req.url === "/health") { @@ -53,13 +136,6 @@ export function createRelayServer(config: NodeRelayServerConfig): RelayServer { res.end("Not found"); }); - // NOTE: We intentionally use `noServer` + a manual `upgrade` handler instead - // of `new WebSocketServer({ server, path })`. - // - // In some environments (notably when running via `tsx -e`), we observed - // websocket upgrade requests falling through to the HTTP handler (404), - // despite `ws` being configured with `{ server, path: "/ws" }`. - // Handling `upgrade` explicitly is more robust and easier to debug. const wss = new WebSocketServer({ noServer: true }); httpServer.on("upgrade", (req, socket: Socket, head) => { @@ -82,11 +158,17 @@ export function createRelayServer(config: NodeRelayServerConfig): RelayServer { const serverId = url.searchParams.get("serverId"); const role = url.searchParams.get("role"); + const clientId = (url.searchParams.get("clientId") ?? "").trim(); if (!serverId || !role || (role !== "server" && role !== "client")) { socket.write("HTTP/1.1 400 Bad Request\r\n\r\n"); socket.destroy(); return; } + if (role === "client" && !clientId) { + socket.write("HTTP/1.1 400 Bad Request\r\n\r\n"); + socket.destroy(); + return; + } wss.handleUpgrade(req, socket, head, (ws) => { wss.emit("connection", ws, req); @@ -101,13 +183,48 @@ export function createRelayServer(config: NodeRelayServerConfig): RelayServer { const url = new URL(req.url ?? "", `http://${req.headers.host}`); const serverId = url.searchParams.get("serverId")!; const role = url.searchParams.get("role") as ConnectionRole; + const clientId = (url.searchParams.get("clientId") ?? "").trim(); + const session = getSession(serverId); - const connection = wrapWebSocket(ws, role); - relay.addConnection(serverId, role, connection); + const isControl = role === "server" && !clientId; + const isServerData = role === "server" && !!clientId; + + if (isControl) { + try { + session.control?.close(1008, "Replaced by new connection"); + } catch { + // ignore + } + session.control = ws; + try { + ws.send(JSON.stringify({ type: "sync", clientIds: listClientIds(session) })); + } catch { + // ignore + } + } + + if (role === "client") { + const set = session.clients.get(clientId) ?? new Set(); + set.add(ws); + session.clients.set(clientId, set); + notifyControl(session, { type: "client_connected", clientId }); + } + + if (isServerData) { + const prev = session.servers.get(clientId); + if (prev && prev !== ws) { + try { + prev.close(1008, "Replaced by new connection"); + } catch { + // ignore + } + } + session.servers.set(clientId, ws); + flushFrames(session, clientId); + } ws.on("message", (data, isBinary) => { - // If this socket was replaced, ignore any late frames. - if (relay.getSession(serverId)?.[role] !== connection) return; + const normalized = normalizeWsMessage(data, isBinary); if (logFrames) { const preview = (() => { @@ -126,20 +243,92 @@ export function createRelayServer(config: NodeRelayServerConfig): RelayServer { } })(); console.log( - `[Relay] frame ${serverId}/${role} binary=${isBinary} len=${len} preview=${JSON.stringify(preview)}` + `[Relay] frame ${serverId}/${role}${clientId ? `(${clientId})` : ""} binary=${isBinary} len=${len} preview=${JSON.stringify(preview)}` ); } - relay.forward(serverId, role, normalizeWsMessageForRelay(data, isBinary)); + + if (role === "client") { + const server = session.servers.get(clientId); + if (!server || server.readyState !== NodeWebSocket.OPEN) { + bufferFrame(session, clientId, normalized); + return; + } + try { + server.send(normalized); + } catch { + // ignore + } + return; + } + + if (isServerData) { + const targets = session.clients.get(clientId); + if (!targets) return; + for (const target of targets) { + if (target.readyState !== NodeWebSocket.OPEN) continue; + try { + target.send(normalized); + } catch { + // ignore + } + } + } }); - ws.on("close", () => { - // Avoid clearing the current connection if this socket was replaced. - if (relay.getSession(serverId)?.[role] !== connection) return; - relay.removeConnection(serverId, role); + ws.on("close", (code, reasonBuf) => { + const reason = reasonBuf?.toString?.() ?? ""; + + if (isControl) { + if (session.control === ws) session.control = null; + } else if (role === "client") { + const set = session.clients.get(clientId); + if (set) { + set.delete(ws); + if (set.size === 0) { + session.clients.delete(clientId); + session.pending.delete(clientId); + const server = session.servers.get(clientId); + if (server) { + try { + server.close(1001, "Client disconnected"); + } catch { + // ignore + } + session.servers.delete(clientId); + } + notifyControl(session, { type: "client_disconnected", clientId }); + } + } + } else if (isServerData) { + if (session.servers.get(clientId) === ws) { + session.servers.delete(clientId); + } + const targets = session.clients.get(clientId); + if (targets) { + for (const target of targets) { + try { + target.close(1012, "Server disconnected"); + } catch { + // ignore + } + } + } + } + + if (logUpgrades) { + console.log( + `[Relay] close ${serverId}/${role}${clientId ? `(${clientId})` : ""} code=${code} reason=${JSON.stringify(reason)}` + ); + } + + if (session.control === null && session.clients.size === 0 && session.servers.size === 0) { + sessions.delete(serverId); + console.log(`[Relay] Session closed: ${serverId}`); + } }); ws.on("error", (error) => { - console.error(`[Relay] WebSocket error for ${serverId}/${role}:`, error); + console.error(`[Relay] WebSocket error for ${serverId}/${role}${clientId ? `(${clientId})` : ""}:`, error); }); }); @@ -156,22 +345,14 @@ export function createRelayServer(config: NodeRelayServerConfig): RelayServer { stop() { return new Promise((resolve) => { - // Stop accepting new connections immediately. Waiting for `wss.close` - // can leave the HTTP server listening while the WS server is closing, - // causing upgrade requests to get 503 responses. try { httpServer.close(() => undefined); } catch { // ignore } - // Close active sessions + force-close any remaining clients. - for (const session of relay.listSessions()) { - relay.closeSession(session.serverId, 1001, "Server shutting down"); - } for (const ws of wss.clients) { try { - // terminate() is the fastest way to ensure `wss.close` completes. (ws as unknown as { terminate?: () => void }).terminate?.(); } catch { try { @@ -201,72 +382,6 @@ export function createRelayServer(config: NodeRelayServerConfig): RelayServer { } }); }, - - getRelay() { - return relay; - }, }; } -function wrapWebSocket(ws: NodeWebSocket, role: ConnectionRole): RelayConnection { - return { - role, - send: (data) => { - if (ws.readyState === NodeWebSocket.OPEN) { - ws.send(data); - } - }, - close: (code, reason) => { - ws.close(code, reason); - }, - }; -} - -function bufferToArrayBuffer(buffer: Buffer): ArrayBuffer { - const out = new Uint8Array(buffer.byteLength); - out.set(buffer); - return out.buffer; -} - -function normalizeWsMessageForRelay(data: unknown, isBinary: boolean): string | ArrayBuffer { - if (isBinary) { - return normalizeWsBinaryMessage(data); - } - return normalizeWsTextMessage(data); -} - -function normalizeWsBinaryMessage(data: unknown): ArrayBuffer { - if (data instanceof ArrayBuffer) { - return data; - } - return bufferToArrayBuffer(bufferFromWsData(data)); -} - -function normalizeWsTextMessage(data: unknown): string { - if (typeof data === "string") { - return data; - } - return bufferFromWsData(data).toString("utf8"); -} - -function bufferFromWsData(data: unknown): Buffer { - if (Buffer.isBuffer(data)) return data; - - if (Array.isArray(data)) { - return Buffer.concat(data.map(bufferFromWsData)); - } - - if (data instanceof ArrayBuffer) { - return Buffer.from(data); - } - - if (ArrayBuffer.isView(data)) { - return Buffer.from(data.buffer, data.byteOffset, data.byteLength); - } - - if (typeof data === "string") { - return Buffer.from(data, "utf8"); - } - - return Buffer.from(String(data), "utf8"); -} diff --git a/packages/relay/src/types.ts b/packages/relay/src/types.ts index 482460b41..aa59ea88d 100644 --- a/packages/relay/src/types.ts +++ b/packages/relay/src/types.ts @@ -26,6 +26,11 @@ export interface RelayConnection { export interface RelaySessionAttachment { serverId: string; role: ConnectionRole; + /** + * v2 relay: unique id for the client connection. Allows the daemon to create + * an independent socket + E2EE channel per connected client. + */ + clientId?: string | null; createdAt: number; } diff --git a/packages/server/src/client/daemon-client.ts b/packages/server/src/client/daemon-client.ts index 3290fa92c..7610892ca 100644 --- a/packages/server/src/client/daemon-client.ts +++ b/packages/server/src/client/daemon-client.ts @@ -294,9 +294,25 @@ export class DaemonClient { >(); private logger: Logger; private pendingSendQueue: PendingSend[] = []; + private relayClientId: string | null = null; constructor(private config: DaemonClientConfig) { this.logger = config.logger ?? consoleLogger; + // Relay v2 requires a clientId so the daemon can create an independent + // socket + E2EE channel per connected client. Generate one per DaemonClient + // instance (stable across reconnects in this tab/app session). + if (isRelayClientWebSocketUrl(this.config.url)) { + try { + const parsed = new URL(this.config.url); + if (!parsed.searchParams.get("clientId")) { + this.relayClientId = `clt_${safeRandomId()}`; + parsed.searchParams.set("clientId", this.relayClientId); + this.config.url = parsed.toString(); + } + } catch { + // ignore - invalid URL will be handled on connect + } + } } // ============================================================================ @@ -2916,6 +2932,17 @@ function describeTransportError(event?: unknown): string { return "Transport error"; } +function safeRandomId(): string { + try { + if (typeof crypto !== "undefined" && typeof crypto.randomUUID === "function") { + return crypto.randomUUID(); + } + } catch { + // ignore + } + return `${Date.now().toString(36)}-${Math.random().toString(36).slice(2)}`; +} + function decodeMessageData(data: unknown): string | null { if (data === null || data === undefined) { return null; diff --git a/packages/server/src/server/daemon-e2e/relay-transport.e2e.test.ts b/packages/server/src/server/daemon-e2e/relay-transport.e2e.test.ts index 1750581b5..0e5e79409 100644 --- a/packages/server/src/server/daemon-e2e/relay-transport.e2e.test.ts +++ b/packages/server/src/server/daemon-e2e/relay-transport.e2e.test.ts @@ -91,11 +91,13 @@ describe("Relay transport (E2EE) - daemon E2E", () => { const offerUrl = parseOfferUrlFromLogs(lines); const { serverId, daemonPublicKeyB64 } = decodeOfferFromFragmentUrl(offerUrl); + const clientId = `clt_test_${Date.now().toString(36)}_${Math.random().toString(36).slice(2)}`; const ws = new WebSocket( buildRelayWebSocketUrl({ endpoint: `127.0.0.1:${relayPort}`, serverId, role: "client", + clientId, }) ); @@ -209,6 +211,7 @@ describe("Relay transport (E2EE) - daemon E2E", () => { endpoint: `127.0.0.1:${relayPort}`, serverId, role: "client", + clientId: `clt_test_${Date.now().toString(36)}_${Math.random().toString(36).slice(2)}`, }) ); diff --git a/packages/server/src/server/relay-transport.ts b/packages/server/src/server/relay-transport.ts index 3abdb5f0e..e7583098f 100644 --- a/packages/server/src/server/relay-transport.ts +++ b/packages/server/src/server/relay-transport.ts @@ -30,6 +30,33 @@ type RelaySocketLike = { once: (event: "close" | "error", listener: (...args: any[]) => void) => void; }; +type ControlMessage = + | { type: "sync"; clientIds: string[] } + | { type: "client_connected"; clientId: string } + | { type: "client_disconnected"; clientId: string }; + +function tryParseControlMessage(raw: unknown): ControlMessage | null { + try { + const text = + typeof raw === "string" ? raw : Buffer.isBuffer(raw) ? raw.toString("utf8") : String(raw); + const parsed = JSON.parse(text) as any; + if (!parsed || typeof parsed !== "object") return null; + if (parsed.type === "sync" && Array.isArray(parsed.clientIds)) { + const clientIds = parsed.clientIds.filter((id: unknown) => typeof id === "string" && id.trim().length > 0); + return { type: "sync", clientIds }; + } + if (parsed.type === "client_connected" && typeof parsed.clientId === "string" && parsed.clientId.trim()) { + return { type: "client_connected", clientId: parsed.clientId.trim() }; + } + if (parsed.type === "client_disconnected" && typeof parsed.clientId === "string" && parsed.clientId.trim()) { + return { type: "client_disconnected", clientId: parsed.clientId.trim() }; + } + return null; + } catch { + return null; + } +} + export function startRelayTransport({ logger, attachSocket, @@ -40,9 +67,10 @@ export function startRelayTransport({ const relayLogger = logger.child({ module: "relay-transport" }); let stopped = false; - let ws: WebSocket | null = null; + let controlWs: WebSocket | null = null; let reconnectTimeout: ReturnType | null = null; let reconnectAttempt = 0; + const dataSockets = new Map(); // clientId -> ws const stop = async (): Promise => { stopped = true; @@ -50,17 +78,25 @@ export function startRelayTransport({ clearTimeout(reconnectTimeout); reconnectTimeout = null; } - if (ws) { + if (controlWs) { + try { + controlWs.close(); + } catch { + // ignore + } + controlWs = null; + } + for (const ws of dataSockets.values()) { try { ws.close(); } catch { // ignore } - ws = null; } + dataSockets.clear(); }; - const connect = (): void => { + const connectControl = (): void => { if (stopped) return; const url = buildRelayWebSocketUrl({ @@ -69,27 +105,17 @@ export function startRelayTransport({ role: "server", }); const socket = new WebSocket(url); - ws = socket; - - let attached = false; + controlWs = socket; socket.on("open", () => { reconnectAttempt = 0; - relayLogger.info({ url }, "relay_connected"); - - if (attached) return; - attached = true; - if (daemonKeyPair) { - void attachEncryptedSocket(socket, daemonKeyPair, relayLogger, attachSocket); - } else { - void attachSocket(socket); - } + relayLogger.info({ url }, "relay_control_connected"); }); socket.on("close", (code, reason) => { relayLogger.warn( { code, reason: reason?.toString?.(), url }, - "relay_disconnected" + "relay_control_disconnected" ); scheduleReconnect(); }); @@ -98,6 +124,32 @@ export function startRelayTransport({ relayLogger.warn({ err, url }, "relay_error"); // close event will schedule reconnect }); + + socket.on("message", (data) => { + const msg = tryParseControlMessage(data); + if (!msg) return; + if (msg.type === "sync") { + for (const clientId of msg.clientIds) { + ensureClientDataSocket(clientId); + } + return; + } + if (msg.type === "client_connected") { + ensureClientDataSocket(msg.clientId); + return; + } + if (msg.type === "client_disconnected") { + const existing = dataSockets.get(msg.clientId); + if (existing) { + try { + existing.close(1001, "Client disconnected"); + } catch { + // ignore + } + dataSockets.delete(msg.clientId); + } + } + }); }; const scheduleReconnect = (): void => { @@ -108,11 +160,53 @@ export function startRelayTransport({ const delayMs = Math.min(30000, 1000 * reconnectAttempt); reconnectTimeout = setTimeout(() => { reconnectTimeout = null; - connect(); + connectControl(); }, delayMs); }; - connect(); + const ensureClientDataSocket = (clientId: string): void => { + if (stopped) return; + if (!clientId) return; + if (dataSockets.has(clientId)) return; + + const url = buildRelayWebSocketUrl({ + endpoint: relayEndpoint, + serverId, + role: "server", + clientId, + }); + const socket = new WebSocket(url); + dataSockets.set(clientId, socket); + + let attached = false; + + socket.on("open", () => { + relayLogger.info({ url, clientId }, "relay_data_connected"); + if (attached) return; + attached = true; + if (daemonKeyPair) { + void attachEncryptedSocket(socket, daemonKeyPair, relayLogger.child({ clientId }), attachSocket); + } else { + void attachSocket(socket); + } + }); + + socket.on("close", (code, reason) => { + relayLogger.warn( + { code, reason: reason?.toString?.(), url, clientId }, + "relay_data_disconnected" + ); + if (dataSockets.get(clientId) === socket) { + dataSockets.delete(clientId); + } + }); + + socket.on("error", (err) => { + relayLogger.warn({ err, url, clientId }, "relay_data_error"); + }); + }; + + connectControl(); return { stop }; } diff --git a/packages/server/src/shared/daemon-endpoints.ts b/packages/server/src/shared/daemon-endpoints.ts index c6b49e323..45f30dac0 100644 --- a/packages/server/src/shared/daemon-endpoints.ts +++ b/packages/server/src/shared/daemon-endpoints.ts @@ -71,6 +71,7 @@ export function buildRelayWebSocketUrl(params: { endpoint: string; serverId: string; role: RelayRole; + clientId?: string; }): string { const { host, port, isIpv6 } = parseHostPort(params.endpoint); const protocol = shouldUseSecureWebSocket(port) ? "wss" : "ws"; @@ -78,6 +79,9 @@ export function buildRelayWebSocketUrl(params: { const url = new URL(`${protocol}://${hostPart}:${port}/ws`); url.searchParams.set("serverId", params.serverId); url.searchParams.set("role", params.role); + if (params.clientId) { + url.searchParams.set("clientId", params.clientId); + } return url.toString(); }