curriculum-project-hub/hub/src/connections/providerProxy.ts

import { randomBytes, timingSafeEqual } from "node:crypto";
import { createServer, type IncomingHttpHeaders, type IncomingMessage, type ServerResponse } from "node:http";
import { once } from "node:events";
import { classifyNetworkFailure, type NetworkFailureCategory } from "./networkFailure.js";

const MAX_PROVIDER_REQUEST_BYTES = 32 * 1024 * 1024;
const HOP_BY_HOP_HEADERS = new Set([
  "connection",
  "keep-alive",
  "proxy-authenticate",
  "proxy-authorization",
  "te",
  "trailer",
  "transfer-encoding",
  "upgrade",
]);
const REPLACED_REQUEST_HEADERS = new Set([
  ...HOP_BY_HOP_HEADERS,
  "authorization",
  "content-length",
  "host",
  "x-api-key",
]);

export interface ProviderUpstreamCredential {
  readonly baseUrl: string;
  readonly authToken: string;
  readonly anthropicApiKey: string;
}

export interface AgentProviderLease {
  readonly sdkEnv: Readonly<Record<string, string | undefined>>;
  readonly sensitiveValues: readonly string[];
  close(): Promise<void>;
}

export interface ProviderProxyDiagnostic {
  readonly code:
    | "provider_proxy_unauthorized"
    | "provider_proxy_request_too_large"
    | "provider_proxy_redirect_refused"
    | "provider_proxy_upstream_failed";
  readonly category: NetworkFailureCategory | "authorization" | "request_limit" | "redirect";
}

export interface ProviderProxyOptions {
  readonly onDiagnostic?: (diagnostic: ProviderProxyDiagnostic) => void;
}

/**
 * Starts one loopback-only proxy for one Agent run. The Agent receives only a
 * random run capability; customer provider credentials stay in this Hub
 * closure and are attached immediately before the upstream request.
 */
export async function openProviderProxyLease(
  upstream: ProviderUpstreamCredential,
  options: ProviderProxyOptions = {},
): Promise<AgentProviderLease> {
  const capability = randomBytes(32).toString("base64url");
  const handler = (request: IncomingMessage, response: ServerResponse): void => {
    void forwardProviderRequest(request, response, capability, upstream, options);
  };
  const server = createServer(handler);
  server.requestTimeout = 5 * 60_000;
  server.headersTimeout = 30_000;
  server.maxRequestsPerSocket = 1_000;
  server.listen(0, "127.0.0.1");
  await Promise.race([
    once(server, "listening"),
    once(server, "error").then(([error]) => Promise.reject(error)),
  ]);
  server.unref();
  const address = server.address();
  if (address === null || typeof address === "string") {
    server.close();
    throw new Error("provider proxy did not bind a TCP address");
  }

  let closed = false;
  return {
    sdkEnv: {
      ANTHROPIC_BASE_URL: `http://127.0.0.1:${address.port}`,
      ANTHROPIC_AUTH_TOKEN: capability,
      ANTHROPIC_API_KEY: "",
    },
    sensitiveValues: [capability],
    async close(): Promise<void> {
      if (closed) return;
      closed = true;
      server.closeAllConnections();
      if (!server.listening) return;
      server.close();
      await once(server, "close");
    },
  };
}

async function forwardProviderRequest(
  request: IncomingMessage,
  response: ServerResponse,
  capability: string,
  upstream: ProviderUpstreamCredential,
  options: ProviderProxyOptions,
): Promise<void> {
  if (!authorized(request.headers.authorization, capability)) {
    options.onDiagnostic?.({ code: "provider_proxy_unauthorized", category: "authorization" });
    response.writeHead(401, { "content-type": "text/plain; charset=utf-8" });
    response.end("unauthorized");
    return;
  }

  try {
    const requestUrl = new URL(request.url ?? "/", "http://provider-proxy.invalid");
    const upstreamBase = upstream.baseUrl.endsWith("/") ? upstream.baseUrl : `${upstream.baseUrl}/`;
    const upstreamUrl = new URL(`${requestUrl.pathname.replace(/^\//, "")}${requestUrl.search}`, upstreamBase);
    const body = await readBoundedBody(request);
    const headers = forwardedRequestHeaders(request.headers);
    headers.set("authorization", `Bearer ${upstream.authToken}`);
    if (upstream.anthropicApiKey !== "") {
      headers.set("x-api-key", upstream.anthropicApiKey);
    }
    headers.set("accept-encoding", "identity");

    const abort = new AbortController();
    request.once("aborted", () => abort.abort());
    const upstreamResponse = await fetch(upstreamUrl, {
      method: request.method ?? "GET",
      headers,
      ...(body.byteLength === 0 ? {} : { body }),
      redirect: "manual",
      signal: abort.signal,
    });
    if (upstreamResponse.status >= 300 && upstreamResponse.status < 400) {
      options.onDiagnostic?.({ code: "provider_proxy_redirect_refused", category: "redirect" });
      await upstreamResponse.body?.cancel();
      response.writeHead(502, { "content-type": "text/plain; charset=utf-8" });
      response.end("provider redirect refused");
      return;
    }
    response.writeHead(upstreamResponse.status, forwardedResponseHeaders(upstreamResponse.headers));
    if (upstreamResponse.body === null) {
      response.end();
      return;
    }
    for await (const chunk of upstreamResponse.body) {
      if (!response.write(Buffer.from(chunk))) {
        await once(response, "drain");
      }
    }
    response.end();
  } catch (error) {
    const diagnostic = error instanceof ProviderProxyRequestError
      ? { code: error.code, category: error.category } as const
      : {
          code: "provider_proxy_upstream_failed",
          category: classifyNetworkFailure(error),
        } as const;
    options.onDiagnostic?.(diagnostic);
    if (!response.headersSent) {
      response.writeHead(502, { "content-type": "text/plain; charset=utf-8" });
      response.end("provider request failed");
    } else {
      response.destroy();
    }
  }
}

function authorized(value: string | undefined, capability: string): boolean {
  if (value === undefined || !value.startsWith("Bearer ")) return false;
  const supplied = Buffer.from(value.slice("Bearer ".length));
  const expected = Buffer.from(capability);
  return supplied.byteLength === expected.byteLength && timingSafeEqual(supplied, expected);
}

function forwardedRequestHeaders(source: IncomingHttpHeaders): Headers {
  const result = new Headers();
  for (const [name, value] of Object.entries(source)) {
    if (value === undefined || REPLACED_REQUEST_HEADERS.has(name.toLowerCase())) continue;
    if (Array.isArray(value)) {
      for (const item of value) result.append(name, item);
    } else {
      result.set(name, value);
    }
  }
  return result;
}

function forwardedResponseHeaders(source: Headers): Record<string, string> {
  const result: Record<string, string> = {};
  for (const [name, value] of source.entries()) {
    const normalized = name.toLowerCase();
    if (HOP_BY_HOP_HEADERS.has(normalized) || normalized === "content-length" ||
        normalized === "content-encoding") continue;
    result[name] = value;
  }
  return result;
}

async function readBoundedBody(request: IncomingMessage): Promise<Buffer> {
  const declared = Number(request.headers["content-length"] ?? 0);
  if (Number.isFinite(declared) && declared > MAX_PROVIDER_REQUEST_BYTES) {
    throw new ProviderProxyRequestError("provider_proxy_request_too_large", "request_limit");
  }
  const chunks: Buffer[] = [];
  let bytes = 0;
  for await (const chunk of request) {
    const buffer = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk as Uint8Array);
    bytes += buffer.byteLength;
    if (bytes > MAX_PROVIDER_REQUEST_BYTES) {
      throw new ProviderProxyRequestError("provider_proxy_request_too_large", "request_limit");
    }
    chunks.push(buffer);
  }
  return Buffer.concat(chunks, bytes);
}

class ProviderProxyRequestError extends Error {
  constructor(
    readonly code: "provider_proxy_request_too_large",
    readonly category: "request_limit",
  ) {
    super(code);
    this.name = "ProviderProxyRequestError";
  }
}