refactor(gateway): unify auth credential resolution

2026-05-08 17:08:27 +00:00 · 2026-02-22 18:21:20 +01:00
parent ded9a59f78
commit 66529c7aa5
13 changed files with 537 additions and 126 deletions
--- a/src/gateway/server/ws-connection/auth-context.ts
+++ b/src/gateway/server/ws-connection/auth-context.ts
@@ -0,0 +1,133 @@
+import type { IncomingMessage } from "node:http";
+import {
+  AUTH_RATE_LIMIT_SCOPE_SHARED_SECRET,
+  type AuthRateLimiter,
+  type RateLimitCheckResult,
+} from "../../auth-rate-limit.js";
+import {
+  authorizeHttpGatewayConnect,
+  authorizeWsControlUiGatewayConnect,
+  type GatewayAuthResult,
+  type ResolvedGatewayAuth,
+} from "../../auth.js";
+
+type HandshakeConnectAuth = {
+  token?: string;
+  deviceToken?: string;
+  password?: string;
+};
+
+export type ConnectAuthState = {
+  authResult: GatewayAuthResult;
+  authOk: boolean;
+  authMethod: GatewayAuthResult["method"];
+  sharedAuthOk: boolean;
+  sharedAuthProvided: boolean;
+  deviceTokenCandidate?: string;
+};
+
+function trimToUndefined(value: string | undefined): string | undefined {
+  if (!value) {
+    return undefined;
+  }
+  const trimmed = value.trim();
+  return trimmed.length > 0 ? trimmed : undefined;
+}
+
+function resolveSharedConnectAuth(
+  connectAuth: HandshakeConnectAuth | null | undefined,
+): { token?: string; password?: string } | undefined {
+  const token = trimToUndefined(connectAuth?.token);
+  const password = trimToUndefined(connectAuth?.password);
+  if (!token && !password) {
+    return undefined;
+  }
+  return { token, password };
+}
+
+function resolveDeviceTokenCandidate(
+  connectAuth: HandshakeConnectAuth | null | undefined,
+): string | undefined {
+  const explicitDeviceToken = trimToUndefined(connectAuth?.deviceToken);
+  if (explicitDeviceToken) {
+    return explicitDeviceToken;
+  }
+  return trimToUndefined(connectAuth?.token);
+}
+
+export async function resolveConnectAuthState(params: {
+  resolvedAuth: ResolvedGatewayAuth;
+  connectAuth: HandshakeConnectAuth | null | undefined;
+  hasDeviceIdentity: boolean;
+  req: IncomingMessage;
+  trustedProxies: string[];
+  allowRealIpFallback: boolean;
+  rateLimiter?: AuthRateLimiter;
+  clientIp?: string;
+}): Promise<ConnectAuthState> {
+  const sharedConnectAuth = resolveSharedConnectAuth(params.connectAuth);
+  const sharedAuthProvided = Boolean(sharedConnectAuth);
+  const deviceTokenCandidate = params.hasDeviceIdentity
+    ? resolveDeviceTokenCandidate(params.connectAuth)
+    : undefined;
+  const hasDeviceTokenCandidate = Boolean(deviceTokenCandidate);
+
+  let authResult: GatewayAuthResult = await authorizeWsControlUiGatewayConnect({
+    auth: params.resolvedAuth,
+    connectAuth: sharedConnectAuth,
+    req: params.req,
+    trustedProxies: params.trustedProxies,
+    allowRealIpFallback: params.allowRealIpFallback,
+    rateLimiter: hasDeviceTokenCandidate ? undefined : params.rateLimiter,
+    clientIp: params.clientIp,
+    rateLimitScope: AUTH_RATE_LIMIT_SCOPE_SHARED_SECRET,
+  });
+
+  if (
+    hasDeviceTokenCandidate &&
+    authResult.ok &&
+    params.rateLimiter &&
+    (authResult.method === "token" || authResult.method === "password")
+  ) {
+    const sharedRateCheck: RateLimitCheckResult = params.rateLimiter.check(
+      params.clientIp,
+      AUTH_RATE_LIMIT_SCOPE_SHARED_SECRET,
+    );
+    if (!sharedRateCheck.allowed) {
+      authResult = {
+        ok: false,
+        reason: "rate_limited",
+        rateLimited: true,
+        retryAfterMs: sharedRateCheck.retryAfterMs,
+      };
+    } else {
+      params.rateLimiter.reset(params.clientIp, AUTH_RATE_LIMIT_SCOPE_SHARED_SECRET);
+    }
+  }
+
+  const sharedAuthResult =
+    sharedConnectAuth &&
+    (await authorizeHttpGatewayConnect({
+      auth: { ...params.resolvedAuth, allowTailscale: false },
+      connectAuth: sharedConnectAuth,
+      req: params.req,
+      trustedProxies: params.trustedProxies,
+      allowRealIpFallback: params.allowRealIpFallback,
+      // Shared-auth probe only; rate-limit side effects are handled in the
+      // primary auth flow (or deferred for device-token candidates).
+      rateLimitScope: AUTH_RATE_LIMIT_SCOPE_SHARED_SECRET,
+    }));
+  const sharedAuthOk =
+    sharedAuthResult?.ok === true &&
+    (sharedAuthResult.method === "token" || sharedAuthResult.method === "password");
+
+  return {
+    authResult,
+    authOk: authResult.ok,
+    authMethod:
+      authResult.method ?? (params.resolvedAuth.mode === "password" ? "password" : "token"),
+    sharedAuthOk,
+    sharedAuthProvided,
+    deviceTokenCandidate,
+  };
+}
--- a/src/gateway/server/ws-connection/auth-messages.ts
+++ b/src/gateway/server/ws-connection/auth-messages.ts
@@ -2,7 +2,7 @@ import { isGatewayCliClient, isWebchatClient } from "../../../utils/message-chan
 import type { ResolvedGatewayAuth } from "../../auth.js";
 import { GATEWAY_CLIENT_IDS } from "../../protocol/client-info.js";

-export type AuthProvidedKind = "token" | "password" | "none";
+export type AuthProvidedKind = "token" | "device-token" | "password" | "none";

 export function formatGatewayAuthFailureMessage(params: {
  authMode: ResolvedGatewayAuth["mode"];
@@ -57,6 +57,9 @@ export function formatGatewayAuthFailureMessage(params: {
  if (authMode === "token" && authProvided === "none") {
    return `unauthorized: gateway token missing (${tokenHint})`;
  }
+  if (authMode === "token" && authProvided === "device-token") {
+    return "unauthorized: device token rejected (pair/repair this device, or provide gateway token)";
+  }
  if (authMode === "password" && authProvided === "none") {
    return `unauthorized: gateway password missing (${passwordHint})`;
  }
--- a/src/gateway/server/ws-connection/message-handler.ts
+++ b/src/gateway/server/ws-connection/message-handler.ts
@@ -24,17 +24,9 @@ import type { createSubsystemLogger } from "../../../logging/subsystem.js";
 import { roleScopesAllow } from "../../../shared/operator-scope-compat.js";
 import { isGatewayCliClient, isWebchatClient } from "../../../utils/message-channel.js";
 import { resolveRuntimeServiceVersion } from "../../../version.js";
-import {
-  AUTH_RATE_LIMIT_SCOPE_DEVICE_TOKEN,
-  AUTH_RATE_LIMIT_SCOPE_SHARED_SECRET,
-  type AuthRateLimiter,
-} from "../../auth-rate-limit.js";
+import { AUTH_RATE_LIMIT_SCOPE_DEVICE_TOKEN, type AuthRateLimiter } from "../../auth-rate-limit.js";
 import type { GatewayAuthResult, ResolvedGatewayAuth } from "../../auth.js";
-import {
-  authorizeHttpGatewayConnect,
-  authorizeWsControlUiGatewayConnect,
-  isLocalDirectRequest,
-} from "../../auth.js";
+import { isLocalDirectRequest } from "../../auth.js";
 import {
  buildCanvasScopedHostUrl,
  CANVAS_CAPABILITY_TTL_MS,
@@ -75,6 +67,7 @@ import {
  refreshGatewayHealthSnapshot,
 } from "../health-state.js";
 import type { GatewayWsClient } from "../ws-types.js";
+import { resolveConnectAuthState } from "./auth-context.js";
 import { formatGatewayAuthFailureMessage, type AuthProvidedKind } from "./auth-messages.js";
 import {
  evaluateMissingDeviceIdentity,
@@ -362,87 +355,40 @@ export function attachGatewayWsMessageHandler(params: {
        });
        const device = controlUiAuthPolicy.device;

-        const resolveAuthState = async () => {
-          const hasDeviceTokenCandidate = Boolean(connectParams.auth?.token && device);
-          let nextAuthResult: GatewayAuthResult = await authorizeWsControlUiGatewayConnect({
-            auth: resolvedAuth,
+        let { authResult, authOk, authMethod, sharedAuthOk, deviceTokenCandidate } =
+          await resolveConnectAuthState({
+            resolvedAuth,
            connectAuth: connectParams.auth,
+            hasDeviceIdentity: Boolean(device),
            req: upgradeReq,
            trustedProxies,
            allowRealIpFallback,
-            rateLimiter: hasDeviceTokenCandidate ? undefined : rateLimiter,
+            rateLimiter,
            clientIp,
-            rateLimitScope: AUTH_RATE_LIMIT_SCOPE_SHARED_SECRET,
          });
-
-          if (
-            hasDeviceTokenCandidate &&
-            nextAuthResult.ok &&
-            rateLimiter &&
-            (nextAuthResult.method === "token" || nextAuthResult.method === "password")
-          ) {
-            const sharedRateCheck = rateLimiter.check(
-              clientIp,
-              AUTH_RATE_LIMIT_SCOPE_SHARED_SECRET,
-            );
-            if (!sharedRateCheck.allowed) {
-              nextAuthResult = {
-                ok: false,
-                reason: "rate_limited",
-                rateLimited: true,
-                retryAfterMs: sharedRateCheck.retryAfterMs,
-              };
-            } else {
-              rateLimiter.reset(clientIp, AUTH_RATE_LIMIT_SCOPE_SHARED_SECRET);
-            }
-          }
-
-          const nextAuthMethod =
-            nextAuthResult.method ?? (resolvedAuth.mode === "password" ? "password" : "token");
-          const sharedAuthResult = hasSharedAuth
-            ? await authorizeHttpGatewayConnect({
-                auth: { ...resolvedAuth, allowTailscale: false },
-                connectAuth: connectParams.auth,
-                req: upgradeReq,
-                trustedProxies,
-                allowRealIpFallback,
-                // Shared-auth probe only; rate-limit side effects are handled in
-                // the primary auth flow (or deferred for device-token candidates).
-                rateLimitScope: AUTH_RATE_LIMIT_SCOPE_SHARED_SECRET,
-              })
-            : null;
-          const nextSharedAuthOk =
-            sharedAuthResult?.ok === true &&
-            (sharedAuthResult.method === "token" || sharedAuthResult.method === "password");
-
-          return {
-            authResult: nextAuthResult,
-            authOk: nextAuthResult.ok,
-            authMethod: nextAuthMethod,
-            sharedAuthOk: nextSharedAuthOk,
-          };
-        };
-
-        let { authResult, authOk, authMethod, sharedAuthOk } = await resolveAuthState();
        const rejectUnauthorized = (failedAuth: GatewayAuthResult) => {
          markHandshakeFailure("unauthorized", {
            authMode: resolvedAuth.mode,
-            authProvided: connectParams.auth?.token
-              ? "token"
-              : connectParams.auth?.password
-                ? "password"
-                : "none",
+            authProvided: connectParams.auth?.password
+              ? "password"
+              : connectParams.auth?.token
+                ? "token"
+                : connectParams.auth?.deviceToken
+                  ? "device-token"
+                  : "none",
            authReason: failedAuth.reason,
            allowTailscale: resolvedAuth.allowTailscale,
          });
          logWsControl.warn(
            `unauthorized conn=${connId} remote=${remoteAddr ?? "?"} client=${clientLabel} ${connectParams.client.mode} v${connectParams.client.version} reason=${failedAuth.reason ?? "unknown"}`,
          );
-          const authProvided: AuthProvidedKind = connectParams.auth?.token
-            ? "token"
-            : connectParams.auth?.password
-              ? "password"
-              : "none";
+          const authProvided: AuthProvidedKind = connectParams.auth?.password
+            ? "password"
+            : connectParams.auth?.token
+              ? "token"
+              : connectParams.auth?.deviceToken
+                ? "device-token"
+                : "none";
          const authMessage = formatGatewayAuthFailureMessage({
            authMode: resolvedAuth.mode,
            authProvided,
@@ -545,7 +491,7 @@ export function attachGatewayWsMessageHandler(params: {
            role,
            scopes,
            signedAtMs: signedAt,
-            token: connectParams.auth?.token ?? null,
+            token: connectParams.auth?.token ?? connectParams.auth?.deviceToken ?? null,
            nonce: providedNonce,
          });
          const rejectDeviceSignatureInvalid = () =>
@@ -562,7 +508,7 @@ export function attachGatewayWsMessageHandler(params: {
          }
        }

-        if (!authOk && connectParams.auth?.token && device) {
+        if (!authOk && device && deviceTokenCandidate) {
          if (rateLimiter) {
            const deviceRateCheck = rateLimiter.check(clientIp, AUTH_RATE_LIMIT_SCOPE_DEVICE_TOKEN);
            if (!deviceRateCheck.allowed) {
@@ -577,7 +523,7 @@ export function attachGatewayWsMessageHandler(params: {
          if (!authResult.rateLimited) {
            const tokenCheck = await verifyDeviceToken({
              deviceId: device.id,
-              token: connectParams.auth.token,
+              token: deviceTokenCandidate,
              role,
              scopes,
            });