init

crates/vpnshare-core/src/nat.rs

@@ -0,0 +1,104 @@
+use crate::lease::PeerId;
+use std::collections::HashMap;
+use std::net::IpAddr;
+use std::time::{Duration, Instant};
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum TransportProtocol {
+    Tcp,
+    Udp,
+    Icmp,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct FlowKey {
+    pub protocol: TransportProtocol,
+    pub source: IpAddr,
+    pub source_port: u16,
+    pub destination: IpAddr,
+    pub destination_port: u16,
+}
+
+#[derive(Debug, Clone)]
+pub struct FlowEntry {
+    pub key: FlowKey,
+    pub peer_id: PeerId,
+    pub created_at: Instant,
+    pub last_seen_at: Instant,
+    pub bytes_from_peer: u64,
+    pub bytes_to_peer: u64,
+}
+
+#[derive(Debug, Default)]
+pub struct FlowTable {
+    entries: HashMap<FlowKey, FlowEntry>,
+}
+
+impl FlowTable {
+    pub fn upsert(&mut self, key: FlowKey, peer_id: PeerId, now: Instant, bytes_from_peer: u64) {
+        self.entries
+            .entry(key.clone())
+            .and_modify(|entry| {
+                entry.last_seen_at = now;
+                entry.bytes_from_peer = entry.bytes_from_peer.saturating_add(bytes_from_peer);
+            })
+            .or_insert(FlowEntry {
+                key,
+                peer_id,
+                created_at: now,
+                last_seen_at: now,
+                bytes_from_peer,
+                bytes_to_peer: 0,
+            });
+    }
+
+    pub fn record_return_bytes(&mut self, key: &FlowKey, bytes_to_peer: u64) {
+        if let Some(entry) = self.entries.get_mut(key) {
+            entry.bytes_to_peer = entry.bytes_to_peer.saturating_add(bytes_to_peer);
+        }
+    }
+
+    pub fn expire_idle(&mut self, now: Instant, idle_timeout: Duration) -> usize {
+        let before = self.entries.len();
+        self.entries
+            .retain(|_, entry| now.duration_since(entry.last_seen_at) <= idle_timeout);
+        before - self.entries.len()
+    }
+
+    pub fn clear(&mut self) {
+        self.entries.clear();
+    }
+
+    pub fn len(&self) -> usize {
+        self.entries.len()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.entries.is_empty()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::net::{IpAddr, Ipv4Addr};
+
+    #[test]
+    fn expires_idle_flows() {
+        let now = Instant::now();
+        let mut table = FlowTable::default();
+        let key = FlowKey {
+            protocol: TransportProtocol::Tcp,
+            source: IpAddr::V4(Ipv4Addr::new(10, 241, 0, 2)),
+            source_port: 1234,
+            destination: IpAddr::V4(Ipv4Addr::new(1, 1, 1, 1)),
+            destination_port: 443,
+        };
+        table.upsert(key, PeerId::from_u128(1), now, 10);
+
+        assert_eq!(table.len(), 1);
+        let expired = table.expire_idle(now + Duration::from_secs(61), Duration::from_secs(60));
+        assert_eq!(expired, 1);
+        assert!(table.is_empty());
+    }
+}