//! v3.1 multi-hop / onion-routing integration test.
//!
//! Drives three actors on loopback in one process:
//!
//! * **Exit** — a vanilla [`UdpServer`] bound on a free UDP port. Its cert SAN is
//!   `"localhost-exit"`. The server's accept task echoes the first three received packets back to
//!   the sender, then drops.
//!
//! * **Relay** — another [`UdpServer`] on a free port, cert SAN `"localhost-relay"`. Its accept
//!   task:
//!     1. accepts one connection (running its own outer Aura mutual-auth handshake with the
//!        client),
//!     2. uses [`crate::relay::rendezvous`] to read the client's `ExtendBridge` envelope and open
//!        a `connect()`ed UDP socket to the exit,
//!     3. spawns [`crate::relay::run_bridge`] to ferry bytes between the client and the bridge.
//!
//! * **Client** — calls [`circuit::dial_circuit_with_relay_name`] with
//!   `relay_server_name = Some("localhost-relay")` and `proto_cfg.server_name = "localhost-exit"`.
//!   The returned [`circuit::CircuitConnection`] should have `peer_id() == Some("localhost-exit")`
//!   — the core multi-hop invariant: the **inner** handshake authenticated the exit's cert
//!   through the relay opaquely, even though the outer hop authenticated the relay's cert.
//!
//! The test then exchanges three packets of varying sizes through the circuit and asserts that
//! every echoed reply matches.

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

use aura_cli::circuit;
use aura_cli::relay::{self, RendezvousOutcome};
use aura_pki::AuraCa;
use aura_proto::{ClientConfig, PacketConnection, ServerConfig};
use aura_transport::{UdpOpts, UdpServer};

const EXIT_SAN: &str = "localhost-exit";
const RELAY_SAN: &str = "localhost-relay";
const CLIENT_ID: &str = "client-multihop";

/// Reserve and immediately release a free UDP port on loopback (the window before re-bind in the
/// same process is negligible on a quiet test).
fn free_udp_port() -> u16 {
    let sock = std::net::UdpSocket::bind("127.0.0.1:0").expect("bind ephemeral udp");
    sock.local_addr().expect("local_addr").port()
}

/// Build a [`ServerConfig`] from one shared CA, with the given SAN.
fn server_cfg(ca: &AuraCa, san: &str) -> ServerConfig {
    let issued = ca.issue_server_cert(san).expect("issue server cert");
    ServerConfig {
        ca_cert_pem: ca.ca_cert_pem(),
        server_cert_pem: issued.cert_pem,
        server_key_pem: issued.key_pem,
    }
}

/// Build a [`ClientConfig`] from one shared CA. `server_name` is used by the **inner** handshake
/// (the exit). The outer handshake's expected SAN is overridden separately at
/// [`circuit::dial_circuit_with_relay_name`] callsite.
fn client_cfg(ca: &AuraCa, server_name: &str) -> ClientConfig {
    let issued = ca.issue_client_cert(CLIENT_ID).expect("issue client cert");
    ClientConfig {
        ca_cert_pem: ca.ca_cert_pem(),
        client_cert_pem: issued.cert_pem,
        client_key_pem: issued.key_pem,
        server_name: server_name.to_string(),
    }
}

/// Spawn the exit server: accept one connection and echo the first three packets back.
async fn spawn_exit(server: UdpServer) {
    let conn = server.accept().await.expect("exit accept");
    // The dropped server keeps the master loop alive via the connection's anchor.
    drop(server);
    let conn: Arc<dyn PacketConnection> = Arc::new(conn);
    for _ in 0..3 {
        match conn.recv_packet().await {
            Ok(pkt) => {
                if conn.send_packet(&pkt).await.is_err() {
                    return;
                }
            }
            Err(_) => return,
        }
    }
}

/// Spawn the relay server: accept one connection, run the rendezvous, and bridge to the exit.
async fn spawn_relay(server: UdpServer, whitelist: Vec<SocketAddr>) {
    let conn = server.accept().await.expect("relay accept");
    drop(server);
    let conn: Arc<dyn PacketConnection> = Arc::new(conn);
    match relay::rendezvous(&conn, &whitelist).await {
        RendezvousOutcome::Bridged { bridge } => {
            relay::run_bridge(conn, bridge).await;
        }
        RendezvousOutcome::Refused => {
            // Test path that exercises whitelist refusal — the relay sent CircuitFailed
            // already; just exit.
        }
        RendezvousOutcome::Fallback { .. } => {
            // The client did not send ExtendBridge — should not happen in the happy path.
            panic!("relay rendezvous fell back unexpectedly");
        }
    }
}

#[tokio::test(flavor = "multi_thread")]
async fn multihop_v3_1_end_to_end() {
    // One shared CA. Each role gets its own server cert with its own SAN.
    let ca = AuraCa::generate("Aura Multi-Hop Test CA").expect("ca");
    let exit_proto = server_cfg(&ca, EXIT_SAN);
    let relay_proto = server_cfg(&ca, RELAY_SAN);
    let client_proto = client_cfg(&ca, EXIT_SAN);

    let exit_port = free_udp_port();
    let relay_port = free_udp_port();
    let exit_addr: SocketAddr = format!("127.0.0.1:{exit_port}").parse().unwrap();
    let relay_addr: SocketAddr = format!("127.0.0.1:{relay_port}").parse().unwrap();

    // Bind both servers BEFORE spawning the client so they are ready to accept.
    let exit_server =
        UdpServer::bind(exit_addr, exit_proto, UdpOpts::default()).expect("bind exit");
    let relay_server =
        UdpServer::bind(relay_addr, relay_proto, UdpOpts::default()).expect("bind relay");
    let exit_actual = exit_server.local_addr().expect("exit addr");
    let relay_actual = relay_server.local_addr().expect("relay addr");

    // Whitelist contains exactly the exit address.
    let whitelist = vec![exit_actual];

    let exit_task = tokio::spawn(spawn_exit(exit_server));
    let relay_task = tokio::spawn(spawn_relay(relay_server, whitelist));

    // Give the servers a beat to enter their accept loops. Not strictly required (accept is
    // resumable) but makes the trace easier to follow on failure.
    tokio::time::sleep(Duration::from_millis(20)).await;

    // Client: dial circuit. proto_cfg.server_name = "localhost-exit" so the inner handshake's
    // verifier checks the exit's SAN; the outer handshake checks the relay's SAN via the explicit
    // override.
    let circuit_conn = tokio::time::timeout(
        Duration::from_secs(30),
        circuit::dial_circuit_with_relay_name(
            &[relay_actual, exit_actual],
            client_proto,
            UdpOpts::default(),
            Some(RELAY_SAN),
        ),
    )
    .await
    .expect("dial_circuit did not finish within 30s")
    .expect("dial_circuit succeeded");

    // The core invariant: the INNER handshake authenticated the EXIT (not the relay).
    assert_eq!(
        circuit_conn.peer_id(),
        Some(EXIT_SAN),
        "circuit.peer_id() must be the exit's SAN — the inner handshake verified the exit's cert"
    );

    // Echo three packets of varying sizes through the circuit.
    let payloads: Vec<Vec<u8>> = vec![
        b"hello multi-hop".to_vec(),
        vec![0xCDu8; 800],
        (0..=255u8).collect(),
    ];
    for pkt in &payloads {
        circuit_conn.send_packet(pkt).await.expect("circuit send");
        let echoed = tokio::time::timeout(Duration::from_secs(5), circuit_conn.recv_packet())
            .await
            .expect("recv timeout")
            .expect("recv from exit through circuit");
        assert_eq!(&echoed, pkt, "echoed payload must match");
    }

    // Clean shutdown — drop the client first, then wait for the actors to finish.
    drop(circuit_conn);
    let _ = tokio::time::timeout(Duration::from_secs(5), exit_task).await;
    let _ = tokio::time::timeout(Duration::from_secs(5), relay_task).await;
}

/// A whitelist that does NOT contain the exit's address must cause `dial_circuit` to fail with an
/// error mentioning "allow_extend_to" (the reason string sent in `CircuitFailed`).
#[tokio::test(flavor = "multi_thread")]
async fn multihop_whitelist_rejects_disallowed_exit() {
    let ca = AuraCa::generate("Aura Multi-Hop Test CA").expect("ca");
    let exit_proto = server_cfg(&ca, EXIT_SAN);
    let relay_proto = server_cfg(&ca, RELAY_SAN);
    let client_proto = client_cfg(&ca, EXIT_SAN);

    let exit_port = free_udp_port();
    let relay_port = free_udp_port();
    let exit_addr: SocketAddr = format!("127.0.0.1:{exit_port}").parse().unwrap();
    let relay_addr: SocketAddr = format!("127.0.0.1:{relay_port}").parse().unwrap();

    let exit_server =
        UdpServer::bind(exit_addr, exit_proto, UdpOpts::default()).expect("bind exit");
    let relay_server =
        UdpServer::bind(relay_addr, relay_proto, UdpOpts::default()).expect("bind relay");
    let exit_actual = exit_server.local_addr().expect("exit addr");
    let relay_actual = relay_server.local_addr().expect("relay addr");

    // Whitelist contains a different (fake) exit; the real exit is NOT allowed.
    let fake: SocketAddr = "10.255.255.1:9".parse().unwrap();
    let whitelist = vec![fake];

    // Exit task: just sit there; we expect the relay never bridges to it.
    let _exit_task = tokio::spawn(async move {
        // Accept may never resolve; exit when test ends.
        let _ = exit_server.accept().await;
    });
    let relay_task = tokio::spawn(spawn_relay(relay_server, whitelist));

    tokio::time::sleep(Duration::from_millis(20)).await;

    // dial_circuit must error with a message mentioning "allow_extend_to".
    let res = tokio::time::timeout(
        Duration::from_secs(15),
        circuit::dial_circuit_with_relay_name(
            &[relay_actual, exit_actual],
            client_proto,
            UdpOpts::default(),
            Some(RELAY_SAN),
        ),
    )
    .await
    .expect("dial_circuit_with_relay_name returned within 15s");

    let err = match res {
        Ok(_) => panic!("dial_circuit must fail when exit is not on the whitelist"),
        Err(e) => e,
    };
    let msg = format!("{err:#}");
    assert!(
        msg.contains("allow_extend_to") || msg.contains("not in"),
        "expected 'allow_extend_to' / 'not in' in error, got: {msg}"
    );

    let _ = tokio::time::timeout(Duration::from_secs(2), relay_task).await;
}

/// When the v3.1 relay path is **disabled** at the server, the server's accept-side never reads
/// the client's ExtendBridge envelope as a control message — instead the server would treat the
/// connection as a normal VPN client. From the client's `dial_circuit` perspective the relay
/// never sends `CircuitReady`, so the client times out (`READY_TIMEOUT_SECS`-bounded).
///
/// This test exercises that exact fallback: we run a `UdpServer` with NO rendezvous task,
/// accept the connection, and just keep it open. The client's `dial_circuit` must return an Err
/// whose message mentions a timeout / CircuitReady.
#[tokio::test(flavor = "multi_thread")]
async fn multihop_back_compat_relay_disabled() {
    let ca = AuraCa::generate("Aura Multi-Hop Test CA").expect("ca");
    let exit_proto = server_cfg(&ca, EXIT_SAN);
    let relay_proto = server_cfg(&ca, RELAY_SAN);
    let client_proto = client_cfg(&ca, EXIT_SAN);

    let exit_port = free_udp_port();
    let relay_port = free_udp_port();
    let exit_addr: SocketAddr = format!("127.0.0.1:{exit_port}").parse().unwrap();
    let relay_addr: SocketAddr = format!("127.0.0.1:{relay_port}").parse().unwrap();

    let exit_server =
        UdpServer::bind(exit_addr, exit_proto, UdpOpts::default()).expect("bind exit");
    let relay_server =
        UdpServer::bind(relay_addr, relay_proto, UdpOpts::default()).expect("bind relay");
    let exit_actual = exit_server.local_addr().expect("exit addr");
    let relay_actual = relay_server.local_addr().expect("relay addr");

    // Exit task: idle.
    let _exit_task = tokio::spawn(async move {
        let _ = exit_server.accept().await;
    });
    // Relay task: just accept and keep the connection alive WITHOUT running the rendezvous. This
    // models a v2 server that does not know about `ExtendBridge`. The client's incoming
    // `ExtendBridge` envelope is just an opaque payload from the server's perspective.
    let relay_task = tokio::spawn(async move {
        let conn = relay_server.accept().await.expect("relay accept");
        // Hold the connection until the test ends.
        tokio::time::sleep(Duration::from_secs(20)).await;
        drop(conn);
    });

    tokio::time::sleep(Duration::from_millis(20)).await;

    // The client must time out waiting for CircuitReady.
    let res = tokio::time::timeout(
        Duration::from_secs(20),
        circuit::dial_circuit_with_relay_name(
            &[relay_actual, exit_actual],
            client_proto,
            UdpOpts::default(),
            Some(RELAY_SAN),
        ),
    )
    .await
    .expect("dial_circuit returned within 20s");

    let err = match res {
        Ok(_) => panic!("dial_circuit must fail when the relay never sends CircuitReady"),
        Err(e) => e,
    };
    let msg = format!("{err:#}");
    assert!(
        msg.contains("timeout") || msg.contains("CircuitReady"),
        "expected timeout / CircuitReady in error, got: {msg}"
    );

    relay_task.abort();
}

// ---- v3.2: 3-hop + per-hop client certs + cell padding -----------------------------------------

use aura_cli::cells::CellPaddingConn;
use aura_cli::circuit::HopConfig;

const ENTRY_SAN: &str = "localhost-entry";
const MIDDLE_SAN: &str = "localhost-middle";
const CLIENT_ID_ENTRY: &str = "client-entry";
const CLIENT_ID_MIDDLE: &str = "client-middle";
const CLIENT_ID_EXIT: &str = "client-exit";

/// Build a [`ClientConfig`] with the given CN and expected server SAN. v3.2: a different cert /
/// CN per hop is the identity-unlinkable design.
fn client_cfg_with_cn(ca: &AuraCa, cn: &str, server_name: &str) -> ClientConfig {
    let issued = ca.issue_client_cert(cn).expect("issue client cert");
    ClientConfig {
        ca_cert_pem: ca.ca_cert_pem(),
        client_cert_pem: issued.cert_pem,
        client_key_pem: issued.key_pem,
        server_name: server_name.to_string(),
    }
}

/// v3.2 3-hop end-to-end: `client → A (entry-relay) → B (middle-relay) → C (exit)`. Each hop is
/// a real Aura UdpServer on loopback. The client uses a **different** client cert per hop
/// (identity-unlinkable). The exit echoes three packets which the client must receive back
/// through three layers of AEAD encryption.
#[tokio::test(flavor = "multi_thread")]
async fn multihop_v3_2_three_hops_end_to_end() {
    let ca = AuraCa::generate("Aura v3.2 3-hop Test CA").expect("ca");

    let entry_proto = server_cfg(&ca, ENTRY_SAN);
    let middle_proto = server_cfg(&ca, MIDDLE_SAN);
    let exit_proto = server_cfg(&ca, EXIT_SAN);

    let entry_port = free_udp_port();
    let middle_port = free_udp_port();
    let exit_port = free_udp_port();
    let entry_addr: SocketAddr = format!("127.0.0.1:{entry_port}").parse().unwrap();
    let middle_addr: SocketAddr = format!("127.0.0.1:{middle_port}").parse().unwrap();
    let exit_addr: SocketAddr = format!("127.0.0.1:{exit_port}").parse().unwrap();

    let entry_server =
        UdpServer::bind(entry_addr, entry_proto, UdpOpts::default()).expect("bind entry");
    let middle_server =
        UdpServer::bind(middle_addr, middle_proto, UdpOpts::default()).expect("bind middle");
    let exit_server =
        UdpServer::bind(exit_addr, exit_proto, UdpOpts::default()).expect("bind exit");
    let entry_actual = entry_server.local_addr().expect("entry addr");
    let middle_actual = middle_server.local_addr().expect("middle addr");
    let exit_actual = exit_server.local_addr().expect("exit addr");

    // Whitelists per hop (CIDR-aware): entry allows middle; middle allows exit. Both can be exact
    // entries here; this test exercises the literal-IP:port path.
    let entry_whitelist = vec![middle_actual];
    let middle_whitelist = vec![exit_actual];

    let exit_task = tokio::spawn(spawn_exit(exit_server));
    let middle_task = tokio::spawn(spawn_relay(middle_server, middle_whitelist));
    let entry_task = tokio::spawn(spawn_relay(entry_server, entry_whitelist));

    tokio::time::sleep(Duration::from_millis(50)).await;

    // Per-hop client configs: distinct CN per hop, distinct server_name per hop.
    let hops = vec![
        HopConfig {
            addr: entry_actual,
            proto_cfg: client_cfg_with_cn(&ca, CLIENT_ID_ENTRY, ENTRY_SAN),
        },
        HopConfig {
            addr: middle_actual,
            proto_cfg: client_cfg_with_cn(&ca, CLIENT_ID_MIDDLE, MIDDLE_SAN),
        },
        HopConfig {
            addr: exit_actual,
            proto_cfg: client_cfg_with_cn(&ca, CLIENT_ID_EXIT, EXIT_SAN),
        },
    ];

    let circuit_conn = tokio::time::timeout(
        Duration::from_secs(60),
        circuit::dial_circuit(&hops, UdpOpts::default()),
    )
    .await
    .expect("dial_circuit did not finish within 60s")
    .expect("dial_circuit succeeded");

    // peer_id is the exit's SAN — the innermost handshake authenticated the exit cert through
    // every relay opaquely.
    assert_eq!(
        circuit_conn.peer_id(),
        Some(EXIT_SAN),
        "circuit.peer_id() must be the exit's SAN through 3 hops"
    );

    // Echo three packets — through THREE AEAD layers.
    let payloads: Vec<Vec<u8>> = vec![
        b"hello 3-hop".to_vec(),
        vec![0x77u8; 600],
        (0..200u8).collect(),
    ];
    for pkt in &payloads {
        circuit_conn.send_packet(pkt).await.expect("circuit send");
        let echoed = tokio::time::timeout(Duration::from_secs(10), circuit_conn.recv_packet())
            .await
            .expect("recv timeout")
            .expect("recv from exit through 3-hop circuit");
        assert_eq!(&echoed, pkt, "echoed payload must match");
    }

    drop(circuit_conn);
    let _ = tokio::time::timeout(Duration::from_secs(5), exit_task).await;
    let _ = tokio::time::timeout(Duration::from_secs(5), middle_task).await;
    let _ = tokio::time::timeout(Duration::from_secs(5), entry_task).await;
}

/// v3.2: smoke-test the [`CellPaddingConn`] wrap around a 2-hop circuit. The exit also wraps its
/// `Accepted.conn` in a `CellPaddingConn`; the bytes the client sends are padded cells, ferried
/// opaquely through the relay, and unwrapped by the exit. We exchange three payloads of varying
/// (small) sizes through the padded layer.
#[tokio::test(flavor = "multi_thread")]
async fn multihop_v3_2_cell_padding_smoke() {
    let ca = AuraCa::generate("Aura v3.2 cell-padding Test CA").expect("ca");
    let exit_proto = server_cfg(&ca, EXIT_SAN);
    let relay_proto = server_cfg(&ca, RELAY_SAN);
    let client_proto = client_cfg(&ca, EXIT_SAN);

    let exit_port = free_udp_port();
    let relay_port = free_udp_port();
    let exit_addr: SocketAddr = format!("127.0.0.1:{exit_port}").parse().unwrap();
    let relay_addr: SocketAddr = format!("127.0.0.1:{relay_port}").parse().unwrap();

    let exit_server =
        UdpServer::bind(exit_addr, exit_proto, UdpOpts::default()).expect("bind exit");
    let relay_server =
        UdpServer::bind(relay_addr, relay_proto, UdpOpts::default()).expect("bind relay");
    let exit_actual = exit_server.local_addr().expect("exit addr");
    let relay_actual = relay_server.local_addr().expect("relay addr");

    let whitelist = vec![exit_actual];

    // Exit echoes three CELL-PADDED packets back. The CellPaddingConn wrap on the exit's side
    // means recv_packet returns the original (unpadded) payload, and send_packet pads it again.
    let cell_size = 512;
    let exit_task = tokio::spawn(async move {
        let conn = exit_server.accept().await.expect("exit accept");
        drop(exit_server);
        let conn: Arc<dyn PacketConnection> = Arc::new(conn);
        let wrapped = Arc::new(CellPaddingConn::new(conn, cell_size));
        for _ in 0..3 {
            match wrapped.recv_packet().await {
                Ok(pkt) => {
                    if wrapped.send_packet(&pkt).await.is_err() {
                        return;
                    }
                }
                Err(_) => return,
            }
        }
    });
    let relay_task = tokio::spawn(spawn_relay(relay_server, whitelist));

    tokio::time::sleep(Duration::from_millis(20)).await;

    let circuit_conn = tokio::time::timeout(
        Duration::from_secs(30),
        circuit::dial_circuit_with_relay_name(
            &[relay_actual, exit_actual],
            client_proto,
            UdpOpts::default(),
            Some(RELAY_SAN),
        ),
    )
    .await
    .expect("dial_circuit did not finish within 30s")
    .expect("dial_circuit succeeded");

    // Wrap the client side in CellPaddingConn so its sends become cells.
    let padded: Arc<dyn PacketConnection> =
        Arc::new(CellPaddingConn::new(circuit_conn.into_dyn(), cell_size));

    let payloads: Vec<Vec<u8>> = vec![
        b"tiny".to_vec(),
        vec![0xEFu8; 100],
        b"another payload that fits inside cell".to_vec(),
    ];
    for pkt in &payloads {
        padded.send_packet(pkt).await.expect("padded send");
        let echoed = tokio::time::timeout(Duration::from_secs(10), padded.recv_packet())
            .await
            .expect("recv timeout")
            .expect("recv from padded exit");
        assert_eq!(&echoed, pkt, "padded roundtrip preserves payload");
    }

    drop(padded);
    let _ = tokio::time::timeout(Duration::from_secs(5), exit_task).await;
    let _ = tokio::time::timeout(Duration::from_secs(5), relay_task).await;
}