maestro/crates/maestro-web/src/proxy.rs

//! Reverse-proxy primitives: small hand-written hyper-client handler for
//! `/api/...` HTTP, and an axum WebSocket pump for `/api/attach/:id` +
//! `/api/events`.
//!
//! Keeps dependencies minimal — reqwest for HTTP, tokio_tungstenite for the
//! upstream WS leg. ~150 LOC total.

use std::sync::Arc;

use axum::body::Body;
use axum::extract::ws::{Message as AxMessage, WebSocket};
use axum::extract::{Request, State};
use axum::http::{header, HeaderMap, HeaderValue, StatusCode, Uri};
use axum::response::{IntoResponse, Response};
use futures_util::{SinkExt, StreamExt};
use tokio_tungstenite::tungstenite::Message as WsMessage;

#[derive(Debug, Clone)]
pub struct UpstreamConfig {
    /// HTTP base of the daemon — `http://127.0.0.1:7070`.
    pub daemon_http_base: Arc<String>,
    /// WS base of the daemon — derived from `daemon_http_base` by replacing
    /// the scheme. Cached so we don't redo it per request.
    pub daemon_ws_base: Arc<String>,
}

impl UpstreamConfig {
    pub fn from_daemon_url(daemon_url: &str) -> Self {
        let http_base = daemon_url.trim_end_matches('/').to_string();
        let ws_base = if let Some(rest) = http_base.strip_prefix("https://") {
            format!("wss://{rest}")
        } else if let Some(rest) = http_base.strip_prefix("http://") {
            format!("ws://{rest}")
        } else {
            // Caller passed something already scheme-less or pre-wrapped; trust
            // it. Validation happens at config-load time.
            http_base.clone()
        };
        Self {
            daemon_http_base: Arc::new(http_base),
            daemon_ws_base: Arc::new(ws_base),
        }
    }
}

/// Shared state plumbed into the proxy handlers.
#[derive(Debug, Clone)]
pub struct ProxyState {
    pub upstream: UpstreamConfig,
    pub client: reqwest::Client,
}

impl ProxyState {
    pub fn new(upstream: UpstreamConfig) -> Self {
        Self {
            upstream,
            client: reqwest::Client::builder()
                .redirect(reqwest::redirect::Policy::none())
                .build()
                .expect("reqwest client"),
        }
    }
}

// ---------------------------------------------------------------------------
// HTTP proxy
// ---------------------------------------------------------------------------

/// Forward an HTTP request to the upstream daemon and stream the response back.
pub async fn proxy_http(State(state): State<ProxyState>, req: Request) -> Response {
    let (parts, body) = req.into_parts();
    let upstream_uri = build_uri(&state.upstream.daemon_http_base, &parts.uri);

    let body_bytes = match axum::body::to_bytes(body, usize::MAX).await {
        Ok(b) => b,
        Err(e) => return bad_gateway(&format!("read body: {e}")),
    };

    let mut upstream_req = state
        .client
        .request(parts.method.clone(), upstream_uri.as_str())
        .body(body_bytes.to_vec());

    for (name, value) in filter_request_headers(&parts.headers).iter() {
        upstream_req = upstream_req.header(name, value);
    }

    let upstream_resp = match upstream_req.send().await {
        Ok(r) => r,
        Err(e) => return bad_gateway(&format!("upstream send: {e}")),
    };

    let status = upstream_resp.status();
    let headers = filter_response_headers(upstream_resp.headers());
    let body = match upstream_resp.bytes().await {
        Ok(b) => b,
        Err(e) => return bad_gateway(&format!("upstream body: {e}")),
    };

    let mut resp = Response::builder().status(status);
    for (name, value) in headers.iter() {
        resp = resp.header(name, value);
    }
    resp.body(Body::from(body))
        .unwrap_or_else(|e| bad_gateway(&format!("build response: {e}")))
}

fn build_uri(base: &str, original: &Uri) -> String {
    let mut out = base.to_string();
    out.push_str(original.path());
    if let Some(q) = original.query() {
        out.push('?');
        out.push_str(q);
    }
    out
}

/// Headers we strip before forwarding upstream / downstream. `host` must be
/// dropped (reqwest will set it correctly for the upstream); `content-length`
/// is regenerated by the client; hop-by-hop headers per RFC 7230 §6.1 should
/// not be relayed.
const HOP_BY_HOP: &[&str] = &[
    "connection",
    "keep-alive",
    "proxy-authenticate",
    "proxy-authorization",
    "te",
    "trailers",
    "transfer-encoding",
    "upgrade",
];

fn filter_request_headers(headers: &HeaderMap) -> HeaderMap {
    let mut out = HeaderMap::new();
    for (name, value) in headers {
        let n = name.as_str();
        if n.eq_ignore_ascii_case("host") || n.eq_ignore_ascii_case("content-length") {
            continue;
        }
        if HOP_BY_HOP.iter().any(|h| n.eq_ignore_ascii_case(h)) {
            continue;
        }
        out.append(name.clone(), value.clone());
    }
    out
}

fn filter_response_headers(headers: &reqwest::header::HeaderMap) -> HeaderMap {
    let mut out = HeaderMap::new();
    for (name, value) in headers {
        let n = name.as_str();
        if HOP_BY_HOP.iter().any(|h| n.eq_ignore_ascii_case(h)) {
            continue;
        }
        if let (Ok(n), Ok(v)) = (
            axum::http::HeaderName::from_bytes(name.as_str().as_bytes()),
            HeaderValue::from_bytes(value.as_bytes()),
        ) {
            out.append(n, v);
        }
    }
    out
}

fn bad_gateway(msg: &str) -> Response {
    tracing::warn!(error = msg, "proxy: bad gateway");
    // Plain JSON without a serde_json dep — the message is short and we control
    // every byte we write, so manual escaping of `"` + `\` is enough here.
    let escaped = msg.replace('\\', "\\\\").replace('"', "\\\"");
    let body = format!(r#"{{"error":{{"code":"proxy_failed","message":"{escaped}"}}}}"#);
    (
        StatusCode::BAD_GATEWAY,
        [(header::CONTENT_TYPE, "application/json")],
        body,
    )
        .into_response()
}

// ---------------------------------------------------------------------------
// WebSocket proxy
// ---------------------------------------------------------------------------

/// Open an upstream WS connection at `daemon_ws_base + path?query` and pump
/// frames in both directions until either side closes. Used for
/// `/api/attach/:id` and `/api/events`.
pub async fn pump_ws(client: WebSocket, upstream_url: String) {
    let (upstream, _) = match tokio_tungstenite::connect_async(&upstream_url).await {
        Ok(pair) => pair,
        Err(e) => {
            tracing::warn!(url = %upstream_url, error = %e, "proxy: upstream WS connect failed");
            // Drop the client socket — there's nothing useful to send it.
            return;
        }
    };
    let (mut client_sink, mut client_stream) = client.split();
    let (mut up_sink, mut up_stream) = upstream.split();

    let client_to_up = async {
        while let Some(msg) = client_stream.next().await {
            let Ok(msg) = msg else { break };
            let frame = match msg {
                AxMessage::Text(t) => WsMessage::Text(t),
                AxMessage::Binary(b) => WsMessage::Binary(b),
                AxMessage::Ping(p) => WsMessage::Ping(p),
                AxMessage::Pong(p) => WsMessage::Pong(p),
                AxMessage::Close(c) => {
                    let _ = up_sink
                        .send(WsMessage::Close(c.map(|cf| {
                            tokio_tungstenite::tungstenite::protocol::CloseFrame {
                                code: cf.code.into(),
                                reason: cf.reason,
                            }
                        })))
                        .await;
                    break;
                }
            };
            if up_sink.send(frame).await.is_err() {
                break;
            }
        }
    };

    let up_to_client = async {
        while let Some(msg) = up_stream.next().await {
            let Ok(msg) = msg else { break };
            let frame = match msg {
                WsMessage::Text(t) => AxMessage::Text(t),
                WsMessage::Binary(b) => AxMessage::Binary(b),
                WsMessage::Ping(p) => AxMessage::Ping(p),
                WsMessage::Pong(p) => AxMessage::Pong(p),
                WsMessage::Close(c) => {
                    let _ = client_sink
                        .send(AxMessage::Close(c.map(|cf| {
                            axum::extract::ws::CloseFrame {
                                code: cf.code.into(),
                                reason: cf.reason.to_string().into(),
                            }
                        })))
                        .await;
                    break;
                }
                WsMessage::Frame(_) => continue,
            };
            if client_sink.send(frame).await.is_err() {
                break;
            }
        }
    };

    // Stop as soon as either side hangs up.
    tokio::select! {
        _ = client_to_up => {}
        _ = up_to_client => {}
    }
}

/// Compose the upstream WS URL from the base + path + query of the inbound
/// request.
pub fn upstream_ws_url(state: &ProxyState, uri: &Uri) -> String {
    let mut out = state.upstream.daemon_ws_base.as_str().to_string();
    out.push_str(uri.path());
    if let Some(q) = uri.query() {
        out.push('?');
        out.push_str(q);
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn ws_base_swaps_http_scheme() {
        let u = UpstreamConfig::from_daemon_url("http://127.0.0.1:7070/");
        assert_eq!(&*u.daemon_http_base, "http://127.0.0.1:7070");
        assert_eq!(&*u.daemon_ws_base, "ws://127.0.0.1:7070");
    }

    #[test]
    fn ws_base_swaps_https_scheme() {
        let u = UpstreamConfig::from_daemon_url("https://daemon.example/");
        assert_eq!(&*u.daemon_ws_base, "wss://daemon.example");
    }
}