AuraVPN/crates/aura-transport/tests/udp_knock.rs

//! Integration tests for the UDP **port-knocking** (probe resistance) feature.
//!
//! These exercise the end-to-end behaviour of [`UdpOpts::knock_required`] / [`UdpOpts::knock_key`]:
//!
//! * [`udp_knock_required_silent_drop_on_missing_or_wrong`] — server requires knocking; client does
//!   not knock → server stays silent (no reply within 1 s, so a passive scanner sees a closed port).
//! * [`udp_knock_required_accepts_valid`] — both sides knock with the same key → handshake completes
//!   like usual; the inner Aura mutual auth still drives the auth decision.
//! * [`udp_knock_disabled_back_compat`] — both sides disabled → exact pre-feature wire behaviour.
//!
//! The clock-skew tolerance test (±1 minute) lives as a unit test inside `src/udp.rs` so it can
//! drive [`validate_and_strip_knock`] directly with a faked "now" — much sharper than racing the
//! wall clock here.

use std::sync::Arc;
use std::time::Duration;

use aura_pki::AuraCa;
use aura_proto::{ClientConfig, PacketConnection, ServerConfig};
use aura_transport::{UdpClient, UdpConnection, UdpOpts, UdpServer};
use tokio::net::UdpSocket;

const SERVER_NAME: &str = "localhost";
const CLIENT_ID: &str = "client-knock";

/// Mint CA + server + client cert/key triples, returning matching handshake configs.
fn make_configs() -> (ServerConfig, ClientConfig) {
    let ca = AuraCa::generate("Aura UDP Knock Test CA").expect("generate CA");
    let server_cert = ca
        .issue_server_cert(SERVER_NAME)
        .expect("issue server cert");
    let client_cert = ca.issue_client_cert(CLIENT_ID).expect("issue client cert");
    let ca_pem = ca.ca_cert_pem();
    let server_cfg = ServerConfig {
        ca_cert_pem: ca_pem.clone(),
        server_cert_pem: server_cert.cert_pem,
        server_key_pem: server_cert.key_pem,
    };
    let client_cfg = ClientConfig {
        ca_cert_pem: ca_pem,
        client_cert_pem: client_cert.cert_pem,
        client_key_pem: client_cert.key_pem,
        server_name: SERVER_NAME.to_string(),
    };
    (server_cfg, client_cfg)
}

/// A 32-byte test knock key; in production this is `SHA-256(CA-cert-DER)` (the CLI computes it),
/// but for the transport-level tests any well-known constant is fine — both sides just need the
/// same bytes.
fn test_knock_key() -> [u8; 32] {
    let mut k = [0u8; 32];
    for (i, b) in k.iter_mut().enumerate() {
        *b = (i as u8).wrapping_mul(13).wrapping_add(7);
    }
    k
}

#[tokio::test]
async fn udp_knock_required_silent_drop_on_missing_or_wrong() {
    let (server_cfg, _client_cfg) = make_configs();

    // Server: require knocking with our test key. Tighten the handshake timeout so the test does
    // not have to wait the default 10 s for the (never-arriving) handshake.
    let server_opts = UdpOpts {
        knock_required: true,
        knock_key: Some(test_knock_key()),
        hs_timeout: Duration::from_secs(2),
        ..UdpOpts::default()
    };
    let server = UdpServer::bind("127.0.0.1:0".parse().unwrap(), server_cfg, server_opts)
        .expect("bind udp server");
    let server_addr = server.local_addr().expect("server local_addr");

    // Bind a raw client socket and send a single *unknocked* HS-shaped datagram. We do NOT run
    // `UdpClient::connect` here because that would inject the proto handshake's ClientHello and we
    // want to assert "the server is silent at the wire level".
    let raw_client = UdpSocket::bind("127.0.0.1:0")
        .await
        .expect("bind raw client");
    raw_client.connect(server_addr).await.expect("raw connect");

    // Wire layout the server expects when knock is OFF: 0x01 (HS) || hs_seq(2) || ack(2) || msg.
    // No knock prefix → the server's master loop must drop this silently.
    let mut unknocked_hs = vec![0x01u8, 0x00, 0x00, 0xFF, 0xFF];
    // Append some plausible-looking handshake-message bytes so the datagram is non-trivially sized.
    unknocked_hs.extend_from_slice(&[0u8; 64]);

    raw_client
        .send(&unknocked_hs)
        .await
        .expect("send unknocked HS");

    // The server must NOT reply. Wait 1 s for any inbound datagram; recv_from must time out.
    let mut buf = [0u8; 1024];
    let res = tokio::time::timeout(Duration::from_secs(1), raw_client.recv(&mut buf)).await;
    assert!(
        res.is_err(),
        "server replied to an unknocked HS datagram (got {} bytes), expected wire silence",
        res.unwrap_or(Ok(0)).unwrap_or(0),
    );

    // Cleanup: drop the server explicitly (also tears down the master loop).
    drop(server);
}

#[tokio::test]
async fn udp_knock_required_accepts_valid() {
    let (server_cfg, client_cfg) = make_configs();
    let key = test_knock_key();

    let opts = UdpOpts {
        knock_required: true,
        knock_key: Some(key),
        ..UdpOpts::default()
    };

    let server =
        UdpServer::bind("127.0.0.1:0".parse().unwrap(), server_cfg, opts).expect("bind udp server");
    let server_addr = server.local_addr().expect("server local_addr");

    let accept_task = tokio::spawn(async move { server.accept().await });
    let connect_task =
        tokio::spawn(async move { UdpClient::connect(server_addr, client_cfg, opts).await });

    let server_conn: UdpConnection = tokio::time::timeout(Duration::from_secs(15), accept_task)
        .await
        .expect("server accept timely")
        .expect("accept join")
        .expect("server accept");
    let client_conn: UdpConnection = tokio::time::timeout(Duration::from_secs(15), connect_task)
        .await
        .expect("client connect timely")
        .expect("connect join")
        .expect("client connect");

    assert_eq!(
        server_conn.peer_id(),
        Some(CLIENT_ID),
        "server learned client CN — handshake completed through knocking",
    );

    // Round-trip a packet both ways to prove the data path also works under knocking.
    let server_conn: Arc<dyn PacketConnection> = Arc::new(server_conn);
    let client_conn: Arc<dyn PacketConnection> = Arc::new(client_conn);
    client_conn
        .send_packet(b"knock knock")
        .await
        .expect("client send");
    let got = tokio::time::timeout(Duration::from_secs(5), server_conn.recv_packet())
        .await
        .expect("server recv timely")
        .expect("server recv");
    assert_eq!(got, b"knock knock");
}

#[tokio::test]
async fn udp_knock_disabled_back_compat() {
    let (server_cfg, client_cfg) = make_configs();
    let opts = UdpOpts::default(); // knock_required: false, knock_key: None

    let server =
        UdpServer::bind("127.0.0.1:0".parse().unwrap(), server_cfg, opts).expect("bind udp server");
    let server_addr = server.local_addr().expect("server local_addr");

    let accept_task = tokio::spawn(async move { server.accept().await });
    let connect_task =
        tokio::spawn(async move { UdpClient::connect(server_addr, client_cfg, opts).await });

    let server_conn: UdpConnection = tokio::time::timeout(Duration::from_secs(15), accept_task)
        .await
        .expect("server accept timely")
        .expect("accept join")
        .expect("server accept");
    let client_conn: UdpConnection = tokio::time::timeout(Duration::from_secs(15), connect_task)
        .await
        .expect("client connect timely")
        .expect("connect join")
        .expect("client connect");

    assert_eq!(server_conn.peer_id(), Some(CLIENT_ID));
    let server_conn: Arc<dyn PacketConnection> = Arc::new(server_conn);
    let client_conn: Arc<dyn PacketConnection> = Arc::new(client_conn);
    client_conn
        .send_packet(b"no-knock")
        .await
        .expect("client send");
    let got = tokio::time::timeout(Duration::from_secs(5), server_conn.recv_packet())
        .await
        .expect("server recv timely")
        .expect("server recv");
    assert_eq!(got, b"no-knock");
}