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

//! End-to-end test for the maestro-web proxy.
//!
//! Stands up a tiny axum app that pretends to be the upstream registry,
//! then a `maestro-web` router pointed at it, and verifies HTTP requests
//! and a WebSocket round-trip both make it through.

use std::net::{Ipv4Addr, SocketAddr};
use std::sync::Arc;
use std::time::Duration;

use axum::extract::ws::{Message, WebSocketUpgrade};
use axum::response::IntoResponse;
use axum::routing::{any, get};
use axum::Router;
use futures_util::{SinkExt, StreamExt};
use maestro_web::proxy::UpstreamConfig;
use maestro_web::router;
use tokio::net::TcpListener;
use tokio::sync::oneshot;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::{connect_async, tungstenite::Message as WsMessage};

struct Server {
    addr: SocketAddr,
    _shutdown: oneshot::Sender<()>,
}

async fn spawn_app(app: Router) -> Server {
    let listener = TcpListener::bind(SocketAddr::from((Ipv4Addr::LOCALHOST, 0)))
        .await
        .unwrap();
    let addr = listener.local_addr().unwrap();
    let (tx, rx) = oneshot::channel::<()>();
    tokio::spawn(async move {
        let _ = axum::serve(listener, app)
            .with_graceful_shutdown(async move {
                let _ = rx.await;
            })
            .await;
    });
    Server {
        addr,
        _shutdown: tx,
    }
}

async fn upstream_mock() -> Server {
    let app = Router::new()
        .route(
            "/api/daemons",
            get(|| async {
                (
                    [("set-cookie", "maestro_session=test; Path=/")],
                    "[{\"daemon_id\":\"echo\"}]",
                )
                    .into_response()
            }),
        )
        .route(
            "/api/echo",
            get(|req: axum::extract::Request| async move {
                // Echo the incoming Cookie header back in the body — proves it
                // forwarded.
                let cookie = req
                    .headers()
                    .get("cookie")
                    .and_then(|v| v.to_str().ok())
                    .unwrap_or("none")
                    .to_string();
                (axum::http::StatusCode::OK, format!("cookie={cookie}"))
            }),
        )
        .route(
            "/api/events",
            any(|ws: WebSocketUpgrade| async move {
                ws.on_upgrade(|mut socket| async move {
                    // Echo every text frame back, then close.
                    while let Some(Ok(msg)) = socket.next().await {
                        if let Message::Text(t) = msg {
                            let _ = socket.send(Message::Text(format!("upstream:{t}"))).await;
                        }
                    }
                })
            }),
        );
    spawn_app(app).await
}

#[tokio::test(flavor = "multi_thread")]
async fn proxies_http_with_cookies_and_set_cookie() {
    let upstream = upstream_mock().await;
    let cfg = Arc::new(UpstreamConfig::new(&format!("http://{}", upstream.addr)).unwrap());
    let proxy = spawn_app(router(cfg)).await;

    let client = reqwest::Client::builder().build().unwrap();
    let resp = client
        .get(format!("http://{}/api/daemons", proxy.addr))
        .send()
        .await
        .unwrap();
    assert!(resp.status().is_success());
    let cookie = resp
        .headers()
        .get("set-cookie")
        .and_then(|v| v.to_str().ok())
        .unwrap_or_default()
        .to_string();
    assert!(
        cookie.contains("maestro_session=test"),
        "expected Set-Cookie passthrough, got {cookie:?}"
    );
    let body = resp.text().await.unwrap();
    assert!(body.contains("\"echo\""));

    // Subsequent request: forward the cookie back upstream.
    let resp = client
        .get(format!("http://{}/api/echo", proxy.addr))
        .header("cookie", "maestro_session=test")
        .send()
        .await
        .unwrap();
    let body = resp.text().await.unwrap();
    assert!(
        body.contains("maestro_session=test"),
        "upstream didn't see the cookie: {body:?}"
    );
}

#[tokio::test(flavor = "multi_thread")]
async fn proxies_websocket_text_round_trip() {
    let upstream = upstream_mock().await;
    let cfg = Arc::new(UpstreamConfig::new(&format!("http://{}", upstream.addr)).unwrap());
    let proxy = spawn_app(router(cfg)).await;

    let url = format!("ws://{}/api/events", proxy.addr);
    let req = url.into_client_request().unwrap();
    let (mut ws, _resp) = connect_async(req).await.expect("ws connect");
    ws.send(WsMessage::Text("hello".into())).await.unwrap();

    let deadline = std::time::Instant::now() + Duration::from_secs(3);
    while std::time::Instant::now() < deadline {
        match tokio::time::timeout(Duration::from_millis(200), ws.next()).await {
            Ok(Some(Ok(WsMessage::Text(t)))) => {
                assert_eq!(t, "upstream:hello");
                return;
            }
            Ok(Some(Err(e))) => panic!("ws error: {e}"),
            Ok(None) => panic!("ws closed before echo"),
            _ => continue,
        }
    }
    panic!("did not receive echo within timeout");
}