AgentHub/server/lib/proxy-gateway.js

/**
 * @file proxy-gateway.js
 * @description Server mode's AI Proxy key gateway (S3 plan, P1b). A pane in
 * server mode is never given a proxy profile's API key: its Claude Code / Codex
 * talks to `http://127.0.0.1:<port>/api/proxy-gw/<profile>/v1/…` with its own
 * pane token as the key, and this handler forwards the request to the
 * profile's `baseUrl` with the real key, streaming the body and the answer
 * (SSE included) through untouched.
 *
 * A request is forwarded only when all hold:
 *  - server mode (in local mode the key is the user's own and the pane gets it
 *    as before — this route answers 404);
 *  - a loopback socket with no forwarding header (a pane runs on this host);
 *  - the bearer / `x-api-key` token is the pane token of a LIVE dashboard run
 *    (`dashboard_runs.pane_token_hash`, tmux session alive), and that run was
 *    started on THIS profile (`dashboard_runs.proxy`) — so a token never opens
 *    a profile its pane was not given, nor outlives its pane;
 *  - the path is under `/v1/` of the profile's own base URL (no `..` escape).
 *
 * The token and any client-sent key are dropped; the upstream sees only the
 * profile's key (both `Authorization: Bearer` and `x-api-key`, as
 * `testProxy` sends them). Nothing here logs a key or a token.
 *
 * Ceiling: an agent in the pane can still SPEND through the profile while its
 * pane lives — it just cannot carry the key away.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

const crypto = require("node:crypto");
const http = require("node:http");
const https = require("node:https");

/** A streamed answer may go quiet while the model thinks; not for this long. */
const IDLE_MS = 10 * 60 * 1000;

/** Request headers never passed upstream: hop-by-hop, the pane's token, cookies. */
/**
 * Set on every request the gateway sends upstream. A personal profile may
 * name ANY http(s) base — a loopback one included — and the gateway sends
 * from the dashboard's own uid, which hook auth trusts. The dashboard refuses
 * every request carrying this mark (server/index.js), so a profile pointed
 * back at it can never reach a route as the host; a real proxy ignores it.
 */
const GATEWAY_MARK = "x-agenthub-via-gateway";

const DROP_REQUEST = new Set([
  GATEWAY_MARK,
  "host",
  "connection",
  "keep-alive",
  "proxy-authorization",
  "proxy-connection",
  "te",
  "trailer",
  "transfer-encoding",
  "upgrade",
  "authorization",
  "x-api-key",
  "cookie",
  "x-agenthub-pane-token",
  "x-forwarded-for",
  "x-forwarded-host",
  "x-forwarded-proto",
  "forwarded",
]);

/** Answer headers node sets itself. */
const DROP_RESPONSE = new Set(["connection", "keep-alive", "transfer-encoding"]);

let listenPort = null;

/** The port this dashboard listens on — set by `startServer` once bound. */
function setListenPort(port) {
  listenPort = port;
}

/** This dashboard's port (the bound one once listening). */
function currentPort() {
  return listenPort || parseInt(process.env.DASHBOARD_PORT || "4820", 10);
}

/** The gateway endpoint a pane uses for `name`: Claude's base URL (Codex adds `/v1`). */
function gatewayBase(name) {
  return `http://127.0.0.1:${currentPort()}/api/proxy-gw/${encodeURIComponent(name)}`;
}

const sha256 = (s) => crypto.createHash("sha256").update(s).digest("hex");

const fail = (res, status, code, message) => res.status(status).json({ error: { code, message } });

function bearer(req) {
  const fromAuth = /^Bearer\s+(\S+)$/i.exec(String(req.headers.authorization || ""))?.[1];
  const fromKey = req.headers["x-api-key"];
  return [fromAuth, fromKey].find((t) => typeof t === "string" && t.startsWith("ahpane_")) || null;
}

/** The live run this pane token belongs to, or null. */
function liveRun(token, req) {
  if (!token || token.length > 200) return null;
  const { db } = require("../db");
  const row = db
    .prepare(
      `SELECT id, proxy, lane_id FROM dashboard_runs
        WHERE pane_token_hash = ? AND status IN ('running', 'spawning') AND instr(id, '@') = 0`
    )
    .get(sha256(token));
  if (!row) return null;
  // S3: only from the account that pane runs as (any local uid could post it).
  if (!require("./hook-auth").peerIsPaneAccount(req, row.lane_id)) return null;
  if (!require("./tmux").hasSession(row.id)) return null;
  return row;
}

/** Express handler, mounted at `/api/proxy-gw` BEFORE the body parsers. */
function proxyGateway(req, res) {
  const { isServerMode } = require("./auth");
  const { isLocalRequest } = require("./same-origin");
  if (!isServerMode()) return fail(res, 404, "ENOTFOUND", "not found");
  if (!isLocalRequest(req)) {
    return fail(res, 403, "ELOCALONLY", "the proxy gateway is for this machine's panes only");
  }

  // `/<name>/v1/<rest>?<query>`, relative to the mount point.
  const m = /^\/([^/?]+)(\/v1\/[^?]*)(\?.*)?$/.exec(req.url);
  if (!m) return fail(res, 400, "EBADPATH", "expected /api/proxy-gw/<profile>/v1/…");
  let name;
  try {
    name = decodeURIComponent(m[1]);
  } catch {
    return fail(res, 400, "EBADPATH", "bad profile name");
  }

  const run = liveRun(bearer(req), req);
  if (!run) return fail(res, 401, "EAUTH", "not a live pane's token");
  if (run.proxy !== name) {
    return fail(res, 403, "EPROXY", "this pane was not started on that AI Proxy");
  }

  let profile;
  try {
    profile = require("./ai-proxies").getProxy(name);
  } catch {
    return fail(res, 404, "ENOTFOUND", `no AI Proxy profile named '${name}'`);
  }
  if (profile.builtin) return fail(res, 404, "ENOTFOUND", "native has no endpoint");

  // The upstream URL stays under the profile's own `/v1/` — a `..` in the
  // path is resolved here and refused if it climbs out.
  const base = new URL(profile.baseUrl);
  const prefix = `${base.pathname.replace(/\/+$/, "")}/v1/`;
  const target = new URL(`${base.origin}${prefix}${m[2].slice("/v1/".length)}${m[3] || ""}`);
  if (target.origin !== base.origin || !target.pathname.startsWith(prefix)) {
    return fail(res, 400, "EBADPATH", "path leaves the profile's /v1/");
  }

  const headers = {};
  for (const [k, v] of Object.entries(req.headers)) {
    if (!DROP_REQUEST.has(k.toLowerCase())) headers[k] = v;
  }
  if (profile.apiKey) {
    headers.authorization = `Bearer ${profile.apiKey}`;
    headers["x-api-key"] = profile.apiKey;
  }
  headers[GATEWAY_MARK] = "1";

  const mod = target.protocol === "https:" ? https : http;
  const up = mod.request(target, { method: req.method, headers }, (answer) => {
    const out = {};
    for (const [k, v] of Object.entries(answer.headers)) {
      if (!DROP_RESPONSE.has(k.toLowerCase())) out[k] = v;
    }
    res.writeHead(answer.statusCode || 502, out);
    answer.pipe(res);
  });
  up.setTimeout(IDLE_MS, () => up.destroy(new Error("upstream went quiet")));
  up.on("error", (e) => {
    if (!res.headersSent) fail(res, 502, "EUPSTREAM", e.message);
    else res.destroy();
  });
  // The pane hung up (Ctrl-C mid-stream): stop the upstream call too.
  res.on("close", () => {
    if (!res.writableFinished) up.destroy();
  });
  req.pipe(up);
}

module.exports = {
  proxyGateway,
  gatewayBase,
  setListenPort,
  listenPort: currentPort,
  GATEWAY_MARK,
};