From e0f80cf0e9150bf4ce7cd887a29a8750401a3362 Mon Sep 17 00:00:00 2001
From: Peter Steinberger <steipete@gmail.com>
Date: Sun, 8 Mar 2026 05:40:17 +0000
Subject: [PATCH] fix(ui): align control-ui device auth token signing

---
 ui/src/ui/gateway.node.test.ts | 77 +++++++++++++++-------------------
 ui/src/ui/gateway.ts           | 13 ++++--
 2 files changed, 42 insertions(+), 48 deletions(-)

diff --git a/ui/src/ui/gateway.node.test.ts b/ui/src/ui/gateway.node.test.ts
index 6dcc8662b47..07c63a7117b 100644
--- a/ui/src/ui/gateway.node.test.ts
+++ b/ui/src/ui/gateway.node.test.ts
@@ -75,30 +75,6 @@ vi.mock("./device-identity.ts", () => ({
 
 const { GatewayBrowserClient } = await import("./gateway.ts");
 
-function createStorageMock(): Storage {
-  const store = new Map<string, string>();
-  return {
-    get length() {
-      return store.size;
-    },
-    clear() {
-      store.clear();
-    },
-    getItem(key: string) {
-      return store.get(key) ?? null;
-    },
-    key(index: number) {
-      return Array.from(store.keys())[index] ?? null;
-    },
-    removeItem(key: string) {
-      store.delete(key);
-    },
-    setItem(key: string, value: string) {
-      store.set(key, String(value));
-    },
-  };
-}
-
 function getLatestWebSocket(): MockWebSocket {
   const ws = wsInstances.at(-1);
   if (!ws) {
@@ -118,23 +94,8 @@ describe("GatewayBrowserClient", () => {
       publicKey: "public-key", // pragma: allowlist secret
     });
 
-    const localStorage = createStorageMock();
+    window.localStorage.clear();
     vi.stubGlobal("WebSocket", MockWebSocket);
-    vi.stubGlobal("localStorage", localStorage);
-    vi.stubGlobal("crypto", {
-      randomUUID: vi.fn(() => "req-1"),
-      subtle: {},
-    });
-    vi.stubGlobal("navigator", {
-      language: "en-GB",
-      platform: "test-platform",
-      userAgent: "test-agent",
-    });
-    vi.stubGlobal("window", {
-      clearTimeout: vi.fn(),
-      localStorage,
-      setTimeout: vi.fn(() => 1),
-    });
 
     storeDeviceAuthToken({
       deviceId: "device-1",
@@ -148,7 +109,7 @@ describe("GatewayBrowserClient", () => {
     vi.unstubAllGlobals();
   });
 
-  it("keeps shared auth token separate from cached device token", async () => {
+  it("prefers explicit shared auth over cached device tokens", async () => {
     const client = new GatewayBrowserClient({
       url: "ws://127.0.0.1:18789",
       token: "shared-auth-token",
@@ -162,19 +123,47 @@ describe("GatewayBrowserClient", () => {
       event: "connect.challenge",
       payload: { nonce: "nonce-1" },
     });
-    await Promise.resolve();
+    await vi.waitFor(() => expect(ws.sent.length).toBeGreaterThan(0));
 
     const connectFrame = JSON.parse(ws.sent.at(-1) ?? "{}") as {
       id?: string;
       method?: string;
       params?: { auth?: { token?: string } };
     };
-    expect(connectFrame.id).toBe("req-1");
+    expect(typeof connectFrame.id).toBe("string");
     expect(connectFrame.method).toBe("connect");
     expect(connectFrame.params?.auth?.token).toBe("shared-auth-token");
     expect(signDevicePayloadMock).toHaveBeenCalledWith("private-key", expect.any(String));
     const signedPayload = signDevicePayloadMock.mock.calls[0]?.[1];
+    expect(signedPayload).toContain("|shared-auth-token|nonce-1");
+    expect(signedPayload).not.toContain("stored-device-token");
+  });
+
+  it("uses cached device tokens only when no explicit shared auth is provided", async () => {
+    const client = new GatewayBrowserClient({
+      url: "ws://127.0.0.1:18789",
+    });
+
+    client.start();
+    const ws = getLatestWebSocket();
+    ws.emitOpen();
+    ws.emitMessage({
+      type: "event",
+      event: "connect.challenge",
+      payload: { nonce: "nonce-1" },
+    });
+    await vi.waitFor(() => expect(ws.sent.length).toBeGreaterThan(0));
+
+    const connectFrame = JSON.parse(ws.sent.at(-1) ?? "{}") as {
+      id?: string;
+      method?: string;
+      params?: { auth?: { token?: string } };
+    };
+    expect(typeof connectFrame.id).toBe("string");
+    expect(connectFrame.method).toBe("connect");
+    expect(connectFrame.params?.auth?.token).toBe("stored-device-token");
+    expect(signDevicePayloadMock).toHaveBeenCalledWith("private-key", expect.any(String));
+    const signedPayload = signDevicePayloadMock.mock.calls[0]?.[1];
     expect(signedPayload).toContain("|stored-device-token|nonce-1");
-    expect(signedPayload).not.toContain("shared-auth-token");
   });
 });
diff --git a/ui/src/ui/gateway.ts b/ui/src/ui/gateway.ts
index e5b747ae523..c5d4bad86a3 100644
--- a/ui/src/ui/gateway.ts
+++ b/ui/src/ui/gateway.ts
@@ -205,17 +205,22 @@ export class GatewayBrowserClient {
     const role = "operator";
     let deviceIdentity: Awaited<ReturnType<typeof loadOrCreateDeviceIdentity>> | null = null;
     let canFallbackToShared = false;
-    let authToken = this.opts.token;
+    const explicitGatewayToken = this.opts.token?.trim() || undefined;
+    let authToken = explicitGatewayToken;
     let deviceToken: string | undefined;
 
     if (isSecureContext) {
       deviceIdentity = await loadOrCreateDeviceIdentity();
-      deviceToken = loadDeviceAuthToken({
+      const storedToken = loadDeviceAuthToken({
         deviceId: deviceIdentity.deviceId,
         role,
       })?.token;
-      canFallbackToShared = Boolean(deviceToken && this.opts.token);
+      deviceToken = !(explicitGatewayToken || this.opts.password?.trim())
+        ? (storedToken ?? undefined)
+        : undefined;
+      canFallbackToShared = Boolean(deviceToken && explicitGatewayToken);
     }
+    authToken = explicitGatewayToken ?? deviceToken;
     const auth =
       authToken || this.opts.password
         ? {
@@ -244,7 +249,7 @@ export class GatewayBrowserClient {
         role,
         scopes,
         signedAtMs,
-        token: deviceToken ?? null,
+        token: authToken ?? null,
         nonce,
       });
       const signature = await signDevicePayload(deviceIdentity.privateKey, payload);