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>

packages/app/src/utils/connection-headers.test.ts

@@ -0,0 +1,56 @@
+import { describe, expect, it } from "vitest";
+import { normalizeConnectionHeadersRecord, prepareConnectionHeaders } from "./connection-headers";
+
+describe("prepareConnectionHeaders", () => {
+  it("trims and returns custom headers", () => {
+    expect(
+      prepareConnectionHeaders([
+        { id: 1, name: " X-Tenant ", value: " acme " },
+        { id: 2, name: "X-Empty", value: "" },
+      ]),
+    ).toEqual({ headers: { "X-Tenant": "acme", "X-Empty": "" } });
+  });
+
+  it("ignores completely empty draft rows", () => {
+    expect(prepareConnectionHeaders([{ id: 1, name: "", value: "" }])).toEqual({});
+  });
+
+  it("requires a name when a value is entered", () => {
+    expect(prepareConnectionHeaders([{ id: 1, name: "", value: "acme" }])).toEqual({
+      issue: { type: "missingName" },
+    });
+  });
+
+  it("rejects invalid and duplicate names case-insensitively", () => {
+    expect(prepareConnectionHeaders([{ id: 1, name: "Bad Header", value: "value" }])).toEqual({
+      issue: { type: "invalidName", name: "Bad Header" },
+    });
+    expect(
+      prepareConnectionHeaders([
+        { id: 1, name: "X-Tenant", value: "one" },
+        { id: 2, name: "x-tenant", value: "two" },
+      ]),
+    ).toEqual({ issue: { type: "duplicateName", name: "x-tenant" } });
+  });
+
+  it("rejects line breaks in values", () => {
+    expect(prepareConnectionHeaders([{ id: 1, name: "X-Test", value: "one\ntwo" }])).toEqual({
+      issue: { type: "invalidValue", name: "X-Test" },
+    });
+  });
+});
+
+describe("normalizeConnectionHeadersRecord", () => {
+  it("accepts a valid persisted string record", () => {
+    expect(normalizeConnectionHeadersRecord({ "X-Tenant": "acme" })).toEqual({
+      "X-Tenant": "acme",
+    });
+  });
+
+  it("rejects malformed persisted header values", () => {
+    expect(normalizeConnectionHeadersRecord({ "X-Tenant": 42 })).toBeUndefined();
+    expect(
+      normalizeConnectionHeadersRecord({ "X-Tenant": "acme\r\nX-Injected: true" }),
+    ).toBeUndefined();
+  });
+});

packages/app/src/utils/connection-headers.ts

@@ -0,0 +1,65 @@
+export interface ConnectionHeaderDraft {
+  id: number;
+  name: string;
+  value: string;
+}
+
+export type ConnectionHeaderValidationIssue =
+  | { type: "missingName" }
+  | { type: "invalidName"; name: string }
+  | { type: "invalidValue"; name: string }
+  | { type: "duplicateName"; name: string };
+
+const HTTP_HEADER_NAME_PATTERN = /^[!#$%&'*+.^_`|~0-9A-Za-z-]+$/;
+
+function isPlainRecord(value: unknown): value is Record<string, unknown> {
+  return Boolean(value) && typeof value === "object" && !Array.isArray(value);
+}
+
+export function prepareConnectionHeaders(drafts: readonly ConnectionHeaderDraft[]): {
+  headers?: Record<string, string>;
+  issue?: ConnectionHeaderValidationIssue;
+} {
+  const headers: Record<string, string> = {};
+  const normalizedNames = new Set<string>();
+
+  for (const draft of drafts) {
+    const name = draft.name.trim();
+    const value = draft.value.trim();
+    if (!name && !value) continue;
+    if (!name) return { issue: { type: "missingName" } };
+    if (!HTTP_HEADER_NAME_PATTERN.test(name)) {
+      return { issue: { type: "invalidName", name } };
+    }
+    if (/\r|\n/.test(value)) {
+      return { issue: { type: "invalidValue", name } };
+    }
+
+    const normalizedName = name.toLowerCase();
+    if (normalizedNames.has(normalizedName)) {
+      return { issue: { type: "duplicateName", name } };
+    }
+    normalizedNames.add(normalizedName);
+    headers[name] = value;
+  }
+
+  return Object.keys(headers).length > 0 ? { headers } : {};
+}
+
+export function normalizeConnectionHeadersRecord(
+  value: unknown,
+): Record<string, string> | undefined {
+  if (!isPlainRecord(value)) {
+    return undefined;
+  }
+
+  const drafts: ConnectionHeaderDraft[] = [];
+  for (const [name, headerValue] of Object.entries(value)) {
+    if (typeof headerValue !== "string") {
+      return undefined;
+    }
+    drafts.push({ id: drafts.length, name, value: headerValue });
+  }
+  const result = prepareConnectionHeaders(drafts);
+  return result.issue ? undefined : result.headers;
+}

packages/app/src/utils/test-daemon-connection.test.ts

@@ -134,6 +134,46 @@ describe("test-daemon-connection connectToDaemon", () => {
     expect(probe.createdConfigs()[0]?.password).toBe("shared-secret");
   });
 
+  it("passes direct TCP custom headers into the initial probe client config", async () => {
+    const { connectToDaemon } = await import("./test-daemon-connection");
+    const result = await connectToDaemon(
+      {
+        id: "direct:lan:6767",
+        type: "directTcp",
+        endpoint: "lan:6767",
+        headers: { "X-Tenant": "acme" },
+      },
+      undefined,
+      probe.deps,
+    );
+    await result.client.close();
+
+    expect(probe.createdConfigs()[0]?.headers).toEqual({ "X-Tenant": "acme" });
+  });
+
+  it("uses the Electron WebSocket transport for the initial probe", async () => {
+    const { connectToDaemon } = await import("./test-daemon-connection");
+    const transportFactory: NonNullable<DaemonClientConfig["transportFactory"]> = () => {
+      throw new Error("The fake probe client should not open the transport.");
+    };
+    const result = await connectToDaemon(
+      {
+        id: "direct:lan:6767",
+        type: "directTcp",
+        endpoint: "lan:6767",
+        headers: { "X-Tenant": "acme" },
+      },
+      undefined,
+      {
+        ...probe.deps,
+        createWebSocketTransportFactory: () => transportFactory,
+      },
+    );
+    await result.client.close();
+
+    expect(probe.createdConfigs()[0]?.transportFactory).toBe(transportFactory);
+  });
+
   it("passes performance tracing into the connected client", async () => {
     const { connectToDaemon } = await import("./test-daemon-connection");
     const trace = {

packages/app/src/utils/test-daemon-connection.ts

@@ -11,6 +11,7 @@ import {
 import {
   buildLocalDaemonTransportUrl,
   createDesktopLocalDaemonTransportFactory,
+  createDesktopWebSocketTransportFactory,
 } from "@/desktop/daemon/desktop-daemon-transport";
 
 export interface DaemonProbeClient {
@@ -29,6 +30,7 @@ export interface DaemonConnectionDependencies<TClient extends DaemonProbeClient>
   getClientId(): Promise<string>;
   resolveAppVersion(): string | null;
   createLocalTransportFactory(): DaemonClientConfig["transportFactory"] | null;
+  createWebSocketTransportFactory?(): DaemonClientConfig["transportFactory"] | null;
   buildLocalTransportUrl(input: LocalTransportUrlInput): string;
   createClient(config: DaemonClientConfig): TClient;
 }
@@ -37,10 +39,17 @@ const defaultDaemonConnectionDependencies: DaemonConnectionDependencies<DaemonCl
   getClientId: getOrCreateClientId,
   resolveAppVersion,
   createLocalTransportFactory: createDesktopLocalDaemonTransportFactory,
+  createWebSocketTransportFactory: createDesktopWebSocketTransportFactory,
   buildLocalTransportUrl: buildLocalDaemonTransportUrl,
   createClient: (config) => new DaemonClient(config),
 };
 
+function resolveWebSocketTransportFactory(
+  deps: Pick<DaemonConnectionDependencies<DaemonProbeClient>, "createWebSocketTransportFactory">,
+): DaemonClientConfig["transportFactory"] {
+  return deps.createWebSocketTransportFactory?.() ?? undefined;
+}
+
 function normalizeNonEmptyString(value: unknown): string | null {
   if (typeof value !== "string") return null;
   const trimmed = value.trim();
@@ -103,7 +112,11 @@ export async function buildClientConfig(
   },
   deps: Pick<
     DaemonConnectionDependencies<DaemonProbeClient>,
-    "getClientId" | "resolveAppVersion" | "createLocalTransportFactory" | "buildLocalTransportUrl"
+    | "getClientId"
+    | "resolveAppVersion"
+    | "createLocalTransportFactory"
+    | "createWebSocketTransportFactory"
+    | "buildLocalTransportUrl"
   > = defaultDaemonConnectionDependencies,
 ): Promise<DaemonClientConfig> {
   const clientId = await deps.getClientId();
@@ -135,8 +148,10 @@ export async function buildClientConfig(
   if (connection.type === "directTcp") {
     return {
       ...base,
+      transportFactory: resolveWebSocketTransportFactory(deps),
       url: buildDaemonWebSocketUrl(connection.endpoint, { useTls: connection.useTls ?? false }),
       ...(connection.password ? { password: connection.password } : {}),
+      ...(connection.headers ? { headers: connection.headers } : {}),
     };
   }