path: root/src/nat.rs

//! NAT type detection (STUN-like echo protocol).
//!
//! Detects whether
//! this node has a public IP, is behind a cone NAT, or is
//! behind a symmetric NAT. The result determines how the
//! node communicates:
//!
//! - **Global**: direct communication.
//! - **ConeNat**: hole-punching via DTUN.
//! - **SymmetricNat**: relay via proxy.
//!
//! ## Detection protocol
//!
//! 1. Send NatEcho to a known node.
//! 2. Receive NatEchoReply with our observed external
//!    address and port.
//! 3. If observed == local → Global.
//! 4. If different → we're behind NAT. Send
//!    NatEchoRedirect asking the peer to have a *third*
//!    node echo us from a different port.
//! 5. Receive NatEchoRedirectReply with observed port
//!    from the third node.
//! 6. If ports match → ConeNat (same external port for
//!    different destinations).
//! 7. If ports differ → SymmetricNat (external port
//!    changes per destination).

use std::collections::HashMap;
use std::net::SocketAddr;
use std::time::{Duration, Instant};

use crate::id::NodeId;

/// NAT detection echo timeout.
pub const ECHO_TIMEOUT: Duration = Duration::from_secs(3);

/// Periodic re-detection interval.
pub const NAT_TIMER_INTERVAL: Duration = Duration::from_secs(30);

/// Node's NAT state (visible to the application).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NatState {
    /// Not yet determined.
    Unknown,

    /// Public IP — direct communication.
    Global,

    /// Behind NAT, type not yet classified.
    Nat,

    /// Cone NAT — hole-punching works.
    ConeNat,

    /// Symmetric NAT — must use a proxy relay.
    SymmetricNat,
}

/// Internal state machine for the detection protocol.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum DetectorState {
    /// Initial / idle.
    Undefined,

    /// Sent first echo, waiting for reply.
    EchoWait1,

    /// Behind NAT, sent redirect request, waiting.
    EchoRedirectWait,

    /// Confirmed global (public IP).
    Global,

    /// Behind NAT, type unknown yet.
    Nat,

    /// Sent second echo (via redirected node), waiting.
    EchoWait2,

    /// Confirmed cone NAT.
    ConeNat,

    /// Confirmed symmetric NAT.
    SymmetricNat,
}

/// Pending echo request tracked by nonce.
struct EchoPending {
    sent_at: Instant,
}

/// NAT type detector.
///
/// Drives the echo protocol state machine. The caller
/// feeds received messages via `recv_*` methods and
/// reads the result via `state()`.
pub struct NatDetector {
    state: DetectorState,

    /// Our detected external address (if behind NAT).
    global_addr: Option<SocketAddr>,

    /// Port observed in the first echo reply.
    echo1_port: Option<u16>,

    /// Pending echo requests keyed by nonce.
    pending: HashMap<u32, EchoPending>,
}

impl NatDetector {
    pub fn new(_local_id: NodeId) -> Self {
        Self {
            state: DetectorState::Undefined,
            global_addr: None,
            echo1_port: None,
            pending: HashMap::new(),
        }
    }

    /// Current NAT state as seen by the application.
    pub fn state(&self) -> NatState {
        match self.state {
            DetectorState::Undefined
            | DetectorState::EchoWait1
            | DetectorState::EchoRedirectWait
            | DetectorState::EchoWait2 => NatState::Unknown,
            DetectorState::Global => NatState::Global,
            DetectorState::Nat => NatState::Nat,
            DetectorState::ConeNat => NatState::ConeNat,
            DetectorState::SymmetricNat => NatState::SymmetricNat,
        }
    }

    /// Our detected global address, if known.
    pub fn global_addr(&self) -> Option<SocketAddr> {
        self.global_addr
    }

    /// Force the NAT state (e.g. from configuration).
    pub fn set_state(&mut self, s: NatState) {
        self.state = match s {
            NatState::Unknown => DetectorState::Undefined,
            NatState::Global => DetectorState::Global,
            NatState::Nat => DetectorState::Nat,
            NatState::ConeNat => DetectorState::ConeNat,
            NatState::SymmetricNat => DetectorState::SymmetricNat,
        };
    }

    /// Start detection: prepare a NatEcho to send to a
    /// known peer.
    ///
    /// Returns the nonce to include in the echo message.
    pub fn start_detect(&mut self, nonce: u32) -> u32 {
        self.state = DetectorState::EchoWait1;
        self.pending.insert(
            nonce,
            EchoPending {
                sent_at: Instant::now(),
            },
        );
        nonce
    }

    /// Handle a NatEchoReply: the peer tells us our
    /// observed external address and port.
    /// Replies in unexpected states are silently ignored
    /// (catch-all returns `EchoReplyAction::Ignore`).
    pub fn recv_echo_reply(
        &mut self,
        nonce: u32,
        observed_addr: SocketAddr,
        local_addr: SocketAddr,
    ) -> EchoReplyAction {
        if self.pending.remove(&nonce).is_none() {
            return EchoReplyAction::Ignore;
        }

        match self.state {
            DetectorState::EchoWait1 => {
                if observed_addr.port() == local_addr.port()
                    && observed_addr.ip() == local_addr.ip()
                {
                    // Our address matches → we have a public IP
                    self.state = DetectorState::Global;
                    self.global_addr = Some(observed_addr);
                    EchoReplyAction::DetectionComplete(NatState::Global)
                } else {
                    // Behind NAT — need redirect test
                    self.state = DetectorState::Nat;
                    self.global_addr = Some(observed_addr);
                    self.echo1_port = Some(observed_addr.port());
                    EchoReplyAction::NeedRedirect
                }
            }
            DetectorState::EchoWait2 => {
                // Reply from redirected third node
                let port2 = observed_addr.port();
                if Some(port2) == self.echo1_port {
                    // Same external port → Cone NAT
                    self.state = DetectorState::ConeNat;
                    EchoReplyAction::DetectionComplete(NatState::ConeNat)
                } else {
                    // Different port → Symmetric NAT
                    self.state = DetectorState::SymmetricNat;
                    EchoReplyAction::DetectionComplete(NatState::SymmetricNat)
                }
            }
            _ => EchoReplyAction::Ignore,
        }
    }

    /// After receiving NeedRedirect, start the redirect
    /// phase with a new nonce.
    pub fn start_redirect(&mut self, nonce: u32) {
        self.state = DetectorState::EchoRedirectWait;
        self.pending.insert(
            nonce,
            EchoPending {
                sent_at: Instant::now(),
            },
        );
    }

    /// The redirect was forwarded and a third node will
    /// send us an echo. Transition to EchoWait2.
    pub fn redirect_sent(&mut self, nonce: u32) {
        self.state = DetectorState::EchoWait2;
        self.pending.insert(
            nonce,
            EchoPending {
                sent_at: Instant::now(),
            },
        );
    }

    /// Expire timed-out echo requests.
    ///
    /// If all pending requests timed out and we haven't
    /// completed detection, reset to Undefined.
    pub fn expire_pending(&mut self) {
        self.pending
            .retain(|_, p| p.sent_at.elapsed() < ECHO_TIMEOUT);

        if self.pending.is_empty() {
            match self.state {
                DetectorState::EchoWait1
                | DetectorState::EchoRedirectWait
                | DetectorState::EchoWait2 => {
                    log::debug!("NAT detection timed out, resetting");
                    self.state = DetectorState::Undefined;
                }
                _ => {}
            }
        }
    }

    /// Check if detection is still in progress.
    pub fn is_detecting(&self) -> bool {
        matches!(
            self.state,
            DetectorState::EchoWait1
                | DetectorState::EchoRedirectWait
                | DetectorState::EchoWait2
        )
    }

    /// Check if detection is complete (any terminal state).
    pub fn is_complete(&self) -> bool {
        matches!(
            self.state,
            DetectorState::Global
                | DetectorState::ConeNat
                | DetectorState::SymmetricNat
        )
    }
}

/// Action to take after processing an echo reply.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EchoReplyAction {
    /// Ignore (unknown nonce or wrong state).
    Ignore,

    /// Detection complete with the given NAT state.
    DetectionComplete(NatState),

    /// Need to send a NatEchoRedirect to determine
    /// NAT type (cone vs symmetric).
    NeedRedirect,
}

#[cfg(test)]
mod tests {
    use super::*;

    fn local() -> SocketAddr {
        "192.168.1.100:5000".parse().unwrap()
    }

    #[test]
    fn detect_global() {
        let mut nd = NatDetector::new(NodeId::from_bytes([0x01; 32]));
        let nonce = nd.start_detect(1);
        assert!(nd.is_detecting());

        let action = nd.recv_echo_reply(nonce, local(), local());
        assert_eq!(
            action,
            EchoReplyAction::DetectionComplete(NatState::Global)
        );
        assert_eq!(nd.state(), NatState::Global);
        assert!(nd.is_complete());
    }

    #[test]
    fn detect_cone_nat() {
        let mut nd = NatDetector::new(NodeId::from_bytes([0x01; 32]));

        // Phase 1: echo shows different address
        let n1 = nd.start_detect(1);
        let observed: SocketAddr = "1.2.3.4:5000".parse().unwrap();
        let action = nd.recv_echo_reply(n1, observed, local());
        assert_eq!(action, EchoReplyAction::NeedRedirect);
        assert_eq!(nd.state(), NatState::Nat);

        // Phase 2: redirect → third node echoes with
        // same port
        nd.redirect_sent(2);
        let observed2: SocketAddr = "1.2.3.4:5000".parse().unwrap();
        let action = nd.recv_echo_reply(2, observed2, local());
        assert_eq!(
            action,
            EchoReplyAction::DetectionComplete(NatState::ConeNat)
        );
        assert_eq!(nd.state(), NatState::ConeNat);
    }

    #[test]
    fn detect_symmetric_nat() {
        let mut nd = NatDetector::new(NodeId::from_bytes([0x01; 32]));

        // Phase 1: behind NAT
        let n1 = nd.start_detect(1);
        let observed: SocketAddr = "1.2.3.4:5000".parse().unwrap();
        nd.recv_echo_reply(n1, observed, local());

        // Phase 2: different port from third node
        nd.redirect_sent(2);
        let observed2: SocketAddr = "1.2.3.4:6000".parse().unwrap();
        let action = nd.recv_echo_reply(2, observed2, local());
        assert_eq!(
            action,
            EchoReplyAction::DetectionComplete(NatState::SymmetricNat)
        );
    }

    #[test]
    fn unknown_nonce_ignored() {
        let mut nd = NatDetector::new(NodeId::from_bytes([0x01; 32]));
        nd.start_detect(1);
        let action = nd.recv_echo_reply(999, local(), local());
        assert_eq!(action, EchoReplyAction::Ignore);
    }

    #[test]
    fn force_state() {
        let mut nd = NatDetector::new(NodeId::from_bytes([0x01; 32]));
        nd.set_state(NatState::Global);
        assert_eq!(nd.state(), NatState::Global);
        assert!(nd.is_complete());
    }

    #[test]
    fn timeout_resets() {
        let mut nd = NatDetector::new(NodeId::from_bytes([0x01; 32]));
        nd.start_detect(1);

        // Clear pending manually to simulate timeout
        nd.pending.clear();
        nd.expire_pending();
        assert_eq!(nd.state(), NatState::Unknown);
    }
}