v2

vnt-core/src/context/nat.rs

@@ -0,0 +1,147 @@
+use parking_lot::RwLock;
+use rust_p2p_core::nat::NatInfo;
+use rust_p2p_core::route::Index;
+use rust_p2p_core::tunnel::udp::UDPIndex;
+use std::collections::HashMap;
+use std::net::{Ipv4Addr, SocketAddr};
+use std::sync::Arc;
+
+#[derive(Clone, Default)]
+pub struct MyNatInfo {
+    nat_info: Arc<RwLock<Option<NatInfo>>>,
+}
+impl MyNatInfo {
+    pub fn get(&self) -> Option<NatInfo> {
+        self.nat_info.read().clone()
+    }
+    pub fn update_public_addr(&self, index: Index, addr: SocketAddr) {
+        let (ip, port) = if let Some(r) = mapping_addr(addr) {
+            r
+        } else {
+            return;
+        };
+        log::debug!("public_addr:{},{},index={index:?}", ip, port);
+        let mut nat_info = self.nat_info.write();
+        let Some(nat_info) = nat_info.as_mut() else {
+            return;
+        };
+        if rust_p2p_core::extend::addr::is_ipv4_global(&ip) {
+            if !nat_info.public_ips.contains(&ip) {
+                nat_info.public_ips.push(ip);
+            }
+            match index {
+                Index::Udp(index) => {
+                    let index = match index {
+                        UDPIndex::MainV4(index) => index,
+                        UDPIndex::MainV6(index) => index,
+                        UDPIndex::SubV4(_) => return,
+                    };
+                    if let Some(p) = nat_info.public_udp_ports.get_mut(index) {
+                        *p = port;
+                    }
+                }
+                Index::Tcp(_) => {
+                    nat_info.public_tcp_port = port;
+                }
+                _ => {}
+            }
+        } else {
+            log::debug!("not public addr: {addr:?}")
+        }
+    }
+    pub fn update_tcp_public_addr(&self, addr: SocketAddr) {
+        let SocketAddr::V4(addr) = addr else {
+            return;
+        };
+        let ip = *addr.ip();
+        let port = addr.port();
+        log::info!("tcp_public_addr, {}:{}", ip, port);
+        let mut nat_info = self.nat_info.write();
+        let Some(nat_info) = nat_info.as_mut() else {
+            return;
+        };
+        if ip.is_unspecified() && port == 0 {
+            nat_info.public_tcp_port = 0;
+            return;
+        }
+        if rust_p2p_core::extend::addr::is_ipv4_global(&ip) {
+            if !nat_info.public_ips.contains(&ip) {
+                nat_info.public_ips.push(ip);
+            }
+            nat_info.public_tcp_port = port;
+        } else {
+            log::debug!("not public addr: {addr:?}")
+        }
+    }
+    pub fn replace_nat_info(&self, nat_info: NatInfo) {
+        self.nat_info.write().replace(nat_info);
+    }
+    pub fn clear(&self) {
+        *self.nat_info.write() = None;
+    }
+}
+fn mapping_addr(addr: SocketAddr) -> Option<(Ipv4Addr, u16)> {
+    match addr {
+        SocketAddr::V4(addr) => Some((*addr.ip(), addr.port())),
+        SocketAddr::V6(addr) => addr.ip().to_ipv4_mapped().map(|ip| (ip, addr.port())),
+    }
+}
+
+#[derive(Copy, Clone, Debug)]
+pub struct PunchState {
+    pub count: i64,
+    pub last_ts: i64,
+}
+
+#[derive(Clone, Default)]
+pub struct PunchBackoff {
+    inner: Arc<RwLock<HashMap<Ipv4Addr, PunchState>>>,
+}
+
+impl PunchBackoff {
+    const MAX_BACKOFF_MS: i64 = 3_600_000; // 1h
+    const BASE_MS: i64 = 3000;
+
+    fn now() -> i64 {
+        crate::utils::time::now_ts_ms()
+    }
+
+    pub fn record(&self, ip: Ipv4Addr) {
+        let mut map = self.inner.write();
+        let entry = map.entry(ip).or_insert(PunchState {
+            count: 0,
+            last_ts: 0,
+        });
+        entry.count += 1;
+        entry.last_ts = Self::now();
+    }
+    #[allow(dead_code)]
+    pub fn reset(&self, ip: Ipv4Addr) {
+        self.inner.write().remove(&ip);
+    }
+    #[allow(dead_code)]
+    pub fn get(&self, ip: &Ipv4Addr) -> Option<PunchState> {
+        self.inner.read().get(ip).copied()
+    }
+
+    pub fn should_punch(&self, ip: Ipv4Addr) -> bool {
+        let map = self.inner.read();
+        let Some(state) = map.get(&ip) else {
+            return true;
+        };
+
+        let now = Self::now();
+        let elapsed = now - state.last_ts;
+
+        let mut backoff = Self::BASE_MS * state.count;
+        if backoff > Self::MAX_BACKOFF_MS {
+            backoff = Self::MAX_BACKOFF_MS;
+        }
+
+        elapsed >= backoff
+    }
+    #[allow(dead_code)]
+    pub fn clear(&self) {
+        self.inner.write().clear();
+    }
+}