hyperhive/hive-agent/src/web_ui/proxy.rs

//! Extra web proxies (`HIVE_EXTRA_WEB_PROXIES` / `services.hyperhive.agent.extraWebProxies`).
//!
//! Each declared entry mounts a transparent reverse-proxy at `/extra/<name>/`
//! in the per-agent web UI, forwarding every request (method, headers, body)
//! to the configured upstream. The `/extra/` namespace keeps user-declared
//! proxies from ever colliding with the native agent endpoints (`/api/*`,
//! `/events/*`, …). Response bodies are buffered whole (MVP): an SSE upstream
//! appears as one large response rather than streaming.
//!
//! Upstreams are either a regular `http(s)://` URL (forwarded via `reqwest`)
//! or a Unix domain socket, spelled `unix:<path>` (e.g.
//! `unix:/run/myapp/http.sock`) — forwarded via a raw hyper/1.1 client
//! dialing the socket directly, since `reqwest` has no UDS transport.

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;

use axum::{
    Router,
    body::Bytes,
    extract::State,
    http::{HeaderMap, HeaderName, HeaderValue, Method, StatusCode, Uri},
    response::{IntoResponse, Response},
};
use http_body_util::{BodyExt, Full};
use hyper_util::rt::TokioIo;

use super::AppState;

/// Hop-by-hop headers stripped on both the request and response side — they
/// describe a single transport hop and must not be forwarded to the upstream
/// or back to the client.
const HOP_BY_HOP: [&str; 6] = [
    "connection",
    "keep-alive",
    "transfer-encoding",
    "te",
    "trailer",
    "upgrade",
];

/// Bound on dialing a `unix:` upstream. The socket is local, so a connect
/// still pending after this long is not going to complete.
const UNIX_CONNECT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);

/// Bound on a `unix:` upstream answering with its response head — the same
/// budget the `http(s)://` path gives a whole request. Only the head is
/// bounded: once the upstream has answered, a long body is the upstream
/// working, not stuck.
const UNIX_RESPONSE_HEAD_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);

/// Nest every proxy declared in `HIVE_EXTRA_WEB_PROXIES` (a JSON object
/// `{"<name>": "<upstream_url>"}`) under `/extra/<name>/` on `app`. Absent /
/// blank / invalid JSON is a no-op (logged). Called from [`super::serve`]
/// before the static-dir fallback so `/extra/*` wins over `ServeDir`.
pub(super) fn mount_extra_proxies(mut app: Router<AppState>) -> Router<AppState> {
    let Ok(json) = std::env::var("HIVE_EXTRA_WEB_PROXIES") else {
        return app;
    };
    let Ok(map) = serde_json::from_str::<HashMap<String, String>>(&json) else {
        tracing::warn!("HIVE_EXTRA_WEB_PROXIES: invalid JSON — ignoring");
        return app;
    };
    for (name, upstream) in &map {
        if let Some(svc) = extra_proxy_service(upstream) {
            let mount = format!("/extra/{}", name.trim_matches('/'));
            tracing::info!(mount, upstream, "mounting extra web proxy");
            app = app.nest_service(mount.as_str(), svc);
        }
    }
    app
}

/// Build a transparent reverse-proxy service forwarding every request to
/// `upstream` — either an `http(s)://` URL or a `unix:<path>` Unix domain
/// socket. The caller nests it at a path prefix; axum strips the prefix
/// before the service sees the request. Returns `None` when `upstream` is
/// blank or the underlying client won't build.
fn extra_proxy_service(upstream: &str) -> Option<Router<()>> {
    let upstream = upstream.trim();
    if upstream.is_empty() {
        return None;
    }
    if let Some(sock_path) = upstream.strip_prefix("unix:") {
        let sock_path = sock_path.trim();
        if sock_path.is_empty() {
            tracing::warn!("extra proxy: `unix:` upstream has an empty socket path — ignoring");
            return None;
        }
        return Some(
            Router::new()
                .fallback(unix_proxy_handler)
                .with_state(Arc::new(PathBuf::from(sock_path))),
        );
    }

    let client = match reqwest::Client::builder()
        .timeout(std::time::Duration::from_secs(30))
        .build()
    {
        Ok(c) => Arc::new(c),
        Err(e) => {
            tracing::warn!("extra proxy: failed to build reqwest client: {e}");
            return None;
        }
    };
    let base = Arc::new(upstream.trim_end_matches('/').to_owned());
    Some(
        Router::new()
            .fallback(proxy_handler)
            .with_state((client, base)),
    )
}

/// Forward the request to `{base}{stripped_path_and_query}`, strip hop-by-hop
/// headers both ways, buffer the full response body, and return it verbatim.
async fn proxy_handler(
    State((client, base)): State<(Arc<reqwest::Client>, Arc<String>)>,
    method: Method,
    uri: Uri,
    headers: HeaderMap,
    body: Bytes,
) -> Response {
    let path_and_query = uri
        .path_and_query()
        .map_or("/", axum::http::uri::PathAndQuery::as_str);
    let url = format!("{base}{path_and_query}");

    let mut req_headers = headers;
    for h in HOP_BY_HOP {
        req_headers.remove(h);
    }

    let upstream_resp = match client
        .request(
            reqwest::Method::from_bytes(method.as_str().as_bytes()).unwrap_or(reqwest::Method::GET),
            &url,
        )
        .headers(
            req_headers
                .iter()
                .filter_map(|(k, v)| {
                    let n = reqwest::header::HeaderName::from_bytes(k.as_str().as_bytes()).ok()?;
                    let v = reqwest::header::HeaderValue::from_bytes(v.as_bytes()).ok()?;
                    Some((n, v))
                })
                .collect(),
        )
        .body(body.to_vec())
        .send()
        .await
    {
        Ok(r) => r,
        Err(e) => {
            tracing::warn!(url, "proxy: upstream request failed: {e}");
            return StatusCode::BAD_GATEWAY.into_response();
        }
    };

    let status =
        StatusCode::from_u16(upstream_resp.status().as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
    let mut resp_headers = HeaderMap::new();
    for (name, value) in upstream_resp.headers() {
        if HOP_BY_HOP.contains(&name.as_str()) {
            continue;
        }
        if let (Ok(n), Ok(v)) = (
            HeaderName::from_bytes(name.as_str().as_bytes()),
            HeaderValue::from_bytes(value.as_bytes()),
        ) {
            resp_headers.append(n, v);
        }
    }
    match upstream_resp.bytes().await {
        Ok(b) => (status, resp_headers, b).into_response(),
        Err(e) => {
            tracing::warn!(url, "proxy: failed to read upstream response: {e}");
            StatusCode::BAD_GATEWAY.into_response()
        }
    }
}

/// Forward the request to `sock_path` over a Unix domain socket, dialing a
/// fresh connection per request (MVP — no connection pooling, matching the
/// simplicity of the HTTP path above). `reqwest` has no UDS transport, so
/// this speaks raw HTTP/1.1 via `hyper`'s low-level client directly over a
/// `tokio::net::UnixStream`.
async fn unix_proxy_handler(
    State(sock_path): State<Arc<PathBuf>>,
    method: Method,
    uri: Uri,
    headers: HeaderMap,
    body: Bytes,
) -> Response {
    let path_and_query = uri
        .path_and_query()
        .map_or("/", axum::http::uri::PathAndQuery::as_str)
        .to_owned();

    match forward_over_unix_socket(&sock_path, method, &path_and_query, headers, body).await {
        Ok(resp) => resp,
        Err(e) => {
            tracing::warn!(
                socket = %sock_path.display(),
                path = path_and_query,
                "extra proxy: unix upstream request failed: {e}"
            );
            StatusCode::BAD_GATEWAY.into_response()
        }
    }
}

/// Dial `sock_path`, send one HTTP/1.1 request built from the given parts,
/// and return the axum `Response` built from whatever comes back. Split out
/// of [`unix_proxy_handler`] so the handler can map every failure mode
/// (connect, handshake, send, body read) to the same `BAD_GATEWAY` fallback
/// with one `?`-chain instead of a matching arm per step.
async fn forward_over_unix_socket(
    sock_path: &Path,
    method: Method,
    path_and_query: &str,
    mut headers: HeaderMap,
    body: Bytes,
) -> anyhow::Result<Response> {
    let stream = tokio::time::timeout(
        UNIX_CONNECT_TIMEOUT,
        tokio::net::UnixStream::connect(sock_path),
    )
    .await
    .map_err(|_| anyhow::anyhow!("connect timed out after {UNIX_CONNECT_TIMEOUT:?}"))??;
    let io = TokioIo::new(stream);

    let (mut sender, conn) = hyper::client::conn::http1::handshake(io).await?;
    // The connection future drives I/O in the background; drop-and-forget is
    // fine here since we only ever send one request per dialed socket.
    tokio::spawn(async move {
        if let Err(e) = conn.await {
            tracing::debug!("extra proxy: unix connection closed: {e}");
        }
    });

    for h in HOP_BY_HOP {
        headers.remove(h);
    }
    // HTTP/1.1 requires a Host header; a UDS peer has no meaningful
    // hostname, so `localhost` is the conventional placeholder (matches
    // what tools like `curl --unix-socket` send by default).
    if !headers.contains_key(hyper::header::HOST) {
        headers.insert(hyper::header::HOST, HeaderValue::from_static("localhost"));
    }

    let mut req_builder = hyper::Request::builder().method(method).uri(path_and_query);
    if let Some(h) = req_builder.headers_mut() {
        *h = headers;
    }
    let req = req_builder.body(Full::new(body))?;

    let upstream_resp = tokio::time::timeout(UNIX_RESPONSE_HEAD_TIMEOUT, sender.send_request(req))
        .await
        .map_err(|_| {
            anyhow::anyhow!(
                "waiting for the response head timed out after {UNIX_RESPONSE_HEAD_TIMEOUT:?}"
            )
        })??;
    let status = StatusCode::from_u16(upstream_resp.status().as_u16())?;
    let mut resp_headers = HeaderMap::new();
    for (name, value) in upstream_resp.headers() {
        if HOP_BY_HOP.contains(&name.as_str()) {
            continue;
        }
        resp_headers.append(name.clone(), value.clone());
    }
    let body_bytes = upstream_resp.into_body().collect().await?.to_bytes();
    Ok((status, resp_headers, body_bytes).into_response())
}

#[cfg(test)]
mod tests {
    use super::*;
    use tokio::io::{AsyncReadExt, AsyncWriteExt};

    /// Raw HTTP/1.1 answer with two `Set-Cookie` headers and a body.
    const TWO_COOKIES: &[u8] = b"HTTP/1.1 200 OK\r\n\
        Set-Cookie: session=abc\r\n\
        Set-Cookie: csrf=def\r\n\
        Content-Length: 2\r\n\
        Connection: close\r\n\r\nok";

    /// Read the request head, then answer with [`TWO_COOKIES`].
    async fn answer_once<S: AsyncReadExt + AsyncWriteExt + Unpin>(mut stream: S) {
        let mut buf = [0u8; 1024];
        let _ = stream.read(&mut buf).await;
        stream.write_all(TWO_COOKIES).await.unwrap();
        stream.shutdown().await.unwrap();
    }

    fn cookies(resp: &Response) -> Vec<&str> {
        resp.headers()
            .get_all("set-cookie")
            .iter()
            .map(|v| v.to_str().unwrap())
            .collect()
    }

    #[tokio::test]
    async fn unix_upstream_keeps_repeated_response_headers() {
        let dir = tempfile::tempdir().unwrap();
        let sock = dir.path().join("up.sock");
        let listener = tokio::net::UnixListener::bind(&sock).unwrap();
        tokio::spawn(async move {
            let (stream, _) = listener.accept().await.unwrap();
            answer_once(stream).await;
        });

        let resp =
            forward_over_unix_socket(&sock, Method::GET, "/", HeaderMap::new(), Bytes::new())
                .await
                .unwrap();
        assert_eq!(cookies(&resp), ["session=abc", "csrf=def"]);
    }

    #[tokio::test]
    async fn http_upstream_keeps_repeated_response_headers() {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let base = format!("http://{}", listener.local_addr().unwrap());
        tokio::spawn(async move {
            let (stream, _) = listener.accept().await.unwrap();
            answer_once(stream).await;
        });

        let client = reqwest::Client::builder().no_proxy().build().unwrap();
        let resp = proxy_handler(
            State((Arc::new(client), Arc::new(base))),
            Method::GET,
            Uri::from_static("/"),
            HeaderMap::new(),
            Bytes::new(),
        )
        .await;
        assert_eq!(cookies(&resp), ["session=abc", "csrf=def"]);
    }
}