maestro/src/mcp/discovery.ts

import { createHash } from 'node:crypto';
import { checkSSRFStrict, pinnedFetch } from './ssrf-strict.js';
import type { McpDiscoveryMetadata, TokenEndpointResponse } from './types.js';

export interface FetchOpts {
  /**
   * When true, skip SSRF + HTTPS enforcement and use plain `fetch`.
   * ONLY set in tests against a 127.0.0.1 mock server.
   */
  insecureLocalTestMode?: boolean;
}

async function safeFetch(urlStr: string, init: RequestInit, opts: FetchOpts): Promise<Response> {
  if (opts.insecureLocalTestMode) {
    return fetch(urlStr, init);
  }
  const ssrf = await checkSSRFStrict(urlStr);
  if (!ssrf.ok) throw new Error(`SSRF check failed: ${ssrf.reason}`);
  return pinnedFetch(urlStr, { ...init, pinnedIp: ssrf.pinnedIp, family: ssrf.family });
}

function sameOrigin(a: string, b: string): boolean {
  try {
    const ua = new URL(a);
    const ub = new URL(b);
    return ua.origin === ub.origin;
  } catch {
    return false;
  }
}

export async function fetchDiscovery(
  mcpUrlStr: string,
  opts: FetchOpts = {},
): Promise<McpDiscoveryMetadata> {
  const mcpUrl = new URL(mcpUrlStr);
  const discoveryUrl = `${mcpUrl.origin}/.well-known/oauth-authorization-server`;
  const res = await safeFetch(discoveryUrl, { method: 'GET' }, opts);
  if (!res.ok) throw new Error(`Discovery fetch failed: ${res.status}`);
  const bodyText = await res.text();
  const meta = JSON.parse(bodyText) as {
    issuer: string;
    authorization_endpoint: string;
    token_endpoint: string;
  };

  if (!sameOrigin(meta.authorization_endpoint, mcpUrlStr)) {
    throw new Error(
      `authorization_endpoint origin must match MCP url origin: got ${meta.authorization_endpoint}`,
    );
  }
  if (!sameOrigin(meta.token_endpoint, mcpUrlStr)) {
    throw new Error(
      `token_endpoint origin must match MCP url origin: got ${meta.token_endpoint}`,
    );
  }

  const fingerprint = createHash('sha256').update(bodyText).digest('hex');
  return {
    issuer: meta.issuer,
    authorizationEndpoint: meta.authorization_endpoint,
    tokenEndpoint: meta.token_endpoint,
    fingerprint,
  };
}

export interface ExchangeInput {
  tokenEndpoint: string;
  clientId: string;
  clientSecret: string;
  code: string;
  redirectUri: string;
  codeVerifier: string;
}

export async function exchangeCode(
  input: ExchangeInput,
  opts: FetchOpts = {},
): Promise<TokenEndpointResponse> {
  const body = new URLSearchParams({
    grant_type: 'authorization_code',
    code: input.code,
    redirect_uri: input.redirectUri,
    client_id: input.clientId,
    client_secret: input.clientSecret,
    code_verifier: input.codeVerifier,
  });
  const res = await safeFetch(
    input.tokenEndpoint,
    {
      method: 'POST',
      headers: { 'content-type': 'application/x-www-form-urlencoded' },
      body: body.toString(),
    },
    opts,
  );
  if (!res.ok) {
    const text = await res.text();
    const err = Object.assign(new Error(`token exchange failed: ${res.status} ${text}`), {
      code: extractOauthError(text),
      status: res.status,
    });
    throw err;
  }
  return (await res.json()) as TokenEndpointResponse;
}

export async function refreshAccessToken(
  input: {
    tokenEndpoint: string;
    clientId: string;
    clientSecret: string;
    refreshToken: string;
  },
  opts: FetchOpts = {},
): Promise<TokenEndpointResponse> {
  const body = new URLSearchParams({
    grant_type: 'refresh_token',
    refresh_token: input.refreshToken,
    client_id: input.clientId,
    client_secret: input.clientSecret,
  });
  const res = await safeFetch(
    input.tokenEndpoint,
    {
      method: 'POST',
      headers: { 'content-type': 'application/x-www-form-urlencoded' },
      body: body.toString(),
    },
    opts,
  );
  if (!res.ok) {
    const text = await res.text();
    const err = Object.assign(new Error(`refresh failed: ${res.status} ${text}`), {
      code: extractOauthError(text),
      status: res.status,
    });
    throw err;
  }
  return (await res.json()) as TokenEndpointResponse;
}

function extractOauthError(text: string): string | undefined {
  try {
    const obj = JSON.parse(text) as { error?: string };
    return obj.error;
  } catch {
    return undefined;
  }
}