Add custom headers to direct connections (#2922)

* Add custom headers to direct connections * Address custom header review feedback * Fix direct connection header integration * Fix desktop relay E2E transport setup --------- Co-authored-by: Claude <claude@Claudes-Mac-mini.local> Co-authored-by: Mohamed Boudra <boudra.moha@gmail.com>
2026-08-06 11:05:40 +02:00
parent 098122be8d
commit 73ba37a07a
38 changed files with 1190 additions and 125 deletions
@@ -0,0 +1,31 @@
+import { expect, test } from "../../app/e2e/support/fixtures";
+import { gotoAppShell, openSettings } from "../../app/e2e/support/helpers/app";
+import { getServerId } from "../../app/e2e/support/helpers/server-id";
+import {
+  openAddHostFlow,
+  selectHostConnectionType,
+  toggleHostAdvanced,
+} from "../../app/e2e/support/helpers/settings";
+import { installDesktopRuntime } from "./support/runtime";
+
+test("Electron users can add and remove direct connection custom headers", async ({ page }) => {
+  await installDesktopRuntime(page, { serverId: getServerId() });
+  await gotoAppShell(page);
+  await openSettings(page);
+  await openAddHostFlow(page);
+  await selectHostConnectionType(page, "direct");
+  await toggleHostAdvanced(page);
+
+  const addHeader = page.getByTestId("direct-header-add");
+  await expect(addHeader).toBeVisible();
+  await addHeader.click();
+
+  await page.getByTestId("direct-header-name-0").fill("X-Tenant");
+  await page.getByTestId("direct-header-value-0").fill("acme");
+  await expect(page.getByTestId("direct-header-name-0")).toHaveValue("X-Tenant");
+  await expect(page.getByTestId("direct-header-value-0")).toHaveValue("acme");
+
+  await page.getByTestId("direct-header-remove-0").click();
+  await expect(page.getByTestId("direct-header-name-0")).toHaveCount(0);
+  await expect(page.getByTestId("direct-header-value-0")).toHaveCount(0);
+});
@@ -1,4 +1,5 @@
 import { expect, test } from "../../app/e2e/support/fixtures";
+import type { Page } from "@playwright/test";
 import {
  startIsolatedHostDaemon,
  type IsolatedHostDaemon,
@@ -18,11 +19,12 @@ import {
  openPairDeviceFromHome,
  openRelaySecurityDocs,
  observePairingOfferRequests,
-  prepareLocalPairingHost,
+  preparePairingHost,
  reloadAndOpenPairDevice,
  retryRelayAndExpectFailure,
  switchPairDeviceToHost,
 } from "../../app/e2e/support/helpers/pair-device";
+import { installDesktopRuntime } from "./support/runtime";

 interface RelayConfigDaemonClient {
  close(): Promise<void>;
@@ -30,6 +32,20 @@ interface RelayConfigDaemonClient {
  patchDaemonConfig(config: { relay: { enabled: boolean } }): Promise<unknown>;
 }

+async function prepareLocalPairingHost(
+  page: Page,
+  daemon: Parameters<typeof preparePairingHost>[1],
+  additionalHosts: Parameters<typeof preparePairingHost>[2] = [],
+): Promise<void> {
+  const endpoint = "port" in daemon ? `127.0.0.1:${daemon.port}` : daemon.endpoint;
+  await installDesktopRuntime(page, {
+    serverId: daemon.serverId,
+    daemonListen: endpoint,
+    manageBuiltInDaemon: false,
+  });
+  await preparePairingHost(page, daemon, additionalHosts);
+}
+
 test.describe("local device relay pairing", () => {
  let relayOffDaemon: IsolatedHostDaemon;
  let relayEnableDaemon: IsolatedHostDaemon;
@@ -1,24 +1,15 @@
 import { test, expect } from "../../app/e2e/support/fixtures";
 import { gotoAppShell, openSettings } from "../../app/e2e/support/helpers/app";
+import { getServerId } from "../../app/e2e/support/helpers/server-id";
+import { installDesktopRuntime } from "./support/runtime";

 test.describe("Settings sidebar scrolling", () => {
  test.use({ viewport: { width: 900, height: 260 } });

  test("desktop drag region does not cover the scroll body", async ({ page }) => {
-    await page.addInitScript(() => {
-      window.paseoDesktop = {
-        platform: "darwin",
-        events: { on: () => () => {} },
-        invoke: async (command: string) => {
-          if (command === "get_desktop_settings") {
-            return {
-              releaseChannel: "stable",
-              daemon: { manageBuiltInDaemon: true, keepRunningAfterQuit: true },
-            };
-          }
-          return null;
-        },
-      };
+    await installDesktopRuntime(page, {
+      serverId: getServerId(),
+      manageBuiltInDaemon: true,
    });

    await gotoAppShell(page);
@@ -139,8 +139,123 @@ export async function installDesktopRuntime(
    let daemonRunning = true;
    let currentPid: number | null = cfg.daemonPid ?? null;
    let startCount = 0;
+    let nextTransportSessionId = 0;
+    const transportSessions = new Map<string, WebSocket>();
+    const desktopEventHandlers = new Map<string, Set<(payload: unknown) => void>>();
    window.__desktopDaemonStartRequested = false;

+    function emitDesktopEvent(event: string, payload: unknown) {
+      for (const handler of desktopEventHandlers.get(event) ?? []) {
+        handler(payload);
+      }
+    }
+
+    function bytesToBase64(bytes: Uint8Array): string {
+      let binary = "";
+      for (const byte of bytes) {
+        binary += String.fromCharCode(byte);
+      }
+      return btoa(binary);
+    }
+
+    function base64ToBytes(base64: string): Uint8Array {
+      const binary = atob(base64);
+      const bytes = new Uint8Array(binary.length);
+      for (let index = 0; index < binary.length; index += 1) {
+        bytes[index] = binary.charCodeAt(index);
+      }
+      return bytes;
+    }
+
+    function parseWebSocketTransportTarget(args?: Record<string, unknown>): {
+      url: string;
+      protocols: string[];
+    } {
+      const url = typeof args?.url === "string" ? args.url : "";
+      const protocols = Array.isArray(args?.protocols)
+        ? args.protocols.filter((protocol): protocol is string => typeof protocol === "string")
+        : [];
+      const headers =
+        args?.headers && typeof args.headers === "object"
+          ? Object.entries(args.headers as Record<string, unknown>)
+          : [];
+      const passwordProtocol = protocols.find((protocol) => protocol.startsWith("paseo.bearer."));
+      const supportsHeaders =
+        headers.length === 0 ||
+        (headers.length === 1 &&
+          headers[0][0].toLowerCase() === "authorization" &&
+          typeof headers[0][1] === "string" &&
+          passwordProtocol !== undefined &&
+          headers[0][1] === `Bearer ${passwordProtocol.slice("paseo.bearer.".length)}`);
+
+      if (!url || !supportsHeaders) {
+        throw new Error("The browser-backed desktop E2E transport cannot attach custom headers.");
+      }
+      return { url, protocols };
+    }
+
+    async function openWebSocketTransport(args?: Record<string, unknown>): Promise<string> {
+      const target = parseWebSocketTransportTarget(args);
+      const sessionId = `e2e-websocket-session-${++nextTransportSessionId}`;
+      const socket = new WebSocket(target.url, target.protocols);
+      socket.binaryType = "arraybuffer";
+      transportSessions.set(sessionId, socket);
+
+      return await new Promise<string>((resolve, reject) => {
+        let opened = false;
+
+        socket.addEventListener("open", () => {
+          opened = true;
+          resolve(sessionId);
+          emitDesktopEvent("local-daemon-transport-event", { sessionId, kind: "open" });
+        });
+        socket.addEventListener("message", async (event) => {
+          if (typeof event.data === "string") {
+            emitDesktopEvent("local-daemon-transport-event", {
+              sessionId,
+              kind: "message",
+              text: event.data,
+            });
+            return;
+          }
+          const buffer =
+            event.data instanceof ArrayBuffer
+              ? event.data
+              : await (event.data as Blob).arrayBuffer();
+          emitDesktopEvent("local-daemon-transport-event", {
+            sessionId,
+            kind: "message",
+            binaryBase64: bytesToBase64(new Uint8Array(buffer)),
+          });
+        });
+        socket.addEventListener("close", (event) => {
+          transportSessions.delete(sessionId);
+          if (!opened) {
+            reject(new Error("Desktop E2E WebSocket closed before opening."));
+            return;
+          }
+          emitDesktopEvent("local-daemon-transport-event", {
+            sessionId,
+            kind: "close",
+            code: event.code,
+            reason: event.reason,
+          });
+        });
+        socket.addEventListener("error", () => {
+          if (!opened) {
+            transportSessions.delete(sessionId);
+            reject(new Error("Desktop E2E WebSocket failed to open."));
+            return;
+          }
+          emitDesktopEvent("local-daemon-transport-event", {
+            sessionId,
+            kind: "error",
+            error: "Desktop E2E WebSocket transport error.",
+          });
+        });
+      });
+    }
+
    function buildDaemonStatus() {
      return {
        serverId: cfg.serverId,
@@ -175,6 +290,55 @@ export async function installDesktopRuntime(
      }
    }

+    function patchDesktopSettings(args?: Record<string, unknown>) {
+      const daemon = args?.daemon;
+      if (
+        daemon !== null &&
+        typeof daemon === "object" &&
+        "manageBuiltInDaemon" in daemon &&
+        typeof daemon.manageBuiltInDaemon === "boolean"
+      ) {
+        manageDaemon = daemon.manageBuiltInDaemon;
+      }
+      return {
+        releaseChannel: "stable",
+        daemon: { manageBuiltInDaemon: manageDaemon, keepRunningAfterQuit: true },
+      };
+    }
+
+    function sendWebSocketTransportMessage(args?: Record<string, unknown>): null {
+      const sessionId = typeof args?.sessionId === "string" ? args.sessionId : "";
+      const socket = transportSessions.get(sessionId);
+      if (!socket || socket.readyState !== WebSocket.OPEN) {
+        throw new Error(`Desktop E2E WebSocket session is not open: ${sessionId}`);
+      }
+      if (typeof args?.text === "string") {
+        socket.send(args.text);
+        return null;
+      }
+      if (typeof args?.binaryBase64 === "string") {
+        socket.send(base64ToBytes(args.binaryBase64));
+        return null;
+      }
+      throw new Error("Desktop E2E WebSocket send requires a payload.");
+    }
+
+    function closeWebSocketTransport(args?: Record<string, unknown>): null {
+      const sessionId = typeof args?.sessionId === "string" ? args.sessionId : "";
+      transportSessions.get(sessionId)?.close();
+      transportSessions.delete(sessionId);
+      return null;
+    }
+
+    const webSocketTransportCommandHandlers: Record<
+      string,
+      (args?: Record<string, unknown>) => unknown
+    > = {
+      open_websocket_daemon_transport: openWebSocketTransport,
+      send_local_daemon_transport_message: sendWebSocketTransportMessage,
+      close_local_daemon_transport: closeWebSocketTransport,
+    };
+
    const desktopBridge: {
      platform: string;
      invoke: (command: string, args?: Record<string, unknown>) => Promise<unknown>;
@@ -183,7 +347,10 @@ export async function installDesktopRuntime(
        open: (options?: Record<string, unknown>) => Promise<string | string[] | null>;
      };
      getPendingOpenProject: () => Promise<string | null>;
-      events: { on: () => Promise<() => void> };
+      opener: { openUrl: (url: string) => Promise<void> };
+      events: {
+        on: (event: string, handler: (payload: unknown) => void) => Promise<() => void>;
+      };
      editor?: {
        listTargets: () => Promise<DesktopEditorTargetConfig[]>;
        openTarget: (input: DesktopEditorOpenRecord) => Promise<void>;
@@ -239,19 +406,7 @@ export async function installDesktopRuntime(
        }

        if (command === "patch_desktop_settings") {
-          const daemon = args?.daemon;
-          if (
-            daemon !== null &&
-            typeof daemon === "object" &&
-            "manageBuiltInDaemon" in daemon &&
-            typeof daemon.manageBuiltInDaemon === "boolean"
-          ) {
-            manageDaemon = daemon.manageBuiltInDaemon;
-          }
-          return {
-            releaseChannel: "stable",
-            daemon: { manageBuiltInDaemon: manageDaemon, keepRunningAfterQuit: true },
-          };
+          return patchDesktopSettings(args);
        }

        if (command === "stop_desktop_daemon") {
@@ -264,6 +419,11 @@ export async function installDesktopRuntime(
          return startDesktopDaemon();
        }

+        const transportCommandHandler = webSocketTransportCommandHandlers[command];
+        if (transportCommandHandler) {
+          return await transportCommandHandler(args);
+        }
+
        return null;
      },
      dialog: {
@@ -280,7 +440,24 @@ export async function installDesktopRuntime(
        },
      },
      getPendingOpenProject: async () => null,
-      events: { on: async () => () => undefined },
+      opener: {
+        openUrl: async (url) => {
+          localStorage.setItem("@paseo:e2e-opened-url", url);
+        },
+      },
+      events: {
+        on: async (event, handler) => {
+          const handlers = desktopEventHandlers.get(event) ?? new Set();
+          handlers.add(handler);
+          desktopEventHandlers.set(event, handlers);
+          return () => {
+            handlers.delete(handler);
+            if (handlers.size === 0) {
+              desktopEventHandlers.delete(event);
+            }
+          };
+        },
+      },
    };

    if (cfg.editorTargets) {
@@ -26,8 +26,10 @@ import {
 import { createSkillsCommandHandlers, getSkillsController } from "../integrations/skills/index.js";
 import {
  openLocalTransportSession,
+  openWebSocketTransportSession,
  sendLocalTransportMessage,
  closeLocalTransportSession,
+  type WebSocketTransportTarget,
 } from "./local-transport.js";
 import { createNodeEntrypointInvocation, resolveDaemonRunnerEntrypoint } from "./runtime-paths.js";
 import { runExternalCliJsonCommand, runExternalCliTextCommand } from "./cli/external.js";
@@ -537,6 +539,9 @@ export function createDaemonCommandHandlers(): Record<string, DesktopCommandHand
      const target = args as { transportType: "socket" | "pipe"; transportPath: string };
      return await openLocalTransportSession(target);
    },
+    open_websocket_daemon_transport: async (args) => {
+      return await openWebSocketTransportSession(args as unknown as WebSocketTransportTarget);
+    },
    send_local_daemon_transport_message: async (args) => {
      await sendLocalTransportMessage(
        args as { sessionId: string; text?: string; binaryBase64?: string },
@@ -0,0 +1,86 @@
+import { createServer } from "node:http";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { WebSocketServer } from "ws";
+
+const send = vi.fn();
+
+vi.mock("electron", () => ({
+  BrowserWindow: {
+    getAllWindows: () => [{ webContents: { send } }],
+  },
+}));
+
+import { closeAllTransportSessions, openWebSocketTransportSession } from "./local-transport";
+
+afterEach(() => {
+  closeAllTransportSessions();
+  send.mockReset();
+});
+
+describe("desktop WebSocket transport", () => {
+  it("attaches custom headers when opening a remote daemon session", async () => {
+    const server = createServer();
+    const webSocketServer = new WebSocketServer({ server });
+    const requestHeaders = new Promise<Record<string, string | string[] | undefined>>((resolve) => {
+      webSocketServer.once("connection", (_socket, request) => resolve(request.headers));
+    });
+
+    await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
+    const address = server.address();
+    if (!address || typeof address === "string") {
+      throw new Error("Expected TCP test server address.");
+    }
+
+    try {
+      const sessionId = await openWebSocketTransportSession({
+        transportType: "websocket",
+        url: `ws://127.0.0.1:${address.port}/ws`,
+        headers: { "X-Tenant": "acme" },
+      });
+
+      expect((await requestHeaders)["x-tenant"]).toBe("acme");
+      expect(sessionId).toMatch(/^local-session-/);
+    } finally {
+      closeAllTransportSessions();
+      await new Promise<void>((resolve, reject) => {
+        webSocketServer.close((error) => (error ? reject(error) : resolve()));
+      });
+      await new Promise<void>((resolve, reject) => {
+        server.close((error) => {
+          if (error) {
+            reject(error);
+            return;
+          }
+          resolve();
+        });
+      });
+    }
+  });
+
+  it("rejects non-WebSocket URLs", () => {
+    expect(() =>
+      openWebSocketTransportSession({
+        transportType: "websocket",
+        url: "https://daemon.example/ws",
+      }),
+    ).toThrow("Unsupported WebSocket transport URL");
+  });
+
+  it("rejects malformed custom headers before opening a session", () => {
+    expect(() =>
+      openWebSocketTransportSession({
+        transportType: "websocket",
+        url: "ws://daemon.example/ws",
+        headers: { "Bad Header": "value" },
+      }),
+    ).toThrow("Invalid WebSocket transport header");
+
+    expect(() =>
+      openWebSocketTransportSession({
+        transportType: "websocket",
+        url: "ws://daemon.example/ws",
+        headers: { "X-Tenant": "acme\r\nX-Injected: true" },
+      }),
+    ).toThrow("Invalid WebSocket transport header");
+  });
+});
@@ -6,6 +6,13 @@ interface TransportTarget {
  transportPath: string;
 }

+export interface WebSocketTransportTarget {
+  transportType: "websocket";
+  url: string;
+  headers?: Record<string, string>;
+  protocols?: string[];
+}
+
 interface TransportEventPayload {
  sessionId: string;
  kind: "open" | "message" | "close" | "error";
@@ -23,6 +30,7 @@ interface Session {
 }

 const WS_ENDPOINT_PATH = "/ws";
+const HTTP_HEADER_NAME_PATTERN = /^[!#$%&'*+\-.^_`|~0-9A-Za-z]+$/;

 let nextSessionId = 0;
 const sessions = new Map<string, Session>();
@@ -52,6 +60,28 @@ function describeTransportTarget(target: TransportTarget): string {
  return target.transportType === "pipe" ? "local daemon pipe" : "local daemon socket";
 }

+function normalizeWebSocketTarget(target: WebSocketTransportTarget): WebSocketTransportTarget {
+  const parsed = new URL(target.url);
+  if (parsed.protocol !== "ws:" && parsed.protocol !== "wss:") {
+    throw new Error(`Unsupported WebSocket transport URL: ${target.url}`);
+  }
+
+  const headers: Record<string, string> = {};
+  for (const [name, value] of Object.entries(target.headers ?? {})) {
+    if (!HTTP_HEADER_NAME_PATTERN.test(name) || typeof value !== "string" || /[\r\n]/.test(value)) {
+      throw new Error("Invalid WebSocket transport header.");
+    }
+    headers[name] = value;
+  }
+
+  return {
+    transportType: "websocket",
+    url: parsed.toString(),
+    ...(Object.keys(headers).length > 0 ? { headers } : {}),
+    ...(target.protocols?.length ? { protocols: [...target.protocols] } : {}),
+  };
+}
+
 function decodeTransportMessage(input: { text?: string; binaryBase64?: string }): string | Buffer {
  if (typeof input.text === "string") {
    return input.text;
@@ -65,11 +95,28 @@ function decodeTransportMessage(input: { text?: string; binaryBase64?: string })
 }

 export function openLocalTransportSession(target: TransportTarget): Promise<string> {
+  return openTransportSession(buildLocalWebSocketUrl(target), describeTransportTarget(target));
+}
+
+export function openWebSocketTransportSession(target: WebSocketTransportTarget): Promise<string> {
+  const normalized = normalizeWebSocketTarget(target);
+  return openTransportSession(normalized.url, "remote daemon", {
+    headers: normalized.headers,
+    protocols: normalized.protocols,
+  });
+}
+
+function openTransportSession(
+  url: string,
+  targetDescription: string,
+  options?: { headers?: Record<string, string>; protocols?: string[] },
+): Promise<string> {
  const sessionId = `local-session-${++nextSessionId}`;
-  const url = buildLocalWebSocketUrl(target);

  return new Promise((resolve, reject) => {
-    const ws = new WebSocket(url);
+    const ws = options?.protocols
+      ? new WebSocket(url, options.protocols, { headers: options.headers })
+      : new WebSocket(url, { headers: options?.headers });
    const session: Session = {
      id: sessionId,
      ws,
@@ -121,9 +168,7 @@ export function openLocalTransportSession(target: TransportTarget): Promise<stri
      sessions.delete(sessionId);

      if (!openSettled) {
-        finalizeOpenFailure(
-          `${describeTransportTarget(target)} closed before the session became ready.`,
-        );
+        finalizeOpenFailure(`${targetDescription} closed before the session became ready.`);
        return;
      }

@@ -139,9 +184,7 @@ export function openLocalTransportSession(target: TransportTarget): Promise<stri

    ws.on("error", (err: Error) => {
      if (!openSettled) {
-        finalizeOpenFailure(
-          `Failed to connect to ${describeTransportTarget(target)}: ${err.message}`,
-        );
+        finalizeOpenFailure(`Failed to connect to ${targetDescription}: ${err.message}`);
        return;
      }