1 files changed, 384 insertions, 0 deletions
diff --git a/src/nat.rs b/src/nat.rs
new file mode 100644
index 0000000..cfd5729
--- /dev/null
+++ b/src/nat.rs
@@ -0,0 +1,384 @@
+//! 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);
+    }
+}