remora/crates/remora-credentials/tests/vault_transport.rs

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

use remora_credentials::{SecretMaterial, VaultClient, VaultConfig};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;

/// Vault client must not follow redirects that would re-issue X-Vault-Token
/// (and secret write bodies) to an alternate origin.
#[tokio::test]
async fn vault_client_does_not_follow_redirects_with_token() {
    let evil_hits: Arc<Mutex<usize>> = Arc::new(Mutex::new(0));
    let evil_hits_bg = Arc::clone(&evil_hits);
    let evil = TcpListener::bind("127.0.0.1:0").await.expect("evil bind");
    let evil_addr = evil.local_addr().expect("evil addr");
    tokio::spawn(async move {
        if let Ok((mut stream, _)) = evil.accept().await {
            let mut buf = vec![0u8; 4096];
            let _ = stream.read(&mut buf).await;
            *evil_hits_bg.lock().expect("lock") += 1;
            let _ = stream
                .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n")
                .await;
        }
    });

    let redirector = TcpListener::bind("127.0.0.1:0")
        .await
        .expect("redirector bind");
    let redirector_addr = redirector.local_addr().expect("redirector addr");
    tokio::spawn(async move {
        if let Ok((mut stream, _)) = redirector.accept().await {
            let mut buf = vec![0u8; 8192];
            let _ = read_http_request(&mut stream, &mut buf).await;
            let location = format!("http://{evil_addr}/v1/secret/data/stolen");
            let response =
                format!("HTTP/1.1 302 Found\r\nLocation: {location}\r\nContent-Length: 0\r\n\r\n");
            let _ = stream.write_all(response.as_bytes()).await;
        }
    });

    let client = VaultClient::new(VaultConfig {
        address: format!("http://{redirector_addr}"),
        token: "synthetic-vault-token-must-not-exfiltrate".into(),
        mount: "secret".into(),
        path_prefix: "remora".into(),
        timeout: Duration::from_secs(2),
    })
    .expect("client");

    let result = client
        .put_version(
            "secret",
            "remora/credentials/probe",
            &SecretMaterial::BearerToken {
                token: "secret-body-must-not-exfiltrate".into(),
            },
        )
        .await;

    assert!(
        result.is_err(),
        "redirect response must fail closed, got {result:?}"
    );
    // Give the evil server a moment if a follow were attempted.
    tokio::time::sleep(Duration::from_millis(50)).await;
    assert_eq!(
        *evil_hits.lock().expect("lock"),
        0,
        "Vault client followed a redirect and contacted the evil origin"
    );
}

/// Ambient HTTP(S)_PROXY must not be consulted for Vault traffic.
#[tokio::test]
async fn vault_client_ignores_ambient_proxy_env() {
    // Safety: test process only; restores prior values at end of scope via scopeguard pattern.
    let prior_http = std::env::var("HTTP_PROXY").ok();
    let prior_https = std::env::var("HTTPS_PROXY").ok();
    let prior_all = std::env::var("ALL_PROXY").ok();

    let proxy_hits: Arc<Mutex<usize>> = Arc::new(Mutex::new(0));
    let proxy_hits_bg = Arc::clone(&proxy_hits);
    let proxy = TcpListener::bind("127.0.0.1:0").await.expect("proxy bind");
    let proxy_addr = proxy.local_addr().expect("proxy addr");
    tokio::spawn(async move {
        if let Ok((mut stream, _)) = proxy.accept().await {
            let mut buf = vec![0u8; 4096];
            let _ = stream.read(&mut buf).await;
            *proxy_hits_bg.lock().expect("lock") += 1;
            let _ = stream.shutdown().await;
        }
    });

    // Point proxies at a listener that records any CONNECT/request attempt.
    // SAFETY: single-threaded test mutation of process env for this test only.
    unsafe {
        std::env::set_var("HTTP_PROXY", format!("http://{proxy_addr}"));
        std::env::set_var("HTTPS_PROXY", format!("http://{proxy_addr}"));
        std::env::set_var("ALL_PROXY", format!("http://{proxy_addr}"));
    }

    let vault = TcpListener::bind("127.0.0.1:0").await.expect("vault bind");
    let vault_addr = vault.local_addr().expect("vault addr");
    tokio::spawn(async move {
        if let Ok((mut stream, _)) = vault.accept().await {
            let mut buf = vec![0u8; 8192];
            let _ = read_http_request(&mut stream, &mut buf).await;
            let body = r#"{"data":{"version":1}}"#;
            let response = format!(
                "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{body}",
                body.len()
            );
            let _ = stream.write_all(response.as_bytes()).await;
        }
    });

    let client = VaultClient::new(VaultConfig {
        address: format!("http://{vault_addr}"),
        token: "synthetic-token".into(),
        mount: "secret".into(),
        path_prefix: "remora".into(),
        timeout: Duration::from_secs(2),
    })
    .expect("client");

    let version = client
        .put_version(
            "secret",
            "remora/credentials/probe",
            &SecretMaterial::BearerToken {
                token: "secret-value".into(),
            },
        )
        .await
        .expect("direct vault write without proxy");
    assert_eq!(version, 1);

    tokio::time::sleep(Duration::from_millis(50)).await;
    assert_eq!(
        *proxy_hits.lock().expect("lock"),
        0,
        "Vault client used ambient proxy environment"
    );

    restore_env("HTTP_PROXY", prior_http);
    restore_env("HTTPS_PROXY", prior_https);
    restore_env("ALL_PROXY", prior_all);
}

#[tokio::test]
async fn exact_read_rejects_missing_or_mismatched_version_metadata() {
    for metadata in ["{}", r#"{"version":2}"#] {
        let vault = TcpListener::bind("127.0.0.1:0").await.expect("vault bind");
        let vault_addr = vault.local_addr().expect("vault addr");
        let body = format!(
            r#"{{"data":{{"data":{{"kind":"bearer_token","token":"synthetic"}},"metadata":{metadata}}}}}"#
        );
        tokio::spawn(async move {
            if let Ok((mut stream, _)) = vault.accept().await {
                let mut buf = vec![0u8; 8192];
                let _ = read_http_request(&mut stream, &mut buf).await;
                let response = format!(
                    "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{body}",
                    body.len()
                );
                let _ = stream.write_all(response.as_bytes()).await;
            }
        });

        let client = VaultClient::new(VaultConfig {
            address: format!("http://{vault_addr}"),
            token: "synthetic-token".into(),
            mount: "secret".into(),
            path_prefix: "remora".into(),
            timeout: Duration::from_secs(2),
        })
        .expect("client");

        assert!(matches!(
            client
                .get_version("secret", "remora/credentials/probe", 1)
                .await,
            Err(remora_credentials::CredentialError::Unresolvable)
        ));
    }
}

fn restore_env(key: &str, prior: Option<String>) {
    // SAFETY: restoring process env after test mutation.
    unsafe {
        match prior {
            Some(value) => std::env::set_var(key, value),
            None => std::env::remove_var(key),
        }
    }
}

async fn read_http_request(
    stream: &mut tokio::net::TcpStream,
    buf: &mut [u8],
) -> io::Result<usize> {
    let mut total = 0;
    loop {
        let n = stream.read(&mut buf[total..]).await?;
        if n == 0 {
            break;
        }
        total += n;
        if buf[..total].windows(4).any(|w| w == b"\r\n\r\n") {
            break;
        }
        if total == buf.len() {
            break;
        }
    }
    Ok(total)
}

#[allow(dead_code)]
fn _socket_addr_fmt(addr: SocketAddr) -> String {
    addr.to_string()
}