CL-6391: unify MCP OAuth state onto ConnectStateStore

packages/connections/src/mcp-oauth-routes.test.ts

@@ -518,6 +518,64 @@ describe("MCP OAuth connect flow", () => {
     }
   });
 
+  test("a replayed callback is rejected as state_expired without a second exchange (one-shot state)", async () => {
+    const as = startStubAuthorizationServer();
+    try {
+      const hub = fakeHub();
+      const routes = createMcpOAuthRoutes({
+        hubUrl: "http://hub.test",
+        requireGrant: allowAll,
+        log: () => {},
+        credentialCipher: createNoopCredentialCipher(),
+        apiCall: hub.apiCall,
+        probe: async (): Promise<McpProbeResult> => ({
+          ok: true,
+          toolCount: 2,
+        }),
+      });
+      const app = mountAs(routes);
+
+      const startResponse = await app.request(
+        `/exa/start?url=${encodeURIComponent(as.resourcePath)}&name=Exa`,
+        { redirect: "manual" },
+      );
+      const cookieHeader = startResponse.headers.get("set-cookie") ?? "";
+      const cookie = cookieHeader.split(";")[0] ?? "";
+      const authorizeResponse = await fetch(
+        startResponse.headers.get("location") ?? "",
+        { redirect: "manual" },
+      );
+      const callbackUrl = new URL(
+        authorizeResponse.headers.get("location") ?? "",
+      );
+      const replayableCallback = `${callbackUrl.pathname}${callbackUrl.search}`;
+
+      const firstResponse = await app.request(replayableCallback, {
+        headers: { cookie },
+        redirect: "manual",
+      });
+      expect(firstResponse.headers.get("location") ?? "").toContain(
+        "outcome=connected",
+      );
+
+      // A browser (or an attacker holding a stolen state cookie)
+      // presenting the exact same callback again: the sealed state was
+      // burned by the first arrival, so the replay dies at the state
+      // check — it never reaches the token endpoint a second time.
+      const replayResponse = await app.request(replayableCallback, {
+        headers: { cookie },
+        redirect: "manual",
+      });
+      expect(replayResponse.status).toBe(302);
+      expect(replayResponse.headers.get("location") ?? "").toContain(
+        "code=state_expired",
+      );
+      expect(hub.credentials).toHaveLength(1);
+    } finally {
+      as.stop();
+    }
+  });
+
   test("a callback with no cookie redirects with a state_expired error", async () => {
     const hub = fakeHub();
     const routes = createMcpOAuthRoutes({

packages/connections/src/mcp-oauth-routes.ts

@@ -31,6 +31,7 @@ import {
 import { createHubAPI } from "@workbench/hub-client";
 import type { OAuthClientInformationMixed } from "@modelcontextprotocol/sdk/shared/auth.js";
 import { createMcpOAuthProvider, type McpOAuthSession } from "./mcp-oauth";
+import { createConnectStateStore, randomToken } from "./pkce";
 import { mcpPresetBySlug } from "./mcp-presets";
 import { probeMcpServer, type McpProbeResult } from "./mcp-probe";
 import {
@@ -51,27 +52,31 @@ const ErrorEnvelope = (code: string, message: string) => ({
 
 const OAUTH_STATE_TTL_MS = 10 * 60 * 1000;
 
+// One provider label for the whole MCP connect surface — the sealed
+// state already carries the target slug/url, and the shared store's AAD
+// separates this flow's states from the fixed-registry connectors'.
+const MCP_OAUTH_STATE_PROVIDER = "mcp-oauth";
+
 const McpOAuthStatePayload = type({
-  principalId: "string > 0",
   slug: "string > 0",
   name: "string > 0",
   url: "string > 0",
   returnPath: "string > 0",
-  nonce: "string > 0",
-  expiresAt: "number",
+  // The `state` value `/start` sent to the authorization server, which
+  // the callback requires echoed back — the CSRF binding, distinct from
+  // the sealed envelope's own single-use replay nonce (the shared
+  // store's concern).
+  oauthState: "string > 0",
   "codeVerifier?": "string",
   "clientInformation?": "unknown",
 });
 type McpOAuthStatePayload = typeof McpOAuthStatePayload.infer;
 
-function randomNonce(): string {
-  const bytes = new Uint8Array(16);
-  crypto.getRandomValues(bytes);
-  return Array.from(bytes, (b) => b.toString(16).padStart(2, "0")).join("");
-}
-
-function stateAad(): string {
-  return JSON.stringify(["mcp-oauth-connect-state"]);
+function parseMcpOAuthStatePayload(
+  value: unknown,
+): McpOAuthStatePayload | undefined {
+  const parsed = McpOAuthStatePayload(value);
+  return parsed instanceof type.errors ? undefined : parsed;
 }
 
 function cookieName(slug: string): string {
@@ -114,6 +119,12 @@ export function createMcpOAuthRoutes(
   const returnPathAllowlist =
     deps.returnPathAllowlist ?? DEFAULT_RETURN_PATH_ALLOWLIST;
   const secureCookies = deps.hubUrl.startsWith("https:");
+  const stateStore = createConnectStateStore({
+    cipher: deps.credentialCipher,
+    provider: MCP_OAUTH_STATE_PROVIDER,
+    parsePayload: parseMcpOAuthStatePayload,
+    ttlMs: OAUTH_STATE_TTL_MS,
+  });
 
   function redirectPath(
     returnPath: string,
@@ -158,7 +169,7 @@ export function createMcpOAuthRoutes(
       // already be on the session by the time `redirectToAuthorization`
       // fires. The same value is what `/callback` requires the provider's
       // `?state=` to match.
-      const nonce = randomNonce();
+      const nonce = randomToken();
       const session: McpOAuthSession = { state: nonce };
       const provider = createMcpOAuthProvider({
         callbackUrl,
@@ -196,24 +207,22 @@ export function createMcpOAuthRoutes(
       }
 
       const payload: McpOAuthStatePayload = {
-        principalId: principal.id,
         slug: target.slug,
         name: target.name,
         url: target.url,
         returnPath,
-        nonce,
-        expiresAt: Date.now() + OAUTH_STATE_TTL_MS,
+        oauthState: nonce,
         ...(session.codeVerifier !== undefined
           ? { codeVerifier: session.codeVerifier }
           : {}),
         ...(session.clientInformation !== undefined
           ? { clientInformation: session.clientInformation }
           : {}),
       };
-      const sealed = await deps.credentialCipher.encrypt(
-        JSON.stringify(payload),
-        stateAad(),
-      );
+      const sealed = await stateStore.issue({
+        userId: principal.id,
+        payload,
+      });
       setCookie(c, cookieName(target.slug), sealed, {
         httpOnly: true,
         sameSite: "Lax",
@@ -249,16 +258,16 @@ export function createMcpOAuthRoutes(
         );
       }
 
-      let payload: McpOAuthStatePayload;
-      try {
-        const plaintext = await deps.credentialCipher.decrypt(
-          sealed,
-          stateAad(),
-        );
-        const parsed = McpOAuthStatePayload(JSON.parse(plaintext));
-        if (parsed instanceof type.errors) throw new Error(parsed.summary);
-        payload = parsed;
-      } catch {
+      // One-shot: the shared store burns the sealed state on this
+      // attempt (decrypt + AAD + TTL + user binding + replay guard all
+      // inside `consume`), so a replayed callback dies here without a
+      // second token exchange.
+      const principal = c.get("principal");
+      const payload = await stateStore.consume({
+        state: sealed,
+        userId: principal.id,
+      });
+      if (payload === undefined) {
         return c.redirect(
           redirectPath(fallbackReturn, {
             mcpOauth: slugParam,
@@ -274,14 +283,8 @@ export function createMcpOAuthRoutes(
         defaultReturnPath,
         returnPathAllowlist,
       );
-      const principal = c.get("principal");
       const code = c.req.query("code");
-      if (
-        payload.expiresAt <= Date.now() ||
-        payload.principalId !== principal.id ||
-        code === undefined ||
-        code === ""
-      ) {
+      if (code === undefined || code === "") {
         return c.redirect(
           redirectPath(returnPath, {
             mcpOauth: payload.slug,
@@ -297,7 +300,7 @@ export function createMcpOAuthRoutes(
       // must echo that exact value back -- never optional-when-absent.
       // A missing or mismatched `state` means this callback did not
       // originate from the authorize redirect this session minted.
-      if (c.req.query("state") !== payload.nonce) {
+      if (c.req.query("state") !== payload.oauthState) {
         return c.redirect(
           redirectPath(returnPath, {
             mcpOauth: payload.slug,

packages/connections/src/oauth-routes.ts

@@ -32,10 +32,15 @@
 // restart-proofing hardened.
 import { Hono, type Context } from "hono";
 import { deleteCookie, getCookie, setCookie } from "hono/cookie";
+import { type } from "arktype";
 import type { AppEnv } from "@intx/hub-api";
 import type { CredentialCipher } from "@intx/types";
 import { cookiesFromHeader } from "@workbench/hub-client";
-import { createConnectStateStore, generatePKCEPair } from "./pkce";
+import {
+  createConnectStateStore,
+  generatePKCEPair,
+  type ConnectStateStore,
+} from "./pkce";
 import type { ConnectorDescriptor } from "./descriptor";
 import { CONNECTOR_REGISTRY } from "./registry";
 
@@ -208,6 +213,22 @@ export type CreateOAuthConnectRoutesDeps<E extends AppEnv = AppEnv> = {
 const CONNECT_STATE_TTL_MS = 10 * 60 * 1000;
 const CONNECT_START_RATE_LIMIT_MS = 10_000;
 
+const VerifierStatePayload = type({
+  // Empty for a non-PKCE flow (GitHub's confidential-client web flow
+  // seals `codeVerifier: ""`), so this must accept the empty string —
+  // `string > 0` here silently expired every non-PKCE callback
+  // (CL-6394).
+  codeVerifier: "string",
+});
+type VerifierStatePayload = typeof VerifierStatePayload.infer;
+
+function parseVerifierStatePayload(
+  value: unknown,
+): VerifierStatePayload | undefined {
+  const parsed = VerifierStatePayload(value);
+  return parsed instanceof type.errors ? undefined : parsed;
+}
+
 export function createOAuthConnectRoutes<E extends AppEnv = AppEnv>(
   deps: CreateOAuthConnectRoutesDeps<E>,
 ): Hono<E> {
@@ -221,14 +242,15 @@ export function createOAuthConnectRoutes<E extends AppEnv = AppEnv>(
 
   const stateStores = new Map<
     string,
-    ReturnType<typeof createConnectStateStore>
+    ConnectStateStore<VerifierStatePayload>
   >();
   function stateStoreFor(connectorId: string) {
     let store = stateStores.get(connectorId);
     if (store === undefined) {
       store = createConnectStateStore({
         cipher: deps.credentialCipher,
         provider: connectorId,
+        parsePayload: parseVerifierStatePayload,
         ttlMs: CONNECT_STATE_TTL_MS,
       });
       stateStores.set(connectorId, store);
@@ -341,7 +363,7 @@ export function createOAuthConnectRoutes<E extends AppEnv = AppEnv>(
       : undefined;
     const state = await stateStoreFor(connectorId).issue({
       userId: user.id,
-      codeVerifier: pkce?.codeVerifier ?? "",
+      payload: { codeVerifier: pkce?.codeVerifier ?? "" },
     });
     setCookie(c, stateCookieName(connectorId), state, {
       httpOnly: true,
@@ -467,11 +489,11 @@ export function createOAuthConnectRoutes<E extends AppEnv = AppEnv>(
       );
     }
 
-    const codeVerifier = await stateStoreFor(connectorId).consume({
+    const statePayload = await stateStoreFor(connectorId).consume({
       state: cookieState,
       userId: user.id,
     });
-    if (codeVerifier === undefined) {
+    if (statePayload === undefined) {
       // Not necessarily a real failure: a browser that fires this exact
       // callback twice burns the state on its first, successful arrival
       // and only ever sees this branch on the second.
@@ -504,7 +526,9 @@ export function createOAuthConnectRoutes<E extends AppEnv = AppEnv>(
     const exchangeArgs: Parameters<typeof descriptor.oauth.exchange>[0] = {
       code,
       redirectUri: callbackUrl,
-      ...(descriptor.oauth.usesPKCE ? { codeVerifier } : {}),
+      ...(descriptor.oauth.usesPKCE
+        ? { codeVerifier: statePayload.codeVerifier }
+        : {}),
       ...(clientId !== undefined ? { clientId } : {}),
       ...(clientSecret !== undefined ? { clientSecret } : {}),
     };