remora/crates/remora-outbound/tests/transport.rs

use std::{
    collections::HashMap,
    io,
    net::IpAddr,
    sync::{Arc, Mutex},
};

use http::{HeaderMap, HeaderValue};
use ipnet::IpNet;
use remora_outbound::{
    OutboundClient, OutboundError, OutboundPolicyBuilder, OutboundRequest, Resolver,
};
use tokio::{
    io::{AsyncRead, AsyncReadExt, AsyncWriteExt},
    net::TcpListener,
    task::JoinHandle,
};

#[derive(Clone)]
struct ScriptedResolver {
    state: Arc<Mutex<ResolverState>>,
}

struct ResolverState {
    answers: HashMap<String, Vec<Vec<IpAddr>>>,
    calls: HashMap<String, usize>,
}

#[derive(Debug)]
struct ObservedRequest {
    path: String,
    host: Option<String>,
    authorization: Option<String>,
    cookie: Option<String>,
    api_key: Option<String>,
}

impl ScriptedResolver {
    fn new(answers: impl IntoIterator<Item = (&'static str, Vec<Vec<IpAddr>>)>) -> Self {
        Self {
            state: Arc::new(Mutex::new(ResolverState {
                answers: answers
                    .into_iter()
                    .map(|(host, answers)| (host.to_owned(), answers))
                    .collect(),
                calls: HashMap::new(),
            })),
        }
    }

    fn calls(&self, host: &str) -> usize {
        self.state
            .lock()
            .expect("resolver state")
            .calls
            .get(host)
            .copied()
            .unwrap_or_default()
    }
}

impl Resolver for ScriptedResolver {
    async fn resolve(&self, host: &str, _port: u16) -> io::Result<Vec<IpAddr>> {
        let mut state = self.state.lock().expect("resolver state");
        let call = state.calls.entry(host.to_owned()).or_default();
        let index = *call;
        *call += 1;
        let sequence = state.answers.get(host).cloned().unwrap_or_default();
        Ok(sequence
            .get(index)
            .or_else(|| sequence.last())
            .cloned()
            .unwrap_or_default())
    }
}

fn ip(input: &str) -> IpAddr {
    input.parse().expect("valid test address")
}

fn local_policy(port: u16, max_redirects: usize) -> remora_outbound::OutboundPolicy {
    OutboundPolicyBuilder::new()
        .allow_scheme_port("http", port)
        .allow_non_global("127.0.0.0/8".parse::<IpNet>().expect("network"))
        .max_redirects(max_redirects)
        .build()
}

async fn serve_http(
    connections: usize,
    handler: impl Fn(&ObservedRequest) -> String + Send + Sync + 'static,
) -> (u16, Arc<Mutex<Vec<ObservedRequest>>>, JoinHandle<()>) {
    let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
    let port = listener.local_addr().expect("address").port();
    let observed = Arc::new(Mutex::new(Vec::new()));
    let task_observed = Arc::clone(&observed);
    let handler = Arc::new(handler);
    let task = tokio::spawn(async move {
        for _ in 0..connections {
            let (mut stream, _) = listener.accept().await.expect("accept");
            let request = read_request(&mut stream).await;
            let response = handler(&request);
            task_observed.lock().expect("observations").push(request);
            stream
                .write_all(response.as_bytes())
                .await
                .expect("response");
        }
    });
    (port, observed, task)
}

async fn read_request(stream: &mut (impl AsyncRead + Unpin)) -> ObservedRequest {
    let mut bytes = Vec::new();
    let mut chunk = [0_u8; 1024];
    while !bytes.windows(4).any(|window| window == b"\r\n\r\n") {
        let read = stream.read(&mut chunk).await.expect("request");
        assert_ne!(read, 0, "request ended before headers");
        bytes.extend_from_slice(&chunk[..read]);
    }
    let request = String::from_utf8(bytes).expect("ASCII request");
    let mut lines = request.split("\r\n");
    let path = lines
        .next()
        .and_then(|line| line.split_whitespace().nth(1))
        .expect("request path")
        .to_owned();
    let headers: HashMap<String, String> = lines
        .take_while(|line| !line.is_empty())
        .filter_map(|line| line.split_once(':'))
        .map(|(name, value)| (name.to_ascii_lowercase(), value.trim().to_owned()))
        .collect();
    ObservedRequest {
        path,
        host: headers.get("host").cloned(),
        authorization: headers.get("authorization").cloned(),
        cookie: headers.get("cookie").cloned(),
        api_key: headers.get("x-api-key").cloned(),
    }
}

fn response(status: &str, headers: &[(&str, &str)], body: &str) -> String {
    let fields = headers
        .iter()
        .map(|(name, value)| format!("{name}: {value}\r\n"))
        .collect::<String>();
    format!(
        "HTTP/1.1 {status}\r\n{fields}Content-Length: {}\r\nConnection: close\r\n\r\n{body}",
        body.len()
    )
}

#[tokio::test]
async fn pins_address_and_preserves_http_authority() {
    let (port, observed, server) = serve_http(1, |_| response("200 OK", &[], "ok")).await;
    let resolver = ScriptedResolver::new([("source.test", vec![vec![ip("127.0.0.1")]])]);
    let client = OutboundClient::builder(resolver.clone(), local_policy(port, 5)).build();
    let result = client
        .execute(
            OutboundRequest::get(&format!("http://source.test:{port}/object")).expect("request"),
        )
        .await
        .expect("response");
    assert_eq!(result.response().status(), 200);
    server.await.expect("server");
    let observed = observed.lock().expect("observations");
    assert_eq!(
        observed[0].host.as_deref(),
        Some(format!("source.test:{port}").as_str())
    );
    assert_eq!(observed[0].path, "/object");
    assert_eq!(resolver.calls("source.test"), 1);
}

#[tokio::test]
async fn redirects_reauthorize_and_strip_cross_origin_credentials() {
    let (port, observed, server) = serve_http(2, |request| {
        if request.path == "/redirect" {
            let port = request
                .host
                .as_deref()
                .and_then(|host| host.rsplit_once(':'))
                .map(|(_, port)| port)
                .expect("authority port");
            response(
                "302 Found",
                &[("Location", &format!("http://other.test:{port}/ok"))],
                "",
            )
        } else {
            response("200 OK", &[], "ok")
        }
    })
    .await;
    let resolver = ScriptedResolver::new([
        ("source.test", vec![vec![ip("127.0.0.1")]]),
        ("other.test", vec![vec![ip("127.0.0.1")]]),
    ]);
    let client = OutboundClient::builder(resolver.clone(), local_policy(port, 5)).build();
    let mut headers = HeaderMap::new();
    headers.insert(
        "authorization",
        HeaderValue::from_static("Bearer synthetic"),
    );
    headers.insert("cookie", HeaderValue::from_static("session=synthetic"));
    headers.insert("x-api-key", HeaderValue::from_static("synthetic-api-key"));
    let request = OutboundRequest::get(&format!("http://source.test:{port}/redirect"))
        .expect("request")
        .with_credential_headers(headers)
        .expect("headers");
    let result = client.execute(request).await.expect("response");
    assert_eq!(result.final_url().host_str(), "other.test");
    server.await.expect("server");
    let observed = observed.lock().expect("observations");
    assert_eq!(
        observed[0].authorization.as_deref(),
        Some("Bearer synthetic")
    );
    assert_eq!(observed[0].cookie.as_deref(), Some("session=synthetic"));
    assert_eq!(observed[0].api_key.as_deref(), Some("synthetic-api-key"));
    assert_eq!(observed[1].authorization, None);
    assert_eq!(observed[1].cookie, None);
    assert_eq!(observed[1].api_key, None);
    assert_eq!(resolver.calls("source.test"), 1);
    assert_eq!(resolver.calls("other.test"), 1);
}

#[test]
fn arbitrary_headers_must_be_classified_as_credentials() {
    let mut headers = HeaderMap::new();
    headers.insert("x-api-key", HeaderValue::from_static("synthetic-api-key"));
    assert_eq!(
        OutboundRequest::get("https://example.com/")
            .expect("request")
            .with_public_headers(headers)
            .expect_err("unclassified header"),
        OutboundError::InvalidHeaders
    );
}

#[test]
fn request_debug_redacts_url_and_credential_values() {
    let mut headers = HeaderMap::new();
    headers.insert("x-api-key", HeaderValue::from_static("synthetic-api-key"));
    let request = OutboundRequest::get("https://example.com/sensitive?token=synthetic")
        .expect("request")
        .with_credential_headers(headers)
        .expect("credential headers");
    let debug = format!("{request:?}");
    assert!(!debug.contains("sensitive"));
    assert!(!debug.contains("synthetic-api-key"));
}

#[tokio::test]
async fn response_debug_redacts_url_and_response_headers() {
    let (port, _, server) = serve_http(1, |_| {
        response(
            "200 OK",
            &[("Set-Cookie", "session=synthetic-secret")],
            "ok",
        )
    })
    .await;
    let resolver = ScriptedResolver::new([("source.test", vec![vec![ip("127.0.0.1")]])]);
    let client = OutboundClient::builder(resolver, local_policy(port, 5)).build();
    let result = client
        .execute(
            OutboundRequest::get(&format!(
                "http://source.test:{port}/sensitive?token=synthetic"
            ))
            .expect("request"),
        )
        .await
        .expect("response");
    let debug = format!("{result:?}");
    assert!(debug.contains("200"));
    assert!(!debug.contains("sensitive"));
    assert!(!debug.contains("synthetic-secret"));
    server.await.expect("server");
}

#[tokio::test]
async fn denied_redirect_is_rejected_before_dial() {
    let (port, _, server) = serve_http(1, |request| {
        let port = request
            .host
            .as_deref()
            .and_then(|host| host.rsplit_once(':'))
            .map(|(_, port)| port)
            .expect("authority port");
        response(
            "302 Found",
            &[(
                "Location",
                &format!("http://metadata.test:{port}/latest/meta-data"),
            )],
            "",
        )
    })
    .await;
    let resolver = ScriptedResolver::new([
        ("source.test", vec![vec![ip("127.0.0.1")]]),
        ("metadata.test", vec![vec![ip("169.254.169.254")]]),
    ]);
    let client = OutboundClient::builder(resolver, local_policy(port, 5)).build();
    let result = client
        .execute(OutboundRequest::get(&format!("http://source.test:{port}/")).expect("request"))
        .await;
    assert_eq!(
        result.expect_err("denied redirect"),
        OutboundError::AddressDenied
    );
    server.await.expect("server");
}

#[tokio::test]
async fn redirect_count_is_bounded() {
    let (port, _, server) =
        serve_http(3, |_| response("302 Found", &[("Location", "/loop")], "")).await;
    let resolver = ScriptedResolver::new([("source.test", vec![vec![ip("127.0.0.1")]])]);
    let client = OutboundClient::builder(resolver.clone(), local_policy(port, 2)).build();
    let result = client
        .execute(OutboundRequest::get(&format!("http://source.test:{port}/loop")).expect("request"))
        .await;
    assert_eq!(
        result.expect_err("redirect bound"),
        OutboundError::RedirectLimit
    );
    server.await.expect("server");
    assert_eq!(resolver.calls("source.test"), 3);
}

#[tokio::test]
async fn a_new_request_cannot_reuse_a_pool_after_dns_changes() {
    let (port, _, server) = serve_http(1, |_| response("200 OK", &[], "ok")).await;
    let resolver = ScriptedResolver::new([(
        "rebind.test",
        vec![vec![ip("127.0.0.1")], vec![ip("169.254.169.254")]],
    )]);
    let client = OutboundClient::builder(resolver.clone(), local_policy(port, 5)).build();
    let url = format!("http://rebind.test:{port}/ok");
    client
        .execute(OutboundRequest::get(&url).expect("first request"))
        .await
        .expect("first response");
    assert_eq!(
        client
            .execute(OutboundRequest::get(&url).expect("second request"))
            .await
            .expect_err("rebound address"),
        OutboundError::AddressDenied
    );
    server.await.expect("server");
    assert_eq!(resolver.calls("rebind.test"), 2);
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn ignores_ambient_proxy_configuration() {
    let (port, _, server) = serve_http(1, |_| response("200 OK", &[], "ok")).await;
    let status = std::process::Command::new(std::env::current_exe().expect("test executable"))
        .args(["--exact", "ambient_proxy_child", "--nocapture"])
        .env("REMORA_PROXY_CHILD_PORT", port.to_string())
        .env("HTTP_PROXY", "http://169.254.169.254:8080")
        .status()
        .expect("proxy child");
    assert!(status.success());
    server.await.expect("server");
}

#[test]
fn ambient_proxy_child() {
    let Ok(port) = std::env::var("REMORA_PROXY_CHILD_PORT") else {
        return;
    };
    let port = port.parse::<u16>().expect("child port");
    tokio::runtime::Runtime::new()
        .expect("runtime")
        .block_on(async move {
            let resolver = ScriptedResolver::new([("source.test", vec![vec![ip("127.0.0.1")]])]);
            let client = OutboundClient::builder(resolver, local_policy(port, 5)).build();
            client
                .execute(
                    OutboundRequest::get(&format!("http://source.test:{port}/")).expect("request"),
                )
                .await
                .expect("direct response");
        });
}

#[tokio::test]
async fn tls_verification_uses_canonical_hostname_not_pinned_address() {
    let certified =
        rcgen::generate_simple_self_signed(vec!["source.test".to_owned()]).expect("certificate");
    let certificate_der = certified.cert.der().clone();
    let key_der = certified.signing_key.serialize_der();
    let tls = rustls::ServerConfig::builder()
        .with_no_client_auth()
        .with_single_cert(
            vec![certificate_der.clone()],
            rustls::pki_types::PrivateKeyDer::Pkcs8(key_der.into()),
        )
        .expect("TLS config");
    let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
    let port = listener.local_addr().expect("address").port();
    let server = tokio::spawn(async move {
        let (stream, _) = listener.accept().await.expect("accept");
        let mut stream = tokio_rustls::TlsAcceptor::from(Arc::new(tls))
            .accept(stream)
            .await
            .expect("TLS handshake");
        let request = read_request(&mut stream).await;
        assert_eq!(
            request.host.as_deref(),
            Some(format!("source.test:{port}").as_str())
        );
        stream
            .write_all(response("200 OK", &[], "ok").as_bytes())
            .await
            .expect("response");
    });
    let policy = OutboundPolicyBuilder::new()
        .allow_scheme_port("https", port)
        .allow_non_global("127.0.0.0/8".parse::<IpNet>().expect("network"))
        .build();
    let resolver = ScriptedResolver::new([("source.test", vec![vec![ip("127.0.0.1")]])]);
    let root = reqwest::Certificate::from_der(certificate_der.as_ref()).expect("root");
    let client = OutboundClient::builder(resolver, policy)
        .add_root_certificate(root)
        .build();
    client
        .execute(OutboundRequest::get(&format!("https://source.test:{port}/")).expect("request"))
        .await
        .expect("TLS response");
    server.await.expect("server");
}

#[tokio::test]
async fn ipv6_pinning_preserves_http_authority_and_tls_server_name() {
    let certified =
        rcgen::generate_simple_self_signed(vec!["source6.test".to_owned()]).expect("certificate");
    let certificate_der = certified.cert.der().clone();
    let key_der = certified.signing_key.serialize_der();
    let tls = rustls::ServerConfig::builder()
        .with_no_client_auth()
        .with_single_cert(
            vec![certificate_der.clone()],
            rustls::pki_types::PrivateKeyDer::Pkcs8(key_der.into()),
        )
        .expect("TLS config");
    let listener = match TcpListener::bind("[::1]:0").await {
        Ok(listener) => listener,
        Err(error)
            if error.kind() == io::ErrorKind::AddrNotAvailable
                && std::env::var_os("REMORA_OUTBOUND_REQUIRE_IPV6").is_none() =>
        {
            return;
        }
        Err(error) => panic!("IPv6 bind: {error}"),
    };
    let port = listener.local_addr().expect("address").port();
    let server = tokio::spawn(async move {
        let (stream, _) = listener.accept().await.expect("accept");
        let mut stream = tokio_rustls::TlsAcceptor::from(Arc::new(tls))
            .accept(stream)
            .await
            .expect("TLS handshake");
        let request = read_request(&mut stream).await;
        assert_eq!(
            request.host.as_deref(),
            Some(format!("source6.test:{port}").as_str())
        );
        stream
            .write_all(response("200 OK", &[], "ok").as_bytes())
            .await
            .expect("response");
    });
    let policy = OutboundPolicyBuilder::new()
        .allow_scheme_port("https", port)
        .allow_non_global("::1/128".parse::<IpNet>().expect("network"))
        .build();
    let resolver = ScriptedResolver::new([("source6.test", vec![vec![ip("::1")]])]);
    let root = reqwest::Certificate::from_der(certificate_der.as_ref()).expect("root");
    let client = OutboundClient::builder(resolver, policy)
        .add_root_certificate(root)
        .build();
    client
        .execute(OutboundRequest::get(&format!("https://source6.test:{port}/")).expect("request"))
        .await
        .expect("IPv6 TLS response");
    server.await.expect("server");
}

#[tokio::test]
async fn https_redirect_cannot_downgrade_to_http() {
    let certified =
        rcgen::generate_simple_self_signed(vec!["source.test".to_owned()]).expect("certificate");
    let certificate_der = certified.cert.der().clone();
    let key_der = certified.signing_key.serialize_der();
    let tls = rustls::ServerConfig::builder()
        .with_no_client_auth()
        .with_single_cert(
            vec![certificate_der.clone()],
            rustls::pki_types::PrivateKeyDer::Pkcs8(key_der.into()),
        )
        .expect("TLS config");
    let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
    let port = listener.local_addr().expect("address").port();
    let server = tokio::spawn(async move {
        let (stream, _) = listener.accept().await.expect("accept");
        let mut stream = tokio_rustls::TlsAcceptor::from(Arc::new(tls))
            .accept(stream)
            .await
            .expect("TLS handshake");
        read_request(&mut stream).await;
        stream
            .write_all(
                response(
                    "302 Found",
                    &[("Location", "http://source.test/object")],
                    "",
                )
                .as_bytes(),
            )
            .await
            .expect("response");
    });
    let policy = OutboundPolicyBuilder::new()
        .allow_scheme_port("https", port)
        .allow_scheme_port("http", 80)
        .allow_non_global("127.0.0.0/8".parse::<IpNet>().expect("network"))
        .build();
    let resolver = ScriptedResolver::new([("source.test", vec![vec![ip("127.0.0.1")]])]);
    let root = reqwest::Certificate::from_der(certificate_der.as_ref()).expect("root");
    let client = OutboundClient::builder(resolver, policy)
        .add_root_certificate(root)
        .build();
    let result = client
        .execute(OutboundRequest::get(&format!("https://source.test:{port}/")).expect("request"))
        .await;
    assert_eq!(
        result.expect_err("HTTPS downgrade"),
        OutboundError::HttpsDowngrade
    );
    server.await.expect("server");
}