relay(v2): support multiple clients per serverId

This commit is contained in:
Mohamed Boudra
2026-02-05 09:48:43 +07:00
parent 98896bfb37
commit 36e105ed30
9 changed files with 663 additions and 169 deletions
+160 -14
View File
@@ -66,16 +66,70 @@ interface CFResponseInit extends ResponseInit {
export class RelayDurableObject { export class RelayDurableObject {
private state: DurableObjectState; private state: DurableObjectState;
private pendingClientFrames = new Map<string, Array<string | ArrayBuffer>>();
constructor(state: DurableObjectState) { constructor(state: DurableObjectState) {
this.state = state; 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<string>();
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<Response> { async fetch(request: Request): Promise<Response> {
const url = new URL(request.url); const url = new URL(request.url);
const role = url.searchParams.get("role") as ConnectionRole | null; const role = url.searchParams.get("role") as ConnectionRole | null;
const serverId = url.searchParams.get("serverId"); 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")) { if (!role || (role !== "server" && role !== "client")) {
return new Response("Missing or invalid role parameter", { status: 400 }); 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 }); 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"); const upgradeHeader = request.headers.get("Upgrade");
if (!upgradeHeader || upgradeHeader.toLowerCase() !== "websocket") { if (!upgradeHeader || upgradeHeader.toLowerCase() !== "websocket") {
return new Response("Expected WebSocket upgrade", { status: 426 }); return new Response("Expected WebSocket upgrade", { status: 426 });
} }
// Close any existing connection with the same role const isServerControl = role === "server" && !clientId;
const existingConnections = this.state.getWebSockets(role); const isServerData = role === "server" && !!clientId;
for (const ws of existingConnections) {
ws.close(1008, "Replaced by new connection"); // 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 // Create WebSocket pair
const pair = new (globalThis as unknown as { WebSocketPair: new () => WebSocketPair }).WebSocketPair(); const pair = new (globalThis as unknown as { WebSocketPair: new () => WebSocketPair }).WebSocketPair();
const [client, server] = [pair[0], pair[1]]; const [client, server] = [pair[0], pair[1]];
// Accept with hibernation support, tagged by role const tags: string[] = [];
this.state.acceptWebSocket(server, [role]); 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 // Store attachment for hibernation recovery
const attachment: RelaySessionAttachment = { const attachment: RelaySessionAttachment = {
serverId, serverId,
role, role,
clientId: clientId || null,
createdAt: Date.now(), createdAt: Date.now(),
}; };
(server as WebSocketWithAttachment).serializeAttachment(attachment); (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, { return new Response(null, {
status: 101, status: 101,
@@ -129,15 +233,35 @@ export class RelayDurableObject {
return; return;
} }
const { role } = attachment; const { role, clientId } = attachment;
const targetRole = role === "server" ? "client" : "server"; if (!clientId) {
const targets = this.state.getWebSockets(targetRole); // 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) { for (const target of targets) {
try { try {
target.send(message); target.send(message);
} catch (error) { } 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; if (!attachment) return;
console.log( 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. if (attachment.role === "client" && attachment.clientId) {
// The client/server will detect the peer is gone when they try to send. 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
}
}
}
} }
/** /**
+84 -16
View File
@@ -1,6 +1,6 @@
import { describe, it, expect, beforeAll, afterAll } from "vitest"; import { describe, it, expect, beforeAll, afterAll } from "vitest";
import { WebSocket } from "ws"; import { WebSocket } from "ws";
import { createRelayServer, RelayServer } from "./node-adapter.js"; import { createRelayServer, type RelayServer } from "./node-adapter.js";
import { import {
generateKeyPair, generateKeyPair,
exportPublicKey, exportPublicKey,
@@ -26,6 +26,7 @@ describe("E2E Relay with E2EE", () => {
it("full flow: daemon and client exchange encrypted messages through relay", async () => { it("full flow: daemon and client exchange encrypted messages through relay", async () => {
const serverId = "test-session-" + Date.now(); const serverId = "test-session-" + Date.now();
const clientId = "clt_test_" + Date.now() + "_" + Math.random().toString(36).slice(2);
// === DAEMON SIDE === // === DAEMON SIDE ===
// Generate keypair (public key goes in QR) // Generate keypair (public key goes in QR)
@@ -34,14 +35,14 @@ describe("E2E Relay with E2EE", () => {
// QR would contain: { serverId, daemonPubKeyB64, relay: { endpoint } } // QR would contain: { serverId, daemonPubKeyB64, relay: { endpoint } }
// Daemon connects to relay as "server" role // Daemon connects to relay as "server" control role
const daemonWs = new WebSocket( const daemonControlWs = new WebSocket(
`ws://127.0.0.1:${TEST_PORT}/ws?serverId=${serverId}&role=server` `ws://127.0.0.1:${TEST_PORT}/ws?serverId=${serverId}&role=server`
); );
await new Promise<void>((resolve, reject) => { await new Promise<void>((resolve, reject) => {
daemonWs.on("open", resolve); daemonControlWs.on("open", resolve);
daemonWs.on("error", reject); daemonControlWs.on("error", reject);
}); });
// === CLIENT SIDE === // === CLIENT SIDE ===
@@ -57,9 +58,9 @@ describe("E2E Relay with E2EE", () => {
daemonPubKeyOnClient daemonPubKeyOnClient
); );
// Client connects to relay as "client" role // Client connects to relay as "client" role (must include clientId in v2)
const clientWs = new WebSocket( 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<void>((resolve, reject) => { await new Promise<void>((resolve, reject) => {
@@ -67,6 +68,30 @@ describe("E2E Relay with E2EE", () => {
clientWs.on("error", reject); clientWs.on("error", reject);
}); });
// Wait for relay to notify daemon control socket, then open the per-client server-data socket.
await new Promise<void>((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<void>((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) // Client sends hello with its public key (this message is NOT encrypted - it's the handshake)
const helloMsg = JSON.stringify({ type: "hello", key: clientPubKeyB64 }); const helloMsg = JSON.stringify({ type: "hello", key: clientPubKeyB64 });
clientWs.send(helloMsg); clientWs.send(helloMsg);
@@ -150,6 +175,7 @@ describe("E2E Relay with E2EE", () => {
it("relay only sees opaque bytes after handshake", async () => { it("relay only sees opaque bytes after handshake", async () => {
const serverId = "opaque-test-" + Date.now(); const serverId = "opaque-test-" + Date.now();
const clientId = "clt_opaque_" + Date.now() + "_" + Math.random().toString(36).slice(2);
// Setup keys // Setup keys
const daemonKeyPair = await generateKeyPair(); const daemonKeyPair = await generateKeyPair();
@@ -170,18 +196,59 @@ describe("E2E Relay with E2EE", () => {
daemonPubKey daemonPubKey
); );
// Connect both const daemonControlWs = new WebSocket(
const daemonWs = new WebSocket(
`ws://127.0.0.1:${TEST_PORT}/ws?serverId=${serverId}&role=server` `ws://127.0.0.1:${TEST_PORT}/ws?serverId=${serverId}&role=server`
); );
const clientWs = new WebSocket( await new Promise<void>((r) => daemonControlWs.on("open", r));
`ws://127.0.0.1:${TEST_PORT}/ws?serverId=${serverId}&role=client`
);
await Promise.all([ const waitForClientSeen = new Promise<void>((resolve, reject) => {
new Promise<void>((r) => daemonWs.on("open", r)), const timeout = setTimeout(
new Promise<void>((r) => clientWs.on("open", r)), () => 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<void>((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<void>((r) => daemonWs.on("open", r));
// Handshake (not encrypted)
clientWs.send(JSON.stringify({ type: "hello", key: clientPubKeyB64 }));
await new Promise<void>((resolve) => {
daemonWs.once("message", () => resolve());
});
// Send encrypted secret // Send encrypted secret
const secret = "This is a secret that relay cannot read"; const secret = "This is a secret that relay cannot read";
@@ -207,6 +274,7 @@ describe("E2E Relay with E2EE", () => {
); );
expect(decrypted).toBe(secret); expect(decrypted).toBe(secret);
daemonControlWs.close();
daemonWs.close(); daemonWs.close();
clientWs.close(); clientWs.close();
}); });
+40 -8
View File
@@ -36,8 +36,14 @@ describe("Live relay (relay.paseo.sh) E2E", () => {
await withRetry( await withRetry(
async () => { async () => {
const serverId = `live-${Date.now()}-${Math.random().toString(16).slice(2)}`; const serverId = `live-${Date.now()}-${Math.random().toString(16).slice(2)}`;
const serverUrl = `${RELAY_BASE_URL}/ws?serverId=${encodeURIComponent(serverId)}&role=server`; const clientId = `clt_live_${Date.now()}_${Math.random().toString(16).slice(2)}`;
const clientUrl = `${RELAY_BASE_URL}/ws?serverId=${encodeURIComponent(serverId)}&role=client`; 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 === // === Key setup ===
const daemonKeyPair = await generateKeyPair(); const daemonKeyPair = await generateKeyPair();
@@ -53,8 +59,9 @@ describe("Live relay (relay.paseo.sh) E2E", () => {
); );
// === Connect === // === Connect ===
const daemonWs = new WebSocket(serverUrl); const daemonControlWs = new WebSocket(serverControlUrl);
const clientWs = new WebSocket(clientUrl); const clientWs = new WebSocket(clientUrl);
let daemonWs: WebSocket | null = null;
const waitOpen = (ws: WebSocket, label: string) => const waitOpen = (ws: WebSocket, label: string) =>
new Promise<void>((resolve, reject) => { new Promise<void>((resolve, reject) => {
@@ -73,7 +80,31 @@ describe("Live relay (relay.paseo.sh) E2E", () => {
}); });
try { try {
await Promise.all([waitOpen(daemonWs, "server"), waitOpen(clientWs, "client")]); await Promise.all([
waitOpen(daemonControlWs, "server-control"),
waitOpen(clientWs, "client"),
]);
await new Promise<void>((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 === // === Handshake ===
// Client sends hello with its public key (not encrypted). // 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")), () => reject(new Error("Timed out waiting for hello")),
10_000 10_000
); );
daemonWs.once("message", (data) => { daemonWs!.once("message", (data) => {
clearTimeout(timeout); clearTimeout(timeout);
resolve(data.toString()); resolve(data.toString());
}); });
@@ -117,7 +148,7 @@ describe("Live relay (relay.paseo.sh) E2E", () => {
() => reject(new Error("Timed out waiting for encrypted message")), () => reject(new Error("Timed out waiting for encrypted message")),
10_000 10_000
); );
daemonWs.once("message", (data) => { daemonWs!.once("message", (data) => {
clearTimeout(timeout); clearTimeout(timeout);
resolve(data as Buffer); resolve(data as Buffer);
}); });
@@ -139,7 +170,7 @@ describe("Live relay (relay.paseo.sh) E2E", () => {
daemonSharedKey, daemonSharedKey,
plaintextFromDaemon plaintextFromDaemon
); );
daemonWs.send(Buffer.from(ciphertextFromDaemon)); daemonWs!.send(Buffer.from(ciphertextFromDaemon));
const clientReceivedCiphertext = await new Promise<Buffer>( const clientReceivedCiphertext = await new Promise<Buffer>(
(resolve, reject) => { (resolve, reject) => {
@@ -164,7 +195,8 @@ describe("Live relay (relay.paseo.sh) E2E", () => {
); );
expect(decryptedOnClient).toBe(plaintextFromDaemon); expect(decryptedOnClient).toBe(plaintextFromDaemon);
} finally { } finally {
daemonWs.close(); daemonControlWs.close();
daemonWs?.close();
clientWs.close(); clientWs.close();
} }
}, },
+227 -112
View File
@@ -1,8 +1,7 @@
import http from "http"; import http from "http";
import type { Socket } from "net"; import type { Socket } from "net";
import { WebSocketServer, WebSocket as NodeWebSocket } from "ws"; import { WebSocketServer, WebSocket as NodeWebSocket } from "ws";
import { Relay } from "./relay.js"; import type { ConnectionRole } from "./types.js";
import type { ConnectionRole, RelayConnection } from "./types.js";
export interface NodeRelayServerConfig { export interface NodeRelayServerConfig {
port: number; port: number;
@@ -12,23 +11,62 @@ export interface NodeRelayServerConfig {
/** /**
* Standalone Node.js relay server for self-hosting. * Standalone Node.js relay server for self-hosting.
* *
* This is a separate process that bridges daemon↔client connections. * v2 protocol:
* Use this when you want to self-host a relay instead of using Cloudflare. * - 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: * This allows multiple independent clients, each with its own E2EE handshake.
* ```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
*/ */
export interface RelayServer { export interface RelayServer {
start(): Promise<void>; start(): Promise<void>;
stop(): Promise<void>; stop(): Promise<void>;
getRelay(): Relay; }
type Session = {
serverId: string;
control: NodeWebSocket | null;
clients: Map<string, Set<NodeWebSocket>>; // clientId -> sockets
servers: Map<string, NodeWebSocket>; // clientId -> server-data socket
pending: Map<string, Array<string | ArrayBuffer>>; // 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 { 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 logFrames = process.env.PASEO_RELAY_LOG_FRAMES === "1";
const logUpgrades = process.env.PASEO_RELAY_LOG_UPGRADES === "1"; const logUpgrades = process.env.PASEO_RELAY_LOG_UPGRADES === "1";
const relay = new Relay({ const sessions = new Map<string, Session>();
onSessionCreated: (id) => console.log(`[Relay] Session created: ${id}`),
onSessionBridged: (id) => console.log(`[Relay] Session bridged: ${id}`), const getSession = (serverId: string): Session => {
onSessionClosed: (id) => console.log(`[Relay] Session closed: ${id}`), let session = sessions.get(serverId);
onError: (id, err) => console.error(`[Relay] Session ${id} error:`, err), 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) => { const httpServer = http.createServer((req, res) => {
if (req.url === "/health") { if (req.url === "/health") {
@@ -53,13 +136,6 @@ export function createRelayServer(config: NodeRelayServerConfig): RelayServer {
res.end("Not found"); 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 }); const wss = new WebSocketServer({ noServer: true });
httpServer.on("upgrade", (req, socket: Socket, head) => { httpServer.on("upgrade", (req, socket: Socket, head) => {
@@ -82,11 +158,17 @@ export function createRelayServer(config: NodeRelayServerConfig): RelayServer {
const serverId = url.searchParams.get("serverId"); const serverId = url.searchParams.get("serverId");
const role = url.searchParams.get("role"); const role = url.searchParams.get("role");
const clientId = (url.searchParams.get("clientId") ?? "").trim();
if (!serverId || !role || (role !== "server" && role !== "client")) { if (!serverId || !role || (role !== "server" && role !== "client")) {
socket.write("HTTP/1.1 400 Bad Request\r\n\r\n"); socket.write("HTTP/1.1 400 Bad Request\r\n\r\n");
socket.destroy(); socket.destroy();
return; 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.handleUpgrade(req, socket, head, (ws) => {
wss.emit("connection", ws, req); 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 url = new URL(req.url ?? "", `http://${req.headers.host}`);
const serverId = url.searchParams.get("serverId")!; const serverId = url.searchParams.get("serverId")!;
const role = url.searchParams.get("role") as ConnectionRole; const role = url.searchParams.get("role") as ConnectionRole;
const clientId = (url.searchParams.get("clientId") ?? "").trim();
const session = getSession(serverId);
const connection = wrapWebSocket(ws, role); const isControl = role === "server" && !clientId;
relay.addConnection(serverId, role, connection); 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) => { ws.on("message", (data, isBinary) => {
// If this socket was replaced, ignore any late frames. const normalized = normalizeWsMessage(data, isBinary);
if (relay.getSession(serverId)?.[role] !== connection) return;
if (logFrames) { if (logFrames) {
const preview = (() => { const preview = (() => {
@@ -126,20 +243,92 @@ export function createRelayServer(config: NodeRelayServerConfig): RelayServer {
} }
})(); })();
console.log( 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", () => { ws.on("close", (code, reasonBuf) => {
// Avoid clearing the current connection if this socket was replaced. const reason = reasonBuf?.toString?.() ?? "";
if (relay.getSession(serverId)?.[role] !== connection) return;
relay.removeConnection(serverId, role); 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) => { 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() { stop() {
return new Promise((resolve) => { 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 { try {
httpServer.close(() => undefined); httpServer.close(() => undefined);
} catch { } catch {
// ignore // 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) { for (const ws of wss.clients) {
try { try {
// terminate() is the fastest way to ensure `wss.close` completes.
(ws as unknown as { terminate?: () => void }).terminate?.(); (ws as unknown as { terminate?: () => void }).terminate?.();
} catch { } catch {
try { 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");
}
+5
View File
@@ -26,6 +26,11 @@ export interface RelayConnection {
export interface RelaySessionAttachment { export interface RelaySessionAttachment {
serverId: string; serverId: string;
role: ConnectionRole; 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; createdAt: number;
} }
@@ -294,9 +294,25 @@ export class DaemonClient {
>(); >();
private logger: Logger; private logger: Logger;
private pendingSendQueue: PendingSend[] = []; private pendingSendQueue: PendingSend[] = [];
private relayClientId: string | null = null;
constructor(private config: DaemonClientConfig) { constructor(private config: DaemonClientConfig) {
this.logger = config.logger ?? consoleLogger; 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"; 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 { function decodeMessageData(data: unknown): string | null {
if (data === null || data === undefined) { if (data === null || data === undefined) {
return null; return null;
@@ -91,11 +91,13 @@ describe("Relay transport (E2EE) - daemon E2E", () => {
const offerUrl = parseOfferUrlFromLogs(lines); const offerUrl = parseOfferUrlFromLogs(lines);
const { serverId, daemonPublicKeyB64 } = decodeOfferFromFragmentUrl(offerUrl); const { serverId, daemonPublicKeyB64 } = decodeOfferFromFragmentUrl(offerUrl);
const clientId = `clt_test_${Date.now().toString(36)}_${Math.random().toString(36).slice(2)}`;
const ws = new WebSocket( const ws = new WebSocket(
buildRelayWebSocketUrl({ buildRelayWebSocketUrl({
endpoint: `127.0.0.1:${relayPort}`, endpoint: `127.0.0.1:${relayPort}`,
serverId, serverId,
role: "client", role: "client",
clientId,
}) })
); );
@@ -209,6 +211,7 @@ describe("Relay transport (E2EE) - daemon E2E", () => {
endpoint: `127.0.0.1:${relayPort}`, endpoint: `127.0.0.1:${relayPort}`,
serverId, serverId,
role: "client", role: "client",
clientId: `clt_test_${Date.now().toString(36)}_${Math.random().toString(36).slice(2)}`,
}) })
); );
+113 -19
View File
@@ -30,6 +30,33 @@ type RelaySocketLike = {
once: (event: "close" | "error", listener: (...args: any[]) => void) => void; 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({ export function startRelayTransport({
logger, logger,
attachSocket, attachSocket,
@@ -40,9 +67,10 @@ export function startRelayTransport({
const relayLogger = logger.child({ module: "relay-transport" }); const relayLogger = logger.child({ module: "relay-transport" });
let stopped = false; let stopped = false;
let ws: WebSocket | null = null; let controlWs: WebSocket | null = null;
let reconnectTimeout: ReturnType<typeof setTimeout> | null = null; let reconnectTimeout: ReturnType<typeof setTimeout> | null = null;
let reconnectAttempt = 0; let reconnectAttempt = 0;
const dataSockets = new Map<string, WebSocket>(); // clientId -> ws
const stop = async (): Promise<void> => { const stop = async (): Promise<void> => {
stopped = true; stopped = true;
@@ -50,17 +78,25 @@ export function startRelayTransport({
clearTimeout(reconnectTimeout); clearTimeout(reconnectTimeout);
reconnectTimeout = null; reconnectTimeout = null;
} }
if (ws) { if (controlWs) {
try {
controlWs.close();
} catch {
// ignore
}
controlWs = null;
}
for (const ws of dataSockets.values()) {
try { try {
ws.close(); ws.close();
} catch { } catch {
// ignore // ignore
} }
ws = null;
} }
dataSockets.clear();
}; };
const connect = (): void => { const connectControl = (): void => {
if (stopped) return; if (stopped) return;
const url = buildRelayWebSocketUrl({ const url = buildRelayWebSocketUrl({
@@ -69,27 +105,17 @@ export function startRelayTransport({
role: "server", role: "server",
}); });
const socket = new WebSocket(url); const socket = new WebSocket(url);
ws = socket; controlWs = socket;
let attached = false;
socket.on("open", () => { socket.on("open", () => {
reconnectAttempt = 0; reconnectAttempt = 0;
relayLogger.info({ url }, "relay_connected"); relayLogger.info({ url }, "relay_control_connected");
if (attached) return;
attached = true;
if (daemonKeyPair) {
void attachEncryptedSocket(socket, daemonKeyPair, relayLogger, attachSocket);
} else {
void attachSocket(socket);
}
}); });
socket.on("close", (code, reason) => { socket.on("close", (code, reason) => {
relayLogger.warn( relayLogger.warn(
{ code, reason: reason?.toString?.(), url }, { code, reason: reason?.toString?.(), url },
"relay_disconnected" "relay_control_disconnected"
); );
scheduleReconnect(); scheduleReconnect();
}); });
@@ -98,6 +124,32 @@ export function startRelayTransport({
relayLogger.warn({ err, url }, "relay_error"); relayLogger.warn({ err, url }, "relay_error");
// close event will schedule reconnect // 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 => { const scheduleReconnect = (): void => {
@@ -108,11 +160,53 @@ export function startRelayTransport({
const delayMs = Math.min(30000, 1000 * reconnectAttempt); const delayMs = Math.min(30000, 1000 * reconnectAttempt);
reconnectTimeout = setTimeout(() => { reconnectTimeout = setTimeout(() => {
reconnectTimeout = null; reconnectTimeout = null;
connect(); connectControl();
}, delayMs); }, 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 }; return { stop };
} }
@@ -71,6 +71,7 @@ export function buildRelayWebSocketUrl(params: {
endpoint: string; endpoint: string;
serverId: string; serverId: string;
role: RelayRole; role: RelayRole;
clientId?: string;
}): string { }): string {
const { host, port, isIpv6 } = parseHostPort(params.endpoint); const { host, port, isIpv6 } = parseHostPort(params.endpoint);
const protocol = shouldUseSecureWebSocket(port) ? "wss" : "ws"; const protocol = shouldUseSecureWebSocket(port) ? "wss" : "ws";
@@ -78,6 +79,9 @@ export function buildRelayWebSocketUrl(params: {
const url = new URL(`${protocol}://${hostPart}:${port}/ws`); const url = new URL(`${protocol}://${hostPart}:${port}/ws`);
url.searchParams.set("serverId", params.serverId); url.searchParams.set("serverId", params.serverId);
url.searchParams.set("role", params.role); url.searchParams.set("role", params.role);
if (params.clientId) {
url.searchParams.set("clientId", params.clientId);
}
return url.toString(); return url.toString();
} }