fix: add auth token modal for browser mode authentication

ThomasK33 · ThomasK33 · commit 175a409abc44 · 2025-11-26T17:50:59.000+01:00
When the server requires authentication (--auth-token), the browser client
now shows a modal prompting the user to enter the auth token. The token is:
- Stored in localStorage for subsequent visits
- Can also be passed via URL query parameter (?token=...)
- Cleared and re-prompted if authentication fails

This replaces the previous server-side injection approach with a cleaner
user-driven authentication flow.

Change-Id: I5599266df30340bcc5ca016a14a67a5d74c52669
Signed-off-by: Thomas Kosiewski &lt;tk@coder.com&gt;
diff --git a/src/browser/components/AuthTokenModal.tsx b/src/browser/components/AuthTokenModal.tsx
@@ -0,0 +1,111 @@
+import { useState, useCallback } from "react";
+import { Modal } from "./Modal";
+
+interface AuthTokenModalProps {
+  isOpen: boolean;
+  onSubmit: (token: string) => void;
+  error?: string | null;
+}
+
+const AUTH_TOKEN_STORAGE_KEY = "mux:auth-token";
+
+export function getStoredAuthToken(): string | null {
+  try {
+    return localStorage.getItem(AUTH_TOKEN_STORAGE_KEY);
+  } catch {
+    return null;
+  }
+}
+
+export function setStoredAuthToken(token: string): void {
+  try {
+    localStorage.setItem(AUTH_TOKEN_STORAGE_KEY, token);
+  } catch {
+    // Ignore storage errors
+  }
+}
+
+export function clearStoredAuthToken(): void {
+  try {
+    localStorage.removeItem(AUTH_TOKEN_STORAGE_KEY);
+  } catch {
+    // Ignore storage errors
+  }
+}
+
+export function AuthTokenModal(props: AuthTokenModalProps) {
+  const [token, setToken] = useState("");
+
+  const { onSubmit } = props;
+  const handleSubmit = useCallback(
+    (e: React.FormEvent) => {
+      e.preventDefault();
+      if (token.trim()) {
+        setStoredAuthToken(token.trim());
+        onSubmit(token.trim());
+      }
+    },
+    [token, onSubmit]
+  );
+
+  return (
+    <Modal isOpen={props.isOpen} onClose={() => undefined} title="Authentication Required">
+      <form onSubmit={handleSubmit} style={{ display: "flex", flexDirection: "column", gap: 16 }}>
+        <p style={{ margin: 0, color: "var(--color-text-secondary)" }}>
+          This server requires an authentication token. Enter the token provided when the server was
+          started.
+        </p>
+
+        {props.error && (
+          <div
+            style={{
+              padding: "8px 12px",
+              borderRadius: 4,
+              backgroundColor: "var(--color-error-background, rgba(255, 0, 0, 0.1))",
+              color: "var(--color-error, #ff6b6b)",
+              fontSize: 13,
+            }}
+          >
+            {props.error}
+          </div>
+        )}
+
+        <input
+          type="password"
+          value={token}
+          onChange={(e) => setToken(e.target.value)}
+          placeholder="Enter auth token"
+          autoFocus
+          style={{
+            padding: "10px 12px",
+            borderRadius: 4,
+            border: "1px solid var(--color-border)",
+            backgroundColor: "var(--color-input-background)",
+            color: "var(--color-text)",
+            fontSize: 14,
+            outline: "none",
+          }}
+        />
+
+        <button
+          type="submit"
+          disabled={!token.trim()}
+          style={{
+            padding: "10px 16px",
+            borderRadius: 4,
+            border: "none",
+            backgroundColor: token.trim()
+              ? "var(--color-primary)"
+              : "var(--color-button-disabled-background)",
+            color: token.trim() ? "white" : "var(--color-text-disabled)",
+            fontSize: 14,
+            fontWeight: 500,
+            cursor: token.trim() ? "pointer" : "not-allowed",
+          }}
+        >
+          Connect
+        </button>
+      </form>
+    </Modal>
+  );
+}
diff --git a/src/browser/orpc/react.tsx b/src/browser/orpc/react.tsx
@@ -1,9 +1,14 @@
-import { createContext, useContext, useEffect, useState } from "react";
+import { createContext, useContext, useEffect, useState, useCallback } from "react";
 import { createClient } from "@/common/orpc/client";
 import { RPCLink as WebSocketLink } from "@orpc/client/websocket";
 import { RPCLink as MessagePortLink } from "@orpc/client/message-port";
 import type { AppRouter } from "@/node/orpc/router";
 import type { RouterClient } from "@orpc/server";
+import {
+  AuthTokenModal,
+  getStoredAuthToken,
+  clearStoredAuthToken,
+} from "@/browser/components/AuthTokenModal";
 
 type ORPCClient = ReturnType<typeof createClient>;
 
@@ -17,73 +22,196 @@ interface ORPCProviderProps {
   client?: ORPCClient;
 }
 
-export const ORPCProvider = (props: ORPCProviderProps) => {
-  const [client, setClient] = useState<ORPCClient | null>(props.client ?? null);
+type ConnectionState =
+  | { status: "connecting" }
+  | { status: "connected"; client: ORPCClient; cleanup: () => void }
+  | { status: "auth_required"; error?: string }
+  | { status: "error"; error: string };
 
-  useEffect(() => {
-    // If client provided externally, use it directly
-    if (props.client) {
-      setClient(() => props.client!);
-      window.__ORPC_CLIENT__ = props.client;
-      return;
-    }
-
-    let cleanup: () => void;
-    let newClient: ORPCClient;
-
-    // Detect Electron mode by checking if window.api exists (exposed by preload script)
-    // window.api.platform contains the actual OS platform (darwin/win32/linux), not "electron"
-    if (window.api) {
-      // Electron Mode: Use MessageChannel
-      const { port1: clientPort, port2: serverPort } = new MessageChannel();
-
-      // Send port to preload/main
-      window.postMessage("start-orpc-client", "*", [serverPort]);
-
-      const link = new MessagePortLink({
-        port: clientPort,
-      });
-      clientPort.start();
-
-      newClient = createClient(link);
-      cleanup = () => {
-        clientPort.close();
-      };
-    } else {
-      // Browser Mode: Use HTTP/WebSocket
-      // Assume server is at same origin or configured via VITE_BACKEND_URL
-      // eslint-disable-next-line @typescript-eslint/ban-ts-comment, @typescript-eslint/prefer-ts-expect-error
-      // @ts-ignore - import.meta is available in Vite
-      const API_BASE = import.meta.env.VITE_BACKEND_URL ?? window.location.origin;
-      const WS_BASE = API_BASE.replace("http://", "ws://").replace("https://", "wss://");
-
-      const ws = new WebSocket(`${WS_BASE}/orpc/ws`);
-      const link = new WebSocketLink({
-        websocket: ws,
+function getApiBase(): string {
+  // eslint-disable-next-line @typescript-eslint/ban-ts-comment, @typescript-eslint/prefer-ts-expect-error
+  // @ts-ignore - import.meta is available in Vite
+  return import.meta.env.VITE_BACKEND_URL ?? window.location.origin;
+}
+
+function createElectronClient(): { client: ORPCClient; cleanup: () => void } {
+  const { port1: clientPort, port2: serverPort } = new MessageChannel();
+  window.postMessage("start-orpc-client", "*", [serverPort]);
+
+  const link = new MessagePortLink({ port: clientPort });
+  clientPort.start();
+
+  return {
+    client: createClient(link),
+    cleanup: () => clientPort.close(),
+  };
+}
+
+function createBrowserClient(authToken: string | null): {
+  client: ORPCClient;
+  cleanup: () => void;
+  ws: WebSocket;
+} {
+  const API_BASE = getApiBase();
+  const WS_BASE = API_BASE.replace("http://", "ws://").replace("https://", "wss://");
+
+  const wsUrl = authToken
+    ? `${WS_BASE}/orpc/ws?token=${encodeURIComponent(authToken)}`
+    : `${WS_BASE}/orpc/ws`;
+
+  const ws = new WebSocket(wsUrl);
+  const link = new WebSocketLink({ websocket: ws });
+
+  return {
+    client: createClient(link),
+    cleanup: () => ws.close(),
+    ws,
+  };
+}
+
+export const ORPCProvider = (props: ORPCProviderProps) => {
+  const [state, setState] = useState<ConnectionState>({ status: "connecting" });
+  const [authToken, setAuthToken] = useState<string | null>(() => {
+    // Check URL param first, then localStorage
+    const urlParams = new URLSearchParams(window.location.search);
+    return urlParams.get("token") ?? getStoredAuthToken();
+  });
+
+  const connect = useCallback(
+    (token: string | null) => {
+      // If client provided externally, use it directly
+      if (props.client) {
+        window.__ORPC_CLIENT__ = props.client;
+        setState({ status: "connected", client: props.client, cleanup: () => undefined });
+        return;
+      }
+
+      // Electron mode - no auth needed
+      if (window.api) {
+        const { client, cleanup } = createElectronClient();
+        window.__ORPC_CLIENT__ = client;
+        setState({ status: "connected", client, cleanup });
+        return;
+      }
+
+      // Browser mode - connect with optional auth token
+      setState({ status: "connecting" });
+      const { client, cleanup, ws } = createBrowserClient(token);
+
+      ws.addEventListener("open", () => {
+        // Connection successful - test with a ping to verify auth
+        client.general
+          .ping("auth-check")
+          .then(() => {
+            window.__ORPC_CLIENT__ = client;
+            setState({ status: "connected", client, cleanup });
+          })
+          .catch((err: unknown) => {
+            cleanup();
+            const errMsg = err instanceof Error ? err.message : String(err);
+            if (errMsg.includes("UNAUTHORIZED") || errMsg.includes("401")) {
+              clearStoredAuthToken();
+              setState({ status: "auth_required", error: token ? "Invalid token" : undefined });
+            } else {
+              setState({ status: "error", error: errMsg });
+            }
+          });
       });
 
-      newClient = createClient(link);
-      cleanup = () => {
-        ws.close();
-      };
-    }
+      ws.addEventListener("error", () => {
+        // WebSocket connection failed - might be auth issue or network
+        cleanup();
+        // If we had a token and failed, likely auth issue
+        if (token) {
+          clearStoredAuthToken();
+          setState({ status: "auth_required", error: "Connection failed - invalid token?" });
+        } else {
+          // Try without token first, server might not require auth
+          // If server requires auth, the ping will fail with UNAUTHORIZED
+          setState({ status: "auth_required" });
+        }
+      });
 
-    // Pass a function to setClient to prevent React from treating the client (which is a callable Proxy)
-    // as a functional state update. Without this, React calls client(prevState), triggering a request to root /.
-    setClient(() => newClient);
+      ws.addEventListener("close", (event) => {
+        // 1008 = Policy Violation (often used for auth failures)
+        // 4401 = Custom unauthorized code
+        if (event.code === 1008 || event.code === 4401) {
+          cleanup();
+          clearStoredAuthToken();
+          setState({ status: "auth_required", error: "Authentication required" });
+        }
+      });
+    },
+    [props.client]
+  );
 
-    window.__ORPC_CLIENT__ = newClient;
+  // Initial connection attempt
+  useEffect(() => {
+    connect(authToken);
 
     return () => {
-      cleanup();
+      if (state.status === "connected") {
+        state.cleanup();
+      }
     };
-  }, [props.client]);
+    // Only run on mount and when authToken changes via handleAuthSubmit
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, []);
+
+  const handleAuthSubmit = useCallback(
+    (token: string) => {
+      setAuthToken(token);
+      connect(token);
+    },
+    [connect]
+  );
+
+  // Show auth modal if auth is required
+  if (state.status === "auth_required") {
+    return (
+      <AuthTokenModal isOpen={true} onSubmit={handleAuthSubmit} error={state.error ?? null} />
+    );
+  }
+
+  // Show error state
+  if (state.status === "error") {
+    return (
+      <div
+        style={{
+          display: "flex",
+          alignItems: "center",
+          justifyContent: "center",
+          height: "100vh",
+          color: "var(--color-error, #ff6b6b)",
+          flexDirection: "column",
+          gap: 16,
+        }}
+      >
+        <div>Failed to connect to server</div>
+        <div style={{ fontSize: 13, color: "var(--color-text-secondary)" }}>{state.error}</div>
+        <button
+          onClick={() => connect(authToken)}
+          style={{
+            padding: "8px 16px",
+            borderRadius: 4,
+            border: "1px solid var(--color-border)",
+            background: "var(--color-button-background)",
+            color: "var(--color-text)",
+            cursor: "pointer",
+          }}
+        >
+          Retry
+        </button>
+      </div>
+    );
+  }
 
-  if (!client) {
+  // Show loading while connecting
+  if (state.status === "connecting") {
     return null; // Or a loading spinner
   }
 
-  return <ORPCContext.Provider value={client}>{props.children}</ORPCContext.Provider>;
+  return <ORPCContext.Provider value={state.client}>{props.children}</ORPCContext.Provider>;
 };
 
 export const useORPC = (): RouterClient<AppRouter> => {
