@runuai/host 0.4.2 → 0.4.3

package/lib/mcp-connections.ts

@@ -0,0 +1,522 @@
+/**
+ * User-authed MCP connections (ADR-057) — the host side of the generic
+ * handshake. Everything secret happens here: OAuth discovery (RFC 9728
+ * protected-resource metadata → RFC 8414 AS metadata), dynamic client
+ * registration (RFC 7591), PKCE, the code exchange, token refresh, and the
+ * sealed secret store. The cloud only ever sees non-secret status/URLs.
+ *
+ * Auth kinds:
+ *   - "oauth" — discovered from a 401 + WWW-Authenticate chain.
+ *   - "token" — a static header the user pasted (write-only over the bridge).
+ *   - "none"  — the URL answered an unauthenticated initialize.
+ */
+
+import { createHash, randomBytes } from "node:crypto";
+import { eq } from "drizzle-orm";
+
+import { getDb, schema } from "./db";
+import { sealAesGcm, openAesGcm } from "./secrets";
+import type { McpOp } from "../src/protocol";
+
+const FETCH_TIMEOUT_MS = 10_000;
+// Refresh an access token this long before its recorded expiry.
+const REFRESH_LEAD_MS = 60_000;
+
+export interface McpAck {
+  status: "connected" | "auth_required" | "disconnected";
+  authorizeUrl?: string;
+  scopes?: string[];
+}
+
+interface OauthSecret {
+  accessToken: string;
+  refreshToken?: string;
+  expiresAt?: number; // epoch ms; absent = treat as non-expiring
+}
+
+interface HeaderSecret {
+  headerName: string;
+  headerValue: string;
+}
+
+// --- sealed single-column packing (same format as host-env.ts) --------------
+
+function pack(value: string): string {
+  const sealed = sealAesGcm(value);
+  return `${sealed.ct.toString("base64")}.${sealed.nonce.toString("base64")}`;
+}
+
+function unpack(enc: string): string {
+  const dot = enc.indexOf(".");
+  const ct = Buffer.from(enc.slice(0, dot), "base64");
+  const nonce = Buffer.from(enc.slice(dot + 1), "base64");
+  return openAesGcm(ct, nonce);
+}
+
+// --- fetch helpers -----------------------------------------------------------
+
+async function timedFetch(url: string, init?: RequestInit): Promise<Response> {
+  return fetch(url, {
+    ...init,
+    signal: AbortSignal.timeout(FETCH_TIMEOUT_MS),
+    redirect: "follow",
+  });
+}
+
+async function fetchJson(url: string): Promise<Record<string, unknown> | null> {
+  try {
+    const res = await timedFetch(url, { headers: { accept: "application/json" } });
+    if (!res.ok) return null;
+    return (await res.json()) as Record<string, unknown>;
+  } catch {
+    return null;
+  }
+}
+
+/** POST an MCP `initialize` (streamable HTTP). Returns the raw Response. */
+async function mcpInitialize(
+  url: string,
+  header?: HeaderSecret | null,
+): Promise<Response> {
+  return timedFetch(url, {
+    method: "POST",
+    headers: {
+      "content-type": "application/json",
+      accept: "application/json, text/event-stream",
+      ...(header ? { [header.headerName]: header.headerValue } : {}),
+    },
+    body: JSON.stringify({
+      jsonrpc: "2.0",
+      id: 0,
+      method: "initialize",
+      params: {
+        protocolVersion: "2025-06-18",
+        capabilities: {},
+        clientInfo: { name: "uai-host", version: "1.0" },
+      },
+    }),
+  });
+}
+
+// --- OAuth discovery (RFC 9728 → RFC 8414) -----------------------------------
+
+/** Candidate well-known URLs, path-aware per spec, root fallback last. */
+function wellKnownCandidates(base: string, suffix: string): string[] {
+  const u = new URL(base);
+  const path = u.pathname.replace(/\/$/, "");
+  const out: string[] = [];
+  if (path && path !== "") {
+    out.push(`${u.origin}/.well-known/${suffix}${path}`);
+  }
+  out.push(`${u.origin}/.well-known/${suffix}`);
+  return out;
+}
+
+/** Parse `resource_metadata="…"` out of a WWW-Authenticate header. */
+function resourceMetadataUrl(wwwAuth: string | null): string | null {
+  if (!wwwAuth) return null;
+  const m = /resource_metadata="([^"]+)"/i.exec(wwwAuth);
+  return m?.[1] ?? null;
+}
+
+interface Discovered {
+  authorizationEndpoint: string;
+  tokenEndpoint: string;
+  registrationEndpoint?: string;
+  scopesSupported?: string[];
+}
+
+/**
+ * From the MCP URL + the 401's WWW-Authenticate, find the authorization
+ * server's metadata. Tolerates servers that skip RFC 9728 (falls back to
+ * treating the MCP origin as its own authorization server).
+ */
+async function discoverAuth(
+  mcpUrl: string,
+  wwwAuth: string | null,
+): Promise<Discovered> {
+  // 1. Protected-resource metadata → authorization_servers[0].
+  const prmCandidates = [
+    ...(resourceMetadataUrl(wwwAuth) ? [resourceMetadataUrl(wwwAuth)!] : []),
+    ...wellKnownCandidates(mcpUrl, "oauth-protected-resource"),
+  ];
+  let asBase: string | null = null;
+  let resourceScopes: string[] | undefined;
+  for (const url of prmCandidates) {
+    const prm = await fetchJson(url);
+    const servers = prm?.authorization_servers;
+    if (Array.isArray(servers) && typeof servers[0] === "string") {
+      asBase = servers[0];
+      if (Array.isArray(prm?.scopes_supported)) {
+        resourceScopes = prm.scopes_supported as string[];
+      }
+      break;
+    }
+  }
+  // Servers that skip RFC 9728 often ARE their own AS.
+  asBase ??= new URL(mcpUrl).origin;
+
+  // 2. AS metadata (RFC 8414, then OIDC discovery).
+  const asCandidates = [
+    ...wellKnownCandidates(asBase, "oauth-authorization-server"),
+    ...wellKnownCandidates(asBase, "openid-configuration"),
+  ];
+  for (const url of asCandidates) {
+    const meta = await fetchJson(url);
+    if (
+      typeof meta?.authorization_endpoint === "string" &&
+      typeof meta?.token_endpoint === "string"
+    ) {
+      return {
+        authorizationEndpoint: meta.authorization_endpoint,
+        tokenEndpoint: meta.token_endpoint,
+        registrationEndpoint:
+          typeof meta.registration_endpoint === "string"
+            ? meta.registration_endpoint
+            : undefined,
+        scopesSupported:
+          resourceScopes ??
+          (Array.isArray(meta.scopes_supported)
+            ? (meta.scopes_supported as string[])
+            : undefined),
+      };
+    }
+  }
+  throw new Error(
+    "the server requires auth but exposes no OAuth metadata (" +
+      "no /.well-known/oauth-protected-resource or …/oauth-authorization-server)",
+  );
+}
+
+/** RFC 7591 dynamic client registration. Returns client id (+secret). */
+async function registerClient(
+  registrationEndpoint: string,
+  redirectUri: string,
+): Promise<{ clientId: string; clientSecret?: string }> {
+  const res = await timedFetch(registrationEndpoint, {
+    method: "POST",
+    headers: { "content-type": "application/json" },
+    body: JSON.stringify({
+      client_name: "Uai",
+      redirect_uris: [redirectUri],
+      grant_types: ["authorization_code", "refresh_token"],
+      response_types: ["code"],
+      token_endpoint_auth_method: "none",
+    }),
+  });
+  if (!res.ok) {
+    throw new Error(`dynamic client registration failed (${res.status})`);
+  }
+  const json = (await res.json()) as Record<string, unknown>;
+  if (typeof json.client_id !== "string") {
+    throw new Error("registration response carried no client_id");
+  }
+  return {
+    clientId: json.client_id,
+    clientSecret:
+      typeof json.client_secret === "string" ? json.client_secret : undefined,
+  };
+}
+
+function b64url(buf: Buffer): string {
+  return buf.toString("base64url");
+}
+
+// --- op handlers --------------------------------------------------------------
+
+export async function handleMcpOp(op: McpOp): Promise<McpAck> {
+  switch (op.kind) {
+    case "probe":
+      return probe(op);
+    case "oauth.complete":
+      return oauthComplete(op.connectionId, op.code);
+    case "disconnect": {
+      getDb()
+        .delete(schema.mcpConnections)
+        .where(eq(schema.mcpConnections.id, op.connectionId))
+        .run();
+      return { status: "disconnected" };
+    }
+  }
+}
+
+async function probe(op: Extract<McpOp, { kind: "probe" }>): Promise<McpAck> {
+  const db = getDb();
+  const now = Date.now();
+
+  // Static-header connection: verify the header works, seal, done.
+  if (op.headerName && op.headerValue) {
+    const header = { headerName: op.headerName, headerValue: op.headerValue };
+    const res = await mcpInitialize(op.url, header);
+    if (res.status === 401 || res.status === 403) {
+      throw new Error(`the server rejected the token (${res.status})`);
+    }
+    if (!res.ok) throw new Error(`unexpected response (${res.status})`);
+    db.insert(schema.mcpConnections)
+      .values({
+        id: op.connectionId,
+        userId: op.userId,
+        url: op.url,
+        authKind: "token",
+        secretEnc: pack(JSON.stringify(header)),
+        status: "connected",
+        updatedAt: now,
+      })
+      .onConflictDoUpdate({
+        target: schema.mcpConnections.id,
+        set: {
+          url: op.url,
+          authKind: "token",
+          secretEnc: pack(JSON.stringify(header)),
+          status: "connected",
+          updatedAt: now,
+        },
+      })
+      .run();
+    return { status: "connected" };
+  }
+
+  // Bare probe.
+  const res = await mcpInitialize(op.url);
+  if (res.ok) {
+    db.insert(schema.mcpConnections)
+      .values({
+        id: op.connectionId,
+        userId: op.userId,
+        url: op.url,
+        authKind: "none",
+        status: "connected",
+        updatedAt: now,
+      })
+      .onConflictDoUpdate({
+        target: schema.mcpConnections.id,
+        set: { url: op.url, authKind: "none", status: "connected", updatedAt: now },
+      })
+      .run();
+    return { status: "connected" };
+  }
+  if (res.status !== 401 && res.status !== 403) {
+    throw new Error(
+      `not an MCP server? initialize returned ${res.status} ${res.statusText}`,
+    );
+  }
+  if (!op.state || !op.redirectUri) {
+    throw new Error("server requires OAuth but no state/redirectUri provided");
+  }
+
+  // OAuth: discovery → client → PKCE → authorize URL.
+  const found = await discoverAuth(op.url, res.headers.get("www-authenticate"));
+  let clientId = op.clientId;
+  let clientSecret = op.clientSecret;
+  if (!clientId) {
+    if (!found.registrationEndpoint) {
+      throw new Error(
+        "the authorization server does not support dynamic client " +
+          "registration — enter a client id/secret manually",
+      );
+    }
+    const reg = await registerClient(found.registrationEndpoint, op.redirectUri);
+    clientId = reg.clientId;
+    clientSecret = reg.clientSecret;
+  }
+
+  const verifier = b64url(randomBytes(32));
+  const challenge = b64url(createHash("sha256").update(verifier).digest());
+  const authorize = new URL(found.authorizationEndpoint);
+  authorize.searchParams.set("response_type", "code");
+  authorize.searchParams.set("client_id", clientId);
+  authorize.searchParams.set("redirect_uri", op.redirectUri);
+  authorize.searchParams.set("state", op.state);
+  authorize.searchParams.set("code_challenge", challenge);
+  authorize.searchParams.set("code_challenge_method", "S256");
+  // RFC 8707 resource indicator — REQUIRED by the MCP auth spec.
+  authorize.searchParams.set("resource", op.url);
+  if (found.scopesSupported?.length) {
+    authorize.searchParams.set("scope", found.scopesSupported.join(" "));
+  }
+
+  const fields = {
+    userId: op.userId,
+    url: op.url,
+    authKind: "oauth",
+    tokenEndpoint: found.tokenEndpoint,
+    redirectUri: op.redirectUri,
+    clientId,
+    clientSecretEnc: clientSecret ? pack(clientSecret) : null,
+    pkceVerifierEnc: pack(verifier),
+    secretEnc: null,
+    status: "pending",
+    scopes: found.scopesSupported?.join(" ") ?? null,
+    updatedAt: now,
+  };
+  db.insert(schema.mcpConnections)
+    .values({ id: op.connectionId, ...fields })
+    .onConflictDoUpdate({ target: schema.mcpConnections.id, set: fields })
+    .run();
+  return { status: "auth_required", authorizeUrl: authorize.toString() };
+}
+
+async function oauthComplete(connectionId: string, code: string): Promise<McpAck> {
+  const db = getDb();
+  const row = db
+    .select()
+    .from(schema.mcpConnections)
+    .where(eq(schema.mcpConnections.id, connectionId))
+    .get();
+  if (!row) throw new Error("unknown connection — probe first");
+  if (!row.tokenEndpoint || !row.clientId || !row.pkceVerifierEnc || !row.redirectUri) {
+    throw new Error("connection has no pending OAuth flow");
+  }
+
+  const body = new URLSearchParams({
+    grant_type: "authorization_code",
+    code,
+    redirect_uri: row.redirectUri,
+    client_id: row.clientId,
+    code_verifier: unpack(row.pkceVerifierEnc),
+    // RFC 8707, mirrored from the authorize request.
+    resource: row.url,
+  });
+  const headers: Record<string, string> = {
+    "content-type": "application/x-www-form-urlencoded",
+    accept: "application/json",
+  };
+  if (row.clientSecretEnc) {
+    headers.authorization =
+      "Basic " +
+      Buffer.from(`${row.clientId}:${unpack(row.clientSecretEnc)}`).toString(
+        "base64",
+      );
+  }
+  const res = await timedFetch(row.tokenEndpoint, { method: "POST", headers, body });
+  const json = (await res.json().catch(() => ({}))) as Record<string, unknown>;
+  if (!res.ok || typeof json.access_token !== "string") {
+    const detail =
+      typeof json.error === "string" ? json.error : `status ${res.status}`;
+    throw new Error(`token exchange failed (${detail})`);
+  }
+
+  const secret: OauthSecret = {
+    accessToken: json.access_token,
+    refreshToken:
+      typeof json.refresh_token === "string" ? json.refresh_token : undefined,
+    expiresAt:
+      typeof json.expires_in === "number"
+        ? Date.now() + json.expires_in * 1000
+        : undefined,
+  };
+  const scopes =
+    typeof json.scope === "string" && json.scope.length > 0
+      ? json.scope
+      : row.scopes;
+  db.update(schema.mcpConnections)
+    .set({
+      secretEnc: pack(JSON.stringify(secret)),
+      pkceVerifierEnc: null,
+      status: "connected",
+      scopes,
+      updatedAt: Date.now(),
+    })
+    .where(eq(schema.mcpConnections.id, connectionId))
+    .run();
+  return {
+    status: "connected",
+    scopes: scopes ? scopes.split(" ").filter(Boolean) : undefined,
+  };
+}
+
+// --- gateway support ----------------------------------------------------------
+
+/** A connected row for `id`, or null. */
+export function getConnection(id: string): schema.McpConnection | null {
+  return (
+    getDb()
+      .select()
+      .from(schema.mcpConnections)
+      .where(eq(schema.mcpConnections.id, id))
+      .get() ?? null
+  );
+}
+
+/** All CONNECTED connections owned by `userId` (task-up + gateway ACL). */
+export function listConnectionsForUser(userId: string): schema.McpConnection[] {
+  return getDb()
+    .select()
+    .from(schema.mcpConnections)
+    .where(eq(schema.mcpConnections.userId, userId))
+    .all()
+    .filter((r) => r.status === "connected");
+}
+
+/**
+ * The Authorization (or custom) header for a connection, refreshing an
+ * expired OAuth access token first (rotating the stored refresh token when
+ * the server issues a new one). Null when the connection needs no header.
+ */
+export async function authHeaderFor(
+  row: schema.McpConnection,
+  forceRefresh = false,
+): Promise<{ name: string; value: string } | null> {
+  if (row.authKind === "none") return null;
+  if (!row.secretEnc) throw new Error("connection has no stored secret");
+  if (row.authKind === "token") {
+    const h = JSON.parse(unpack(row.secretEnc)) as HeaderSecret;
+    return { name: h.headerName, value: h.headerValue };
+  }
+  let secret = JSON.parse(unpack(row.secretEnc)) as OauthSecret;
+  if (
+    forceRefresh ||
+    (secret.expiresAt && secret.expiresAt - REFRESH_LEAD_MS < Date.now())
+  ) {
+    secret = await refreshTokens(row, secret);
+  }
+  return { name: "authorization", value: `Bearer ${secret.accessToken}` };
+}
+
+async function refreshTokens(
+  row: schema.McpConnection,
+  secret: OauthSecret,
+): Promise<OauthSecret> {
+  if (!secret.refreshToken || !row.tokenEndpoint || !row.clientId) {
+    throw new Error("access token expired and no refresh token is stored");
+  }
+  const body = new URLSearchParams({
+    grant_type: "refresh_token",
+    refresh_token: secret.refreshToken,
+    client_id: row.clientId,
+    resource: row.url,
+  });
+  const headers: Record<string, string> = {
+    "content-type": "application/x-www-form-urlencoded",
+    accept: "application/json",
+  };
+  if (row.clientSecretEnc) {
+    headers.authorization =
+      "Basic " +
+      Buffer.from(`${row.clientId}:${unpack(row.clientSecretEnc)}`).toString(
+        "base64",
+      );
+  }
+  const res = await timedFetch(row.tokenEndpoint, { method: "POST", headers, body });
+  const json = (await res.json().catch(() => ({}))) as Record<string, unknown>;
+  if (!res.ok || typeof json.access_token !== "string") {
+    throw new Error(`token refresh failed (status ${res.status})`);
+  }
+  const next: OauthSecret = {
+    accessToken: json.access_token,
+    refreshToken:
+      typeof json.refresh_token === "string"
+        ? json.refresh_token
+        : secret.refreshToken,
+    expiresAt:
+      typeof json.expires_in === "number"
+        ? Date.now() + json.expires_in * 1000
+        : undefined,
+  };
+  getDb()
+    .update(schema.mcpConnections)
+    .set({ secretEnc: pack(JSON.stringify(next)), updatedAt: Date.now() })
+    .where(eq(schema.mcpConnections.id, row.id))
+    .run();
+  return next;
+}