wgcloak/tests/http3_runtime.rs

use std::{
    fs, net::SocketAddr, os::unix::fs::PermissionsExt, path::Path, process::Command, sync::Arc,
    time::Duration,
};

use rustls::pki_types::pem::PemObject;
use tempfile::TempDir;
use tokio::{
    net::UdpSocket,
    sync::mpsc,
    time::{sleep, timeout},
};
use wgcloak::{
    client,
    config::{self, Config},
    control, server, stats, transport,
    transport::{Http3ClientConnector, TransportConnector, TransportSession},
};

const SECRET_ONE: &str = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";
const SECRET_TWO: &str = "BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB";
const WRONG_SECRET: &str = "CCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCC";

fn free_udp_port() -> u16 {
    let socket = std::net::UdpSocket::bind("127.0.0.1:0").unwrap();
    socket.local_addr().unwrap().port()
}

fn private_runtime_dir(root: &Path) -> std::path::PathBuf {
    let path = root.join("runtime");
    fs::create_dir(&path).unwrap();
    fs::set_permissions(&path, fs::Permissions::from_mode(0o700)).unwrap();
    path
}

fn server_config(
    directory: &TempDir,
    backend: SocketAddr,
    direct_port: u16,
    http3_port: u16,
) -> config::ServerConfig {
    let certificate = directory.path().join("server-cert.pem");
    let private_key = directory.path().join("server-key.pem");
    let generated = Command::new("openssl")
        .args(["req", "-x509", "-newkey", "rsa:2048", "-nodes", "-keyout"])
        .arg(&private_key)
        .args(["-out"])
        .arg(&certificate)
        .args([
            "-subj",
            "/CN=127.0.0.1",
            "-addext",
            "subjectAltName=IP:127.0.0.1",
            "-addext",
            "basicConstraints=critical,CA:FALSE",
            "-addext",
            "keyUsage=critical,digitalSignature,keyEncipherment",
            "-addext",
            "extendedKeyUsage=serverAuth",
            "-days",
            "1",
        ])
        .output()
        .expect("openssl is required for the HTTP/3 runtime integration test");
    assert!(
        generated.status.success(),
        "openssl certificate generation failed: {}",
        String::from_utf8_lossy(&generated.stderr)
    );

    Config::from_str(&format!(
        "[server]\nlisten='127.0.0.1'\n[backend]\ntarget='{backend}'\n[transports.direct_udp]\nenabled=true\nport={direct_port}\n[transports.quic]\nenabled=false\nport={}\n[transports.http3]\nenabled=true\nport={http3_port}\npath='/tunnel'\n[transports.tls_tcp]\nenabled=true\nport={direct_port}\n[tls]\ncertificate='{}'\nprivate_key='{}'\n[[auth.clients]]\nid='one'\nsecret='{SECRET_ONE}'\n[[auth.clients]]\nid='two'\nsecret='{SECRET_TWO}'\n[limits]\nhandshake_timeout_ms=3000\n",
        free_udp_port(),
        certificate.display(),
        private_key.display()
    ))
    .unwrap()
    .into_server()
    .unwrap()
}

fn client_config(
    id: &str,
    secret: &str,
    http3_port: u16,
    direct_probe_port: u16,
    local_port: u16,
) -> config::ClientConfig {
    Config::from_str(&format!(
        "[client]\nserver='127.0.0.1'\nid='{id}'\nsecret='{secret}'\n[local]\nlisten='127.0.0.1:{local_port}'\n[transport]\nmode='auto'\n[transport.auto]\norder=['direct_udp','http3']\nprobe_timeout_ms=700\nretries=1\n[transport.direct_udp]\nport={direct_probe_port}\n[transport.http3]\nport={http3_port}\npath='/tunnel'\n[network]\nmax_datagram_size=1200\n"
    ))
    .unwrap()
    .into_client()
    .unwrap()
}

fn test_client_tls(certificate: &[u8]) -> quinn::ClientConfig {
    let _ = rustls::crypto::ring::default_provider().install_default();
    let certs: Vec<_> = rustls::pki_types::CertificateDer::pem_slice_iter(certificate)
        .collect::<Result<_, _>>()
        .unwrap();
    let mut roots = rustls::RootCertStore::empty();
    roots.add(certs[0].clone()).unwrap();
    let mut tls = rustls::ClientConfig::builder_with_provider(Arc::new(
        rustls::crypto::ring::default_provider(),
    ))
    .with_protocol_versions(&[&rustls::version::TLS13])
    .unwrap()
    .with_root_certificates(roots)
    .with_no_client_auth();
    tls.alpn_protocols = vec![b"h3".to_vec()];
    let crypto = quinn::crypto::rustls::QuicClientConfig::try_from(tls).unwrap();
    quinn::ClientConfig::new(Arc::new(crypto))
}

async fn wait_control_contains(path: &std::path::Path, command: &str, needle: &str) {
    let mut last_response = None;
    let result = timeout(Duration::from_secs(6), async {
        loop {
            if let Ok(response) = control::send_command(path, command).await {
                if response.contains(needle) {
                    return;
                }
                last_response = Some(response);
            }
            sleep(Duration::from_millis(40)).await;
        }
    })
    .await;
    assert!(
        result.is_ok(),
        "control {command} did not contain {needle:?}; latest response: {last_response:?}"
    );
}

#[tokio::test]
async fn auto_http3_runtime_relays_isolated_clients_and_rejects_bad_secret() {
    if Command::new("openssl").arg("version").output().is_err() {
        return;
    }
    let directory = tempfile::tempdir().unwrap();
    let runtime_dir = private_runtime_dir(directory.path());
    let backend = UdpSocket::bind("127.0.0.1:0").await.unwrap();
    let backend_addr = backend.local_addr().unwrap();
    let (peer_tx, mut peer_rx) = mpsc::unbounded_channel();
    let backend_task = tokio::spawn(async move {
        let mut packet = vec![0u8; u16::MAX as usize];
        while let Ok((len, peer)) = backend.recv_from(&mut packet).await {
            let _ = peer_tx.send(peer);
            let _ = backend.send_to(&packet[..len], peer).await;
        }
    });

    let certificate = directory.path().join("server-cert.pem");
    let direct_server_port = free_udp_port();
    let http3_port = free_udp_port();
    let server = server_config(&directory, backend_addr, direct_server_port, http3_port);
    let server_control = runtime_dir.join("server.sock");
    let server_control_for_task = server_control.clone();
    let server_task = tokio::spawn(server::run_with_control_path(
        server,
        server_control_for_task,
    ));
    wait_control_contains(
        &server_control,
        "status",
        "transports=direct_udp,http3,tls_tcp",
    )
    .await;

    let certificate_bytes = fs::read(certificate).unwrap();
    let client_tls = test_client_tls(&certificate_bytes);
    let direct_probe_port = free_udp_port();
    assert_ne!(direct_probe_port, direct_server_port);

    let local_one = free_udp_port();
    let local_two = free_udp_port();
    let app_one = UdpSocket::bind("127.0.0.1:0").await.unwrap();
    let app_two = UdpSocket::bind("127.0.0.1:0").await.unwrap();
    let client_one_config =
        client_config("one", SECRET_ONE, http3_port, direct_probe_port, local_one);
    let client_two_config =
        client_config("two", SECRET_TWO, http3_port, direct_probe_port, local_two);
    let client_one_control = runtime_dir.join("client-one.sock");
    let client_two_control = runtime_dir.join("client-two.sock");
    let client_one_task = tokio::spawn(client::run_with_http3_connector_and_control_path(
        client_one_config.clone(),
        client_one_control.clone(),
        Some(Arc::new(Http3ClientConnector::with_client_config(
            client_one_config,
            client_tls.clone(),
        )) as Arc<dyn TransportConnector>),
    ));
    let client_two_task = tokio::spawn(client::run_with_http3_connector_and_control_path(
        client_two_config.clone(),
        client_two_control.clone(),
        Some(Arc::new(Http3ClientConnector::with_client_config(
            client_two_config,
            client_tls.clone(),
        )) as Arc<dyn TransportConnector>),
    ));

    wait_control_contains(&server_control, "status", "active_sessions=2").await;
    app_one
        .send_to(b"client-one", ("127.0.0.1", local_one))
        .await
        .unwrap();
    app_two
        .send_to(b"client-two", ("127.0.0.1", local_two))
        .await
        .unwrap();
    let mut response = [0u8; 64];
    let (len_one, _) = timeout(Duration::from_secs(4), app_one.recv_from(&mut response))
        .await
        .expect("first HTTP/3 echo timed out")
        .unwrap();
    assert_eq!(&response[..len_one], b"client-one");
    let (len_two, _) = timeout(Duration::from_secs(4), app_two.recv_from(&mut response))
        .await
        .expect("second HTTP/3 echo timed out")
        .unwrap();
    assert_eq!(&response[..len_two], b"client-two");
    let peer_one = timeout(Duration::from_secs(2), peer_rx.recv())
        .await
        .unwrap()
        .unwrap();
    let peer_two = timeout(Duration::from_secs(2), peer_rx.recv())
        .await
        .unwrap()
        .unwrap();
    assert_ne!(
        peer_one, peer_two,
        "each tunnel must own its backend UDP socket"
    );

    let bad_local = free_udp_port();
    let bad_client_config = client_config(
        "one",
        WRONG_SECRET,
        http3_port,
        direct_probe_port,
        bad_local,
    );
    let bad_control = runtime_dir.join("client-bad.sock");
    let bad_task = tokio::spawn(client::run_with_http3_connector_and_control_path(
        bad_client_config.clone(),
        bad_control.clone(),
        Some(Arc::new(Http3ClientConnector::with_client_config(
            bad_client_config,
            client_tls,
        )) as Arc<dyn TransportConnector>),
    ));
    let bad_result = timeout(Duration::from_secs(6), bad_task)
        .await
        .expect("bad-secret client should finish its rejected startup")
        .unwrap();
    assert!(bad_result.is_err(), "wrong project secret must be rejected");
    wait_control_contains(&server_control, "stats", "\"authentication_rejections\":1").await;
    let aggregate = control::send_command(&server_control, "stats")
        .await
        .unwrap();
    assert!(aggregate.contains("\"transport\":\"multi\""));
    assert!(aggregate.contains("\"direct_udp\":{"));
    assert!(aggregate.contains("\"http3\":{"));
    assert!(aggregate.contains("\"tls_tcp\":{"));

    for path in [&client_one_control, &client_two_control] {
        control::send_command(path, "shutdown").await.unwrap();
    }
    timeout(Duration::from_secs(4), client_one_task)
        .await
        .unwrap()
        .unwrap()
        .unwrap();
    timeout(Duration::from_secs(4), client_two_task)
        .await
        .unwrap()
        .unwrap()
        .unwrap();
    control::send_command(&server_control, "shutdown")
        .await
        .unwrap();
    timeout(Duration::from_secs(5), server_task)
        .await
        .unwrap()
        .unwrap()
        .unwrap();
    backend_task.abort();
}

#[tokio::test]
async fn conflicting_udp_listeners_are_rejected() {
    let parsed = Config::from_str(&format!(
        "[server]\nlisten='127.0.0.1'\n[backend]\ntarget='127.0.0.1:9'\n[transports.direct_udp]\nenabled=true\nport=54001\n[transports.quic]\nenabled=true\nport=54001\n[transports.http3]\nenabled=false\n[transports.tls_tcp]\nenabled=true\nport=54001\n[tls]\ncertificate='placeholder-cert'\nprivate_key='placeholder-key'\n[[auth.clients]]\nid='one'\nsecret='{SECRET_ONE}'\n"
    ))
    .unwrap()
    .into_server()
    .unwrap();
    let directory = tempfile::tempdir().unwrap();
    let runtime_dir = private_runtime_dir(directory.path());
    // Validation happens before binding any listener and reports only the
    // incompatible same-protocol UDP bind; TCP is allowed to reuse the number.
    let result = server::run_with_control_path(parsed, runtime_dir.join("conflict.sock")).await;
    assert!(result.unwrap_err().to_string().contains("UDP listeners"));
}

#[tokio::test]
async fn control_stats_and_status_merge_multiple_transport_snapshots() {
    let directory = tempfile::tempdir().unwrap();
    let path = private_runtime_dir(directory.path()).join("aggregate.sock");
    let state = control::ControlState::default();
    state
        .set_status("listening; transports=direct_udp,http3")
        .await;
    let direct = stats::StatsSnapshot {
        datagrams_sent: 2,
        bytes_sent: 40,
        ..stats::StatsSnapshot::default()
    };
    let http3 = stats::StatsSnapshot {
        datagrams_received: 3,
        authentication_rejections: 1,
        ..stats::StatsSnapshot::default()
    };
    state.set_transport_stats("direct_udp", 1, direct).await;
    state.set_transport_stats("http3", 2, http3).await;
    let server = control::ControlServer::bind(&path, state.clone())
        .await
        .unwrap();
    let task = tokio::spawn(server.run());

    let status = control::send_command(&path, "status").await.unwrap();
    assert!(status.contains("active_sessions=3"));
    assert!(status.contains("active_sessions_by_transport=direct_udp:1,http3:2"));
    let response = control::send_command(&path, "stats").await.unwrap();
    assert!(response.contains("\"datagrams_sent\":2"));
    assert!(response.contains("\"bytes_sent\":40"));
    assert!(response.contains("\"datagrams_received\":3"));
    assert!(response.contains("\"authentication_rejections\":1"));
    assert!(response.contains("\"direct_udp\":{"));
    assert!(response.contains("\"http3\":{"));

    state.shutdown.cancel();
    timeout(Duration::from_secs(1), task)
        .await
        .unwrap()
        .unwrap()
        .unwrap();
}

#[tokio::test]
async fn direct_udp_runtime_cancellation_runs_session_cleanup() {
    let port = free_udp_port();
    let server = Config::from_str(&format!(
        "[server]\nlisten='127.0.0.1'\n[backend]\ntarget='127.0.0.1:9'\n[transports.direct_udp]\nenabled=true\nport={port}\n[transports.quic]\nenabled=false\n[transports.http3]\nenabled=false\n[transports.tls_tcp]\nenabled=false\n[[auth.clients]]\nid='lifecycle'\nsecret='{SECRET_ONE}'\n"
    ))
    .unwrap()
    .into_server()
    .unwrap();
    let client = Config::from_str(&format!(
        "[client]\nserver='127.0.0.1'\nid='lifecycle'\nsecret='{SECRET_ONE}'\n[transport]\nmode='direct_udp'\n[transport.direct_udp]\nport={port}\n"
    ))
    .unwrap()
    .into_client()
    .unwrap();

    let listener = transport::DirectUdpServer::bind(server).await.unwrap();
    let observer = listener.handle();
    let shutdown = tokio_util::sync::CancellationToken::new();
    let run = tokio::spawn(listener.run_with_shutdown(shutdown.clone()));
    let session = timeout(
        Duration::from_secs(3),
        transport::direct_udp::DirectUdpSession::connect(&client),
    )
    .await
    .expect("direct UDP client authentication timed out")
    .unwrap();
    timeout(Duration::from_secs(2), async {
        while observer.active_sessions().await != 1 {
            sleep(Duration::from_millis(10)).await;
        }
    })
    .await
    .expect("server did not register the authenticated direct session");

    shutdown.cancel();
    timeout(Duration::from_secs(2), run)
        .await
        .expect("direct listener did not complete graceful cleanup")
        .unwrap()
        .unwrap();
    assert_eq!(observer.active_sessions().await, 0);
    session.close().await.unwrap();
}