wgcloak/tests/quic.rs

use std::{collections::HashMap, sync::Arc, time::Duration};

use base64::Engine;
use bytes::Bytes;
use quinn::{ClientConfig as QuinnClientConfig, Endpoint, ServerConfig as QuinnServerConfig};
use rustls::pki_types::{CertificateDer, PrivatePkcs8KeyDer};
use tokio::time::timeout;
use wgcloak::{
    config,
    transport::{
        quic::{QuicClientConnector, QuicServer},
        AdmissionController, TransportError, TransportSession,
    },
};

const SECRET: &str = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";

// A test-only localhost certificate, valid for ten years.  Keeping the DER
// here avoids adding a certificate-generation crate to the application.
const CERT_DER_B64: &str = "MIIDHDCCAgSgAwIBAgIUJnI/VlHMWu6j0dXr/t1OmMQrLgowDQYJKoZIhvcNAQELBQAwFDESMBAGA1UEAwwJbG9jYWxob3N0MB4XDTI2MDkxNDIwNDEyMFoXDTM2MDkxMTIwNDEyMFowFDESMBAGA1UEAwwJbG9jYWxob3N0MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAm4d4r5Py4GdfQrLit2oytLv20L+y7oM/VTDzzJfFMPmewOh0XyS0BI4lyRbCgOuUh686UPvF/BU8uJq24OO7wt/2CrE3FfwYX7BDjIWMuFevHuv896imabbX4IJ7lzVugI3862QSW09QSEQk3ggYs3z/F16cnSRZm9aIURgdPQsLmMzOwcPtDGwxp2jf0Rky0kFRxhUnqfIWko4hny7GjWL8VG0WzYJY6aUevjPlN/Oq/akr9xUAr2QJenEesP77do4X31LIZmsvB6GqKwaSQ0AbwNHW7rDiaE4dkZMIk5l95YgYnrzyOC8IhUwaSZt0wODxk/euXPYr+RFJv2FuOQIDAQABo2YwZDAdBgNVHQ4EFgQUhViBEXLkp7YUVZaWzLi2SfRmfmYwHwYDVR0jBBgwFoAUhViBEXLkp7YUVZaWzLi2SfRmfmYwFAYDVR0RBA0wC4IJbG9jYWxob3N0MAwGA1UdEwEB/wQCMAAwDQYJKoZIhvcNAQELBQADggEBAEFC7da0bK78IUL0k6lZ82pqFImX2bNteMh/r1IJG0xqB/O2dOLUDHybpF7HToommALgRu6zhlT0d1eHg0cMYEgrKyq4BAAXGHqOt3OdKKL5nAWmI04dlWsrG9aTU3hTQhe0++hnTbVm+YJnYAp+ufOxx858FeothjLWu1Lr81Nm01FnmZbXcFno3nosyl1GPxSRVZl2XH9NsfMSZJUbSU6mHXTO1pEH6mjYmgD362CeJc+0DOzevdLsDozxUoWmdEuqlldauESW6XVK0WxRSJ1i2/B/mtYViHr5R/j5NYgG5STnNnwaNwjYK+0AhxvTtH3jAp0Bsluw4d+xKdsV/YI=";
const KEY_DER_B64: &str = "MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQCbh3ivk/LgZ19CsuK3ajK0u/bQv7Lugz9VMPPMl8Uw+Z7A6HRfJLQEjiXJFsKA65SHrzpQ+8X8FTy4mrbg47vC3/YKsTcV/BhfsEOMhYy4V68e6/z3qKZpttfggnuXNW6AjfzrZBJbT1BIRCTeCBizfP8XXpydJFmb1ohRGB09CwuYzM7Bw+0MbDGnaN/RGTLSQVHGFSep8haSjiGfLsaNYvxUbRbNgljppR6+M+U386r9qSv3FQCvZAl6cR6w/vt2jhffUshmay8HoaorBpJDQBvA0dbusOJoTh2RkwiTmX3liBievPI4LwiFTBpJm3TA4PGT965c9iv5EUm/YW45AgMBAAECggEAMs0EtCD4t6i8Ol6Q5EUZa+RrgZLkILkZEjwLnjvCqrUWUD99Z2hnLoe5ut2myZJpcQpg23vpluXe8BOFAZ8RQN2DkcmByzoNQhAieD9VEA+AAXMmrURWmvggulIUiXbn8bCvOLn2rwGQoRI1EpsFanc9IUaDlGZ/qIHiA88DGeFDuchbcgClpXAX4lBzqhLKv/CCIhhGxmyWc1RPlZUef7uTO1Zh7q2HxnO5iBKobvPfMq5VfvzE4rXeSQIHtW+yJFLU2sNjaYry2JTLEqGIZhWMROmzlkW8x/cj0GvjLbY4ltkIpLNSb+WRbY3zZNhQ81kosBZtiYRj3vgr/FzvSwKBgQDQ/dxghryGDqIIOfyZ59+2tpPrrLX4vy7tmAQg6zJIiQVHMZ/Av60QvY7lA0dvDfQmfJuNzTOXtwhgnYkHbg6jlFqpT7P1CZQOzbF0UuFCNvwzE/jgCEKUExckDa+Du7g/VOglB9B4XKqJnQLpZkzz4Hwmuvz4oUtZJCl796HTMwKBgQC+gyPKrNp6AthxZ1gAu+aIjf18vgGZjANQx3cRFQl5BqSyTwZpxIfuiKrT+sU65N2VCWpGH6oTzGruaXBXzG0rlNCMAL5XaJqJehvLAnejUqmXhyi8ouH174Q4azDngPOCFaSjuYJwpX4nUXwMExqs9hwkH9pV3tn6qzEpKJA44wKBgBa+SalHgGSmWwPCK7ckALiFjftrxDHZvGWm0WwzrDkwkB57TANR2KeluVM61B9FyXgWv2bjGeY8yPCatDrWw0kR7/vU7KSR8kbaI6GhE1p9WVzFo6PdwkETAsWtp5gSaa6asti0zP8MNpLbP2He/CjeGSSt3LHB+80rnkgbXsT5AoGAVaTytCxc0zjwQJrqzvmpu33JlQh+A58T5n6lIkLIzWWVKKuqAGZfjXZXBZDYlY9JpbRj2D5n4Rqmf6Na47183ni6ABZde0IFJwSD9hGX4w5IO1yBjbvTxRH16L8XbxtjXPmJ2CiPYUCko9xX3b3A4jyJjhvbtXrHl47aPNyFaW0CgYEAxDOcW5r5pXXkCF/mIqSklGa1jDO9TeMavLo1AguHaUlJEQU8DgGM9Flx3senFASN2EjwwTB3aScVuxe3X29LeZnCNsuTIzCVGpSNyFLUjZVU0fRaf6R9JIcI7dVoflPAPCdrqBVkezK5Eyu8f3NCwnVD+zv+iWA+RE1zUyKy4sM=";

fn endpoint_and_client(
    max_datagram: usize,
) -> (Endpoint, QuinnClientConfig, CertificateDer<'static>) {
    let cert_bytes = base64::engine::general_purpose::STANDARD
        .decode(CERT_DER_B64)
        .unwrap();
    let key_bytes = base64::engine::general_purpose::STANDARD
        .decode(KEY_DER_B64)
        .unwrap();
    let cert = CertificateDer::from(cert_bytes);
    let key = PrivatePkcs8KeyDer::from(key_bytes);
    let mut server_config =
        QuinnServerConfig::with_single_cert(vec![cert.clone()], key.into()).unwrap();
    let mut transport_config = quinn::TransportConfig::default();
    transport_config.keep_alive_interval(Some(Duration::from_secs(1)));
    transport_config.datagram_send_buffer_size(max_datagram * 8);
    server_config.transport_config(Arc::new(transport_config));
    let endpoint = Endpoint::server(server_config, "127.0.0.1:0".parse().unwrap()).unwrap();
    let mut roots = rustls::RootCertStore::empty();
    roots.add(cert.clone()).unwrap();
    let client = QuinnClientConfig::with_root_certificates(Arc::new(roots)).unwrap();
    (endpoint, client, cert)
}

fn client_config(port: u16, secret: &str) -> config::ClientConfig {
    config::Config::from_str(&format!(
        "[client]\nserver = \"localhost\"\nid = \"test\"\nsecret = \"{secret}\"\n[transport.quic]\nport = {port}\n[transport.auto]\nprobe_timeout_ms = 2000\n[network]\nmax_datagram_size = 1200\n"
    ))
    .unwrap()
    .into_client()
    .unwrap()
}

#[tokio::test]
async fn authenticated_datagram_round_trip_and_limit() {
    let (endpoint, client_tls, _cert) = endpoint_and_client(1200);
    let mut clients = HashMap::new();
    clients.insert(
        "test".to_owned(),
        base64::engine::general_purpose::STANDARD
            .decode(SECRET)
            .unwrap(),
    );
    let admission = AdmissionController::new(&config::LimitsConfig::default());
    let server = QuicServer::from_endpoint_with_admission(
        endpoint,
        clients,
        1200,
        Duration::from_secs(2),
        admission.clone(),
    );
    let server_addr = server.local_addr().unwrap();
    let port = server_addr.port();
    let client = QuicClientConnector::with_client_config(client_config(port, SECRET), client_tls);
    let server_task = tokio::spawn(async move { server.accept().await });
    let client_session = client.connect_session().await.unwrap();
    let server_session = server_task.await.unwrap().unwrap();
    let remote_addr = client_session.remote_addr().unwrap();
    assert_eq!(remote_addr.port(), server_addr.port());
    assert!(match remote_addr.ip() {
        std::net::IpAddr::V4(address) => address.is_loopback(),
        std::net::IpAddr::V6(address) => {
            address.is_loopback()
                || address
                    .to_ipv4_mapped()
                    .map(|mapped| mapped.is_loopback())
                    .unwrap_or(false)
        }
    });
    assert!(client_session.rtt().is_some());
    assert_eq!(client_session.authenticated_keepalive_interval(), None);
    client_session
        .send_datagram(Bytes::from_static(b"opaque"))
        .await
        .unwrap();
    assert_eq!(
        timeout(Duration::from_secs(2), server_session.recv_datagram())
            .await
            .unwrap()
            .unwrap(),
        Bytes::from_static(b"opaque")
    );
    assert_eq!(admission.active_sessions(), 1);
    assert!(matches!(
        client_session
            .send_datagram(Bytes::from(vec![0; 1201]))
            .await,
        Err(TransportError::Oversized { maximum: 1200 })
    ));
    client_session.close().await.unwrap();
    server_session.close().await.unwrap();
    assert_eq!(admission.active_sessions(), 0);
}

#[tokio::test]
async fn bad_secret_is_rejected_as_authentication_error() {
    let (endpoint, client_tls, _cert) = endpoint_and_client(1200);
    let mut clients = HashMap::new();
    clients.insert(
        "test".to_owned(),
        base64::engine::general_purpose::STANDARD
            .decode(SECRET)
            .unwrap(),
    );
    let server = QuicServer::from_endpoint(endpoint, clients, 1200, Duration::from_secs(2));
    let port = server.local_addr().unwrap().port();
    let wrong = "BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB";
    let client = QuicClientConnector::with_client_config(client_config(port, wrong), client_tls);
    let server_task = tokio::spawn(async move { server.accept().await });
    let result = client.connect_session().await;
    assert!(matches!(
        result,
        Err(TransportError::AuthenticationRejected(_))
    ));
    assert!(matches!(
        server_task.await.unwrap(),
        Err(TransportError::AuthenticationRejected(_))
    ));
}

#[tokio::test]
async fn certificate_hostname_mismatch_is_not_accepted() {
    let (endpoint, client_tls, _cert) = endpoint_and_client(1200);
    let mut clients = HashMap::new();
    clients.insert(
        "test".to_owned(),
        base64::engine::general_purpose::STANDARD
            .decode(SECRET)
            .unwrap(),
    );
    let server = QuicServer::from_endpoint(endpoint, clients, 1200, Duration::from_secs(2));
    let port = server.local_addr().unwrap().port();
    let server_task = tokio::spawn(async move { server.accept().await });
    // The connector's configured hostname is localhost, which matches.  This
    // test changes it to an unrelated name while retaining the localhost root.
    let mut mismatch = client_config(port, SECRET);
    mismatch.client.server = "not-localhost".to_owned();
    let mismatch = QuicClientConnector::with_client_config(mismatch, client_tls);
    let result = mismatch.connect_session().await;
    assert!(result.is_err());
    server_task.abort();
}