wgcloak/tests/direct_udp.rs

#[path = "../src/config.rs"]
mod config;
#[path = "../src/datagram.rs"]
mod datagram;
#[path = "../src/protocol/mod.rs"]
mod protocol;
#[path = "../src/stats.rs"]
mod stats;
#[path = "../src/transport/mod.rs"]
mod transport;

use std::time::{Duration, SystemTime, UNIX_EPOCH};

use base64::Engine;
use bytes::Bytes;
use config::Config;
use tokio::{net::UdpSocket, time::timeout};
use transport::{
    direct_udp::{DirectUdpServer, DirectUdpSession},
    TransportError, TransportSession,
};

const SECRET: &str = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";

fn configs(port: u16, backend: u16) -> (config::ServerConfig, config::ClientConfig) {
    let server = Config::from_str(&format!(
        "[server]\nlisten = \"127.0.0.1\"\n[backend]\ntarget = \"127.0.0.1:{backend}\"\n[transports.direct_udp]\nport = {port}\n[[auth.clients]]\nid = \"test\"\nsecret = \"{SECRET}\"\n"
    )).unwrap().into_server().unwrap();
    let client = Config::from_str(&format!(
        "[client]\nserver = \"127.0.0.1\"\nid = \"test\"\nsecret = \"{SECRET}\"\n[transport.direct_udp]\nport = {port}\n[network]\nmax_datagram_size = 1200\n"
    )).unwrap().into_client().unwrap();
    (server, client)
}

fn multi_configs(
    port: u16,
    backend: u16,
) -> (
    config::ServerConfig,
    config::ClientConfig,
    config::ClientConfig,
) {
    let server = Config::from_str(&format!(
        "[server]\nlisten = \"127.0.0.1\"\n[backend]\ntarget = \"127.0.0.1:{backend}\"\n[transports.direct_udp]\nport = {port}\n[[auth.clients]]\nid = \"one\"\nsecret = \"{SECRET}\"\n[[auth.clients]]\nid = \"two\"\nsecret = \"{SECRET}\"\n"
    )).unwrap().into_server().unwrap();
    let client = |id: &str| {
        Config::from_str(&format!(
            "[client]\nserver = \"127.0.0.1\"\nid = \"{id}\"\nsecret = \"{SECRET}\"\n[transport.direct_udp]\nport = {port}\n[network]\nmax_datagram_size = 1200\n"
        )).unwrap().into_client().unwrap()
    };
    (server, client("one"), client("two"))
}

#[tokio::test]
async fn direct_udp_round_trip_and_auth_rejection() {
    let backend = UdpSocket::bind("127.0.0.1:0").await.unwrap();
    let backend_addr = backend.local_addr().unwrap();
    let port = 45781;
    let (server_config, client_config) = configs(port, backend_addr.port());
    let server = DirectUdpServer::bind(server_config).await.unwrap();
    let stats = server.stats();
    let server_task = tokio::spawn(server.run());
    let backend_task = tokio::spawn(async move {
        let mut buf = [0u8; 2048];
        let (len, peer) = backend.recv_from(&mut buf).await.unwrap();
        backend.send_to(&buf[..len], peer).await.unwrap();
    });
    let session = DirectUdpSession::connect(&client_config).await.unwrap();
    session
        .send_datagram(Bytes::from_static(b"opaque"))
        .await
        .unwrap();
    assert_eq!(
        timeout(Duration::from_secs(2), session.recv_datagram())
            .await
            .unwrap()
            .unwrap(),
        Bytes::from_static(b"opaque")
    );
    let mut wrong = client_config.clone();
    wrong.client.secret.secret =
        Some("BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB".to_owned());
    let result = DirectUdpSession::connect(&wrong).await;
    assert!(matches!(
        result,
        Err(TransportError::AuthenticationRejected(_))
    ));
    session.close().await.unwrap();
    backend_task.await.unwrap();
    tokio::time::sleep(Duration::from_millis(20)).await;
    assert_eq!(stats.snapshot().expired_sessions, 0);
    server_task.abort();
}

#[tokio::test]
async fn direct_udp_enforces_payload_limit_and_expires_idle_session() {
    let backend = UdpSocket::bind("127.0.0.1:0").await.unwrap();
    let backend_addr = backend.local_addr().unwrap();
    let (mut server_config, client_config) = configs(45782, backend_addr.port());
    server_config.session.idle_timeout_secs = 1;
    let server = DirectUdpServer::bind(server_config).await.unwrap();
    let stats = server.stats();
    let server_task = tokio::spawn(server.run());
    let session = DirectUdpSession::connect(&client_config).await.unwrap();
    assert!(matches!(
        session.send_datagram(Bytes::from(vec![0u8; 1201])).await,
        Err(TransportError::Oversized { maximum: 1200 })
    ));
    tokio::time::sleep(Duration::from_secs(3)).await;
    assert_eq!(stats.snapshot().expired_sessions, 1);
    session.close().await.unwrap();
    server_task.abort();
}

#[tokio::test]
async fn concurrent_clients_use_independent_backend_sockets() {
    let backend = UdpSocket::bind("127.0.0.1:0").await.unwrap();
    let backend_addr = backend.local_addr().unwrap();
    let (server_config, mut client_one, mut client_two) = multi_configs(45783, backend_addr.port());
    let server = DirectUdpServer::bind(server_config).await.unwrap();
    let stats = server.stats();
    let server_task = tokio::spawn(server.run());
    client_one.transport.direct_udp.port = 45783;
    client_two.transport.direct_udp.port = 45783;
    let backend_task = tokio::spawn(async move {
        let mut peers = Vec::new();
        let mut buf = [0u8; 2048];
        for _ in 0..2 {
            let (len, peer) = backend.recv_from(&mut buf).await.unwrap();
            peers.push(peer);
            backend.send_to(&buf[..len], peer).await.unwrap();
        }
        peers
    });
    let first = DirectUdpSession::connect(&client_one).await.unwrap();
    let second = DirectUdpSession::connect(&client_two).await.unwrap();
    assert!(stats.snapshot().handshake_attempts >= 2);
    first
        .send_datagram(Bytes::from_static(b"first"))
        .await
        .unwrap();
    second
        .send_datagram(Bytes::from_static(b"second"))
        .await
        .unwrap();
    assert_eq!(
        timeout(Duration::from_secs(2), first.recv_datagram())
            .await
            .unwrap()
            .unwrap(),
        Bytes::from_static(b"first")
    );
    assert_eq!(
        timeout(Duration::from_secs(2), second.recv_datagram())
            .await
            .unwrap()
            .unwrap(),
        Bytes::from_static(b"second")
    );
    let peers = backend_task.await.unwrap();
    assert_ne!(
        peers[0], peers[1],
        "each authenticated client must have its own backend socket"
    );
    first.close().await.unwrap();
    second.close().await.unwrap();
    server_task.abort();
}

#[tokio::test]
async fn malformed_and_unauthenticated_packets_do_not_kill_server() {
    let backend = UdpSocket::bind("127.0.0.1:0").await.unwrap();
    let backend_addr = backend.local_addr().unwrap();
    let (server_config, client_config) = configs(45784, backend_addr.port());
    let server = DirectUdpServer::bind(server_config).await.unwrap();
    let stats = server.stats();
    let server_task = tokio::spawn(server.run());
    let raw = UdpSocket::bind("127.0.0.1:0").await.unwrap();
    raw.send_to(&[1, 2, 3], "127.0.0.1:45784").await.unwrap();
    raw.send_to(&[0u8; 64], "127.0.0.1:45784").await.unwrap();
    let mut wrong = client_config.clone();
    wrong.client.secret.secret = Some("B".repeat(64));
    assert!(matches!(
        DirectUdpSession::connect(&wrong).await,
        Err(TransportError::AuthenticationRejected(_))
    ));
    let backend_task = tokio::spawn(async move {
        let mut buf = [0u8; 2048];
        let (len, peer) = backend.recv_from(&mut buf).await.unwrap();
        backend.send_to(&buf[..len], peer).await.unwrap();
    });
    let session = DirectUdpSession::connect(&client_config).await.unwrap();
    session
        .send_datagram(Bytes::from_static(b"still-alive"))
        .await
        .unwrap();
    assert_eq!(
        timeout(Duration::from_secs(2), session.recv_datagram())
            .await
            .unwrap()
            .unwrap(),
        Bytes::from_static(b"still-alive")
    );
    assert!(stats.snapshot().datagrams_dropped >= 3);
    session.close().await.unwrap();
    backend_task.await.unwrap();
    server_task.abort();
}

#[tokio::test]
async fn authenticated_source_port_change_updates_return_peer() {
    let backend = UdpSocket::bind("127.0.0.1:0").await.unwrap();
    let backend_addr = backend.local_addr().unwrap();
    let (server_config, _) = configs(45785, backend_addr.port());
    let server = DirectUdpServer::bind(server_config).await.unwrap();
    let server_task = tokio::spawn(server.run());
    let first = UdpSocket::bind("127.0.0.1:0").await.unwrap();
    let second = UdpSocket::bind("127.0.0.1:0").await.unwrap();
    let secret = base64::engine::general_purpose::STANDARD
        .decode(SECRET)
        .unwrap();
    let mut nonce = [0u8; 16];
    nonce[0] = 1;
    let hello = protocol::HelloPacket::authenticated(
        &secret,
        Bytes::from_static(b"test"),
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs(),
        nonce,
    )
    .unwrap();
    first
        .send_to(&hello.encode().unwrap(), "127.0.0.1:45785")
        .await
        .unwrap();
    let mut wire = [0u8; 2048];
    let (len, _) = timeout(Duration::from_secs(2), first.recv_from(&mut wire))
        .await
        .unwrap()
        .unwrap();
    let accept = match protocol::parse_packet(&wire[..len]).unwrap() {
        protocol::Packet::Accept(packet) => packet,
        other => panic!("unexpected handshake response: {other:?}"),
    };
    let key =
        protocol::derive_session_key(&secret, &nonce, &accept.server_nonce, accept.session_id);
    let data = protocol::DataPacket::authenticated(
        &key,
        accept.session_id,
        1,
        Bytes::from_static(b"rebound"),
    );
    second
        .send_to(&data.encode().unwrap(), "127.0.0.1:45785")
        .await
        .unwrap();
    let backend_task = tokio::spawn(async move {
        let mut buf = [0u8; 2048];
        let (len, peer) = backend.recv_from(&mut buf).await.unwrap();
        backend.send_to(&buf[..len], peer).await.unwrap();
    });
    let (len, _) = timeout(Duration::from_secs(2), second.recv_from(&mut wire))
        .await
        .unwrap()
        .unwrap();
    match protocol::parse_packet(&wire[..len]).unwrap() {
        protocol::Packet::Data(packet) => {
            assert_eq!(packet.payload, Bytes::from_static(b"rebound"))
        }
        other => panic!("unexpected rebound response: {other:?}"),
    }
    backend_task.await.unwrap();
    server_task.abort();
}

#[tokio::test]
async fn unavailable_backend_does_not_terminate_listener() {
    let (server_config, client_config) = configs(45786, 9);
    let server = DirectUdpServer::bind(server_config).await.unwrap();
    let server_task = tokio::spawn(server.run());
    let session = DirectUdpSession::connect(&client_config).await.unwrap();
    session
        .send_datagram(Bytes::from_static(b"backend-may-be-down"))
        .await
        .unwrap();
    tokio::time::sleep(Duration::from_millis(100)).await;
    assert!(
        !server_task.is_finished(),
        "backend failure must not terminate the public listener"
    );
    session.close().await.unwrap();
    server_task.abort();
}