@phi-code-admin/phi-code 0.82.3 → 0.84.0

package/extensions/phi/mcp/oauth-provider.ts

@@ -0,0 +1,367 @@
+/**
+ * File-based OAuthClientProvider for MCP servers.
+ *
+ * Implements the full OAuth2 Authorization Code flow with PKCE and
+ * Dynamic Client Registration (RFC 7591) as required by the MCP spec.
+ *
+ * Token and client state are persisted per-server under ~/.pi/agent/mcp-auth/
+ * so they survive pi restarts without requiring re-authorization.
+ *
+ * Usage: config adds `auth: { ... }` to a server config. This module
+ * constructs an OAuthClientProvider and the transport factory wires it in.
+ */
+
+import type { OAuthClientProvider, OAuthDiscoveryState } from "@modelcontextprotocol/sdk/client/auth.js";
+import type { OAuthClientMetadata, OAuthClientInformationMixed, OAuthTokens } from "@modelcontextprotocol/sdk/shared/auth.js";
+import { readFile, writeFile, mkdir, unlink } from "node:fs/promises";
+import { join } from "node:path";
+import { homedir } from "node:os";
+import { createHash } from "node:crypto";
+
+// ─── Types ────────────────────────────────────────────────────────────────────
+
+/** Auth config in mcp.json server entry. Matches the Zod AuthConfigSchema in config.ts. */
+export interface AuthConfig {
+  /** Auth type. Currently only "oauth" is supported. */
+  type?: "oauth" | undefined;
+  /**
+   * Callback URL the OAuth server redirects to after authorization.
+   * Default: auto-detected local callback server.
+   */
+  redirectUrl?: string | undefined;
+  /**
+   * Optional scope to request during authorization.
+   */
+  scope?: string | undefined;
+  /**
+   * Pre-registered client_id (skip dynamic client registration).
+   */
+  clientId?: string | undefined;
+  /**
+   * Pre-registered client_secret.
+   */
+  clientSecret?: string | undefined;
+}
+
+// ─── Callback Server Port ───────────────────────────────────────────────────────
+
+let callbackPort = 19876;
+
+/** Set the callback server port. Called by the auth flow when the server starts. */
+export function setCallbackPort(port: number): void {
+  callbackPort = port;
+}
+
+/** Get the current callback server port. */
+export function getCallbackPort(): number {
+  return callbackPort;
+}
+
+// ─── Persistent State Types ───────────────────────────────────────────────────
+
+interface StoredClientInfo {
+  client_id: string;
+  client_secret: string | undefined;
+}
+
+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;
+}
+
+interface StoredState {
+  clientInfo: StoredClientInfo | undefined;
+  tokens: StoredTokens | undefined;
+  codeVerifier: string | undefined;
+  discoveryState: OAuthDiscoveryState | undefined;
+}
+
+// ─── File helpers ─────────────────────────────────────────────────────────────
+
+const AUTH_DIR = join(homedir(), ".pi", "agent", "mcp-auth");
+
+function statePath(serverName: string): string {
+  // Hash the server name to avoid filesystem issues with special chars
+  const hash = createHash("sha256").update(serverName).digest("hex").slice(0, 16);
+  return join(AUTH_DIR, `${hash}.json`);
+}
+
+async function loadState(serverName: string): Promise<StoredState> {
+  try {
+    const raw = await readFile(statePath(serverName), "utf8");
+    const parsed = JSON.parse(raw) as StoredState;
+    return {
+      clientInfo: parsed.clientInfo ?? undefined,
+      tokens: parsed.tokens ?? undefined,
+      codeVerifier: parsed.codeVerifier ?? undefined,
+      discoveryState: parsed.discoveryState ?? undefined,
+    };
+  } catch {
+    return {
+      clientInfo: undefined,
+      tokens: undefined,
+      codeVerifier: undefined,
+      discoveryState: undefined,
+    };
+  }
+}
+
+async function saveState(serverName: string, state: StoredState): Promise<void> {
+  await mkdir(AUTH_DIR, { recursive: true });
+  // Only write defined fields
+  const toWrite: Record<string, unknown> = {};
+  if (state.clientInfo !== undefined) toWrite.clientInfo = state.clientInfo;
+  if (state.tokens !== undefined) toWrite.tokens = state.tokens;
+  if (state.codeVerifier !== undefined) toWrite.codeVerifier = state.codeVerifier;
+  if (state.discoveryState !== undefined) toWrite.discoveryState = state.discoveryState;
+  await writeFile(statePath(serverName), JSON.stringify(toWrite, null, 2), "utf8");
+}
+
+// ─── OAuthClientProvider Implementation ────────────────────────────────────────
+
+export class McpOAuthProvider implements OAuthClientProvider {
+  private serverName: string;
+  private authConfig: AuthConfig;
+  private _redirectUrl: string | undefined;
+  private _onAuthRequired: ((url: URL) => void | Promise<void>) | undefined;
+  private _oauthState: string | undefined;
+
+  constructor(
+    serverName: string,
+    authConfig: AuthConfig,
+    onAuthRequired?: (url: URL) => void | Promise<void>,
+  ) {
+    this.serverName = serverName;
+    this.authConfig = authConfig;
+    this._redirectUrl = authConfig.redirectUrl;
+    this._onAuthRequired = onAuthRequired;
+  }
+
+  // --- redirectUrl ---
+
+  get redirectUrl(): string | URL {
+    // Use configured redirect URL if provided, otherwise use the callback server
+    // Use 127.0.0.1 explicitly (IPv4) to match the callback server binding
+    return this._redirectUrl || `http://127.0.0.1:${callbackPort}/callback`;
+  }
+
+  // --- clientMetadata ---
+
+  get clientMetadata(): OAuthClientMetadata {
+    const redirectUrl = String(this.redirectUrl);
+    return {
+      client_name: `pi-mcp/${this.serverName}`,
+      redirect_uris: [redirectUrl],
+      grant_types: ["authorization_code", "refresh_token"],
+      response_types: ["code"],
+      token_endpoint_auth_method: this.authConfig.clientSecret ? "client_secret_basic" : "none",
+    };
+  }
+
+  // --- clientInformation ---
+
+  async clientInformation(): Promise<OAuthClientInformationMixed | undefined> {
+    // If static credentials provided, use those
+    if (this.authConfig.clientId) {
+      return {
+        client_id: this.authConfig.clientId,
+        ...(this.authConfig.clientSecret && { client_secret: this.authConfig.clientSecret }),
+      };
+    }
+    // Otherwise load from persisted DCR state
+    const state = await loadState(this.serverName);
+    if (state.clientInfo) {
+      return {
+        client_id: state.clientInfo.client_id,
+        ...(state.clientInfo.client_secret && { client_secret: state.clientInfo.client_secret }),
+      };
+    }
+    return undefined;
+  }
+
+  // --- saveClientInformation (DCR) ---
+
+  async saveClientInformation(clientInformation: OAuthClientInformationMixed): Promise<void> {
+    const state = await loadState(this.serverName);
+    state.clientInfo = {
+      client_id: clientInformation.client_id,
+      client_secret: clientInformation.client_secret,
+    };
+    await saveState(this.serverName, state);
+  }
+
+  // --- tokens ---
+
+  async tokens(): Promise<OAuthTokens | undefined> {
+    const state = await loadState(this.serverName);
+    if (!state.tokens) return undefined;
+
+    // Always return stored tokens — even if expired.
+    // The SDK's auth() function checks tokens?.refresh_token and attempts
+    // silent refresh before falling back to a new authorization flow.
+    // Returning undefined for expired tokens would prevent that silent refresh
+    // and force the user to re-authenticate via browser every time.
+    //
+    // Flow when we return expired tokens:
+    //   transport sends expired access_token → 401
+    //   → auth() sees refresh_token → silent refresh → success → retry
+    //
+    // Flow when we return undefined (WRONG):
+    //   transport has no token → auth() → no refresh possible → REDIRECT
+    //   → user must re-authenticate in browser
+
+    // Build OAuthTokens — only include defined fields
+    const tokens: Record<string, string> = {
+      access_token: state.tokens.access_token,
+    };
+    if (state.tokens.token_type !== undefined) tokens.token_type = state.tokens.token_type;
+    if (state.tokens.refresh_token !== undefined) tokens.refresh_token = state.tokens.refresh_token;
+    return tokens as unknown as OAuthTokens;
+  }
+
+  // --- saveTokens ---
+
+  async saveTokens(tokens: OAuthTokens): Promise<void> {
+    const state = await loadState(this.serverName);
+    state.tokens = {
+      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: new Date().toISOString(),
+    };
+    await saveState(this.serverName, state);
+  }
+
+  // --- redirectToAuthorization ---
+
+  async redirectToAuthorization(authorizationUrl: URL): Promise<void> {
+    if (this._onAuthRequired) {
+      await this._onAuthRequired(authorizationUrl);
+    } else {
+      // Fallback: just log it
+      console.error(
+        `[pi-mcp] OAuth authorization required for "${this.serverName}".`,
+        `\n  Open this URL in your browser: ${authorizationUrl.toString()}`,
+      );
+    }
+  }
+
+  // --- PKCE code verifier ---
+
+  async saveCodeVerifier(codeVerifier: string): Promise<void> {
+    const state = await loadState(this.serverName);
+    state.codeVerifier = codeVerifier;
+    await saveState(this.serverName, state);
+  }
+
+  async codeVerifier(): Promise<string> {
+    const state = await loadState(this.serverName);
+    if (!state.codeVerifier) {
+      throw new Error(`[pi-mcp] No PKCE code verifier found for "${this.serverName}"`);
+    }
+    return state.codeVerifier;
+  }
+
+  // --- Discovery state caching ---
+
+  async saveDiscoveryState(discState: OAuthDiscoveryState): Promise<void> {
+    const state = await loadState(this.serverName);
+    state.discoveryState = discState;
+    await saveState(this.serverName, state);
+  }
+
+  async discoveryState(): Promise<OAuthDiscoveryState | undefined> {
+    const state = await loadState(this.serverName);
+    return state.discoveryState;
+  }
+
+  // --- OAuth state parameter (CSRF protection) ---
+
+  /**
+   * Set the OAuth state parameter before calling auth().
+   * This should be called with a cryptographically random value.
+   */
+  setState(state: string): void {
+    this._oauthState = state;
+  }
+
+  /**
+   * Returns the OAuth state parameter for CSRF protection.
+   * This is called by the SDK's auth() function when building the authorization URL.
+   * Returns empty string if no state has been set (no CSRF protection).
+   */
+  async state(): Promise<string> {
+    return this._oauthState || "";
+  }
+
+  // --- Credential invalidation ---
+
+  async invalidateCredentials(scope: "all" | "client" | "tokens" | "verifier" | "discovery"): Promise<void> {
+    const state = await loadState(this.serverName);
+    switch (scope) {
+      case "all":
+        state.clientInfo = undefined;
+        state.tokens = undefined;
+        state.codeVerifier = undefined;
+        state.discoveryState = undefined;
+        break;
+      case "client":
+        state.clientInfo = undefined;
+        break;
+      case "tokens":
+        state.tokens = undefined;
+        break;
+      case "verifier":
+        state.codeVerifier = undefined;
+        break;
+      case "discovery":
+        state.discoveryState = undefined;
+        break;
+    }
+    await saveState(this.serverName, state);
+  }
+}
+
+// ─── Public Helpers ────────────────────────────────────────────────────────────
+
+/**
+ * Get auth status info for a server — whether tokens exist, when they were saved, etc.
+ * Returns null if no auth state file exists at all.
+ */
+export async function getAuthStatus(serverName: string): Promise<{
+  hasTokens: boolean;
+  hasClientInfo: boolean;
+  savedAt: string | undefined;
+  scope: string | undefined;
+} | null> {
+  const state = await loadState(serverName);
+  if (
+    state.clientInfo === undefined &&
+    state.tokens === undefined &&
+    state.codeVerifier === undefined &&
+    state.discoveryState === undefined
+  ) {
+    return null;
+  }
+  return {
+    hasTokens: state.tokens !== undefined,
+    hasClientInfo: state.clientInfo !== undefined,
+    savedAt: state.tokens?.saved_at,
+    scope: state.tokens?.scope,
+  };
+}
+
+/**
+ * Reset all OAuth state for a server (tokens, client info, PKCE verifier, discovery).
+ * Used to force re-authorization on next connection.
+ */
+export async function resetAuth(serverName: string): Promise<void> {
+  await unlink(statePath(serverName)).catch(() => {});
+}