Fix relay E2EE reconnect races (#973)

* Fix relay E2EE reconnect races * fix(relay): use WebSocket protocol ping for daemon control keepalive The daemon's control socket pinged the relay every 10s with a JSON `{type:"ping"}` message. Every JSON message wakes a hibernated Cloudflare Durable Object via `webSocketMessage`, so an idle daemon was billing relay DO CPU 24/7 just to echo `pong`. Switch to `socket.ping()` (RFC 6455 protocol ping frames). Cloudflare's runtime auto-responds at the edge without waking hibernated DOs, so an idle daemon-to-relay connection now costs zero DO CPU. Liveness still works: the `ws` library emits a `pong` event on the daemon when the edge replies, which refreshes `controlLastSeenAt`. `markControlReady` stays tied to DO-originated messages (sync, connected, disconnected), so we still detect a broken DO at connect. The DO keeps its JSON ping handler tagged `COMPAT(relay-json-ping)` so old daemons (< v0.1.76) continue to work. * test(relay): assert protocol-ping pongs arrive in e2e idle test Adds a `relay_control_pong_received` debug log to the daemon's pong handler, then asserts in the existing idle-relay e2e test that: - at least one pong arrived during the 12s idle wait (proving the Cloudflare edge auto-responds to RFC 6455 protocol pings) - no stale terminations fired (proving the keepalive path is healthy) If the daemon ever regresses to JSON-level keepalive pings (which wake the hibernated relay Durable Object and incur CPU billing), this test will catch it: the `pong` event would never fire because the DO's reply would arrive via `message`, not as a protocol pong frame. The test's capturing logger is bumped to `debug` level so the assertion can read the new debug log; production daemons stay at `info`. Confirmed locally against wrangler-hosted relay: 2 pongs / 0 stale per 12s window. * test(relay): catch dual-ping regression via DO-side legacy ping log The existing idle-test assertion catches "revert to JSON ping only" via the missing pong event, but misses the dual-ping regression where someone adds `socket.send(JSON.stringify({type:"ping"}))` alongside `socket.ping()`. Both paths would work, the pong assertion still passes, and DO CPU billing silently returns. Add a console.log in the DO's handleControlKeepalive (which only runs when a JSON ping reaches the DO), then assert in the idle e2e test that no such log appeared during the 12s idle window. Verified the assertion catches the regression: temporarily injected a `socket.send(JSON.stringify({type:"ping"}))` next to socket.ping(), test failed with `expected 1 to be 0` and the relay logged `legacy_json_ping_received`. Removed the injection. Log noise during the COMPAT window doubles as useful telemetry: it counts how many old daemons are still pinging, which informs the removal date for the COMPAT shim.
2026-05-13 13:44:11 +08:00
parent 15cab6014c
commit d24087c10a
6 changed files with 301 additions and 11 deletions
@@ -302,6 +302,58 @@ function makeProbeMap(
 }

 describe("HostRuntimeController", () => {
+  it("replaces the active relay client when re-pairing changes the daemon public key", async () => {
+    const oldRelay: HostConnection = {
+      id: "relay:wss:relay.paseo.sh:443",
+      type: "relay",
+      relayEndpoint: "relay.paseo.sh:443",
+      useTls: true,
+      daemonPublicKeyB64: "pk_old",
+    };
+    const newRelay: HostConnection = {
+      ...oldRelay,
+      daemonPublicKeyB64: "pk_new",
+    };
+    const createdClients: Array<{ client: FakeDaemonClient; connection: HostConnection }> = [];
+    const controller = new HostRuntimeController({
+      host: makeHost({
+        connections: [oldRelay],
+        preferredConnectionId: oldRelay.id,
+      }),
+      deps: {
+        createClient: ({ connection }) => {
+          const client = new FakeDaemonClient();
+          createdClients.push({ client, connection });
+          return client as unknown as DaemonClient;
+        },
+        connectToDaemon: async ({ host, connection }) => ({
+          client: makeConnectedProbeClient(5) as unknown as DaemonClient,
+          serverId: host.serverId,
+          hostname: connection.id,
+        }),
+        getClientId: async () => "cid_test_runtime",
+      },
+    });
+
+    await (
+      controller as unknown as {
+        switchToConnection: (input: { connectionId: string }) => Promise<void>;
+      }
+    ).switchToConnection({ connectionId: oldRelay.id });
+    expect(controller.getSnapshot().client).toBe(createdClients[0]?.client);
+
+    await controller.updateHost(
+      makeHost({
+        connections: [newRelay],
+        preferredConnectionId: newRelay.id,
+      }),
+    );
+
+    expect(createdClients.map((entry) => entry.connection)).toEqual([oldRelay, newRelay]);
+    expect(createdClients[0]?.client.closeCalls).toBe(1);
+    expect(controller.getSnapshot().client).toBe(createdClients[1]?.client);
+  });
+
  it("keeps known hosts in connecting when a created client reports idle during connect", async () => {
    const host = makeHost({
      connections: [
@@ -602,8 +602,20 @@ export class HostRuntimeController {
  }

  async updateHost(host: HostProfile): Promise<void> {
+    const activeConnectionId = this.snapshot.activeConnectionId;
+    const previousActiveConnection = findConnectionById(this.host, activeConnectionId);
    this.host = host;
    this.trackConnectionFirstSeen();
+    const nextActiveConnection = findConnectionById(this.host, activeConnectionId);
+    if (
+      activeConnectionId &&
+      previousActiveConnection &&
+      nextActiveConnection &&
+      !equal(previousActiveConnection, nextActiveConnection)
+    ) {
+      this.connectionLastProbedAt.delete(activeConnectionId);
+      await this.switchToConnection({ connectionId: activeConnectionId });
+    }
    await this.runProbeCycleNow();
  }

@@ -192,11 +192,18 @@ export class RelayDurableObject {
    }
  }

+  // COMPAT(relay-json-ping): Old daemons (< v0.1.76) send JSON {type:"ping"} on the control
+  // socket and rely on a JSON {type:"pong"} reply to keep controlLastSeenAt fresh. New daemons
+  // use WebSocket protocol pings (auto-answered at the edge, DO stays hibernated). Remove this
+  // handler once the supported-daemon floor is >= v0.1.76 (target: 2026-11-13).
  private handleControlKeepalive(ws: WebSocket, message: string): void {
    try {
      const parsed: unknown = JSON.parse(message);
      const parsedRecord = isRecord(parsed) ? parsed : null;
      if (parsedRecord?.type !== "ping") return;
+      // Logged so the daemon-side e2e idle test can assert no JSON ping reached the DO
+      // (which would indicate a regression to app-level pings that wake the DO).
+      console.log("[Relay DO] legacy_json_ping_received");
      try {
        ws.send(JSON.stringify({ type: "pong", ts: Date.now() }));
      } catch {
@@ -10,6 +10,14 @@ import { Buffer } from "node:buffer";
 import { generateLocalPairingOffer } from "../pairing-offer.js";
 import { createTestPaseoDaemon } from "../test-utils/paseo-daemon.js";
 import { createClientChannel, type Transport } from "@getpaseo/relay/e2ee";
+import {
+  deriveSharedKey,
+  decrypt,
+  encrypt,
+  exportPublicKey,
+  generateKeyPair,
+  importPublicKey,
+} from "@getpaseo/relay";
 import { buildRelayWebSocketUrl } from "../../shared/daemon-endpoints.js";
 import { ConnectionOfferSchema } from "../../shared/connection-offer.js";
 import { WSOutboundMessageSchema } from "../../shared/messages.js";
@@ -25,7 +33,7 @@ function createCapturingLogger() {
      cb();
    },
  });
-  const logger = pino({ level: "info" }, stream);
+  const logger = pino({ level: "debug" }, stream);
  return { logger, lines };
 }

@@ -65,6 +73,23 @@ function decodeOfferFromFragmentUrl(url: string): {
  return { serverId: offer.serverId, daemonPublicKeyB64: offer.daemonPublicKeyB64 };
 }

+function encodeCiphertext(ciphertext: ArrayBuffer): string {
+  return Buffer.from(new Uint8Array(ciphertext)).toString("base64");
+}
+
+function decodeCiphertext(text: string): ArrayBuffer {
+  const buffer = Buffer.from(text, "base64");
+  return buffer.buffer.slice(buffer.byteOffset, buffer.byteOffset + buffer.byteLength);
+}
+
+function parseEncryptedJson(sharedKey: Uint8Array, text: string): unknown {
+  const plaintext = decrypt(sharedKey, decodeCiphertext(text));
+  if (typeof plaintext !== "string") {
+    throw new Error("Expected encrypted relay frame to contain UTF-8 JSON");
+  }
+  return JSON.parse(plaintext);
+}
+
 async function getAvailablePort(): Promise<number> {
  return await new Promise((resolve, reject) => {
    const server = net.createServer();
@@ -142,8 +167,10 @@ async function waitForRelayWebSocketReady(port: number, timeout = 60000): Promis
 (shouldRunRelayE2e ? describe : describe.skip)("Relay transport (E2EE) - daemon E2E", () => {
  let relayPort: number;
  let relayProcess: ChildProcess | null = null;
+  let relayStdoutLines: string[] = [];

  const startRelay = async () => {
+    relayStdoutLines = [];
    relayPort = await getAvailablePort();
    const relayDir = path.resolve(process.cwd(), "../relay");
    relayProcess = spawn(
@@ -173,6 +200,7 @@ async function waitForRelayWebSocketReady(port: number, timeout = 60000): Promis
        .split("\n")
        .filter((l) => l.trim());
      for (const line of lines) {
+        relayStdoutLines.push(line);
        // eslint-disable-next-line no-console
        console.log(`[relay] ${line}`);
      }
@@ -366,6 +394,22 @@ async function waitForRelayWebSocketReady(port: number, timeout = 60000): Promis
      );
      expect(handshakeFailures.length).toBe(0);

+      // Protocol pings (RFC 6455 control frames) are auto-answered at the Cloudflare edge
+      // without waking the hibernated relay Durable Object. After 12s, the keepalive interval
+      // (10s) must have fired at least once and received a pong. If this assertion fails, the
+      // daemon may have regressed to app-level JSON pings (which wake the DO and cost CPU).
+      const pongLines = lines.filter((l) => l.includes("relay_control_pong_received"));
+      expect(pongLines.length).toBeGreaterThan(0);
+      const staleLines = lines.filter((l) => l.includes("relay_control_stale_terminating"));
+      expect(staleLines.length).toBe(0);
+      // Guard against the dual-ping regression: if a JSON {type:"ping"} reaches the DO during
+      // the idle window, the DO logs `legacy_json_ping_received`. The current daemon code must
+      // not send JSON pings on the control socket — only protocol pings via socket.ping().
+      const legacyPingLines = relayStdoutLines.filter((l) =>
+        l.includes("legacy_json_ping_received"),
+      );
+      expect(legacyPingLines.length).toBe(0);
+
      const ws = new WebSocket(
        buildRelayWebSocketUrl({
          endpoint: `127.0.0.1:${relayPort}`,
@@ -458,4 +502,144 @@ async function waitForRelayWebSocketReady(port: number, timeout = 60000): Promis
      await stopRelay();
    }
  }, 90000);
+
+  test("daemon accepts a relay client that pipelines app hello after E2EE hello", async () => {
+    process.env.PASEO_PRIMARY_LAN_IP = "192.168.1.12";
+
+    const { logger, lines } = createCapturingLogger();
+    await startRelay();
+
+    const daemon = await createTestPaseoDaemon({
+      listen: "127.0.0.1",
+      logger,
+      relayEnabled: true,
+      relayEndpoint: `127.0.0.1:${relayPort}`,
+    });
+
+    try {
+      const offerUrl = await getPairingOfferUrl({
+        paseoHome: daemon.paseoHome,
+        relayEnabled: daemon.config.relayEnabled,
+        relayEndpoint: daemon.config.relayEndpoint,
+        relayPublicEndpoint: daemon.config.relayPublicEndpoint,
+        appBaseUrl: daemon.config.appBaseUrl,
+      });
+      const { serverId, daemonPublicKeyB64 } = decodeOfferFromFragmentUrl(offerUrl);
+      const clientKeyPair = generateKeyPair();
+      const sharedKey = deriveSharedKey(
+        clientKeyPair.secretKey,
+        importPublicKey(daemonPublicKeyB64),
+      );
+
+      const ws = new WebSocket(
+        buildRelayWebSocketUrl({
+          endpoint: `127.0.0.1:${relayPort}`,
+          useTls: false,
+          serverId,
+          role: "client",
+        }),
+      );
+
+      const received = await new Promise<unknown>((resolve, reject) => {
+        const timeout = setTimeout(() => {
+          ws.close();
+          reject(new Error("timed out waiting for server_info"));
+        }, 20000);
+
+        const settleResolve = (value: unknown) => {
+          clearTimeout(timeout);
+          resolve(value);
+          ws.close();
+        };
+        const settleReject = (reason: unknown) => {
+          clearTimeout(timeout);
+          reject(reason);
+          ws.close();
+        };
+
+        ws.on("open", () => {
+          ws.send(
+            JSON.stringify({
+              type: "e2ee_hello",
+              key: exportPublicKey(clientKeyPair.publicKey),
+            }),
+          );
+          ws.send(
+            encodeCiphertext(
+              encrypt(
+                sharedKey,
+                JSON.stringify({
+                  type: "hello",
+                  clientId: "cid_relay_pipelined_hello",
+                  clientType: "cli",
+                  protocolVersion: 1,
+                }),
+              ),
+            ),
+          );
+        });
+
+        ws.on("message", (data) => {
+          try {
+            const text = typeof data === "string" ? data : data.toString();
+            const maybePlaintext = JSON.parse(text) as unknown;
+            if (
+              maybePlaintext &&
+              typeof maybePlaintext === "object" &&
+              "type" in maybePlaintext &&
+              maybePlaintext.type === "e2ee_ready"
+            ) {
+              return;
+            }
+          } catch {
+            // encrypted frame; parse below
+          }
+
+          try {
+            const parsed = WSOutboundMessageSchema.parse(
+              parseEncryptedJson(sharedKey, data.toString()),
+            );
+            if (
+              parsed.type === "session" &&
+              parsed.message.type === "status" &&
+              parsed.message.payload?.status === "server_info"
+            ) {
+              settleResolve({
+                type: parsed.type,
+                message: {
+                  type: parsed.message.type,
+                  payload: { status: parsed.message.payload.status },
+                },
+              });
+            }
+          } catch (error) {
+            settleReject(error);
+          }
+        });
+
+        ws.on("close", (code, reason) => {
+          settleReject(new Error(`relay client closed before server_info: ${code} ${reason}`));
+        });
+        ws.on("error", (err) => {
+          settleReject(err);
+        });
+      });
+
+      expect(received).toEqual({
+        type: "session",
+        message: {
+          type: "status",
+          payload: { status: "server_info" },
+        },
+      });
+    } catch (err) {
+      const tail = lines.slice(-50).join("");
+      // eslint-disable-next-line no-console
+      console.error("daemon logs (tail):\n", tail);
+      throw err;
+    } finally {
+      await daemon.close();
+      await stopRelay();
+    }
+  }, 90000);
 });
@@ -31,6 +31,7 @@ class FakeRelayWebSocket {
  readyState = FakeRelayWebSocket.CONNECTING;
  sent: Array<string | Uint8Array | ArrayBuffer> = [];
  terminateCalls = 0;
+  pingCalls = 0;
  private readonly listeners = new Map<string, Array<(...args: unknown[]) => void>>();

  constructor(readonly url: string) {}
@@ -67,6 +68,13 @@ class FakeRelayWebSocket {
    this.sent.push(data);
  }

+  ping() {
+    if (this.readyState !== FakeRelayWebSocket.OPEN) {
+      throw new Error(`WebSocket not open (readyState=${this.readyState})`);
+    }
+    this.pingCalls += 1;
+  }
+
  open() {
    this.readyState = FakeRelayWebSocket.OPEN;
    this.emit("open");
@@ -76,6 +84,10 @@ class FakeRelayWebSocket {
    this.emit("message", data);
  }

+  pong() {
+    this.emit("pong");
+  }
+
  private off(event: string, listener: (...args: unknown[]) => void) {
    const handlers = this.listeners.get(event) ?? [];
    this.listeners.set(
@@ -135,9 +147,9 @@ describe("relay-transport control lifecycle", () => {

    control.open();
    expect(hasLogMessage(logger, "info", "relay_control_connected")).toBe(false);
-    expect(control.sent.length).toBeGreaterThan(0);
+    expect(control.pingCalls).toBeGreaterThan(0);

-    control.message(JSON.stringify({ type: "pong", ts: Date.now() }));
+    control.message(JSON.stringify({ type: "sync", connectionIds: [] }));
    expect(hasLogMessage(logger, "info", "relay_control_connected")).toBe(true);
  });

@@ -180,7 +192,7 @@ describe("relay-transport control lifecycle", () => {

    const control = relay.sockets[0];
    control.open();
-    control.message(JSON.stringify({ type: "pong", ts: Date.now() }));
+    control.message(JSON.stringify({ type: "sync", connectionIds: [] }));
    logger.messages.length = 0;

    vi.advanceTimersByTime(40_000);
@@ -208,7 +220,7 @@ describe("relay-transport control lifecycle", () => {

    const control = relay.sockets[0];
    control.open();
-    control.message(JSON.stringify({ type: "pong", ts: Date.now() }));
+    control.message(JSON.stringify({ type: "sync", connectionIds: [] }));
    control.message(JSON.stringify({ type: "connected", connectionId: "clt_test" }));

    const dataSocket = relay.sockets[1];
@@ -240,7 +252,7 @@ describe("relay-transport control lifecycle", () => {

    const control = relay.sockets[0];
    control.open();
-    control.message(JSON.stringify({ type: "pong", ts: Date.now() }));
+    control.message(JSON.stringify({ type: "sync", connectionIds: [] }));
    control.message(JSON.stringify({ type: "connected", connectionId: "clt_test" }));

    expect(relay.sockets[0]?.url).toMatch(/^wss:\/\/\[::1\]\/ws\?/);
@@ -35,8 +35,9 @@ interface RelaySocketLike {

 interface RelayWebSocketLike extends RelaySocketLike {
  terminate: () => void;
+  ping: () => void;
  on: (
-    event: "open" | "message" | "close" | "error",
+    event: "open" | "message" | "close" | "error" | "pong",
    listener: (...args: unknown[]) => void,
  ) => void;
 }
@@ -228,7 +229,9 @@ export function startRelayTransport({
        const now = Date.now();
        const staleForMs = now - controlLastSeenAt;
        // If the control socket is half-open or silently dropped, ws may never emit "close".
-        // Use app-level ping/pong to detect staleness and force a reconnect.
+        // Use a WebSocket protocol ping to detect staleness and force a reconnect.
+        // Cloudflare's runtime auto-responds to protocol pings at the edge without waking the
+        // hibernated relay Durable Object, so this keepalive does not incur DO CPU billing.
        if (staleForMs > CONTROL_STALE_TIMEOUT_MS) {
          relayLogger.warn(
            { url, staleForMs, connectionId, staleTimeoutMs: CONTROL_STALE_TIMEOUT_MS },
@@ -243,7 +246,7 @@ export function startRelayTransport({
        }

        try {
-          socket.send(JSON.stringify({ type: "ping", ts: now }));
+          socket.ping();
        } catch (error) {
          relayLogger.warn({ err: error, connectionId }, "relay_control_ping_send_failed");
          try {
@@ -254,7 +257,7 @@ export function startRelayTransport({
        }
      }, CONTROL_PING_INTERVAL_MS);
      try {
-        socket.send(JSON.stringify({ type: "ping", ts: Date.now() }));
+        socket.ping();
      } catch (error) {
        relayLogger.warn({ err: error, connectionId }, "relay_control_ping_send_failed");
        try {
@@ -290,6 +293,12 @@ export function startRelayTransport({
      // close event will schedule reconnect
    });

+    socket.on("pong", () => {
+      if (controlWs !== socket) return;
+      controlLastSeenAt = Date.now();
+      relayLogger.debug({ connectionId }, "relay_control_pong_received");
+    });
+
    socket.on("message", (data) => {
      if (controlWs !== socket) return;
      controlLastSeenAt = Date.now();
@@ -423,8 +432,17 @@ async function attachEncryptedSocket(
  try {
    const relayTransport = createRelayTransportAdapter(socket, logger);
    const emitter = new EventEmitter();
+    const pendingMessages: Array<string | ArrayBuffer> = [];
+    let attached = false;
+    const emitMessage = (data: string | ArrayBuffer) => {
+      if (attached) {
+        emitter.emit("message", data);
+        return;
+      }
+      pendingMessages.push(data);
+    };
    const channel = await createDaemonChannel(relayTransport, daemonKeyPair, {
-      onmessage: (data) => emitter.emit("message", data),
+      onmessage: emitMessage,
      onclose: (code, reason) => emitter.emit("close", code, reason),
      onerror: (error) => {
        logger.warn({ err: error }, "relay_e2ee_error");
@@ -433,6 +451,11 @@ async function attachEncryptedSocket(
    });
    const encryptedSocket = createEncryptedSocket(channel, emitter);
    await attachSocket(encryptedSocket, metadata);
+    attached = true;
+    for (const message of pendingMessages) {
+      emitter.emit("message", message);
+    }
+    pendingMessages.length = 0;
  } catch (error) {
    logger.warn({ err: error }, "relay_e2ee_handshake_failed");
    try {