Add experimental profiling support

Cargo.toml

@@ -12,10 +12,13 @@ license = "GPL-3.0"
 keywords = []
 categories = []
 
-
+[features]
+default = []
+profiler = ["profiling/profile-with-puffin", "puffin_http", "puffin"]
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 [dependencies]
+protomask-tun = { path = "protomask-tun", version = "0.1.0" }
 tokio = { version = "1.29.1", features = [
     "macros",
     "rt-multi-thread",
@@ -25,6 +28,9 @@ tokio = { version = "1.29.1", features = [
 clap = { version = "4.3.11", features = ["derive"] }
 serde = { version = "1.0.171", features = ["derive"] }
 ipnet = { version = "2.8.0", features = ["serde"] }
+puffin_http = { version = "^0.13.0", optional = true }
+puffin = { version = "0.12.1", optional = true }
+profiling = "1.0.8"
 toml = "0.7.6"
 log = "0.4.19"
 fern = "0.6.2"
@@ -36,7 +42,7 @@ bimap = "0.6.3"
 pnet_packet = "0.33.0"
 rtnetlink = "0.13.0"
 futures = "0.3.28"
-protomask-tun = { path = "protomask-tun", version = "0.1.0" }
+cfg-if = "1.0.0"
 
 [[bin]]
 name = "protomask"

src/cli/main.rs

@@ -1,5 +1,6 @@
 //! This is the entrypoint for `protomask` from the command line.
 
+use cfg_if::cfg_if;
 use clap::Parser;
 use config::Config;
 use logging::enable_logger;
@@ -26,6 +27,16 @@ pub async fn main() {
         std::process::exit(1);
     }
 
+    // Enable the profiler if enabled by build flags
+    cfg_if! {
+        if #[cfg(feature = "profiler")] {
+            let puffin_listen_addr = format!("[::]:{}", puffin_http::DEFAULT_PORT);
+            log::info!("Puffin HTTP server listening on: {}", puffin_listen_addr);
+            let _puffin_server = puffin_http::Server::new(&puffin_listen_addr).unwrap();
+            puffin::set_scopes_on(true);
+        }
+    }
+
     // Create the NAT64 instance
     let mut nat64 = Nat64::new(
         config.nat64_prefix,

src/nat/error.rs

@@ -0,0 +1,17 @@
+use super::table;
+
+#[derive(Debug, thiserror::Error)]
+pub enum Nat64Error {
+    #[error(transparent)]
+    TableError(#[from] table::TableError),
+    #[error(transparent)]
+    TunError(#[from] protomask_tun::Error),
+    #[error(transparent)]
+    IoError(#[from] std::io::Error),
+    #[error(transparent)]
+    PacketHandlingError(#[from] crate::packet::error::PacketError),
+    #[error(transparent)]
+    PacketReceiveError(#[from] tokio::sync::broadcast::error::RecvError),
+    #[error(transparent)]
+    PacketSendError(#[from] tokio::sync::mpsc::error::SendError<Vec<u8>>),
+}

src/nat/mod.rs

@@ -4,6 +4,7 @@ use crate::packet::{
 };
 
 use self::{
+    error::Nat64Error,
     table::Nat64Table,
     utils::{embed_address, extract_address},
 };
@@ -15,27 +16,10 @@ use std::{
 };
 use tokio::sync::{broadcast, mpsc};
 
+mod error;
 mod table;
 mod utils;
 
-#[derive(Debug, thiserror::Error)]
-pub enum Nat64Error {
-    #[error(transparent)]
-    TableError(#[from] table::TableError),
-    #[error(transparent)]
-    TunError(#[from] protomask_tun::Error),
-    #[error(transparent)]
-    IoError(#[from] std::io::Error),
-    // #[error(transparent)]
-    // XlatError(#[from] xlat::PacketTranslationError),
-    #[error(transparent)]
-    PacketHandlingError(#[from] crate::packet::error::PacketError),
-    #[error(transparent)]
-    PacketReceiveError(#[from] broadcast::error::RecvError),
-    #[error(transparent)]
-    PacketSendError(#[from] mpsc::error::SendError<Vec<u8>>),
-}
-
 pub struct Nat64 {
     table: Nat64Table,
     interface: TunDevice,
@@ -81,6 +65,9 @@ impl Nat64 {
 
         // Process packets in a loop
         loop {
+            // Start a new profiler frame
+            profiling::finish_frame!();
+
             // Try to read a packet
             match rx.recv().await {
                 Ok(packet) => {
@@ -90,6 +77,8 @@ impl Nat64 {
                     // Separate logic is needed for handling IPv4 vs IPv6 packets, so a check must be done here
                     match packet[0] >> 4 {
                         4 => {
+                            profiling::scope!("Ipv4 to IPv6");
+
                             // Parse the packet
                             let packet: Ipv4Packet<Vec<u8>> = packet.try_into()?;
 
@@ -113,14 +102,37 @@ impl Nat64 {
                             });
                         }
                         6 => {
+                            profiling::scope!("Ipv6 to IPv4");
+
                             // Parse the packet
                             let packet: Ipv6Packet<Vec<u8>> = packet.try_into()?;
 
+                            // Drop packets "coming from" the NAT64 prefix
+                            if self.ipv6_nat_prefix.contains(&packet.source_address) {
+                                log::warn!(
+                                    "Dropping packet \"from\" NAT64 prefix: {} -> {}",
+                                    packet.source_address,
+                                    packet.destination_address
+                                );
+                                continue;
+                            }
+
                             // Get the new source and dest addresses
                             let new_source =
                                 self.table.get_or_assign_ipv4(packet.source_address)?;
                             let new_destination = extract_address(packet.destination_address);
 
+                            // Drop packets destined for private IPv4 addresses
+                            if new_destination.is_private() {
+                                log::warn!(
+                                    "Dropping packet destined for private IPv4 address: {} -> {} ({})",
+                                    packet.source_address,
+                                    packet.destination_address,
+                                    new_destination
+                                );
+                                continue;
+                            }
+
                             // Spawn a task to process the packet
                             tokio::spawn(async move {
                                 let output =

src/nat/table.rs

@@ -47,6 +47,7 @@ impl Nat64Table {
     }
 
     /// Make a reservation for an IP address pair for eternity
+    #[profiling::function]
     pub fn add_infinite_reservation(
         &mut self,
         ipv6: Ipv6Addr,
@@ -68,6 +69,7 @@ impl Nat64Table {
     }
 
     /// Check if a given address exists in the table
+    #[profiling::function]
     pub fn contains(&self, address: &IpAddr) -> bool {
         match address {
             IpAddr::V4(ipv4) => self.reservations.contains_right(ipv4),
@@ -76,6 +78,7 @@ impl Nat64Table {
     }
 
     /// Get or assign an IPv4 address for the given IPv6 address
+    #[profiling::function]
     pub fn get_or_assign_ipv4(&mut self, ipv6: Ipv6Addr) -> Result<Ipv4Addr, TableError> {
         // Prune old reservations
         self.prune();
@@ -110,6 +113,7 @@ impl Nat64Table {
     }
 
     /// Try to find an IPv6 address for the given IPv4 address
+    #[profiling::function]
     pub fn get_reverse(&mut self, ipv4: Ipv4Addr) -> Result<Ipv6Addr, TableError> {
         // Prune old reservations
         self.prune();
@@ -129,11 +133,13 @@ impl Nat64Table {
     }
 
     /// Check if an address is within the IPv4 pool
+    #[profiling::function]
     pub fn is_address_within_pool(&self, address: &Ipv4Addr) -> bool {
         self.ipv4_pool.iter().any(|net| net.contains(address))
     }
 
     /// Calculate the translated version of any address
+    #[profiling::function]
     pub fn calculate_xlat_addr(
         &mut self,
         input: &IpAddr,
@@ -186,6 +192,7 @@ impl Nat64Table {
 
 impl Nat64Table {
     /// Prune old reservations
+    #[profiling::function]
     pub fn prune(&mut self) {
         let now = Instant::now();
 

src/packet/protocols/icmp.rs

@@ -29,6 +29,7 @@ where
 {
     type Error = PacketError;
 
+    #[profiling::function]
     fn try_from(bytes: Vec<u8>) -> Result<Self, Self::Error> {
         // Parse the packet
         let packet =
@@ -43,11 +44,11 @@ where
     }
 }
 
-
 impl<T> Into<Vec<u8>> for IcmpPacket<T>
 where
     T: Into<Vec<u8>>,
 {
+    #[profiling::function]
     fn into(self) -> Vec<u8> {
         // Convert the payload into raw bytes
         let payload: Vec<u8> = self.payload.into();

src/packet/protocols/icmpv6.rs

@@ -42,6 +42,7 @@ where
     T: From<Vec<u8>>,
 {
     /// Construct a new ICMPv6 packet from raw bytes
+    #[profiling::function]
     pub fn new_from_bytes(
         bytes: &[u8],
         source_address: Ipv6Addr,
@@ -64,6 +65,7 @@ where
 
 impl Icmpv6Packet<RawBytes> {
     /// Construct a new ICMPv6 packet with a raw payload from raw bytes
+    #[profiling::function]
     pub fn new_from_bytes_raw_payload(
         bytes: &[u8],
         source_address: Ipv6Addr,
@@ -88,6 +90,7 @@ impl<T> Into<Vec<u8>> for Icmpv6Packet<T>
 where
     T: Into<Vec<u8>>,
 {
+    #[profiling::function]
     fn into(self) -> Vec<u8> {
         // Convert the payload into raw bytes
         let payload: Vec<u8> = self.payload.into();

src/packet/protocols/ipv4.rs

@@ -53,6 +53,7 @@ impl<T> Ipv4Packet<T> {
         }
     }
 
+    #[profiling::function]
     fn options_length_words(&self) -> u8 {
         self.options
             .iter()
@@ -68,6 +69,7 @@ where
 {
     type Error = PacketError;
 
+    #[profiling::function]
     fn try_from(bytes: Vec<u8>) -> Result<Self, Self::Error> {
         // Parse the packet
         let packet =
@@ -94,6 +96,7 @@ impl<T> Into<Vec<u8>> for Ipv4Packet<T>
 where
     T: Into<Vec<u8>> + Clone,
 {
+    #[profiling::function]
     fn into(self) -> Vec<u8> {
         // Convert the payload into raw bytes
         let payload: Vec<u8> = self.payload.clone().into();

src/packet/protocols/ipv6.rs

@@ -44,6 +44,7 @@ where
 {
     type Error = PacketError;
 
+    #[profiling::function]
     fn try_from(bytes: Vec<u8>) -> Result<Self, Self::Error> {
         // Parse the packet
         let packet =
@@ -66,6 +67,7 @@ impl<T> Into<Vec<u8>> for Ipv6Packet<T>
 where
     T: Into<Vec<u8>>,
 {
+    #[profiling::function]
     fn into(self) -> Vec<u8> {
         // Convert the payload into raw bytes
         let payload: Vec<u8> = self.payload.into();

src/packet/protocols/tcp.rs

@@ -113,6 +113,7 @@ where
     T: From<Vec<u8>>,
 {
     /// Construct a new TCP packet from bytes
+    #[profiling::function]
     pub fn new_from_bytes(
         bytes: &[u8],
         source_address: IpAddr,
@@ -147,6 +148,7 @@ where
 
 impl TcpPacket<RawBytes> {
     /// Construct a new TCP packet with a raw payload from bytes
+    #[profiling::function]
     pub fn new_from_bytes_raw_payload(
         bytes: &[u8],
         source_address: IpAddr,
@@ -183,6 +185,7 @@ impl<T> Into<Vec<u8>> for TcpPacket<T>
 where
     T: Into<Vec<u8>>,
 {
+    #[profiling::function]
     fn into(self) -> Vec<u8> {
         // Get the options length in words
         let options_length = self.options_length();

src/packet/protocols/udp.rs

@@ -84,6 +84,7 @@ where
     T: From<Vec<u8>>,
 {
     /// Construct a new UDP packet from bytes
+    #[profiling::function]
     pub fn new_from_bytes(
         bytes: &[u8],
         source_address: IpAddr,
@@ -112,6 +113,7 @@ where
 
 impl UdpPacket<RawBytes> {
     /// Construct a new UDP packet with a raw payload from bytes
+    #[profiling::function]
     pub fn new_from_bytes_raw_payload(
         bytes: &[u8],
         source_address: IpAddr,
@@ -142,6 +144,7 @@ impl<T> Into<Vec<u8>> for UdpPacket<T>
 where
     T: Into<Vec<u8>>,
 {
+    #[profiling::function]
     fn into(self) -> Vec<u8> {
         // Convert the payload into raw bytes
         let payload: Vec<u8> = self.payload.into();

src/packet/xlat/icmp/mod.rs

@@ -12,6 +12,7 @@ use super::ip::{translate_ipv4_to_ipv6, translate_ipv6_to_ipv4};
 mod type_code;
 
 /// Translates an ICMP packet to an ICMPv6 packet
+#[profiling::function]
 pub fn translate_icmp_to_icmpv6(
     input: IcmpPacket<RawBytes>,
     new_source: Ipv6Addr,
@@ -55,6 +56,7 @@ pub fn translate_icmp_to_icmpv6(
 }
 
 /// Translates an ICMPv6 packet to an ICMP packet
+#[profiling::function]
 pub fn translate_icmpv6_to_icmp(
     input: Icmpv6Packet<RawBytes>,
     new_source: Ipv4Addr,
@@ -83,7 +85,7 @@ pub fn translate_icmpv6_to_icmp(
                 buffer.extend_from_slice(&inner_payload);
                 buffer
             })
-        },
+        }
         _ => input.payload,
     };
 

src/packet/xlat/icmp/type_code.rs

@@ -8,6 +8,7 @@ use pnet_packet::{
 use crate::packet::error::PacketError;
 
 /// Best effort translation from an ICMP type and code to an ICMPv6 type and code
+#[profiling::function]
 pub fn translate_type_and_code_4_to_6(
     icmp_type: IcmpType,
     icmp_code: IcmpCode,
@@ -55,6 +56,7 @@ pub fn translate_type_and_code_4_to_6(
 }
 
 /// Best effort translation from an ICMPv6 type and code to an ICMP type and code
+#[profiling::function]
 pub fn translate_type_and_code_6_to_4(
     icmp_type: Icmpv6Type,
     icmp_code: Icmpv6Code,

src/packet/xlat/ip.rs

@@ -17,6 +17,7 @@ use super::{
 };
 
 /// Translates an IPv4 packet to an IPv6 packet
+#[profiling::function]
 pub fn translate_ipv4_to_ipv6(
     input: Ipv4Packet<Vec<u8>>,
     new_source: Ipv6Addr,
@@ -69,6 +70,7 @@ pub fn translate_ipv4_to_ipv6(
 }
 
 /// Translates an IPv6 packet to an IPv4 packet
+#[profiling::function]
 pub fn translate_ipv6_to_ipv4(
     input: Ipv6Packet<Vec<u8>>,
     new_source: Ipv4Addr,

src/packet/xlat/tcp.rs

@@ -6,6 +6,7 @@ use crate::packet::{
 };
 
 /// Translates an IPv4 TCP packet to an IPv6 TCP packet
+#[profiling::function]
 pub fn translate_tcp4_to_tcp6(
     input: TcpPacket<RawBytes>,
     new_source_addr: Ipv6Addr,
@@ -26,6 +27,7 @@ pub fn translate_tcp4_to_tcp6(
 }
 
 /// Translates an IPv6 TCP packet to an IPv4 TCP packet
+#[profiling::function]
 pub fn translate_tcp6_to_tcp4(
     input: TcpPacket<RawBytes>,
     new_source_addr: Ipv4Addr,

src/packet/xlat/udp.rs

@@ -6,6 +6,7 @@ use crate::packet::{
 };
 
 /// Translates an IPv4 UDP packet to an IPv6 UDP packet
+#[profiling::function]
 pub fn translate_udp4_to_udp6(
     input: UdpPacket<RawBytes>,
     new_source_addr: Ipv6Addr,
@@ -20,6 +21,7 @@ pub fn translate_udp4_to_udp6(
 }
 
 /// Translates an IPv6 UDP packet to an IPv4 UDP packet
+#[profiling::function]
 pub fn translate_udp6_to_udp4(
     input: UdpPacket<RawBytes>,
     new_source_addr: Ipv4Addr,