AuraVPN/crates/aura-tunnel/tests/routes.rs

//! Deterministic tests for the tunnel data plane: routing-table classification, the `dst_ip`
//! parser, DNS host-route registration (no live query), and the router run-loop driven by a mock
//! [`PacketConnection`] and a mock TUN. None of these touch the network or require root.

use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
use std::sync::Arc;

use async_trait::async_trait;
use aura_proto::PacketConnection;
use aura_tunnel::router::dst_ip;
use aura_tunnel::tun::PacketIo;
use aura_tunnel::{AuraDns, AuraRouter, PacketCounters, RouteAction, RouteTable};
use tokio::sync::{mpsc, RwLock};

// ---- §8.4 RouteTable classification --------------------------------------------------------------

/// `192.168.1.1` matches the `192.168.0.0/16 -> Direct` rule under a `Vpn` default.
#[test]
fn test_route_classify_cidr() {
    let mut table = RouteTable::new(RouteAction::Vpn);
    table.add_cidr("192.168.0.0/16".parse().unwrap(), RouteAction::Direct);

    let ip: IpAddr = "192.168.1.1".parse().unwrap();
    assert_eq!(table.classify(ip), RouteAction::Direct);
}

/// `8.8.8.8` matches no rule and falls through to the `Vpn` default.
#[test]
fn test_route_classify_vpn() {
    let mut table = RouteTable::new(RouteAction::Vpn);
    table.add_cidr("192.168.0.0/16".parse().unwrap(), RouteAction::Direct);

    let ip: IpAddr = "8.8.8.8".parse().unwrap();
    assert_eq!(table.classify(ip), RouteAction::Vpn);
}

/// Longest-prefix wins: a more specific `/24 -> Vpn` overrides a less specific `/16 -> Direct`.
#[test]
fn test_route_priority() {
    let mut table = RouteTable::new(RouteAction::Direct);
    table.add_cidr("10.0.0.0/8".parse().unwrap(), RouteAction::Direct);
    table.add_cidr("10.1.2.0/24".parse().unwrap(), RouteAction::Vpn);

    // Inside the /24 -> the most specific rule (Vpn) wins.
    assert_eq!(
        table.classify("10.1.2.5".parse().unwrap()),
        RouteAction::Vpn
    );
    // Inside the /8 but outside the /24 -> the /8 rule (Direct) applies.
    assert_eq!(
        table.classify("10.9.9.9".parse().unwrap()),
        RouteAction::Direct
    );
    // Outside both -> default (Direct).
    assert_eq!(
        table.classify("8.8.8.8".parse().unwrap()),
        RouteAction::Direct
    );
}

/// A host route (`/32`) is the most specific possible match and overrides any broader rule.
#[test]
fn test_route_host_route_overrides() {
    let mut table = RouteTable::new(RouteAction::Vpn);
    table.add_cidr("0.0.0.0/0".parse().unwrap(), RouteAction::Vpn);
    table.add_host_route("1.2.3.4".parse().unwrap(), RouteAction::Direct);

    assert_eq!(
        table.classify("1.2.3.4".parse().unwrap()),
        RouteAction::Direct
    );
    assert_eq!(table.classify("1.2.3.5".parse().unwrap()), RouteAction::Vpn);
}

// ---- dst_ip parser ------------------------------------------------------------------------------

/// IPv4 header: version nibble 4, destination at bytes 16..20.
#[test]
fn test_dst_ip_v4() {
    let mut pkt = [0u8; 20];
    pkt[0] = 0x45; // version 4, IHL 5
    pkt[16] = 8;
    pkt[17] = 8;
    pkt[18] = 4;
    pkt[19] = 4;
    assert_eq!(dst_ip(&pkt), Some(IpAddr::V4(Ipv4Addr::new(8, 8, 4, 4))));
}

/// IPv6 header: version nibble 6, destination at bytes 24..40.
#[test]
fn test_dst_ip_v6() {
    let mut pkt = [0u8; 40];
    pkt[0] = 0x60; // version 6
    let dst = Ipv6Addr::new(0x2001, 0x4860, 0x4860, 0, 0, 0, 0, 0x8888);
    pkt[24..40].copy_from_slice(&dst.octets());
    assert_eq!(dst_ip(&pkt), Some(IpAddr::V6(dst)));
}

/// Too-short and unknown-version packets parse to `None`.
#[test]
fn test_dst_ip_invalid() {
    assert_eq!(dst_ip(&[]), None);
    assert_eq!(dst_ip(&[0x45, 0, 0]), None); // v4 but truncated
    let short_v6 = [0x60u8; 39];
    assert_eq!(dst_ip(&short_v6), None); // v6 but truncated
    let weird = [0x35u8; 64];
    assert_eq!(dst_ip(&weird), None); // version nibble 3
}

// ---- §8.5 AuraDns::register_ips (no live query) -------------------------------------------------

/// `register_ips` inserts each address as a host route in the shared table and caches the set —
/// validated without any DNS query.
#[tokio::test]
async fn test_dns_register_ips_no_query() {
    let routes = Arc::new(RwLock::new(RouteTable::new(RouteAction::Vpn)));
    let mut dns = AuraDns::new(Arc::clone(&routes)).await.unwrap();

    let ips = vec![
        IpAddr::V4(Ipv4Addr::new(93, 184, 216, 34)),
        IpAddr::V6(Ipv6Addr::new(
            0x2606, 0x2800, 0x220, 1, 0x248, 0x1893, 0x25c8, 0x1946,
        )),
    ];
    dns.register_ips("example.com", &ips, RouteAction::Direct)
        .await;

    // Both addresses now classify as Direct host routes.
    let table = routes.read().await;
    assert_eq!(table.classify(ips[0]), RouteAction::Direct);
    assert_eq!(table.classify(ips[1]), RouteAction::Direct);
    drop(table);

    // And the resolution is cached.
    assert_eq!(dns.cached("example.com"), Some(ips.as_slice()));
}

// ---- §8.6 AuraRouter run-loop with mock PacketConnection + mock TUN -----------------------------

/// In-memory fake TUN: `read_packet` drains an injected queue (and parks when empty), `write_packet`
/// forwards to a channel the test observes.
struct MockTun {
    inbound: mpsc::Receiver<Vec<u8>>,
    written: mpsc::Sender<Vec<u8>>,
}

#[async_trait]
impl PacketIo for MockTun {
    async fn read_packet(&mut self) -> std::io::Result<Vec<u8>> {
        match self.inbound.recv().await {
            Some(pkt) => Ok(pkt),
            // Channel closed: surface EOF so the router can stop cleanly.
            None => Err(std::io::Error::new(
                std::io::ErrorKind::UnexpectedEof,
                "mock TUN closed",
            )),
        }
    }

    async fn write_packet(&mut self, packet: &[u8]) -> std::io::Result<()> {
        self.written
            .send(packet.to_vec())
            .await
            .map_err(|_| std::io::Error::new(std::io::ErrorKind::BrokenPipe, "test dropped"))
    }
}

/// Mock encrypted connection backed by mpsc: `send_packet` forwards to a channel the test reads;
/// `recv_packet` drains a channel the test feeds.
struct MockConn {
    sent: mpsc::Sender<Vec<u8>>,
    to_recv: tokio::sync::Mutex<mpsc::Receiver<Vec<u8>>>,
}

#[async_trait]
impl PacketConnection for MockConn {
    async fn send_packet(&self, packet: &[u8]) -> anyhow::Result<()> {
        self.sent.send(packet.to_vec()).await?;
        Ok(())
    }

    async fn recv_packet(&self) -> anyhow::Result<Vec<u8>> {
        let mut rx = self.to_recv.lock().await;
        match rx.recv().await {
            Some(pkt) => Ok(pkt),
            None => Err(anyhow::anyhow!("mock conn closed")),
        }
    }
}

/// Build a minimal valid IPv4 packet whose destination is `dst`.
fn ipv4_packet_to(dst: Ipv4Addr) -> Vec<u8> {
    let mut pkt = vec![0u8; 20];
    pkt[0] = 0x45;
    let o = dst.octets();
    pkt[16..20].copy_from_slice(&o);
    pkt
}

#[tokio::test]
async fn test_router_vpn_outbound_and_inbound() {
    // Channels wiring the mocks to the test.
    let (tun_in_tx, tun_in_rx) = mpsc::channel::<Vec<u8>>(8); // test -> TUN read
    let (tun_out_tx, mut tun_out_rx) = mpsc::channel::<Vec<u8>>(8); // TUN write -> test
    let (conn_sent_tx, mut conn_sent_rx) = mpsc::channel::<Vec<u8>>(8); // conn.send -> test
    let (conn_recv_tx, conn_recv_rx) = mpsc::channel::<Vec<u8>>(8); // test -> conn.recv

    let tun = MockTun {
        inbound: tun_in_rx,
        written: tun_out_tx,
    };
    let conn: Arc<dyn PacketConnection> = Arc::new(MockConn {
        sent: conn_sent_tx,
        to_recv: tokio::sync::Mutex::new(conn_recv_rx),
    });

    // Default Vpn so a packet to 8.8.8.8 is routed through the connection.
    let routes = Arc::new(RwLock::new(RouteTable::new(RouteAction::Vpn)));
    let router = AuraRouter::new(tun, Arc::clone(&routes), conn);

    let handle = tokio::spawn(router.run());

    // (a) Outbound: emit a packet to 8.8.8.8 from the TUN -> it must reach the connection.
    let out_pkt = ipv4_packet_to(Ipv4Addr::new(8, 8, 8, 8));
    tun_in_tx.send(out_pkt.clone()).await.unwrap();
    let got = tokio::time::timeout(std::time::Duration::from_secs(2), conn_sent_rx.recv())
        .await
        .expect("router did not forward outbound packet to connection in time")
        .expect("connection sent channel closed");
    assert_eq!(got, out_pkt, "VPN-routed packet should be sent verbatim");

    // (b) Inbound: feed a packet into the connection -> it must be written to the TUN.
    let in_pkt = ipv4_packet_to(Ipv4Addr::new(10, 0, 0, 9));
    conn_recv_tx.send(in_pkt.clone()).await.unwrap();
    let written = tokio::time::timeout(std::time::Duration::from_secs(2), tun_out_rx.recv())
        .await
        .expect("router did not write inbound packet to TUN in time")
        .expect("TUN write channel closed");
    assert_eq!(written, in_pkt, "inbound packet should be written verbatim");

    // Shut the router down by closing the TUN read channel.
    drop(tun_in_tx);
    let _ = tokio::time::timeout(std::time::Duration::from_secs(2), handle).await;
}

/// A Direct-routed outbound packet must NOT be forwarded to the VPN connection (it goes to the v1
/// `send_direct` stub instead).
#[tokio::test]
async fn test_router_direct_not_sent_to_vpn() {
    let (tun_in_tx, tun_in_rx) = mpsc::channel::<Vec<u8>>(8);
    let (tun_out_tx, _tun_out_rx) = mpsc::channel::<Vec<u8>>(8);
    let (conn_sent_tx, mut conn_sent_rx) = mpsc::channel::<Vec<u8>>(8);
    let (_conn_recv_tx, conn_recv_rx) = mpsc::channel::<Vec<u8>>(8);

    let tun = MockTun {
        inbound: tun_in_rx,
        written: tun_out_tx,
    };
    let conn: Arc<dyn PacketConnection> = Arc::new(MockConn {
        sent: conn_sent_tx,
        to_recv: tokio::sync::Mutex::new(conn_recv_rx),
    });

    // Default Vpn, but a /16 -> Direct rule makes 192.168.x.x bypass the tunnel.
    let routes = Arc::new(RwLock::new(RouteTable::new(RouteAction::Vpn)));
    routes
        .write()
        .await
        .add_cidr("192.168.0.0/16".parse().unwrap(), RouteAction::Direct);
    let router = AuraRouter::new(tun, routes, conn);
    let handle = tokio::spawn(router.run());

    tun_in_tx
        .send(ipv4_packet_to(Ipv4Addr::new(192, 168, 1, 1)))
        .await
        .unwrap();

    // The connection must receive nothing within a short window.
    let res =
        tokio::time::timeout(std::time::Duration::from_millis(300), conn_sent_rx.recv()).await;
    assert!(
        res.is_err(),
        "Direct-routed packet must not be sent to the VPN connection"
    );

    drop(tun_in_tx);
    let _ = tokio::time::timeout(std::time::Duration::from_secs(2), handle).await;
}

// ---- PacketCounters wiring through AuraRouter ----------------------------------------------------

/// A VPN-routed outbound packet bumps `tx`; a DIRECT-routed outbound packet *also* bumps `tx`
/// (the v1 stub still counts as "tx from the TUN"); a packet pumped through the connection and
/// successfully written to the TUN bumps `rx`.
#[tokio::test]
async fn test_router_packet_counters_increment_for_tx_and_rx() {
    let (tun_in_tx, tun_in_rx) = mpsc::channel::<Vec<u8>>(8);
    let (tun_out_tx, mut tun_out_rx) = mpsc::channel::<Vec<u8>>(8);
    let (conn_sent_tx, mut conn_sent_rx) = mpsc::channel::<Vec<u8>>(8);
    let (conn_recv_tx, conn_recv_rx) = mpsc::channel::<Vec<u8>>(8);

    let tun = MockTun {
        inbound: tun_in_rx,
        written: tun_out_tx,
    };
    let conn: Arc<dyn PacketConnection> = Arc::new(MockConn {
        sent: conn_sent_tx,
        to_recv: tokio::sync::Mutex::new(conn_recv_rx),
    });

    // Default Vpn with a /16 -> Direct override so we can exercise both classifier branches.
    let mut table = RouteTable::new(RouteAction::Vpn);
    table.add_cidr("192.168.0.0/16".parse().unwrap(), RouteAction::Direct);
    let routes = Arc::new(RwLock::new(table));
    let counters = PacketCounters::new();
    let router = AuraRouter::with_stats(tun, routes, conn, Some(counters.clone()));
    let handle = tokio::spawn(router.run());

    // (a) VPN packet -> reaches connection and bumps tx to 1.
    let vpn_pkt = ipv4_packet_to(Ipv4Addr::new(8, 8, 8, 8));
    tun_in_tx.send(vpn_pkt.clone()).await.unwrap();
    let got = tokio::time::timeout(std::time::Duration::from_secs(2), conn_sent_rx.recv())
        .await
        .expect("router did not forward to connection")
        .expect("conn sent closed");
    assert_eq!(got, vpn_pkt);

    // (b) DIRECT packet -> goes to send_direct stub but tx still counts.
    let direct_pkt = ipv4_packet_to(Ipv4Addr::new(192, 168, 1, 1));
    tun_in_tx.send(direct_pkt).await.unwrap();

    // (c) Inbound packet -> written to TUN, bumps rx to 1.
    let in_pkt = ipv4_packet_to(Ipv4Addr::new(10, 0, 0, 9));
    conn_recv_tx.send(in_pkt.clone()).await.unwrap();
    let written = tokio::time::timeout(std::time::Duration::from_secs(2), tun_out_rx.recv())
        .await
        .expect("router did not write inbound packet to TUN")
        .expect("TUN write closed");
    assert_eq!(written, in_pkt);

    // Wait until both tx events have been observed (the DIRECT path doesn't surface anywhere
    // externally — poll the counter).
    let mut waited_ms = 0u64;
    while counters.tx_count() < 2 && waited_ms < 2000 {
        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
        waited_ms += 10;
    }
    assert_eq!(
        counters.tx_count(),
        2,
        "both VPN- and DIRECT-routed packets must bump tx"
    );
    assert_eq!(
        counters.rx_count(),
        1,
        "one packet was written to the TUN, so rx must be 1"
    );

    drop(tun_in_tx);
    let _ = tokio::time::timeout(std::time::Duration::from_secs(2), handle).await;
}

/// `AuraRouter::new` (no counters) must not panic and must not blow up on packets — verifies the
/// `None` branch of `with_stats` short-circuits safely.
#[tokio::test]
async fn test_router_no_counters_still_routes() {
    let (tun_in_tx, tun_in_rx) = mpsc::channel::<Vec<u8>>(8);
    let (tun_out_tx, _tun_out_rx) = mpsc::channel::<Vec<u8>>(8);
    let (conn_sent_tx, mut conn_sent_rx) = mpsc::channel::<Vec<u8>>(8);
    let (_conn_recv_tx, conn_recv_rx) = mpsc::channel::<Vec<u8>>(8);

    let tun = MockTun {
        inbound: tun_in_rx,
        written: tun_out_tx,
    };
    let conn: Arc<dyn PacketConnection> = Arc::new(MockConn {
        sent: conn_sent_tx,
        to_recv: tokio::sync::Mutex::new(conn_recv_rx),
    });
    let routes = Arc::new(RwLock::new(RouteTable::new(RouteAction::Vpn)));

    let router = AuraRouter::new(tun, routes, conn);
    let handle = tokio::spawn(router.run());

    let pkt = ipv4_packet_to(Ipv4Addr::new(1, 1, 1, 1));
    tun_in_tx.send(pkt.clone()).await.unwrap();
    let got = tokio::time::timeout(std::time::Duration::from_secs(2), conn_sent_rx.recv())
        .await
        .expect("router did not forward without counters")
        .expect("conn sent closed");
    assert_eq!(got, pkt);

    drop(tun_in_tx);
    let _ = tokio::time::timeout(std::time::Duration::from_secs(2), handle).await;
}