feat: implement daemon connectivity updates and relay protocol versioning

2026-02-20 19:27:56 +07:00
parent 6dad3212da
commit 0350f00840
49 changed files with 4146 additions and 1351 deletions
@@ -1,5 +1,5 @@
 import { afterEach, describe, expect, it, vi } from "vitest";
-import { RelayDurableObject } from "./cloudflare-adapter.js";
+import relayWorker, { RelayDurableObject } from "./cloudflare-adapter.js";

 type MockSocket = WebSocket & {
  send: ReturnType<typeof vi.fn>;
@@ -42,6 +42,57 @@ function createMockState() {
  };
 }

+async function withMockWebSocketPair(
+  run: (sockets: { clientWs: MockSocket; serverWs: MockSocket }) => Promise<void> | void
+): Promise<void> {
+  const serverWs = createMockSocket();
+  const clientWs = createMockSocket();
+  const WebSocketPairMock = class {
+    [index: number]: WebSocket;
+    constructor() {
+      this[0] = clientWs as unknown as WebSocket;
+      this[1] = serverWs as unknown as WebSocket;
+    }
+  };
+
+  const previousPair = (globalThis as unknown as { WebSocketPair?: unknown }).WebSocketPair;
+  (globalThis as unknown as { WebSocketPair: unknown }).WebSocketPair = WebSocketPairMock;
+  try {
+    await run({ clientWs, serverWs });
+  } finally {
+    if (previousPair === undefined) {
+      delete (globalThis as unknown as { WebSocketPair?: unknown }).WebSocketPair;
+    } else {
+      (globalThis as unknown as { WebSocketPair: unknown }).WebSocketPair = previousPair;
+    }
+  }
+}
+
+describe("RelayDurableObject versioning", () => {
+  it("accepts legacy v1 client sockets without clientId", async () => {
+    const { state } = createMockState();
+    await withMockWebSocketPair(async () => {
+      const relay = new RelayDurableObject(state as any);
+      const req = new Request("https://relay.test/ws?role=client&serverId=srv_test&v=1", {
+        headers: {
+          Upgrade: "websocket",
+        },
+      });
+      await relay.fetch(req).catch(() => undefined);
+      expect(state.acceptWebSocket).toHaveBeenCalled();
+    });
+  });
+
+  it("rejects v2 client sockets when clientId is missing", async () => {
+    const { state } = createMockState();
+    const relay = new RelayDurableObject(state as any);
+    const req = new Request("https://relay.test/ws?role=client&serverId=srv_test&v=2");
+    const response = await relay.fetch(req);
+    expect(response.status).toBe(400);
+    await expect(response.text()).resolves.toBe("Missing clientId parameter");
+  });
+});
+
 describe("RelayDurableObject control nudge/reset behavior", () => {
  afterEach(() => {
    vi.useRealTimers();
@@ -93,4 +144,119 @@ describe("RelayDurableObject control nudge/reset behavior", () => {
    vi.advanceTimersByTime(5_000);
    expect(control.close).toHaveBeenCalledWith(1011, "Control unresponsive");
  });
+
+  it("does not replace existing client sockets for the same clientId", async () => {
+    const existingClient = createMockSocket({
+      version: "2",
+      role: "client",
+      clientId: "clt_same_session",
+      serverId: "srv_test",
+      createdAt: Date.now(),
+    });
+    const { state, setTagSockets } = createMockState();
+    setTagSockets("client:clt_same_session", [existingClient]);
+    setTagSockets("client", [existingClient]);
+
+    await withMockWebSocketPair(async () => {
+      const relay = new RelayDurableObject(state as any);
+      const req = new Request(
+        "https://relay.test/ws?role=client&serverId=srv_test&clientId=clt_same_session&v=2",
+        {
+          headers: {
+            Upgrade: "websocket",
+          },
+        }
+      );
+
+      await relay.fetch(req).catch(() => undefined);
+      expect(existingClient.close).not.toHaveBeenCalled();
+    });
+  });
+
+  it("keeps server data socket alive while at least one client socket remains", () => {
+    const clientId = "clt_multi";
+    const disconnectedClient = createMockSocket({
+      version: "2",
+      role: "client",
+      clientId,
+      serverId: "srv_test",
+      createdAt: Date.now(),
+    });
+    const stillConnectedClient = createMockSocket({
+      version: "2",
+      role: "client",
+      clientId,
+      serverId: "srv_test",
+      createdAt: Date.now(),
+    });
+    const serverData = createMockSocket();
+    const control = createMockSocket();
+    const { state, setTagSockets } = createMockState();
+
+    setTagSockets("server-control", [control]);
+    setTagSockets(`server:${clientId}`, [serverData]);
+    setTagSockets("client", [stillConnectedClient]);
+    setTagSockets(`client:${clientId}`, [stillConnectedClient]);
+
+    const relay = new RelayDurableObject(state as any);
+    relay.webSocketClose(
+      disconnectedClient as unknown as WebSocket,
+      1001,
+      "Client disconnected",
+      true
+    );
+
+    expect(serverData.close).not.toHaveBeenCalled();
+    expect(control.send).not.toHaveBeenCalledWith(
+      JSON.stringify({ type: "client_disconnected", clientId })
+    );
+  });
+});
+
+describe("relay worker endpoint routing", () => {
+  it("routes missing v to legacy v1 isolated DO ids", async () => {
+    const fetch = vi.fn(async (request: Request) => new Response(`ok:${new URL(request.url).searchParams.get("v")}`));
+    const get = vi.fn(() => ({ fetch }));
+    const idFromName = vi.fn(() => ({ toString: () => "id" }));
+
+    const response = await relayWorker.fetch(
+      new Request("https://relay.test/ws?serverId=srv_test&role=server"),
+      { RELAY: { idFromName, get } } as any
+    );
+
+    expect(idFromName).toHaveBeenCalledWith("relay-v1:srv_test");
+    expect(fetch).toHaveBeenCalledTimes(1);
+    await expect(response.text()).resolves.toBe("ok:1");
+  });
+
+  it("routes v=2 to v2 isolated DO ids", async () => {
+    const fetch = vi.fn(async (request: Request) => new Response(`ok:${new URL(request.url).searchParams.get("v")}`));
+    const get = vi.fn(() => ({ fetch }));
+    const idFromName = vi.fn(() => ({ toString: () => "id" }));
+
+    const response = await relayWorker.fetch(
+      new Request("https://relay.test/ws?serverId=srv_test&role=server&v=2"),
+      { RELAY: { idFromName, get } } as any
+    );
+
+    expect(idFromName).toHaveBeenCalledWith("relay-v2:srv_test");
+    expect(fetch).toHaveBeenCalledTimes(1);
+    await expect(response.text()).resolves.toBe("ok:2");
+  });
+
+  it("rejects invalid v values", async () => {
+    const fetch = vi.fn();
+    const get = vi.fn(() => ({ fetch }));
+    const idFromName = vi.fn(() => ({ toString: () => "id" }));
+
+    const response = await relayWorker.fetch(
+      new Request("https://relay.test/ws?serverId=srv_test&role=server&v=nope"),
+      { RELAY: { idFromName, get } } as any
+    );
+
+    expect(response.status).toBe(400);
+    await expect(response.text()).resolves.toBe("Invalid v parameter (expected 1 or 2)");
+    expect(idFromName).not.toHaveBeenCalled();
+    expect(fetch).not.toHaveBeenCalled();
+  });
 });
@@ -19,6 +19,21 @@

 import type { ConnectionRole, RelaySessionAttachment } from "./types.js";

+type RelayProtocolVersion = "1" | "2";
+
+const LEGACY_RELAY_VERSION: RelayProtocolVersion = "1";
+const CURRENT_RELAY_VERSION: RelayProtocolVersion = "2";
+
+function resolveRelayVersion(rawValue: string | null): RelayProtocolVersion | null {
+  if (rawValue == null) return LEGACY_RELAY_VERSION;
+  const value = rawValue.trim();
+  if (!value) return LEGACY_RELAY_VERSION;
+  if (value === LEGACY_RELAY_VERSION || value === CURRENT_RELAY_VERSION) {
+    return value;
+  }
+  return null;
+}
+
 type WebSocketPair = {
  0: WebSocket;
  1: WebSocket;
@@ -54,13 +69,14 @@ interface DurableObjectStub {
 /**
 * Durable Object that handles WebSocket relay for a single session.
 *
- * WebSockets connect to this DO in three shapes:
+ * v1 WebSockets connect in two shapes:
+ * - role=server: daemon socket
+ * - role=client: app/client socket
+ *
+ * v2 WebSockets connect in three shapes:
 * - role=server (no clientId): daemon control socket (one per serverId)
 * - role=server&clientId=...: daemon per-client data socket (one per clientId)
- * - role=client&clientId=...: app/client socket (one per clientId)
- *
- * Messages are forwarded between the per-client data sockets and their matching
- * client sockets. The DO hibernates when idle.
+ * - role=client&clientId=...: app/client socket (many per clientId)
 */
 interface CFResponseInit extends ResponseInit {
  webSocket?: WebSocket;
@@ -72,7 +88,26 @@ export class RelayDurableObject {

  constructor(state: DurableObjectState) {
    this.state = state;
+  }

+  private createWebSocketPair(): [WebSocket, WebSocket] {
+    const pair = new (globalThis as unknown as { WebSocketPair: new () => WebSocketPair }).WebSocketPair();
+    return [pair[0], pair[1]];
+  }
+
+  private requireWebSocketUpgrade(request: Request): Response | null {
+    const upgradeHeader = request.headers.get("Upgrade");
+    if (!upgradeHeader || upgradeHeader.toLowerCase() !== "websocket") {
+      return new Response("Expected WebSocket upgrade", { status: 426 });
+    }
+    return null;
+  }
+
+  private asSwitchingProtocolsResponse(client: WebSocket): Response {
+    return new Response(null, {
+      status: 101,
+      webSocket: client,
+    } as CFResponseInit);
  }

  private hasServerDataSocket(clientId: string): boolean {
@@ -180,39 +215,53 @@ export class RelayDurableObject {
    }
  }

-  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() : "";
+  private fetchV1(request: Request, role: ConnectionRole, serverId: string): Response {
+    const upgradeError = this.requireWebSocketUpgrade(request);
+    if (upgradeError) return upgradeError;

-    if (!role || (role !== "server" && role !== "client")) {
-      return new Response("Missing or invalid role parameter", { status: 400 });
+    for (const ws of this.state.getWebSockets(role)) {
+      ws.close(1008, "Replaced by new connection");
    }

-    if (!serverId) {
-      return new Response("Missing serverId parameter", { status: 400 });
-    }
+    const [client, server] = this.createWebSocketPair();
+    this.state.acceptWebSocket(server, [role]);

+    const attachment: RelaySessionAttachment = {
+      serverId,
+      role,
+      version: LEGACY_RELAY_VERSION,
+      clientId: null,
+      createdAt: Date.now(),
+    };
+    (server as WebSocketWithAttachment).serializeAttachment(attachment);
+
+    console.log(`[Relay DO] v1:${role} connected to session ${serverId}`);
+
+    return this.asSwitchingProtocolsResponse(client);
+  }
+
+  private fetchV2(
+    request: Request,
+    role: ConnectionRole,
+    serverId: string,
+    clientId: string
+  ): Response {
    // 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 });
-    }
+    const upgradeError = this.requireWebSocketUpgrade(request);
+    if (upgradeError) return upgradeError;

    const isServerControl = role === "server" && !clientId;
    const isServerData = role === "server" && !!clientId;

-    // Close any existing connection with the same identity.
+    // Close any existing server-side connection with the same identity.
    // - server-control: single per serverId
    // - server-data: single per clientId
-    // - client: single per clientId
+    // - client: many sockets per clientId are allowed
    if (isServerControl) {
      for (const ws of this.state.getWebSockets("server-control")) {
        ws.close(1008, "Replaced by new connection");
@@ -221,15 +270,9 @@ export class RelayDurableObject {
      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]];
+    const [client, server] = this.createWebSocketPair();

    const tags: string[] = [];
    if (role === "client") {
@@ -240,20 +283,19 @@ export class RelayDurableObject {
      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,
+      version: CURRENT_RELAY_VERSION,
      clientId: clientId || null,
      createdAt: Date.now(),
    };
    (server as WebSocketWithAttachment).serializeAttachment(attachment);

    console.log(
-      `[Relay DO] ${role}${isServerControl ? "(control)" : ""}${isServerData ? `(data:${clientId})` : role === "client" ? `(${clientId})` : ""} connected to session ${serverId}`
+      `[Relay DO] v2:${role}${isServerControl ? "(control)" : ""}${isServerData ? `(data:${clientId})` : role === "client" ? `(${clientId})` : ""} connected to session ${serverId}`
    );

    if (role === "client") {
@@ -274,10 +316,34 @@ export class RelayDurableObject {
      this.flushClientFrames(clientId, server);
    }

-    return new Response(null, {
-      status: 101,
-      webSocket: client,
-    } as CFResponseInit);
+    return this.asSwitchingProtocolsResponse(client);
+  }
+
+  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() : "";
+    const version = resolveRelayVersion(url.searchParams.get("v"));
+
+    if (!role || (role !== "server" && role !== "client")) {
+      return new Response("Missing or invalid role parameter", { status: 400 });
+    }
+
+    if (!serverId) {
+      return new Response("Missing serverId parameter", { status: 400 });
+    }
+
+    if (!version) {
+      return new Response("Invalid v parameter (expected 1 or 2)", { status: 400 });
+    }
+
+    if (version === LEGACY_RELAY_VERSION) {
+      return this.fetchV1(request, role, serverId);
+    }
+
+    return this.fetchV2(request, role, serverId, clientId);
  }

  /**
@@ -290,12 +356,27 @@ export class RelayDurableObject {
      return;
    }

+    const version = attachment.version ?? LEGACY_RELAY_VERSION;
+
+    if (version === LEGACY_RELAY_VERSION) {
+      const targetRole = attachment.role === "server" ? "client" : "server";
+      const targets = this.state.getWebSockets(targetRole);
+      for (const target of targets) {
+        try {
+          target.send(message);
+        } catch (error) {
+          console.error(`[Relay DO] Failed to forward to ${targetRole}:`, error);
+        }
+      }
+      return;
+    }
+
    const { role, clientId } = attachment;
    if (!clientId) {
      // Control channel: support simple app-level keepalive.
      if (typeof message === "string") {
        try {
-          const parsed = JSON.parse(message) as any;
+          const parsed = JSON.parse(message) as unknown as { type?: unknown };
          if (parsed?.type === "ping") {
            try {
              ws.send(JSON.stringify({ type: "pong", ts: Date.now() }));
@@ -349,13 +430,25 @@ export class RelayDurableObject {
    const attachment = (ws as WebSocketWithAttachment).deserializeAttachment() as RelaySessionAttachment | null;
    if (!attachment) return;

+    const version = attachment.version ?? LEGACY_RELAY_VERSION;
    console.log(
-      `[Relay DO] ${attachment.role}${attachment.clientId ? `(${attachment.clientId})` : ""} disconnected from session ${attachment.serverId} (${code}: ${reason})`
+      `[Relay DO] v${version}:${attachment.role}${attachment.clientId ? `(${attachment.clientId})` : ""} disconnected from session ${attachment.serverId} (${code}: ${reason})`
    );

+    if (version === LEGACY_RELAY_VERSION) {
+      return;
+    }
+
    if (attachment.role === "client" && attachment.clientId) {
+      const remainingClientSockets = this.state
+        .getWebSockets(`client:${attachment.clientId}`)
+        .some((socket) => socket !== ws);
+      if (remainingClientSockets) {
+        return;
+      }
+
      this.pendingClientFrames.delete(attachment.clientId);
-      // Close the matching server-data socket so the daemon can clean up quickly.
+      // Last socket for this session closed: now clean up matching server-data socket.
      for (const serverWs of this.state.getWebSockets(`server:${attachment.clientId}`)) {
        try {
          serverWs.close(1001, "Client disconnected");
@@ -412,10 +505,19 @@ export default {
        return new Response("Missing serverId parameter", { status: 400 });
      }

-      // Route to Durable Object instance for this session
-      const id = env.RELAY.idFromName(serverId);
+      const version = resolveRelayVersion(url.searchParams.get("v"));
+      if (!version) {
+        return new Response("Invalid v parameter (expected 1 or 2)", { status: 400 });
+      }
+
+      // Route to a version-isolated Durable Object instance.
+      const id = env.RELAY.idFromName(`relay-v${version}:${serverId}`);
      const stub = env.RELAY.get(id);
-      return stub.fetch(request);
+
+      const normalizedUrl = new URL(request.url);
+      normalizedUrl.searchParams.set("v", version);
+      const normalizedRequest = new Request(normalizedUrl.toString(), request);
+      return stub.fetch(normalizedRequest);
    }

    return new Response("Not found", { status: 404 });
@@ -53,7 +53,7 @@ async function waitForRelayWebSocketReady(port: number, timeout = 60000): Promis
  const start = Date.now();
  while (Date.now() - start < timeout) {
    const serverId = `probe-${Date.now().toString(36)}-${Math.random().toString(36).slice(2)}`;
-    const probeUrl = `ws://127.0.0.1:${port}/ws?serverId=${serverId}&role=server`;
+    const probeUrl = `ws://127.0.0.1:${port}/ws?serverId=${serverId}&role=server&v=2`;
    const opened = await new Promise<boolean>((resolve) => {
      const ws = new WebSocket(probeUrl);
      const timer = setTimeout(() => {
@@ -144,7 +144,7 @@ async function waitForRelayWebSocketReady(port: number, timeout = 60000): Promis

    // Daemon connects to relay as "server" control role
    const daemonControlWs = new WebSocket(
-      `ws://127.0.0.1:${relayPort}/ws?serverId=${serverId}&role=server`
+      `ws://127.0.0.1:${relayPort}/ws?serverId=${serverId}&role=server&v=2`
    );

    await new Promise<void>((resolve, reject) => {
@@ -199,7 +199,7 @@ async function waitForRelayWebSocketReady(port: number, timeout = 60000): Promis

    // Client connects to relay as "client" role (must include clientId)
    const clientWs = new WebSocket(
-      `ws://127.0.0.1:${relayPort}/ws?serverId=${serverId}&role=client&clientId=${clientId}`
+      `ws://127.0.0.1:${relayPort}/ws?serverId=${serverId}&role=client&clientId=${clientId}&v=2`
    );

    await new Promise<void>((resolve, reject) => {
@@ -210,7 +210,7 @@ async function waitForRelayWebSocketReady(port: number, timeout = 60000): Promis
    await waitForClientSeen;

    const daemonWs = new WebSocket(
-      `ws://127.0.0.1:${relayPort}/ws?serverId=${serverId}&role=server&clientId=${clientId}`
+      `ws://127.0.0.1:${relayPort}/ws?serverId=${serverId}&role=server&clientId=${clientId}&v=2`
    );
    await new Promise<void>((resolve, reject) => {
      daemonWs.on("open", resolve);
@@ -322,7 +322,7 @@ async function waitForRelayWebSocketReady(port: number, timeout = 60000): Promis
    );

    const daemonControlWs = new WebSocket(
-      `ws://127.0.0.1:${relayPort}/ws?serverId=${serverId}&role=server`
+      `ws://127.0.0.1:${relayPort}/ws?serverId=${serverId}&role=server&v=2`
    );
    await new Promise<void>((r) => daemonControlWs.on("open", r));

@@ -359,13 +359,13 @@ async function waitForRelayWebSocketReady(port: number, timeout = 60000): Promis
    });

    const clientWs = new WebSocket(
-      `ws://127.0.0.1:${relayPort}/ws?serverId=${serverId}&role=client&clientId=${clientId}`
+      `ws://127.0.0.1:${relayPort}/ws?serverId=${serverId}&role=client&clientId=${clientId}&v=2`
    );
    await new Promise<void>((r) => clientWs.on("open", r));
    await waitForClientSeen;

    const daemonWs = new WebSocket(
-      `ws://127.0.0.1:${relayPort}/ws?serverId=${serverId}&role=server&clientId=${clientId}`
+      `ws://127.0.0.1:${relayPort}/ws?serverId=${serverId}&role=server&clientId=${clientId}&v=2`
    );
    await new Promise<void>((r) => daemonWs.on("open", r));

@@ -37,13 +37,13 @@ describe("Live relay (relay.paseo.sh) E2E", () => {
      async () => {
        const serverId = `live-${Date.now()}-${Math.random().toString(16).slice(2)}`;
        const clientId = `clt_live_${Date.now()}_${Math.random().toString(16).slice(2)}`;
-        const serverControlUrl = `${RELAY_BASE_URL}/ws?serverId=${encodeURIComponent(serverId)}&role=server`;
+        const serverControlUrl = `${RELAY_BASE_URL}/ws?serverId=${encodeURIComponent(serverId)}&role=server&v=2`;
        const serverDataUrl = `${RELAY_BASE_URL}/ws?serverId=${encodeURIComponent(
          serverId
-        )}&role=server&clientId=${encodeURIComponent(clientId)}`;
+        )}&role=server&clientId=${encodeURIComponent(clientId)}&v=2`;
        const clientUrl = `${RELAY_BASE_URL}/ws?serverId=${encodeURIComponent(
          serverId
-        )}&role=client&clientId=${encodeURIComponent(clientId)}`;
+        )}&role=client&clientId=${encodeURIComponent(clientId)}&v=2`;

        // === Key setup ===
        const daemonKeyPair = await generateKeyPair();
@@ -13,6 +13,12 @@ export type ConnectionRole = "server" | "client";
 export interface RelaySessionAttachment {
  serverId: string;
  role: ConnectionRole;
+  /**
+   * Relay protocol version carried by this socket.
+   * v1: single server/client socket pair
+   * v2: control + per-client data sockets
+   */
+  version?: "1" | "2";
  /**
   * Unique id for the client connection. Allows the daemon to create an
   * independent socket + E2EE channel per connected client.