fix(gateway): harden token fallback/reconnect behavior and docs (#42507)

* fix(gateway): harden token fallback and auth reconnect handling * docs(gateway): clarify auth retry and token-drift recovery * fix(gateway): tighten auth reconnect gating across clients * fix: harden gateway token retry (#42507) (thanks @joshavant)
2026-03-12 07:20:45 +00:00 · 2026-03-10 17:05:57 -05:00
parent ff2e7a2945
commit a76e810193
21 changed files with 1188 additions and 80 deletions
--- a/ui/src/ui/gateway.node.test.ts
+++ b/ui/src/ui/gateway.node.test.ts
@@ -1,5 +1,5 @@
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
-import { storeDeviceAuthToken } from "./device-auth.ts";
+import { loadDeviceAuthToken, storeDeviceAuthToken } from "./device-auth.ts";
 import type { DeviceIdentity } from "./device-identity.ts";

 const wsInstances = vi.hoisted((): MockWebSocket[] => []);
@@ -54,6 +54,12 @@ class MockWebSocket {
    this.readyState = 3;
  }

+  emitClose(code = 1000, reason = "") {
+    for (const handler of this.handlers.close) {
+      handler({ code, reason });
+    }
+  }
+
  emitOpen() {
    for (const handler of this.handlers.open) {
      handler();
@@ -106,6 +112,7 @@ describe("GatewayBrowserClient", () => {
  });

  afterEach(() => {
+    vi.useRealTimers();
    vi.unstubAllGlobals();
  });

@@ -166,4 +173,212 @@ describe("GatewayBrowserClient", () => {
    const signedPayload = signDevicePayloadMock.mock.calls[0]?.[1];
    expect(signedPayload).toContain("|stored-device-token|nonce-1");
  });
+
+  it("retries once with device token after token mismatch when shared token is explicit", async () => {
+    vi.useFakeTimers();
+    const client = new GatewayBrowserClient({
+      url: "ws://127.0.0.1:18789",
+      token: "shared-auth-token",
+    });
+
+    client.start();
+    const ws1 = getLatestWebSocket();
+    ws1.emitOpen();
+    ws1.emitMessage({
+      type: "event",
+      event: "connect.challenge",
+      payload: { nonce: "nonce-1" },
+    });
+    await vi.waitFor(() => expect(ws1.sent.length).toBeGreaterThan(0));
+    const firstConnect = JSON.parse(ws1.sent.at(-1) ?? "{}") as {
+      id: string;
+      params?: { auth?: { token?: string; deviceToken?: string } };
+    };
+    expect(firstConnect.params?.auth?.token).toBe("shared-auth-token");
+    expect(firstConnect.params?.auth?.deviceToken).toBeUndefined();
+
+    ws1.emitMessage({
+      type: "res",
+      id: firstConnect.id,
+      ok: false,
+      error: {
+        code: "INVALID_REQUEST",
+        message: "unauthorized",
+        details: { code: "AUTH_TOKEN_MISMATCH", canRetryWithDeviceToken: true },
+      },
+    });
+    await vi.waitFor(() => expect(ws1.readyState).toBe(3));
+    ws1.emitClose(4008, "connect failed");
+
+    await vi.advanceTimersByTimeAsync(800);
+    const ws2 = getLatestWebSocket();
+    expect(ws2).not.toBe(ws1);
+    ws2.emitOpen();
+    ws2.emitMessage({
+      type: "event",
+      event: "connect.challenge",
+      payload: { nonce: "nonce-2" },
+    });
+    await vi.waitFor(() => expect(ws2.sent.length).toBeGreaterThan(0));
+    const secondConnect = JSON.parse(ws2.sent.at(-1) ?? "{}") as {
+      id: string;
+      params?: { auth?: { token?: string; deviceToken?: string } };
+    };
+    expect(secondConnect.params?.auth?.token).toBe("shared-auth-token");
+    expect(secondConnect.params?.auth?.deviceToken).toBe("stored-device-token");
+
+    ws2.emitMessage({
+      type: "res",
+      id: secondConnect.id,
+      ok: false,
+      error: {
+        code: "INVALID_REQUEST",
+        message: "unauthorized",
+        details: { code: "AUTH_TOKEN_MISMATCH" },
+      },
+    });
+    await vi.waitFor(() => expect(ws2.readyState).toBe(3));
+    ws2.emitClose(4008, "connect failed");
+    expect(loadDeviceAuthToken({ deviceId: "device-1", role: "operator" })?.token).toBe(
+      "stored-device-token",
+    );
+    await vi.advanceTimersByTimeAsync(30_000);
+    expect(wsInstances).toHaveLength(2);
+
+    vi.useRealTimers();
+  });
+
+  it("treats IPv6 loopback as trusted for bounded device-token retry", async () => {
+    vi.useFakeTimers();
+    const client = new GatewayBrowserClient({
+      url: "ws://[::1]:18789",
+      token: "shared-auth-token",
+    });
+
+    client.start();
+    const ws1 = getLatestWebSocket();
+    ws1.emitOpen();
+    ws1.emitMessage({
+      type: "event",
+      event: "connect.challenge",
+      payload: { nonce: "nonce-1" },
+    });
+    await vi.waitFor(() => expect(ws1.sent.length).toBeGreaterThan(0));
+    const firstConnect = JSON.parse(ws1.sent.at(-1) ?? "{}") as {
+      id: string;
+      params?: { auth?: { token?: string; deviceToken?: string } };
+    };
+    expect(firstConnect.params?.auth?.token).toBe("shared-auth-token");
+    expect(firstConnect.params?.auth?.deviceToken).toBeUndefined();
+
+    ws1.emitMessage({
+      type: "res",
+      id: firstConnect.id,
+      ok: false,
+      error: {
+        code: "INVALID_REQUEST",
+        message: "unauthorized",
+        details: { code: "AUTH_TOKEN_MISMATCH", canRetryWithDeviceToken: true },
+      },
+    });
+    await vi.waitFor(() => expect(ws1.readyState).toBe(3));
+    ws1.emitClose(4008, "connect failed");
+
+    await vi.advanceTimersByTimeAsync(800);
+    const ws2 = getLatestWebSocket();
+    expect(ws2).not.toBe(ws1);
+    ws2.emitOpen();
+    ws2.emitMessage({
+      type: "event",
+      event: "connect.challenge",
+      payload: { nonce: "nonce-2" },
+    });
+    await vi.waitFor(() => expect(ws2.sent.length).toBeGreaterThan(0));
+    const secondConnect = JSON.parse(ws2.sent.at(-1) ?? "{}") as {
+      params?: { auth?: { token?: string; deviceToken?: string } };
+    };
+    expect(secondConnect.params?.auth?.token).toBe("shared-auth-token");
+    expect(secondConnect.params?.auth?.deviceToken).toBe("stored-device-token");
+
+    client.stop();
+    vi.useRealTimers();
+  });
+
+  it("continues reconnecting on first token mismatch when no retry was attempted", async () => {
+    vi.useFakeTimers();
+    window.localStorage.clear();
+
+    const client = new GatewayBrowserClient({
+      url: "ws://127.0.0.1:18789",
+      token: "shared-auth-token",
+    });
+
+    client.start();
+    const ws1 = getLatestWebSocket();
+    ws1.emitOpen();
+    ws1.emitMessage({
+      type: "event",
+      event: "connect.challenge",
+      payload: { nonce: "nonce-1" },
+    });
+    await vi.waitFor(() => expect(ws1.sent.length).toBeGreaterThan(0));
+    const firstConnect = JSON.parse(ws1.sent.at(-1) ?? "{}") as { id: string };
+
+    ws1.emitMessage({
+      type: "res",
+      id: firstConnect.id,
+      ok: false,
+      error: {
+        code: "INVALID_REQUEST",
+        message: "unauthorized",
+        details: { code: "AUTH_TOKEN_MISMATCH" },
+      },
+    });
+    await vi.waitFor(() => expect(ws1.readyState).toBe(3));
+    ws1.emitClose(4008, "connect failed");
+
+    await vi.advanceTimersByTimeAsync(800);
+    expect(wsInstances).toHaveLength(2);
+
+    client.stop();
+    vi.useRealTimers();
+  });
+
+  it("does not auto-reconnect on AUTH_TOKEN_MISSING", async () => {
+    vi.useFakeTimers();
+    window.localStorage.clear();
+
+    const client = new GatewayBrowserClient({
+      url: "ws://127.0.0.1:18789",
+    });
+
+    client.start();
+    const ws1 = getLatestWebSocket();
+    ws1.emitOpen();
+    ws1.emitMessage({
+      type: "event",
+      event: "connect.challenge",
+      payload: { nonce: "nonce-1" },
+    });
+    await vi.waitFor(() => expect(ws1.sent.length).toBeGreaterThan(0));
+    const connect = JSON.parse(ws1.sent.at(-1) ?? "{}") as { id: string };
+
+    ws1.emitMessage({
+      type: "res",
+      id: connect.id,
+      ok: false,
+      error: {
+        code: "INVALID_REQUEST",
+        message: "unauthorized",
+        details: { code: "AUTH_TOKEN_MISSING" },
+      },
+    });
+    await vi.waitFor(() => expect(ws1.readyState).toBe(3));
+    ws1.emitClose(4008, "connect failed");
+
+    await vi.advanceTimersByTimeAsync(30_000);
+    expect(wsInstances).toHaveLength(1);
+
+    vi.useRealTimers();
+  });
 });
--- a/ui/src/ui/gateway.ts
+++ b/ui/src/ui/gateway.ts
@@ -7,6 +7,7 @@ import {
 } from "../../../src/gateway/protocol/client-info.js";
 import {
  ConnectErrorDetailCodes,
+  readConnectErrorRecoveryAdvice,
  readConnectErrorDetailCode,
 } from "../../../src/gateway/protocol/connect-error-details.js";
 import { clearDeviceAuthToken, loadDeviceAuthToken, storeDeviceAuthToken } from "./device-auth.ts";
@@ -57,11 +58,9 @@ export function resolveGatewayErrorDetailCode(
 * Auth errors that won't resolve without user action — don't auto-reconnect.
 *
 * NOTE: AUTH_TOKEN_MISMATCH is intentionally NOT included here because the
- * browser client has a device-token fallback flow: a stale cached device token
- * triggers a mismatch, sendConnect() clears it, and the next reconnect retries
- * with opts.token (the shared gateway token). Blocking reconnect on mismatch
- * would break that fallback. The rate limiter still catches persistent wrong
- * tokens after N failures → AUTH_RATE_LIMITED stops the loop.
+ * browser client supports a bounded one-time retry with a cached device token
+ * when the endpoint is trusted. Reconnect suppression for mismatch is handled
+ * with client state (after retry budget is exhausted).
 */
 export function isNonRecoverableAuthError(error: GatewayErrorInfo | undefined): boolean {
  if (!error) {
@@ -72,10 +71,30 @@ export function isNonRecoverableAuthError(error: GatewayErrorInfo | undefined):
    code === ConnectErrorDetailCodes.AUTH_TOKEN_MISSING ||
    code === ConnectErrorDetailCodes.AUTH_PASSWORD_MISSING ||
    code === ConnectErrorDetailCodes.AUTH_PASSWORD_MISMATCH ||
-    code === ConnectErrorDetailCodes.AUTH_RATE_LIMITED
+    code === ConnectErrorDetailCodes.AUTH_RATE_LIMITED ||
+    code === ConnectErrorDetailCodes.PAIRING_REQUIRED ||
+    code === ConnectErrorDetailCodes.CONTROL_UI_DEVICE_IDENTITY_REQUIRED ||
+    code === ConnectErrorDetailCodes.DEVICE_IDENTITY_REQUIRED
  );
 }

+function isTrustedRetryEndpoint(url: string): boolean {
+  try {
+    const gatewayUrl = new URL(url, window.location.href);
+    const host = gatewayUrl.hostname.trim().toLowerCase();
+    const isLoopbackHost =
+      host === "localhost" || host === "::1" || host === "[::1]" || host === "127.0.0.1";
+    const isLoopbackIPv4 = host.startsWith("127.");
+    if (isLoopbackHost || isLoopbackIPv4) {
+      return true;
+    }
+    const pageUrl = new URL(window.location.href);
+    return gatewayUrl.host === pageUrl.host;
+  } catch {
+    return false;
+  }
+}
+
 export type GatewayHelloOk = {
  type: "hello-ok";
  protocol: number;
@@ -127,6 +146,8 @@ export class GatewayBrowserClient {
  private connectTimer: number | null = null;
  private backoffMs = 800;
  private pendingConnectError: GatewayErrorInfo | undefined;
+  private pendingDeviceTokenRetry = false;
+  private deviceTokenRetryBudgetUsed = false;

  constructor(private opts: GatewayBrowserClientOptions) {}

@@ -140,6 +161,8 @@ export class GatewayBrowserClient {
    this.ws?.close();
    this.ws = null;
    this.pendingConnectError = undefined;
+    this.pendingDeviceTokenRetry = false;
+    this.deviceTokenRetryBudgetUsed = false;
    this.flushPending(new Error("gateway client stopped"));
  }

@@ -161,6 +184,14 @@ export class GatewayBrowserClient {
      this.ws = null;
      this.flushPending(new Error(`gateway closed (${ev.code}): ${reason}`));
      this.opts.onClose?.({ code: ev.code, reason, error: connectError });
+      const connectErrorCode = resolveGatewayErrorDetailCode(connectError);
+      if (
+        connectErrorCode === ConnectErrorDetailCodes.AUTH_TOKEN_MISMATCH &&
+        this.deviceTokenRetryBudgetUsed &&
+        !this.pendingDeviceTokenRetry
+      ) {
+        return;
+      }
      if (!isNonRecoverableAuthError(connectError)) {
        this.scheduleReconnect();
      }
@@ -215,9 +246,20 @@ export class GatewayBrowserClient {
        deviceId: deviceIdentity.deviceId,
        role,
      })?.token;
-      deviceToken = !(explicitGatewayToken || this.opts.password?.trim())
-        ? (storedToken ?? undefined)
-        : undefined;
+      const shouldUseDeviceRetryToken =
+        this.pendingDeviceTokenRetry &&
+        !deviceToken &&
+        Boolean(explicitGatewayToken) &&
+        Boolean(storedToken) &&
+        isTrustedRetryEndpoint(this.opts.url);
+      if (shouldUseDeviceRetryToken) {
+        deviceToken = storedToken ?? undefined;
+        this.pendingDeviceTokenRetry = false;
+      } else {
+        deviceToken = !(explicitGatewayToken || this.opts.password?.trim())
+          ? (storedToken ?? undefined)
+          : undefined;
+      }
      canFallbackToShared = Boolean(deviceToken && explicitGatewayToken);
    }
    authToken = explicitGatewayToken ?? deviceToken;
@@ -225,6 +267,7 @@ export class GatewayBrowserClient {
      authToken || this.opts.password
        ? {
            token: authToken,
+            deviceToken,
            password: this.opts.password,
          }
        : undefined;
@@ -282,6 +325,8 @@ export class GatewayBrowserClient {

    void this.request<GatewayHelloOk>("connect", params)
      .then((hello) => {
+        this.pendingDeviceTokenRetry = false;
+        this.deviceTokenRetryBudgetUsed = false;
        if (hello?.auth?.deviceToken && deviceIdentity) {
          storeDeviceAuthToken({
            deviceId: deviceIdentity.deviceId,
@@ -294,6 +339,33 @@ export class GatewayBrowserClient {
        this.opts.onHello?.(hello);
      })
      .catch((err: unknown) => {
+        const connectErrorCode =
+          err instanceof GatewayRequestError ? resolveGatewayErrorDetailCode(err) : null;
+        const recoveryAdvice =
+          err instanceof GatewayRequestError ? readConnectErrorRecoveryAdvice(err.details) : {};
+        const retryWithDeviceTokenRecommended =
+          recoveryAdvice.recommendedNextStep === "retry_with_device_token";
+        const canRetryWithDeviceTokenHint =
+          recoveryAdvice.canRetryWithDeviceToken === true ||
+          retryWithDeviceTokenRecommended ||
+          connectErrorCode === ConnectErrorDetailCodes.AUTH_TOKEN_MISMATCH;
+        const shouldRetryWithDeviceToken =
+          !this.deviceTokenRetryBudgetUsed &&
+          !deviceToken &&
+          Boolean(explicitGatewayToken) &&
+          Boolean(deviceIdentity) &&
+          Boolean(
+            loadDeviceAuthToken({
+              deviceId: deviceIdentity?.deviceId ?? "",
+              role,
+            })?.token,
+          ) &&
+          canRetryWithDeviceTokenHint &&
+          isTrustedRetryEndpoint(this.opts.url);
+        if (shouldRetryWithDeviceToken) {
+          this.pendingDeviceTokenRetry = true;
+          this.deviceTokenRetryBudgetUsed = true;
+        }
        if (err instanceof GatewayRequestError) {
          this.pendingConnectError = {
            code: err.gatewayCode,
@@ -303,7 +375,11 @@ export class GatewayBrowserClient {
        } else {
          this.pendingConnectError = undefined;
        }
-        if (canFallbackToShared && deviceIdentity) {
+        if (
+          canFallbackToShared &&
+          deviceIdentity &&
+          connectErrorCode === ConnectErrorDetailCodes.AUTH_DEVICE_TOKEN_MISMATCH
+        ) {
          clearDeviceAuthToken({ deviceId: deviceIdentity.deviceId, role });
        }
        this.ws?.close(CONNECT_FAILED_CLOSE_CODE, "connect failed");