@bastani/atomic 0.7.17 → 0.8.0

package/dist/builtin/mcp/mcp-auth.ts

@@ -0,0 +1,297 @@
+/**
+ * MCP Auth Storage Module
+ * 
+ * Handles secure storage of OAuth credentials, tokens, client information,
+ * and PKCE state for MCP servers. Maintains backward compatibility with
+ * per-server directory structure.
+ * 
+ * Token storage location: $MCP_OAUTH_DIR/<server>/tokens.json when set,
+ * otherwise <Pi agent dir>/mcp-oauth/<server>/tokens.json
+ */
+
+import { mkdirSync, readFileSync, writeFileSync, existsSync, rmSync } from 'fs';
+import { join } from 'path';
+import { getAgentPath } from './agent-dir.ts';
+
+/** OAuth token storage format */
+export interface StoredTokens {
+  accessToken: string;
+  refreshToken?: string;
+  expiresAt?: number; // Unix timestamp in seconds
+  scope?: string;
+}
+
+/** OAuth client information from dynamic or static registration */
+export interface StoredClientInfo {
+  clientId: string;
+  clientSecret?: string;
+  clientIdIssuedAt?: number;
+  clientSecretExpiresAt?: number;
+}
+
+/** Complete auth entry for a server */
+export interface AuthEntry {
+  tokens?: StoredTokens;
+  clientInfo?: StoredClientInfo;
+  codeVerifier?: string;
+  oauthState?: string;
+  serverUrl?: string; // Track the URL these credentials are for
+}
+
+// Base directory for auth storage - can be overridden via env var for testing
+function getAuthBaseDir(): string {
+  const override = process.env.MCP_OAUTH_DIR?.trim();
+  return override ? override : getAgentPath('mcp-oauth');
+}
+
+/**
+ * Get the server-specific directory path.
+ */
+function getServerDir(serverName: string): string {
+  return join(getAuthBaseDir(), serverName);
+}
+
+/**
+ * Get the tokens file path for a server.
+ */
+function getTokensFilePath(serverName: string): string {
+  return join(getServerDir(serverName), 'tokens.json');
+}
+
+/**
+ * Ensure the server directory exists with secure permissions.
+ */
+function ensureServerDir(serverName: string): void {
+  const dir = getServerDir(serverName);
+  if (!existsSync(dir)) {
+    mkdirSync(dir, { recursive: true, mode: 0o700 });
+  }
+}
+
+/**
+ * Read the auth entry for a server from disk.
+ * Returns undefined if file doesn't exist.
+ */
+function readAuthEntry(serverName: string): AuthEntry | undefined {
+  try {
+    const filePath = getTokensFilePath(serverName);
+    if (!existsSync(filePath)) {
+      return undefined;
+    }
+    const data = readFileSync(filePath, 'utf-8');
+    return JSON.parse(data) as AuthEntry;
+  } catch (error) {
+    console.error(`Failed to read auth entry for ${serverName}:`, error);
+    return undefined;
+  }
+}
+
+/**
+ * Write the auth entry for a server to disk with secure permissions.
+ */
+function writeAuthEntry(serverName: string, entry: AuthEntry): void {
+  ensureServerDir(serverName);
+  const filePath = getTokensFilePath(serverName);
+  writeFileSync(filePath, JSON.stringify(entry, null, 2), { mode: 0o600 });
+}
+
+/**
+ * Get auth entry for a server.
+ */
+export function getAuthEntry(serverName: string): AuthEntry | undefined {
+  return readAuthEntry(serverName);
+}
+
+/**
+ * Get auth entry and validate it's for the correct URL.
+ * Returns undefined if URL has changed (credentials are invalid).
+ */
+export function getAuthForUrl(serverName: string, serverUrl: string): AuthEntry | undefined {
+  const entry = getAuthEntry(serverName);
+  if (!entry) return undefined;
+
+  // If no serverUrl is stored, this is from an old version - consider it invalid
+  if (!entry.serverUrl) return undefined;
+
+  // If URL has changed, credentials are invalid
+  if (entry.serverUrl !== serverUrl) return undefined;
+
+  return entry;
+}
+
+/**
+ * Save auth entry for a server.
+ */
+export function saveAuthEntry(serverName: string, entry: AuthEntry, serverUrl?: string): void {
+  // Always update serverUrl if provided
+  if (serverUrl) {
+    entry.serverUrl = serverUrl;
+  }
+  writeAuthEntry(serverName, entry);
+}
+
+/**
+ * Remove auth entry for a server.
+ * Also removes the server directory if empty.
+ */
+export function removeAuthEntry(serverName: string): void {
+  try {
+    const filePath = getTokensFilePath(serverName);
+    if (existsSync(filePath)) {
+      writeFileSync(filePath, '{}', { mode: 0o600 });
+    }
+    // Try to remove the directory
+    const dir = getServerDir(serverName);
+    if (existsSync(dir)) {
+      try {
+        rmSync(dir, { recursive: true });
+      } catch {
+        // Directory may not be empty, ignore
+      }
+    }
+  } catch (error) {
+    console.error(`Failed to remove auth entry for ${serverName}:`, error);
+  }
+}
+
+/**
+ * Update tokens for a server.
+ */
+export function updateTokens(
+  serverName: string, 
+  tokens: StoredTokens, 
+  serverUrl?: string
+): void {
+  const entry = getAuthEntry(serverName) ?? {};
+  if (serverUrl && entry.serverUrl !== serverUrl) {
+    delete entry.clientInfo;
+    delete entry.codeVerifier;
+    delete entry.oauthState;
+  }
+  entry.tokens = tokens;
+  saveAuthEntry(serverName, entry, serverUrl);
+}
+
+/**
+ * Update client info for a server.
+ */
+export function updateClientInfo(
+  serverName: string, 
+  clientInfo: StoredClientInfo, 
+  serverUrl?: string
+): void {
+  const entry = getAuthEntry(serverName) ?? {};
+  if (serverUrl && entry.serverUrl !== serverUrl) {
+    delete entry.tokens;
+    delete entry.codeVerifier;
+    delete entry.oauthState;
+  }
+  entry.clientInfo = clientInfo;
+  saveAuthEntry(serverName, entry, serverUrl);
+}
+
+/**
+ * Update code verifier for a server.
+ */
+export function updateCodeVerifier(serverName: string, codeVerifier: string, serverUrl?: string): void {
+  const entry = getAuthEntry(serverName) ?? {};
+  if (serverUrl && entry.serverUrl !== serverUrl) {
+    delete entry.tokens;
+    delete entry.clientInfo;
+    delete entry.oauthState;
+  }
+  entry.codeVerifier = codeVerifier;
+  saveAuthEntry(serverName, entry, serverUrl);
+}
+
+/**
+ * Clear code verifier for a server.
+ */
+export function clearCodeVerifier(serverName: string): void {
+  const entry = getAuthEntry(serverName);
+  if (entry) {
+    delete entry.codeVerifier;
+    saveAuthEntry(serverName, entry);
+  }
+}
+
+/**
+ * Update OAuth state for a server.
+ */
+export function updateOAuthState(serverName: string, state: string, serverUrl?: string): void {
+  const entry = getAuthEntry(serverName) ?? {};
+  if (serverUrl && entry.serverUrl !== serverUrl) {
+    delete entry.tokens;
+    delete entry.clientInfo;
+    delete entry.codeVerifier;
+  }
+  entry.oauthState = state;
+  saveAuthEntry(serverName, entry, serverUrl);
+}
+
+/**
+ * Get OAuth state for a server.
+ */
+export function getOAuthState(serverName: string): string | undefined {
+  const entry = getAuthEntry(serverName);
+  return entry?.oauthState;
+}
+
+/**
+ * Clear OAuth state for a server.
+ */
+export function clearOAuthState(serverName: string): void {
+  const entry = getAuthEntry(serverName);
+  if (entry) {
+    delete entry.oauthState;
+    saveAuthEntry(serverName, entry);
+  }
+}
+
+/**
+ * Check if stored tokens are expired.
+ * Returns null if no tokens exist, false if no expiry or not expired, true if expired.
+ */
+export function isTokenExpired(serverName: string): boolean | null {
+  const entry = getAuthEntry(serverName);
+  if (!entry?.tokens) return null;
+  if (!entry.tokens.expiresAt) return false;
+  return entry.tokens.expiresAt < Date.now() / 1000;
+}
+
+/**
+ * Check if a server has stored tokens.
+ */
+export function hasStoredTokens(serverName: string): boolean {
+  const entry = getAuthEntry(serverName);
+  return !!entry?.tokens;
+}
+
+/**
+ * Clear all credentials for a server.
+ */
+export function clearAllCredentials(serverName: string): void {
+  removeAuthEntry(serverName);
+}
+
+/**
+ * Clear only client info for a server.
+ */
+export function clearClientInfo(serverName: string): void {
+  const entry = getAuthEntry(serverName);
+  if (entry) {
+    delete entry.clientInfo;
+    saveAuthEntry(serverName, entry);
+  }
+}
+
+/**
+ * Clear only tokens for a server.
+ */
+export function clearTokens(serverName: string): void {
+  const entry = getAuthEntry(serverName);
+  if (entry) {
+    delete entry.tokens;
+    saveAuthEntry(serverName, entry);
+  }
+}

package/dist/builtin/mcp/mcp-callback-server.ts

@@ -0,0 +1,284 @@
+/**
+ * MCP OAuth Callback Server
+ * 
+ * HTTP server that handles OAuth callbacks from the authorization server.
+ * Uses Node.js http module for compatibility.
+ */
+
+import { createServer, type Server, type IncomingMessage, type ServerResponse } from "http"
+import {
+  OAUTH_CALLBACK_PATH,
+  getConfiguredOAuthCallbackPort,
+  getOAuthCallbackPort,
+  setOAuthCallbackPort,
+} from "./mcp-oauth-provider.ts"
+
+// HTML templates for callback responses
+const HTML_SUCCESS = `<!DOCTYPE html>
+<html>
+<head>
+  <title>Pi - Authorization Successful</title>
+  <style>
+    body { font-family: system-ui, -apple-system, sans-serif; display: flex; justify-content: center; align-items: center; height: 100vh; margin: 0; background: #1a1a2e; color: #eee; }
+    .container { text-align: center; padding: 2rem; }
+    h1 { color: #4ade80; margin-bottom: 1rem; }
+    p { color: #aaa; }
+  </style>
+</head>
+<body>
+  <div class="container">
+    <h1>Authorization Successful</h1>
+    <p>You can close this window and return to Pi.</p>
+  </div>
+  <script>setTimeout(() => window.close(), 2000);</script>
+</body>
+</html>`
+
+const HTML_ERROR = () => `<!DOCTYPE html>
+<html>
+<head>
+  <title>Pi - Authorization Failed</title>
+  <style>
+    body { font-family: system-ui, -apple-system, sans-serif; display: flex; justify-content: center; align-items: center; height: 100vh; margin: 0; background: #1a1a2e; color: #eee; }
+    .container { text-align: center; padding: 2rem; }
+    h1 { color: #f87171; margin-bottom: 1rem; }
+    p { color: #aaa; }
+    .error { color: #fca5a5; font-family: monospace; margin-top: 1rem; padding: 1rem; background: rgba(248,113,113,0.1); border-radius: 0.5rem; }
+  </style>
+</head>
+<body>
+  <div class="container">
+    <h1>Authorization Failed</h1>
+    <p>An error occurred during authorization.</p>
+    <div class="error">Authorization request failed.</div>
+  </div>
+</body>
+</html>`
+
+/** Pending authorization request */
+interface PendingAuth {
+  resolve: (code: string) => void
+  reject: (error: Error) => void
+  timeout: ReturnType<typeof setTimeout>
+}
+
+/** Server singleton state */
+let server: Server | undefined
+const pendingAuths = new Map<string, PendingAuth>()
+
+/** Timeout for callback completion (5 minutes) */
+const CALLBACK_TIMEOUT_MS = 5 * 60 * 1000
+
+const MAX_PORT_SCAN_ATTEMPTS = 25
+
+interface EnsureCallbackServerOptions {
+  strictPort?: boolean
+}
+
+/**
+ * Handle incoming HTTP requests to the callback server.
+ */
+function handleRequest(req: IncomingMessage, res: ServerResponse): void {
+  const url = new URL(req.url || "/", `http://${req.headers.host}`)
+
+  // Only handle the callback path
+  if (url.pathname !== OAUTH_CALLBACK_PATH) {
+    res.writeHead(404, { "Content-Type": "text/plain" })
+    res.end("Not found")
+    return
+  }
+
+  const code = url.searchParams.get("code")
+  const state = url.searchParams.get("state")
+  const error = url.searchParams.get("error")
+  const errorDescription = url.searchParams.get("error_description")
+
+  // Enforce state parameter presence for CSRF protection
+  if (!state) {
+    const errorMsg = "Missing required state parameter - potential CSRF attack"
+    res.writeHead(400, { "Content-Type": "text/html" })
+    res.end(HTML_ERROR(errorMsg))
+    return
+  }
+
+  // Handle OAuth errors
+  if (error) {
+    const errorMsg = errorDescription || error
+    // Send HTTP response first before rejecting promise
+    res.writeHead(200, { "Content-Type": "text/html" })
+    res.end(HTML_ERROR(errorMsg))
+    // Reject promise after response is sent (defer to allow test to attach handler)
+    if (pendingAuths.has(state)) {
+      const pending = pendingAuths.get(state)!
+      clearTimeout(pending.timeout)
+      pendingAuths.delete(state)
+      setTimeout(() => pending.reject(new Error(errorMsg)), 0)
+    }
+    return
+  }
+
+  // Require authorization code
+  if (!code) {
+    res.writeHead(400, { "Content-Type": "text/html" })
+    res.end(HTML_ERROR("No authorization code provided"))
+    return
+  }
+
+  // Validate state parameter
+  if (!pendingAuths.has(state)) {
+    const errorMsg = "Invalid or expired state parameter - potential CSRF attack"
+    res.writeHead(400, { "Content-Type": "text/html" })
+    res.end(HTML_ERROR(errorMsg))
+    return
+  }
+
+  const pending = pendingAuths.get(state)!
+
+  // Clear timeout and resolve the pending promise
+  clearTimeout(pending.timeout)
+  pendingAuths.delete(state)
+  pending.resolve(code)
+
+  res.writeHead(200, { "Content-Type": "text/html" })
+  res.end(HTML_SUCCESS)
+}
+
+/**
+ * Ensure the callback server is running.
+ * If strictPort is true, requires binding on the configured callback port.
+ * If strictPort is false, scans forward for an available local port.
+ */
+export async function ensureCallbackServer(options: EnsureCallbackServerOptions = {}): Promise<void> {
+  const configuredPort = getConfiguredOAuthCallbackPort()
+  const strictPort = options.strictPort === true
+
+  if (server) {
+    if (!strictPort || getOAuthCallbackPort() === configuredPort) return
+
+    if (pendingAuths.size > 0) {
+      throw new Error(
+        `OAuth callback server is running on port ${getOAuthCallbackPort()}, but strict callback port ${configuredPort} is required and cannot be switched while authorizations are pending`
+      )
+    }
+
+    await stopCallbackServer()
+  }
+
+  const preferredPort = configuredPort
+  const maxAttempts = strictPort ? 1 : MAX_PORT_SCAN_ATTEMPTS
+  let lastError: Error | undefined
+
+  for (let offset = 0; offset < maxAttempts; offset++) {
+    const candidatePort = preferredPort + offset
+    const candidateServer = createServer(handleRequest)
+
+    try {
+      await new Promise<void>((resolve, reject) => {
+        candidateServer.once("error", (err) => {
+          reject(err)
+        })
+
+        candidateServer.listen(candidatePort, "localhost", () => {
+          resolve()
+        })
+      })
+
+      server = candidateServer
+      server.unref()
+      setOAuthCallbackPort(candidatePort)
+      return
+    } catch (error) {
+      const nodeError = error as NodeJS.ErrnoException
+      await new Promise<void>((resolve) => {
+        candidateServer.close(() => resolve())
+      })
+
+      if (nodeError.code !== "EADDRINUSE") {
+        throw error
+      }
+
+      lastError = error instanceof Error ? error : new Error(String(error))
+    }
+  }
+
+  if (strictPort) {
+    throw new Error(
+      `OAuth callback port ${preferredPort} is already in use. Pre-registered OAuth clients require an exact redirect URI; set MCP_OAUTH_CALLBACK_PORT to your registered port or free port ${preferredPort}`,
+      { cause: lastError }
+    )
+  }
+
+  throw new Error(
+    `OAuth callback port ${preferredPort} is already in use and no free port was found in range ${preferredPort}-${preferredPort + MAX_PORT_SCAN_ATTEMPTS - 1}`,
+    { cause: lastError }
+  )
+}
+
+/**
+ * Wait for a callback with the given OAuth state.
+ * Returns a promise that resolves with the authorization code.
+ */
+export function waitForCallback(oauthState: string): Promise<string> {
+  return new Promise((resolve, reject) => {
+    const timeout = setTimeout(() => {
+      if (pendingAuths.has(oauthState)) {
+        pendingAuths.delete(oauthState)
+        reject(new Error("OAuth callback timeout - authorization took too long"))
+      }
+    }, CALLBACK_TIMEOUT_MS)
+
+    pendingAuths.set(oauthState, { resolve, reject, timeout })
+  })
+}
+
+/**
+ * Cancel a pending authorization by state.
+ */
+export function cancelPendingCallback(oauthState: string): void {
+  const pending = pendingAuths.get(oauthState)
+  if (pending) {
+    clearTimeout(pending.timeout)
+    pendingAuths.delete(oauthState)
+    pending.reject(new Error("Authorization cancelled"))
+  }
+}
+
+/**
+ * Stop the callback server and reject all pending authorizations.
+ */
+export async function stopCallbackServer(): Promise<void> {
+  if (server) {
+    await new Promise<void>((resolve) => {
+      server!.close(() => {
+        resolve()
+      })
+    })
+    server = undefined
+  }
+
+  setOAuthCallbackPort(getConfiguredOAuthCallbackPort())
+
+  // Reject all pending auths (defer to allow any pending operations to complete)
+  const pendingList = Array.from(pendingAuths.entries())
+  pendingAuths.clear()
+  setTimeout(() => {
+    for (const [, pending] of pendingList) {
+      clearTimeout(pending.timeout)
+      pending.reject(new Error("OAuth callback server stopped"))
+    }
+  }, 0)
+}
+
+/**
+ * Check if the callback server is running.
+ */
+export function isCallbackServerRunning(): boolean {
+  return server !== undefined
+}
+
+/**
+ * Get the number of pending authorizations.
+ */
+export function getPendingAuthCount(): number {
+  return pendingAuths.size
+}