pi-mcp-extension/src/oauth-state-store.ts

Full-featured, lightweight MCP extension for the Pi coding agent: stdio, SSE & Streamable HTTP transports, OAuth 2.1 (PKCE + DCR), lazy/eager lifecycles — minimal dependencies, zero bloat.

/**
 * Persisted OAuth state schemas, origin-binding helpers, and atomic file I/O.
 */

import type { OAuthDiscoveryState } from "@modelcontextprotocol/sdk/client/auth.js";
import { readFile, writeFile, mkdir, unlink, chmod, rename } from "node:fs/promises";
import { readFileSync } from "node:fs";
import { randomUUID } from "node:crypto";
import { z } from "zod";
import { AuthRequiredError, McpError } from "./errors.js";
import { authDir, statePath } from "./auth-lock.js";
import {
  canonicalServerUrl,
  isResourceCompatibleWithServerUrl,
  urlsMatch,
} from "./oauth-policy.js";

export { canonicalServerUrl, isResourceCompatibleWithServerUrl, urlsMatch };

// ─── Persistent State Types & Schema ──────────────────────────────────────────

export interface StoredClientInfo {
  client_id: string;
  client_secret: string | undefined;
  client_id_issued_at: number | undefined;
  client_secret_expires_at: number | undefined;
  token_endpoint_auth_method: string | undefined;
  redirect_uris: string[] | undefined;
  issuer: string | undefined;
}

export interface StoredTokens {
  access_token: string;
  token_type: string | undefined;
  refresh_token: string | undefined;
  expires_in: number | undefined;
  scope: string | undefined;
  /** ISO timestamp when tokens were saved (for expiry estimation). */
  saved_at: string | undefined;
  issuer: string | undefined;
}

export interface StoredState {
  serverUrl?: string | undefined;
  clientInfo: StoredClientInfo | undefined;
  tokens: StoredTokens | undefined;
  discoveryState: OAuthDiscoveryState | undefined;
}

// WHY: Credentials persisted before SDK 1.31 (GHSA-6qxp-vccf-f47h) have no `issuer`
// property. We allow `issuer: z.string().optional()` so legacy state loads gracefully;
// the SDK binds them to the discovered authorization server on first successful use.
// Unknown keys are stripped by z.object, so explicit schema inclusion is required to
// persist and return the SDK's issuer stamp.
const StoredClientInfoSchema = z.object({
  client_id: z.string(),
  client_secret: z.string().optional(),
  client_id_issued_at: z.number().optional(),
  client_secret_expires_at: z.number().optional(),
  token_endpoint_auth_method: z.string().optional(),
  redirect_uris: z.array(z.string()).optional(),
  issuer: z.string().optional(),
});

const StoredTokensSchema = z.object({
  access_token: z.string(),
  token_type: z.string().optional(),
  refresh_token: z.string().optional(),
  expires_in: z.number().optional(),
  scope: z.string().optional(),
  saved_at: z.string().optional(),
  issuer: z.string().optional(),
});

export const StoredStateSchema = z.object({
  serverUrl: z.string().optional(),
  clientInfo: StoredClientInfoSchema.optional(),
  tokens: StoredTokensSchema.optional(),
  discoveryState: z.record(z.string(), z.unknown()).optional(),
});

// ─── File helpers ─────────────────────────────────────────────────────────────

type ParsedStoredState = z.output<typeof StoredStateSchema>;

function emptyStoredState(expectedServerUrl?: string): StoredState {
  return {
    serverUrl: expectedServerUrl ? canonicalServerUrl(expectedServerUrl) : undefined,
    clientInfo: undefined,
    tokens: undefined,
    discoveryState: undefined,
  };
}

async function readStoredStateFile(
  path: string,
  serverName: string,
): Promise<ParsedStoredState | null> {
  let raw: string;
  try {
    raw = await readFile(path, "utf8");
  } catch (err) {
    if ((err as NodeJS.ErrnoException).code === "ENOENT") return null;
    throw err;
  }

  try {
    const json: unknown = JSON.parse(raw);
    return StoredStateSchema.parse(json);
  } catch (err) {
    const backupPath = `${path}.corrupt-${Date.now()}`;
    await rename(path, backupPath);
    throw new McpError(
      `Corrupt OAuth state file for "${serverName}" backed up to ${backupPath}`,
      serverName,
      "config",
      err,
    );
  }
}

function cachedDiscoveryResource(parsed: ParsedStoredState): unknown {
  const discoveryState = parsed.discoveryState as OAuthDiscoveryState | undefined;
  const metadata = discoveryState?.resourceMetadata as { resource?: unknown } | undefined;
  return metadata?.resource;
}

function hasCompatibleCachedResource(parsed: ParsedStoredState, expectedServerUrl: string): boolean {
  const resource = cachedDiscoveryResource(parsed);
  return typeof resource !== "string" || isResourceCompatibleWithServerUrl(resource, expectedServerUrl);
}

function shouldDiscardStoredState(parsed: ParsedStoredState, expectedServerUrl: string): boolean {
  // WHY: Bind OAuth state to the canonical server URL so a project config with the same
  // server name cannot inherit global Bearer tokens for a different origin/path.
  if (parsed.serverUrl !== undefined && !urlsMatch(parsed.serverUrl, expectedServerUrl)) {
    return true;
  }
  // Legacy state file without top-level serverUrl whose cached resourceMetadata belongs to another server URL.
  return parsed.serverUrl === undefined && !hasCompatibleCachedResource(parsed, expectedServerUrl);
}

function cachedDiscoveryStateFor(
  parsed: ParsedStoredState,
  expectedServerUrl: string,
): OAuthDiscoveryState | undefined {
  if (!hasCompatibleCachedResource(parsed, expectedServerUrl)) {
    // WHY: Discard mismatched discoveryState so the MCP SDK's selectResourceURL()
    // never throws "Protected resource ... does not match expected".
    return undefined;
  }
  return parsed.discoveryState as OAuthDiscoveryState | undefined;
}

function toStoredClientInfo(info: NonNullable<ParsedStoredState["clientInfo"]>): StoredClientInfo {
  return {
    client_id: info.client_id,
    client_secret: info.client_secret,
    client_id_issued_at: info.client_id_issued_at,
    client_secret_expires_at: info.client_secret_expires_at,
    token_endpoint_auth_method: info.token_endpoint_auth_method,
    redirect_uris: info.redirect_uris,
    issuer: info.issuer,
  };
}

function toStoredTokens(tokens: NonNullable<ParsedStoredState["tokens"]>): StoredTokens {
  return {
    access_token: tokens.access_token,
    token_type: tokens.token_type,
    refresh_token: tokens.refresh_token,
    expires_in: tokens.expires_in,
    scope: tokens.scope,
    saved_at: tokens.saved_at,
    issuer: tokens.issuer,
  };
}

function normalizeStoredState(
  parsed: ParsedStoredState,
  discoveryState: OAuthDiscoveryState | undefined,
  expectedServerUrl?: string,
): StoredState {
  return {
    serverUrl: parsed.serverUrl ?? (expectedServerUrl ? canonicalServerUrl(expectedServerUrl) : undefined),
    clientInfo: parsed.clientInfo ? toStoredClientInfo(parsed.clientInfo) : undefined,
    tokens: parsed.tokens ? toStoredTokens(parsed.tokens) : undefined,
    discoveryState,
  };
}

export async function loadState(
  serverName: string,
  storageDir?: string,
  expectedServerUrl?: string,
): Promise<StoredState> {
  const path = statePath(serverName, storageDir);
  const parsed = await readStoredStateFile(path, serverName);
  if (parsed === null) return emptyStoredState(expectedServerUrl);
  if (expectedServerUrl === undefined) {
    return normalizeStoredState(parsed, parsed.discoveryState as OAuthDiscoveryState | undefined);
  }
  if (shouldDiscardStoredState(parsed, expectedServerUrl)) {
    return emptyStoredState(expectedServerUrl);
  }
  return normalizeStoredState(parsed, cachedDiscoveryStateFor(parsed, expectedServerUrl), expectedServerUrl);
}

export async function saveState(
  serverName: string,
  state: StoredState,
  storageDir?: string,
  isCancelled: () => boolean = () => false,
): Promise<void> {
  const directory = authDir(storageDir);
  await mkdir(directory, { recursive: true, mode: 0o700 });
  await chmod(directory, 0o700);
  // Only write defined fields (never write codeVerifier to shared state)
  const toWrite: Record<string, unknown> = {};
  if (state.serverUrl !== undefined) toWrite.serverUrl = state.serverUrl;
  if (state.clientInfo !== undefined) toWrite.clientInfo = state.clientInfo;
  if (state.tokens !== undefined) toWrite.tokens = state.tokens;
  if (state.discoveryState !== undefined) toWrite.discoveryState = state.discoveryState;
  if (isCancelled()) throw new AuthRequiredError(serverName, "Authorization cancelled");
  const path = statePath(serverName, storageDir);
  const tmpPath = `${path}.tmp-${randomUUID()}`;
  try {
    await writeFile(tmpPath, JSON.stringify(toWrite, null, 2), { encoding: "utf8", mode: 0o600 });
    await chmod(tmpPath, 0o600);
    if (isCancelled()) {
      await unlink(tmpPath).catch(() => {
        // WHY: Best-effort cleanup of temporary state file when cancelled before rename.
      });
      throw new AuthRequiredError(serverName, "Authorization cancelled");
    }
    await rename(tmpPath, path);
    await chmod(path, 0o600);
  } catch (err) {
    await unlink(tmpPath).catch(() => {
      // WHY: Best-effort cleanup of temporary state file after write or rename failure.
    });
    throw err;
  }
}

function readAndParseStoredStateSync(path: string): z.output<typeof StoredStateSchema> | undefined {
  try {
    const raw = readFileSync(path, "utf8");
    const json: unknown = JSON.parse(raw);
    const parsed = StoredStateSchema.safeParse(json);
    return parsed.success ? parsed.data : undefined;
  } catch {
    // WHY: State file may not exist yet (ENOENT), directory may be inaccessible, or JSON is malformed; return undefined.
    return undefined;
  }
}

/**
 * Synchronously load stored DCR redirect URIs from disk if valid and origin-compatible.
 * Used by silent McpOAuthProvider to report registered redirect_uris synchronously.
 */
export function loadStoredRedirectUrisSync(
  serverName: string,
  storageDir?: string,
  serverUrl?: string,
): string[] | undefined {
  const data = readAndParseStoredStateSync(statePath(serverName, storageDir));
  if (!data) return undefined;
  if (serverUrl !== undefined && shouldDiscardStoredState(data, serverUrl)) {
    return undefined;
  }
  const uris = data.clientInfo?.redirect_uris;
  if (Array.isArray(uris) && uris.length > 0) {
    return uris.map(String);
  }
  return undefined;
}