citadel_wire 0.14.0

//! UDP Hole Punching Implementation
//!
//! This module implements the core UDP hole punching algorithm for NAT traversal,
//! enabling peer-to-peer connections between nodes behind different NATs. It handles
//! the complexities of coordinated connection establishment and socket management.
//!
//! # Features
//!
//! - Asynchronous hole punching implementation
//! - Dual-stack IPv4/IPv6 support
//! - Automatic retry mechanism
//! - Configurable timeouts
//! - NAT-aware socket binding
//! - Encrypted configuration exchange
//! - Network endpoint integration
//!
//! # Important Notes
//!
//! - Requires coordinated peer configuration
//! - Default timeout includes NAT identification time
//! - Multiple retries on failure (max 3)
//! - Socket binding optimized for NAT type
//! - IPv6 support depends on peer capability
//!
//! # Related Components
//!
//! - [`crate::nat_identification`] - NAT behavior analysis
//! - [`crate::udp_traversal::hole_punch_config`] - Configuration
//! - [`crate::standard::socket_helpers`] - Socket utilities
//! - [`crate::udp_traversal::hole_punched_socket`] - Socket management
//!

use crate::nat_identification::{NatType, IDENTIFY_TIMEOUT};
use crate::udp_traversal::hole_punch_config::HolePunchConfig;
use crate::udp_traversal::hole_punched_socket::HolePunchedUdpSocket;
use crate::udp_traversal::linear::encrypted_config_container::HolePunchConfigContainer;
use crate::udp_traversal::multi::DualStackUdpHolePuncher;
use citadel_io::tokio::net::UdpSocket;
use futures::Future;
use netbeam::reliable_conn::ReliableOrderedStreamToTargetExt;
use netbeam::sync::network_endpoint::NetworkEndpoint;
use netbeam::sync::subscription::Subscribable;
use serde::{Deserialize, Serialize};
use std::net::SocketAddr;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Duration;

/// The address candidates each peer publishes during the synchronized pre-process exchange.
/// `internal_bind_addrs` are the local bind addresses (loopback/same-LAN reachability);
/// `reflexive_addrs` are the *observed* server-reflexive (external) addresses of the actual
/// hole-punch sockets (ICE-style srflx candidates). Reflexive addrs are exact rather than
/// predicted, so they enable direct traversal for NATs the predictive model cannot infer.
#[derive(Serialize, Deserialize)]
struct HolePunchCandidates {
    internal_bind_addrs: Vec<SocketAddr>,
    reflexive_addrs: Vec<SocketAddr>,
}

pub struct UdpHolePuncher<'a> {
    driver: Pin<Box<dyn Future<Output = Result<HolePunchedUdpSocket, anyhow::Error>> + Send + 'a>>,
}

const DEFAULT_TIMEOUT: Duration =
    Duration::from_millis((IDENTIFY_TIMEOUT.as_millis() + 17000) as u64);

impl<'a> UdpHolePuncher<'a> {
    pub fn new(
        conn: &'a NetworkEndpoint,
        encrypted_config_container: HolePunchConfigContainer,
    ) -> Self {
        Self::new_timeout(conn, encrypted_config_container, DEFAULT_TIMEOUT)
    }

    pub fn new_timeout(
        conn: &'a NetworkEndpoint,
        encrypted_config_container: HolePunchConfigContainer,
        timeout: Duration,
    ) -> Self {
        Self {
            driver: Box::pin(
                async move { driver(conn, encrypted_config_container, timeout).await },
            ),
        }
    }
}

impl Future for UdpHolePuncher<'_> {
    type Output = Result<HolePunchedUdpSocket, anyhow::Error>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        self.driver.as_mut().poll(cx)
    }
}

const MAX_RETRIES: usize = 3;

#[cfg_attr(
    feature = "localhost-testing",
    tracing::instrument(level = "trace", target = "citadel", skip_all, ret, err(Debug))
)]
async fn driver(
    conn: &NetworkEndpoint,
    encrypted_config_container: HolePunchConfigContainer,
    timeout: Duration,
) -> Result<HolePunchedUdpSocket, anyhow::Error> {
    let mut retries = 0;
    loop {
        log::trace!(target: "citadel", "[driver] Attempt {}/{} starting (timeout: {:?})", retries + 1, MAX_RETRIES, timeout);
        let task = citadel_io::time::timeout(
            timeout,
            driver_inner(conn, encrypted_config_container.clone()),
        );
        match task.await {
            Ok(Ok(res)) => {
                log::trace!(target: "citadel", "[driver] Attempt {} succeeded!", retries + 1);
                return Ok(res);
            }
            Ok(Err(err)) => {
                log::warn!(target: "citadel", "[driver] Attempt {}/{} failed with error: {err:?}", retries + 1, MAX_RETRIES);
            }
            Err(_) => {
                log::warn!(target: "citadel", "[driver] Attempt {}/{} timed-out after {:?}", retries + 1, MAX_RETRIES, timeout);
            }
        }

        retries += 1;

        if retries >= MAX_RETRIES {
            log::error!(target: "citadel", "[driver] All {} attempts exhausted, giving up", MAX_RETRIES);
            return Err(anyhow::Error::msg("Max retries reached for UDP Traversal"));
        }
        log::trace!(target: "citadel", "[driver] Retrying... ({} attempts remaining)", MAX_RETRIES - retries);
    }
}

async fn driver_inner(
    conn: &NetworkEndpoint,
    mut encrypted_config_container: HolePunchConfigContainer,
) -> Result<HolePunchedUdpSocket, anyhow::Error> {
    log::trace!(target: "citadel", "[driver] Starting hole puncher ...");

    // Step 1: Initiate subscription - this can fail if NetworkEndpoint is stale/closed
    log::trace!(target: "citadel", "[driver] Step 1: Initiating subscription...");
    let stream = match conn.initiate_subscription().await {
        Ok(s) => {
            log::trace!(target: "citadel", "[driver] Step 1: Subscription successful");
            s
        }
        Err(e) => {
            log::error!(target: "citadel", "[driver] Step 1 FAILED: Subscription error: {e:?}");
            return Err(e);
        }
    };
    let stream = &stream;
    let stun_servers = encrypted_config_container.take_stun_servers();

    // Step 2: NAT type identification
    log::trace!(target: "citadel", "[driver] Step 2: Identifying local NAT type...");
    let local_nat_type = match NatType::identify(stun_servers.clone()).await {
        Ok(nat) => {
            log::trace!(target: "citadel", "[driver] Step 2: NAT identification successful: {nat:?}");
            nat
        }
        Err(e) => {
            log::error!(target: "citadel", "[driver] Step 2 FAILED: NAT identification error: {e:?}");
            return Err(anyhow::Error::msg(e.to_string()));
        }
    };
    let local_nat_type = &local_nat_type;

    // Step 3: Exchange NAT types with peer
    log::trace!(target: "citadel", "[driver] Step 3: Exchanging NAT types with peer...");
    if let Err(e) = stream.send_serialized(local_nat_type).await {
        log::error!(target: "citadel", "[driver] Step 3 FAILED: Send NAT type error: {e:?}");
        return Err(anyhow::Error::from(e));
    }
    log::trace!(target: "citadel", "[driver] Step 3a: Sent local NAT type, waiting for peer...");
    let peer_nat_type = match stream.recv_serialized::<NatType>().await {
        Ok(nat) => {
            log::trace!(target: "citadel", "[driver] Step 3: Received peer NAT type: {nat:?}");
            nat
        }
        Err(e) => {
            log::error!(target: "citadel", "[driver] Step 3 FAILED: Receive NAT type error: {e:?}");
            return Err(e.into());
        }
    };
    let peer_nat_type = &peer_nat_type;

    let local_initial_socket = get_optimal_bind_socket(local_nat_type, peer_nat_type)?;
    let internal_bind_addr_optimal = local_initial_socket.local_addr()?;
    let mut sockets = vec![local_initial_socket];
    let mut internal_addresses = vec![internal_bind_addr_optimal];
    if internal_bind_addr_optimal.is_ipv6() {
        let additional_socket = crate::socket_helpers::get_udp_socket("0.0.0.0:0")?;
        internal_addresses.push(additional_socket.local_addr()?);
        sockets.push(additional_socket);
    }

    // Step 3b: ICE-style — observe the *actual* external mapping of each hole-punch socket
    // via STUN, rather than relying solely on the predictive NAT model. This lets traversal
    // succeed for endpoint-independent NATs that randomize the external port. Failure to probe
    // is non-fatal: we simply fall back to the predicted bands.
    let local_reflexive_addrs = probe_reflexive_addrs(&sockets, stun_servers.as_deref()).await;
    log::info!(target: "citadel", "[driver] Local reflexive (srflx) addrs: {local_reflexive_addrs:?}");

    // Step 4: Exchange address candidates, synchronizing the beginning of the hole punch process
    log::trace!(target: "citadel", "[driver] Step 4: Exchanging address candidates...");
    let local_candidates = HolePunchCandidates {
        internal_bind_addrs: internal_addresses,
        reflexive_addrs: local_reflexive_addrs,
    };
    let peer_candidates = match conn.sync_exchange_payload(local_candidates).await {
        Ok(candidates) => {
            log::trace!(target: "citadel", "[driver] Step 4: Sync exchange successful, peer internal: {:?}, peer reflexive: {:?}", candidates.internal_bind_addrs, candidates.reflexive_addrs);
            candidates
        }
        Err(e) => {
            log::error!(target: "citadel", "[driver] Step 4 FAILED: Sync exchange error: {e:?}");
            return Err(e);
        }
    };
    let peer_internal_bind_addrs = peer_candidates.internal_bind_addrs;
    log::info!(target: "citadel", "\n~~~~~~~~~~~~\n [driver] Local NAT type: {local_nat_type:?}\n Peer NAT type: {peer_nat_type:?}");
    log::info!(target: "citadel", "[driver] Local internal bind addr: {internal_bind_addr_optimal:?}\nPeer internal bind addr: {peer_internal_bind_addrs:?}");
    log::info!(target: "citadel", "\n~~~~~~~~~~~~\n");
    // the next functions takes everything insofar obtained into account without causing collisions with any existing
    // connections (e.g., no conflicts with the primary stream existing in conn)
    let hole_punch_config = HolePunchConfig::new(
        peer_nat_type,
        &peer_internal_bind_addrs,
        &peer_candidates.reflexive_addrs,
        sockets,
    );

    // Step 5: Execute DualStackUdpHolePuncher
    let conn = conn.clone();
    log::trace!(target: "citadel", "[driver] Step 5: Creating and executing DualStackUdpHolePuncher...");
    let hole_puncher = match DualStackUdpHolePuncher::new(
        conn.node_type(),
        encrypted_config_container,
        hole_punch_config,
        conn,
    ) {
        Ok(hp) => {
            log::trace!(target: "citadel", "[driver] Step 5a: DualStackUdpHolePuncher created successfully");
            hp
        }
        Err(e) => {
            log::error!(target: "citadel", "[driver] Step 5 FAILED: DualStackUdpHolePuncher creation error: {e:?}");
            return Err(e);
        }
    };

    log::trace!(target: "citadel", "[driver] Step 5b: Awaiting hole punch result...");
    let res = hole_puncher.await;

    log::info!(target: "citadel", "Hole Punch Status: {res:?}");

    res.map_err(|err| {
        anyhow::Error::msg(format!(
            "**HOLE-PUNCH-ERR**: {err:?} | local_nat_type: {local_nat_type:?} | peer_nat_type: {peer_nat_type:?}",
        ))
    })
}

/// Concurrently STUN-probes each hole-punch socket for its server-reflexive (external)
/// address. Probes run in parallel and degrade gracefully: any socket that fails to obtain
/// a reflexive address (e.g. IPv6, or no STUN server reachable) is simply omitted.
async fn probe_reflexive_addrs(
    sockets: &[UdpSocket],
    stun_servers: Option<&[String]>,
) -> Vec<SocketAddr> {
    // Under localhost-testing there are no reachable STUN servers; skip probing entirely
    // (the configured/default servers would be unreachable and just add latency).
    if cfg!(feature = "localhost-testing") {
        return Vec::new();
    }

    let probes = sockets
        .iter()
        .map(|socket| NatType::get_reflexive_addr(socket, stun_servers));
    futures::future::join_all(probes)
        .await
        .into_iter()
        .flatten()
        .collect()
}

/// since the NAT traversal process always ensures that both public-facing and loopback
/// cases are covered, we can start by binding to 0.0.0.0, knowing that 127.0.0.1 will
/// also be covered automatically
///
/// Suppose A binds to ipv6 addr, and B binds to ipv4 addr, then B cannot send packets to
/// A. Only A can send to B via ipv4-mapped-v6 addrs. In order for B to send packets back to A,
/// B will need the ipv4 address of A.
pub fn get_optimal_bind_socket(
    local_nat_info: &NatType,
    peer_nat_info: &NatType,
) -> Result<UdpSocket, anyhow::Error> {
    let mut local_has_an_external_ipv6_addr = false;
    let mut peer_has_an_external_ipv6_addr = false;

    if let Some(other_info) = local_nat_info.ip_addr_info() {
        if other_info.external_ipv6.is_some() {
            local_has_an_external_ipv6_addr = true;
        }
    }

    if let Some(other_info) = peer_nat_info.ip_addr_info() {
        if other_info.external_ipv6.is_some() {
            peer_has_an_external_ipv6_addr = true;
        }
    }

    let local_allows_ipv6 = local_nat_info.is_ipv6_compatible();
    let peer_allows_ipv6 = peer_nat_info.is_ipv6_compatible();

    // only bind to ipv6 if v6 is enabled locally, and, both nodes have an external ipv6 addr,
    // AND, the peer allows ipv6, then go with ipv6
    if local_allows_ipv6
        && local_has_an_external_ipv6_addr
        && peer_has_an_external_ipv6_addr
        && peer_allows_ipv6
    {
        // bind to IN_ADDR6_ANY. Allows both conns from loopback and public internet
        crate::socket_helpers::get_udp_socket("[::]:0")
    } else {
        crate::socket_helpers::get_udp_socket("0.0.0.0:0")
    }
}

pub trait EndpointHolePunchExt {
    fn begin_udp_hole_punch(
        &self,
        encrypted_config_container: HolePunchConfigContainer,
    ) -> UdpHolePuncher<'_>;
}

impl EndpointHolePunchExt for NetworkEndpoint {
    fn begin_udp_hole_punch(
        &self,
        encrypted_config_container: HolePunchConfigContainer,
    ) -> UdpHolePuncher<'_> {
        UdpHolePuncher::new(self, encrypted_config_container)
    }
}

#[cfg(test)]
mod tests {
    use crate::udp_traversal::udp_hole_puncher::EndpointHolePunchExt;
    use citadel_io::tokio;
    use netbeam::sync::test_utils::create_streams_with_addrs_and_lag;
    use rstest::rstest;

    #[rstest]
    #[case(0)]
    #[case(50)]
    #[case(70)]
    #[tokio::test]
    async fn test_dual_hole_puncher(#[case] lag: usize) {
        citadel_logging::setup_log();

        let (server_stream, client_stream) = create_streams_with_addrs_and_lag(lag).await;

        let server = async move {
            let res = server_stream.begin_udp_hole_punch(Default::default()).await;
            log::trace!(target: "citadel", "Server res: {res:?}");
            res.unwrap()
        };

        let client = async move {
            let res = client_stream.begin_udp_hole_punch(Default::default()).await;
            log::trace!(target: "citadel", "Client res: {res:?}");
            res.unwrap()
        };

        let server = citadel_io::tokio::task::spawn(server);
        let client = citadel_io::tokio::task::spawn(client);
        let (res0, res1) = citadel_io::tokio::join!(server, client);
        log::trace!(target: "citadel", "JOIN complete! {res0:?} | {res1:?}");
        let (_res0, _res1) = (res0.unwrap(), res1.unwrap());

        // due to the error "An existing connection was forcibly closed by the remote host",
        // we must gate this test on localhost-testing + NOT windows
        #[cfg(not(target_os = "windows"))]
        {
            let dummy_bytes = b"Hello, world!";

            log::trace!(target: "citadel", "A");
            _res0
                .send_to(dummy_bytes as &[u8], _res0.addr.send_address)
                .await
                .unwrap();
            log::trace!(target: "citadel", "B");
            let buf = &mut [0u8; 4096];
            let (len, _addr) = _res1.recv_from(buf).await.unwrap();
            //assert_eq!(res1.addr.receive_address, addr);
            log::trace!(target: "citadel", "C");
            assert_ne!(len, 0);
            _res1
                .send_to(dummy_bytes, _res1.addr.send_address)
                .await
                .unwrap();
            let (len, _addr) = _res0.recv_from(buf).await.unwrap();
            assert_ne!(len, 0);
            //assert_eq!(res0.addr.receive_address, addr);
            log::trace!(target: "citadel", "D");
        }
    }

    /// Regression test for the hole-punch consensus under high coordination-channel lag — and, since
    /// the retry-desync fix, for the RETRY-RECOVERY path itself. `NetworkConnSimulator` adds a random
    /// 1x..2x delay per message, so 450 ms here means up to 900 ms/msg. That puts a single consensus
    /// attempt right at the production per-attempt timeout (`DEFAULT_TIMEOUT = IDENTIFY_TIMEOUT + 17s
    /// = 20s`), so the unlucky-tail runs exceed it and trigger a RETRY.
    ///
    /// Previously a retry was fatal: re-running the driver re-created the multiplex subscriptions,
    /// and because each attempt consumes a *variable* number of ids, a cancelled attempt left the two
    /// peers' subscription counters skewed — every later subscription then paired on mismatched ids,
    /// the consensus never agreed, the punched paths mismatched, and the exchange deadlocked. The fix
    /// (netbeam: hole-punch conns disable the local pre-reserve fast-path + `subscribe` catches up to
    /// the peer-driven id instead of panicking on a gap) lets a retry RECOVER: the Initiator simply
    /// re-pairs on the Receiver's next id. So at 450 ms this now passes by *succeeding on retry* —
    /// which is exactly what we want to guard. The budget below is sized to absorb a 20s attempt
    /// timeout + the recovering retry.
    #[cfg(not(target_os = "windows"))]
    // Skipped under coverage: llvm-cov instrumentation slows the consensus past the
    // (production) per-attempt timeout, deadlocking this lag-calibrated test in the
    // instrumented environment only. It still runs uninstrumented in the core_libs CI job.
    #[cfg_attr(coverage, ignore)]
    #[tokio::test]
    async fn test_dual_hole_puncher_high_lag_consensus() {
        use std::time::Duration;
        citadel_logging::setup_log();
        const LAG_MS: usize = 450;
        const ITERS: usize = 5;
        // Big enough to absorb a 20s attempt-timeout followed by a recovering retry (the retry now
        // re-pairs instead of deadlocking). A genuine failure (consensus never agreeing) still trips
        // it within MAX_RETRIES × DEFAULT_TIMEOUT.
        const PER_ITER_TIMEOUT: Duration = Duration::from_secs(75);

        for i in 0..ITERS {
            let (server_stream, client_stream) = create_streams_with_addrs_and_lag(LAG_MS).await;

            let iteration = async move {
                let server = citadel_io::tokio::task::spawn(async move {
                    server_stream
                        .begin_udp_hole_punch(Default::default())
                        .await
                        .map_err(|e| e.to_string())
                });
                let client = citadel_io::tokio::task::spawn(async move {
                    client_stream
                        .begin_udp_hole_punch(Default::default())
                        .await
                        .map_err(|e| e.to_string())
                });
                let (res0, res1) = citadel_io::tokio::join!(server, client);
                let s0 = res0
                    .expect("server task panicked")
                    .expect("server punch err");
                let s1 = res1
                    .expect("client task panicked")
                    .expect("client punch err");

                // Exchange a datagram both ways. If the consensus failed and the two sides
                // picked mismatched candidate sockets, this hangs and the per-iter timeout fires.
                let dummy = b"Hello, world!";
                s0.send_to(dummy as &[u8], s0.addr.send_address)
                    .await
                    .unwrap();
                let buf = &mut [0u8; 4096];
                let (len, _) = s1.recv_from(buf).await.unwrap();
                assert_ne!(len, 0);
                s1.send_to(dummy as &[u8], s1.addr.send_address)
                    .await
                    .unwrap();
                let (len, _) = s0.recv_from(buf).await.unwrap();
                assert_ne!(len, 0);
            };

            match citadel_io::tokio::time::timeout(PER_ITER_TIMEOUT, iteration).await {
                Ok(()) => {
                    log::info!(target: "citadel", "[hole-punch-consensus] iter {i}/{ITERS} OK")
                }
                Err(_) => {
                    panic!("iter {i}: hole-punch consensus deadlocked at lag {LAG_MS}ms (per-attempt timeout too short?)")
                }
            }
        }
    }
}