relay(v2): support multiple clients per serverId

2026-02-05 09:48:43 +07:00
parent 98896bfb37
commit 36e105ed30
9 changed files with 663 additions and 169 deletions
@@ -66,16 +66,70 @@ interface CFResponseInit extends ResponseInit {

 export class RelayDurableObject {
  private state: DurableObjectState;
+  private pendingClientFrames = new Map<string, Array<string | ArrayBuffer>>();

  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<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> {
    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
+        }
+      }
+    }
  }

  /**