fast_nat/
cpnat.rs

1use std::{
2    net::{Ipv4Addr, Ipv6Addr},
3    time::Duration,
4};
5
6use ipnet::Ipv4Net;
7use rustc_hash::FxHashMap;
8
9use crate::{bimap::BiHashMap, error::Error, timeout::MaybeTimeout};
10
11/// A table of network address mappings across IPv4 and IPv6
12#[derive(Debug)]
13pub struct CrossProtocolNetworkAddressTable {
14    /// Internal address map
15    addr_map: BiHashMap<u32, u128>,
16    /// Secondary map used to keep track of timeouts
17    timeouts: FxHashMap<(u32, u128), MaybeTimeout>,
18}
19
20impl CrossProtocolNetworkAddressTable {
21    /// Construct a new empty `CrossProtocolNetworkAddressTable`
22    #[must_use]
23    pub fn new() -> Self {
24        Self::default()
25    }
26
27    /// Prune all old mappings
28    #[profiling::function]
29    pub fn prune(&mut self) {
30        log::trace!("Pruning old network address mappings");
31
32        // Compare all mappings against a common timestamp
33        let now = std::time::Instant::now();
34
35        // Remove all old mappings from both the bimap and the timeouts map
36        self.timeouts.retain(|(left, right), timeout| {
37            match timeout {
38                // Retain all indefinite mappings
39                MaybeTimeout::Never => true,
40                // Only retain mappings that haven't timed out yet
41                MaybeTimeout::After { duration, start } => {
42                    let should_retain = now.duration_since(*start) < *duration;
43                    if !should_retain {
44                        log::trace!(
45                            "Mapping {:?} -> {:?} has timed out and will be removed",
46                            left,
47                            right
48                        );
49                        self.addr_map.remove(left, right);
50                    }
51                    should_retain
52                }
53            }
54        });
55    }
56
57    /// Insert a new indefinite mapping
58    #[profiling::function]
59    pub fn insert_indefinite(&mut self, ipv4: Ipv4Addr, ipv6: Ipv6Addr) {
60        self.prune();
61        let (ipv4, ipv6) = (ipv4.into(), ipv6.into());
62        self.addr_map.insert(ipv4, ipv6);
63        self.timeouts.insert((ipv4, ipv6), MaybeTimeout::Never);
64    }
65
66    /// Insert a new mapping with a finite time-to-live
67    #[profiling::function]
68    pub fn insert(&mut self, ipv4: Ipv4Addr, ipv6: Ipv6Addr, duration: Duration) {
69        self.prune();
70        let (ipv4, ipv6) = (ipv4.into(), ipv6.into());
71        self.addr_map.insert(ipv4, ipv6);
72        self.timeouts.insert(
73            (ipv4, ipv6),
74            MaybeTimeout::After {
75                duration,
76                start: std::time::Instant::now(),
77            },
78        );
79    }
80
81    /// Get the IPv6 address for a given IPv4 address
82    #[must_use]
83    #[profiling::function]
84    pub fn get_ipv6(&self, ipv4: &Ipv4Addr) -> Option<Ipv6Addr> {
85        self.addr_map
86            .get_right(&(*ipv4).into())
87            .map(|addr| (*addr).into())
88    }
89
90    /// Get the IPv4 address for a given IPv6 address
91    #[must_use]
92    #[profiling::function]
93    pub fn get_ipv4(&self, ipv6: &Ipv6Addr) -> Option<Ipv4Addr> {
94        self.addr_map
95            .get_left(&(*ipv6).into())
96            .map(|addr| (*addr).into())
97    }
98
99    /// Get the number of mappings in the table
100    #[must_use]
101    #[profiling::function]
102    pub fn len(&self) -> usize {
103        self.addr_map.len()
104    }
105
106    /// Check if the table is empty
107    #[must_use]
108    #[profiling::function]
109    pub fn is_empty(&self) -> bool {
110        self.addr_map.is_empty()
111    }
112}
113
114impl Default for CrossProtocolNetworkAddressTable {
115    fn default() -> Self {
116        Self {
117            addr_map: BiHashMap::new(),
118            timeouts: FxHashMap::default(),
119        }
120    }
121}
122
123#[derive(Debug)]
124pub struct CrossProtocolNetworkAddressTableWithIpv4Pool {
125    /// Internal table
126    table: CrossProtocolNetworkAddressTable,
127    /// Internal pool of IPv4 prefixes to assign new mappings from
128    pool: Vec<Ipv4Net>,
129    /// The timeout to use for new entries
130    timeout: Duration,
131}
132
133impl CrossProtocolNetworkAddressTableWithIpv4Pool {
134    /// Construct a new Cross-protocol network address table with a given IPv4 pool
135    #[must_use]
136    pub fn new(pool: &[Ipv4Net], timeout: Duration) -> Self {
137        Self {
138            table: CrossProtocolNetworkAddressTable::default(),
139            pool: pool.to_vec(),
140            timeout,
141        }
142    }
143
144    /// Insert a new static mapping
145    #[profiling::function]
146    pub fn insert_static(&mut self, ipv4: Ipv4Addr, ipv6: Ipv6Addr) -> Result<(), Error> {
147        if !self.pool.iter().any(|prefix| prefix.contains(&ipv4)) {
148            return Err(Error::InvalidIpv4Address(ipv4));
149        }
150        self.table.insert_indefinite(ipv4, ipv6);
151        Ok(())
152    }
153
154    /// Gets the IPv4 address for a given IPv6 address or inserts a new mapping if one does not exist (if possible)
155    #[profiling::function]
156    pub fn get_or_create_ipv4(&mut self, ipv6: &Ipv6Addr) -> Result<Ipv4Addr, Error> {
157        // Return the known mapping if it exists
158        if let Some(ipv4) = self.table.get_ipv4(ipv6) {
159            return Ok(ipv4);
160        }
161
162        // Find the next available IPv4 address in the pool
163        let new_address = self
164            .pool
165            .iter()
166            .flat_map(Ipv4Net::hosts)
167            .find(|addr| self.table.get_ipv6(addr).is_none())
168            .ok_or(Error::Ipv4PoolExhausted)?;
169
170        // Insert the new mapping
171        self.table.insert(new_address, *ipv6, self.timeout);
172        log::info!(
173            "New cross-protocol address mapping: {} -> {}",
174            ipv6,
175            new_address
176        );
177
178        // Return the new address
179        Ok(new_address)
180    }
181
182    /// Gets the IPv6 address for a given IPv4 address if it exists
183    #[must_use]
184    #[profiling::function]
185    pub fn get_ipv6(&self, ipv4: &Ipv4Addr) -> Option<Ipv6Addr> {
186        self.table.get_ipv6(ipv4)
187    }
188}
fast_nat/cpnat.rs

fast_nat/
cpnat.rs
