From 198b7bcd92ca59c53104d0de8e10816baa264652 Mon Sep 17 00:00:00 2001 From: Evan Pratten Date: Wed, 2 Aug 2023 13:14:51 -0400 Subject: [PATCH] Implement fast NAT --- Cargo.toml | 2 +- libs/easy-tun/src/lib.rs | 1 + libs/fast-nat/Cargo.toml | 17 ++++++ libs/fast-nat/README.md | 0 libs/fast-nat/src/bimap.rs | 111 +++++++++++++++++++++++++++++++++++ libs/fast-nat/src/cpnat.rs | 92 +++++++++++++++++++++++++++++ libs/fast-nat/src/lib.rs | 9 +++ libs/fast-nat/src/nat.rs | 87 +++++++++++++++++++++++++++ libs/fast-nat/src/timeout.rs | 13 ++++ 9 files changed, 331 insertions(+), 1 deletion(-) create mode 100644 libs/fast-nat/Cargo.toml create mode 100644 libs/fast-nat/README.md create mode 100644 libs/fast-nat/src/bimap.rs create mode 100644 libs/fast-nat/src/cpnat.rs create mode 100644 libs/fast-nat/src/lib.rs create mode 100644 libs/fast-nat/src/nat.rs create mode 100644 libs/fast-nat/src/timeout.rs diff --git a/Cargo.toml b/Cargo.toml index 4033a35..5904fc1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -56,4 +56,4 @@ # systemd-units = { enable = false } [workspace] -members=["libs/easy-tun"] \ No newline at end of file +members = ["libs/easy-tun", "libs/fast-nat"] diff --git a/libs/easy-tun/src/lib.rs b/libs/easy-tun/src/lib.rs index 2a78190..87bad45 100644 --- a/libs/easy-tun/src/lib.rs +++ b/libs/easy-tun/src/lib.rs @@ -1,3 +1,4 @@ +#![doc = include_str!("../README.md")] #![deny(clippy::pedantic)] #![allow(clippy::module_name_repetitions)] #![allow(clippy::missing_errors_doc)] diff --git a/libs/fast-nat/Cargo.toml b/libs/fast-nat/Cargo.toml new file mode 100644 index 0000000..00f4927 --- /dev/null +++ b/libs/fast-nat/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "fast-nat" +version = "0.1.0" +authors = ["Evan Pratten "] +edition = "2021" +description = "An OSI layer 3 Network Address Table built for speed" +readme = "README.md" +homepage = "https://github.com/ewpratten/protomask/tree/master/libs/fast-nat" +documentation = "https://docs.rs/fast-nat" +repository = "https://github.com/ewpratten/protomask" +license = "GPL-3.0" +keywords = [] +categories = [] + +[dependencies] +log = "^0.4" +rustc-hash = "1.1.0" diff --git a/libs/fast-nat/README.md b/libs/fast-nat/README.md new file mode 100644 index 0000000..e69de29 diff --git a/libs/fast-nat/src/bimap.rs b/libs/fast-nat/src/bimap.rs new file mode 100644 index 0000000..367f5c7 --- /dev/null +++ b/libs/fast-nat/src/bimap.rs @@ -0,0 +1,111 @@ +use std::hash::Hash; + +use rustc_hash::FxHashMap; + +/// A bi-directional hash map +#[derive(Debug, Clone)] +pub struct BiHashMap { + /// Mapping from a left value to a right value + left_to_right: FxHashMap, + /// Mapping from a right value to a left value + right_to_left: FxHashMap, +} + +impl BiHashMap +where + Left: Eq + Hash + Clone, + Right: Eq + Hash + Clone, +{ + /// Construct a new empty `BiHashMap` + pub fn new() -> Self { + Self::default() + } + + /// Insert a new mapping into the `BiHashMap` + pub fn insert(&mut self, left: Left, right: Right) { + self.left_to_right.insert(left.clone(), right.clone()); + self.right_to_left.insert(right, left); + } + + /// Get the right value for a given left value + pub fn get_right(&self, left: &Left) -> Option<&Right> { + self.left_to_right.get(left) + } + + /// Get the left value for a given right value + pub fn get_left(&self, right: &Right) -> Option<&Left> { + self.right_to_left.get(right) + } + + /// Remove a mapping from the `BiHashMap` + pub fn remove(&mut self, left: &Left, right: &Right) { + self.left_to_right.remove(left); + self.right_to_left.remove(right); + } + + /// Remove a mapping from the `BiHashMap` by left value + #[allow(dead_code)] + pub fn remove_left(&mut self, left: &Left) { + if let Some(right) = self.left_to_right.remove(left) { + self.right_to_left.remove(&right); + } + } + + /// Remove a mapping from the `BiHashMap` by right value + #[allow(dead_code)] + pub fn remove_right(&mut self, right: &Right) { + if let Some(left) = self.right_to_left.remove(right) { + self.left_to_right.remove(&left); + } + } +} + +impl Default for BiHashMap { + fn default() -> Self { + Self { + left_to_right: FxHashMap::default(), + right_to_left: FxHashMap::default(), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_insert() { + let mut bimap = BiHashMap::new(); + bimap.insert(1, "one"); + bimap.insert(2, "two"); + bimap.insert(3, "three"); + + assert_eq!(bimap.get_right(&1), Some(&"one")); + assert_eq!(bimap.get_right(&2), Some(&"two")); + assert_eq!(bimap.get_right(&3), Some(&"three")); + + assert_eq!(bimap.get_left(&"one"), Some(&1)); + assert_eq!(bimap.get_left(&"two"), Some(&2)); + assert_eq!(bimap.get_left(&"three"), Some(&3)); + } + + #[test] + fn test_remove() { + let mut bimap = BiHashMap::new(); + bimap.insert(1, "one"); + bimap.insert(2, "two"); + bimap.insert(3, "three"); + + bimap.remove(&1, &"one"); + assert_eq!(bimap.get_right(&1), None); + assert_eq!(bimap.get_left(&"one"), None); + + bimap.remove_left(&2); + assert_eq!(bimap.get_right(&2), None); + assert_eq!(bimap.get_left(&"two"), None); + + bimap.remove_right(&"three"); + assert_eq!(bimap.get_right(&3), None); + assert_eq!(bimap.get_left(&"three"), None); + } +} \ No newline at end of file diff --git a/libs/fast-nat/src/cpnat.rs b/libs/fast-nat/src/cpnat.rs new file mode 100644 index 0000000..7d67f4b --- /dev/null +++ b/libs/fast-nat/src/cpnat.rs @@ -0,0 +1,92 @@ +use std::time::Duration; + +use rustc_hash::FxHashMap; + +use crate::{bimap::BiHashMap, timeout::MaybeTimeout}; + +/// A table of network address mappings across IPv4 and IPv6 +#[derive(Debug)] +pub struct CrossProtocolNetworkAddressTable { + /// Internal address map + addr_map: BiHashMap, + /// Secondary map used to keep track of timeouts + timeouts: FxHashMap<(u32, u128), MaybeTimeout>, +} + +impl CrossProtocolNetworkAddressTable { + /// Construct a new empty `CrossProtocolNetworkAddressTable` + pub fn new() -> Self { + Self { + addr_map: BiHashMap::new(), + timeouts: FxHashMap::default(), + } + } + + /// Prune all old mappings + pub fn prune(&mut self) { + log::trace!("Pruning old network address mappings"); + + // Compare all mappings against a common timestamp + let now = std::time::Instant::now(); + + // Remove all old mappings from both the bimap and the timeouts map + self.timeouts.retain(|(left, right), timeout| { + match timeout { + // Retain all indefinite mappings + MaybeTimeout::Never => true, + // Only retain mappings that haven't timed out yet + MaybeTimeout::After { duration, start } => { + let should_retain = now.duration_since(*start) < *duration; + if !should_retain { + log::trace!( + "Mapping {:?} -> {:?} has timed out and will be removed", + left, + right + ); + self.addr_map.remove(left, right); + } + should_retain + } + } + }); + } + + /// Insert a new indefinite mapping + pub fn insert_indefinite, T6: Into>(&mut self, ipv4: T4, ipv6: T6) { + self.prune(); + let (ipv4, ipv6) = (ipv4.into(), ipv6.into()); + self.addr_map.insert(ipv4, ipv6); + self.timeouts.insert((ipv4, ipv6), MaybeTimeout::Never); + } + + /// Insert a new mapping with a finite time-to-live + pub fn insert, T6: Into>( + &mut self, + ipv4: T4, + ipv6: T6, + duration: Duration, + ) { + self.prune(); + let (ipv4, ipv6) = (ipv4.into(), ipv6.into()); + self.addr_map.insert(ipv4, ipv6); + self.timeouts.insert( + (ipv4, ipv6), + MaybeTimeout::After { + duration, + start: std::time::Instant::now(), + }, + ); + } + + /// Get the IPv6 address for a given IPv4 address + #[must_use] + pub fn get_ipv6>(&self, ipv4: T) -> Option { + self.addr_map.get_right(&ipv4.into()).copied() + } + + /// Get the IPv4 address for a given IPv6 address + #[must_use] + pub fn get_ipv4>(&self, ipv6: T) -> Option { + self.addr_map.get_left(&ipv6.into()).copied() + } +} diff --git a/libs/fast-nat/src/lib.rs b/libs/fast-nat/src/lib.rs new file mode 100644 index 0000000..e0bc1f4 --- /dev/null +++ b/libs/fast-nat/src/lib.rs @@ -0,0 +1,9 @@ +#![doc = include_str!("../README.md")] + +mod bimap; +mod cpnat; +mod nat; +mod timeout; + +pub use cpnat::CrossProtocolNetworkAddressTable; +pub use nat::NetworkAddressTable; diff --git a/libs/fast-nat/src/nat.rs b/libs/fast-nat/src/nat.rs new file mode 100644 index 0000000..d6bc175 --- /dev/null +++ b/libs/fast-nat/src/nat.rs @@ -0,0 +1,87 @@ +use std::time::Duration; + +use rustc_hash::FxHashMap; + +use crate::{bimap::BiHashMap, timeout::MaybeTimeout}; + +/// A table of network address mappings +#[derive(Debug)] +pub struct NetworkAddressTable { + /// Internal address map + addr_map: BiHashMap, + /// Secondary map used to keep track of timeouts + timeouts: FxHashMap<(u32, u32), MaybeTimeout>, +} + +impl NetworkAddressTable { + /// Construct a new empty `NetworkAddressTable` + pub fn new() -> Self { + Self { + addr_map: BiHashMap::new(), + timeouts: FxHashMap::default(), + } + } + + /// Prune all old mappings + pub fn prune(&mut self) { + log::trace!("Pruning old network address mappings"); + + // Compare all mappings against a common timestamp + let now = std::time::Instant::now(); + + // Remove all old mappings from both the bimap and the timeouts map + self.timeouts.retain(|(left, right), timeout| { + match timeout { + // Retain all indefinite mappings + MaybeTimeout::Never => true, + // Only retain mappings that haven't timed out yet + MaybeTimeout::After { duration, start } => { + let should_retain = now.duration_since(*start) < *duration; + if !should_retain { + log::trace!( + "Mapping {:?} -> {:?} has timed out and will be removed", + left, + right + ); + self.addr_map.remove(left, right); + } + should_retain + } + } + }); + } + + /// Insert a new indefinite mapping + pub fn insert_indefinite>(&mut self, left: T, right: T) { + self.prune(); + let (left, right) = (left.into(), right.into()); + self.addr_map.insert(left, right); + self.timeouts.insert((left, right), MaybeTimeout::Never); + } + + /// Insert a new mapping with a finite time-to-live + pub fn insert>(&mut self, left: T, right: T, duration: Duration) { + self.prune(); + let (left, right) = (left.into(), right.into()); + self.addr_map.insert(left, right); + self.timeouts.insert( + (left, right), + MaybeTimeout::After { + duration, + start: std::time::Instant::now(), + }, + ); + } + + /// Get the right value for a given left value + #[must_use] + pub fn get_right>(&self, left: T) -> Option { + self.addr_map.get_right(&left.into()).copied() + } + + /// Get the left value for a given right value + #[must_use] + pub fn get_left>(&self, right: T) -> Option { + self.addr_map.get_left(&right.into()).copied() + } +} diff --git a/libs/fast-nat/src/timeout.rs b/libs/fast-nat/src/timeout.rs new file mode 100644 index 0000000..9691f6f --- /dev/null +++ b/libs/fast-nat/src/timeout.rs @@ -0,0 +1,13 @@ +use std::time::Duration; + +/// Describes a possible timeout for a mapping +#[derive(Debug, Clone, Copy)] +pub enum MaybeTimeout { + /// Indicates that a mapping should never time out + Never, + /// Indicates that a mapping should time out after a given duration + After { + duration: Duration, + start: std::time::Instant, + }, +} \ No newline at end of file