fast_nat/
nat.rs

1use crate::{bimap::BiHashMap, timeout::MaybeTimeout};
2use rustc_hash::FxHashMap;
3use std::{net::Ipv4Addr, time::Duration};
4
5/// A table of network address mappings
6#[derive(Debug)]
7pub struct NetworkAddressTable {
8    /// Internal address map
9    addr_map: BiHashMap<u32, u32>,
10    /// Secondary map used to keep track of timeouts
11    timeouts: FxHashMap<(u32, u32), MaybeTimeout>,
12}
13
14impl NetworkAddressTable {
15    /// Construct a new empty `NetworkAddressTable`
16    #[must_use]
17    pub fn new() -> Self {
18        Self::default()
19    }
20
21    /// Prune all old mappings
22    #[profiling::function]
23    pub fn prune(&mut self) {
24        log::trace!("Pruning old network address mappings");
25
26        // Compare all mappings against a common timestamp
27        let now = std::time::Instant::now();
28
29        // Remove all old mappings from both the bimap and the timeouts map
30        self.timeouts.retain(|(left, right), timeout| {
31            match timeout {
32                // Retain all indefinite mappings
33                MaybeTimeout::Never => true,
34                // Only retain mappings that haven't timed out yet
35                MaybeTimeout::After { duration, start } => {
36                    let should_retain = now.duration_since(*start) < *duration;
37                    if !should_retain {
38                        log::trace!(
39                            "Mapping {:?} -> {:?} has timed out and will be removed",
40                            left,
41                            right
42                        );
43                        self.addr_map.remove(left, right);
44                    }
45                    should_retain
46                }
47            }
48        });
49    }
50
51    /// Insert a new indefinite mapping
52    #[profiling::function]
53    pub fn insert_indefinite(&mut self, left: Ipv4Addr, right: Ipv4Addr) {
54        self.prune();
55        let (left, right) = (left.into(), right.into());
56        self.addr_map.insert(left, right);
57        self.timeouts.insert((left, right), MaybeTimeout::Never);
58    }
59
60    /// Insert a new mapping with a finite time-to-live
61    #[profiling::function]
62    pub fn insert(&mut self, left: Ipv4Addr, right: Ipv4Addr, duration: Duration) {
63        self.prune();
64        let (left, right) = (left.into(), right.into());
65        self.addr_map.insert(left, right);
66        self.timeouts.insert(
67            (left, right),
68            MaybeTimeout::After {
69                duration,
70                start: std::time::Instant::now(),
71            },
72        );
73    }
74
75    /// Get the right value for a given left value
76    #[must_use]
77    #[profiling::function]
78    pub fn get_right(&self, left: &Ipv4Addr) -> Option<Ipv4Addr> {
79        self.addr_map
80            .get_right(&(*left).into())
81            .map(|addr| (*addr).into())
82    }
83
84    /// Get the left value for a given right value
85    #[must_use]
86    #[profiling::function]
87    pub fn get_left(&self, right: &Ipv4Addr) -> Option<Ipv4Addr> {
88        self.addr_map
89            .get_left(&(*right).into())
90            .map(|addr| (*addr).into())
91    }
92}
93
94impl Default for NetworkAddressTable {
95    fn default() -> Self {
96        Self {
97            addr_map: BiHashMap::new(),
98            timeouts: FxHashMap::default(),
99        }
100    }
101}
fast_nat/nat.rs

fast_nat/
nat.rs
