Implement optional profiling

Cargo.toml

@@ -13,33 +13,6 @@ keywords = []
 categories = []
 exclude = ["/.github/", "/.vscode/"]
 
-# # 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",
-#     # "process",
-#     "sync"
-# ] }
-# clap = { version = "4.3.11", features = ["derive"] }
-# serde = { version = "1.0.171", features = ["derive"] }
-# ipnet = { version = "2.8.0", features = ["serde"] }
-# hyper = { version = "0.14.27", features = ["server", "http1", "tcp"] }
-# owo-colors = { version = "3.5.0", features = ["supports-colors"] }
-# toml = "0.7.6"
-# log = "0.4.19"
-# fern = "0.6.2"
-# serde_path_to_error = "0.1.13"
-# thiserror = "1.0.43"
-# tun-tap = "0.1.3"
-# bimap = "0.6.3"
-# pnet_packet = "0.34.0"
-# rtnetlink = "0.13.0"
-# futures = "0.3.28"
-# prometheus = "0.13.3"
-# lazy_static = "1.4.0"
-
 [workspace]
 members = [
     "libs/easy-tun",
@@ -50,6 +23,16 @@ members = [
     "libs/protomask-metrics",
 ]
 
+[features]
+default = []
+profiler = [
+    "puffin",
+    "puffin_http",
+    "easy-tun/profile-puffin",
+    "fast-nat/profile-puffin",
+    "interproto/profile-puffin",
+]
+
 [[bin]]
 name = "protomask"
 path = "src/protomask.rs"
@@ -76,12 +59,16 @@ tokio = { version = "1.29.1", features = ["macros", "rt-multi-thread"] }
 owo-colors = { version = "3.5.0", features = ["supports-colors"] }
 clap = { version = "4.3.11", features = ["derive"] }
 ipnet = { version = "2.8.0", features = ["serde"] }
+puffin_http = { version = "0.13.0", optional = true }
+puffin = { version = "0.16.0", optional = true }
 serde = { version = "^1.0", features = ["derive"] }
 serde_json = "^1.0"
 log = "0.4.19"
 fern = "0.6.2"
 nix = "0.26.2"
 thiserror = "1.0.44"
+cfg-if = "1.0.0"
+profiling = "1.0.9"
 
 [package.metadata.deb]
 section = "network"

libs/easy-tun/Cargo.toml

@@ -12,6 +12,10 @@ license = "GPL-3.0"
 keywords = []
 categories = []
 
+[features]
+default = []
+profile-puffin = ["profiling/profile-with-puffin"]
+
 [dependencies]
 log = "^0.4"
 libc = "^0.2"

libs/easy-tun/src/tun.rs

@@ -96,6 +96,12 @@ impl Tun {
     pub fn name(&self) -> &str {
         &self.name
     }
+
+    /// Get the underlying file descriptor
+    #[must_use]
+    pub fn fd(&self) -> &File {
+        &self.fd
+    }
 }
 
 impl AsRawFd for Tun {

libs/fast-nat/Cargo.toml

@@ -12,6 +12,10 @@ license = "GPL-3.0"
 keywords = []
 categories = []
 
+[features]
+default = []
+profile-puffin = ["profiling/profile-with-puffin"]
+
 [dependencies]
 log = "^0.4"
 rustc-hash = "1.1.0"

libs/interproto/Cargo.toml

@@ -15,6 +15,7 @@ categories = []
 [features]
 default = []
 metrics = ["protomask-metrics"]
+profile-puffin = ["profiling/profile-with-puffin"]
 
 [dependencies]
 protomask-metrics = { path = "../protomask-metrics", optional = true }

src/args/mod.rs

@@ -1,4 +1,22 @@
 //! This module contains the definitions for each binary's CLI arguments and config file structure for the sake of readability.
 
+use cfg_if::cfg_if;
+
 pub mod protomask;
 pub mod protomask_clat;
+
+
+// Used to trick the build process into including a CLI argument based on a feature flag
+cfg_if! {
+    if #[cfg(feature = "profiler")] {
+        #[derive(Debug, clap::Args)]
+        pub struct ProfilerArgs {
+            /// Expose the puffin HTTP server on this endpoint
+            #[clap(long)]
+            pub puffin_endpoint: Option<std::net::SocketAddr>,
+        }
+    } else {
+        #[derive(Debug, clap::Args)]
+        pub struct ProfilerArgs;
+    }
+}

src/args/protomask.rs

@@ -7,6 +7,8 @@ use ipnet::{Ipv4Net, Ipv6Net};
 
 use crate::common::rfc6052::parse_network_specific_prefix;
 
+use super::ProfilerArgs;
+
 #[derive(clap::Parser)]
 #[clap(author, version, about="Fast and simple NAT64", long_about = None)]
 pub struct Args {
@@ -21,6 +23,9 @@ pub struct Args {
     #[clap(short, long, default_value_t = ("nat%d").to_string())]
     pub interface: String,
 
+    #[command(flatten)]
+    pub profiler_args: ProfilerArgs,
+
     /// Enable verbose logging
     #[clap(short, long)]
     pub verbose: bool,
@@ -102,3 +107,4 @@ impl From<StaticMap> for (Ipv4Addr, Ipv6Addr) {
         (val.ipv4, val.ipv6)
     }
 }
+

src/args/protomask_clat.rs

@@ -1,5 +1,6 @@
 //! Commandline arguments and config file definitions for `protomask-clat`
 
+use super::ProfilerArgs;
 use crate::common::rfc6052::parse_network_specific_prefix;
 use ipnet::{Ipv4Net, Ipv6Net};
 use std::{net::SocketAddr, path::PathBuf};
@@ -18,6 +19,9 @@ pub struct Args {
     #[clap(short, long, default_value_t = ("clat%d").to_string())]
     pub interface: String,
 
+    #[command(flatten)]
+    pub profiler_args: ProfilerArgs,
+
     /// Enable verbose logging
     #[clap(short, long)]
     pub verbose: bool,

src/common/mod.rs

@@ -4,3 +4,4 @@ pub mod logging;
 pub mod packet_handler;
 pub mod permissions;
 pub mod rfc6052;
+pub mod profiler;

src/common/packet_handler.rs

@@ -8,16 +8,9 @@ pub enum PacketHandlingError {
     FastNatError(#[from] fast_nat::error::Error),
 }
 
-/// Handles checking the version number of an IP packet and calling the correct handler with needed data
+
-pub fn handle_packet<Ipv4Handler, Ipv6Handler>(
+/// Get the layer 3 protocol of a packet
-    packet: &[u8],
+pub fn get_layer_3_proto(packet: &[u8]) -> Option<u8> {
-    mut ipv4_handler: Ipv4Handler,
-    mut ipv6_handler: Ipv6Handler,
-) -> Option<Vec<u8>>
-where
-    Ipv4Handler: FnMut(&[u8], &Ipv4Addr, &Ipv4Addr) -> Result<Option<Vec<u8>>, PacketHandlingError>,
-    Ipv6Handler: FnMut(&[u8], &Ipv6Addr, &Ipv6Addr) -> Result<Option<Vec<u8>>, PacketHandlingError>,
-{
     // If the packet is empty, return nothing
     if packet.is_empty() {
         return None;
@@ -26,40 +19,31 @@ where
     // Switch on the layer 3 protocol number to call the correct handler
     let layer_3_proto = packet[0] >> 4;
     log::trace!("New packet with layer 3 protocol: {}", layer_3_proto);
-    let handler_response = match layer_3_proto {
+    Some(layer_3_proto)
-        // IPv4
+}
-        4 => {
-            // Extract the source and destination addresses
-            let source_addr =
-                Ipv4Addr::from(u32::from_be_bytes(packet[12..16].try_into().unwrap()));
-            let destination_addr =
-                Ipv4Addr::from(u32::from_be_bytes(packet[16..20].try_into().unwrap()));
 
-            // Call the handler
+/// Get the source and destination addresses of an IPv4 packet
-            ipv4_handler(packet, &source_addr, &destination_addr)
+pub fn get_ipv4_src_dst(packet: &[u8]) -> (Ipv4Addr, Ipv4Addr) {
-        }
+    let source_addr = Ipv4Addr::from(u32::from_be_bytes(packet[12..16].try_into().unwrap()));
+    let destination_addr = Ipv4Addr::from(u32::from_be_bytes(packet[16..20].try_into().unwrap()));
 
-        // IPv6
+    (source_addr, destination_addr)
-        6 => {
+}
-            // Extract the source and destination addresses
-            let source_addr =
-                Ipv6Addr::from(u128::from_be_bytes(packet[8..24].try_into().unwrap()));
-            let destination_addr =
-                Ipv6Addr::from(u128::from_be_bytes(packet[24..40].try_into().unwrap()));
 
-            // Call the handler
+/// Get the source and destination addresses of an IPv6 packet
-            ipv6_handler(packet, &source_addr, &destination_addr)
+pub fn get_ipv6_src_dst(packet: &[u8]) -> (Ipv6Addr, Ipv6Addr) {
-        }
+    let source_addr = Ipv6Addr::from(u128::from_be_bytes(packet[8..24].try_into().unwrap()));
+    let destination_addr = Ipv6Addr::from(u128::from_be_bytes(packet[24..40].try_into().unwrap()));
 
-        // Unknown protocol numbers can't be handled
+    (source_addr, destination_addr)
-        proto => {
+}
-            log::warn!("Unknown Layer 3 protocol: {}", proto);
-            return None;
-        }
-    };
 
-    // The response from the handler may or may not be a warn-able error
+/// Appropriately handle a translation error
-    match handler_response {
+pub fn handle_translation_error(
+    result: Result<Option<Vec<u8>>, PacketHandlingError>,
+) -> Option<Vec<u8>> {
+    // We may or may not have a warn-able error
+    match result {
         // If we get data, return it
         Ok(data) => data,
         // If we get an error, handle it and return None
@@ -101,3 +85,97 @@ where
         },
     }
 }
+
+// /// Handles checking the version number of an IP packet and calling the correct handler with needed data
+// pub fn handle_packet<Ipv4Handler, Ipv6Handler>(
+//     packet: &[u8],
+//     mut ipv4_handler: Ipv4Handler,
+//     mut ipv6_handler: Ipv6Handler,
+// ) -> Option<Vec<u8>>
+// where
+//     Ipv4Handler: FnMut(&[u8], &Ipv4Addr, &Ipv4Addr) -> Result<Option<Vec<u8>>, PacketHandlingError>,
+//     Ipv6Handler: FnMut(&[u8], &Ipv6Addr, &Ipv6Addr) -> Result<Option<Vec<u8>>, PacketHandlingError>,
+// {
+//     // If the packet is empty, return nothing
+//     if packet.is_empty() {
+//         return None;
+//     }
+
+//     // Switch on the layer 3 protocol number to call the correct handler
+//     let layer_3_proto = packet[0] >> 4;
+//     log::trace!("New packet with layer 3 protocol: {}", layer_3_proto);
+//     let handler_response = match layer_3_proto {
+//         // IPv4
+//         4 => {
+//             // Extract the source and destination addresses
+//             let source_addr =
+//                 Ipv4Addr::from(u32::from_be_bytes(packet[12..16].try_into().unwrap()));
+//             let destination_addr =
+//                 Ipv4Addr::from(u32::from_be_bytes(packet[16..20].try_into().unwrap()));
+
+//             // Call the handler
+//             ipv4_handler(packet, &source_addr, &destination_addr)
+//         }
+
+//         // IPv6
+//         6 => {
+//             // Extract the source and destination addresses
+//             let source_addr =
+//                 Ipv6Addr::from(u128::from_be_bytes(packet[8..24].try_into().unwrap()));
+//             let destination_addr =
+//                 Ipv6Addr::from(u128::from_be_bytes(packet[24..40].try_into().unwrap()));
+
+//             // Call the handler
+//             ipv6_handler(packet, &source_addr, &destination_addr)
+//         }
+
+//         // Unknown protocol numbers can't be handled
+//         proto => {
+//             log::warn!("Unknown Layer 3 protocol: {}", proto);
+//             return None;
+//         }
+//     };
+
+//     // The response from the handler may or may not be a warn-able error
+//     match handler_response {
+//         // If we get data, return it
+//         Ok(data) => data,
+//         // If we get an error, handle it and return None
+//         Err(error) => match error {
+//             PacketHandlingError::InterprotoError(interproto::error::Error::PacketTooShort {
+//                 expected,
+//                 actual,
+//             }) => {
+//                 log::warn!(
+//                     "Got packet with length {} when expecting at least {} bytes",
+//                     actual,
+//                     expected
+//                 );
+//                 None
+//             }
+//             PacketHandlingError::InterprotoError(
+//                 interproto::error::Error::UnsupportedIcmpType(icmp_type),
+//             ) => {
+//                 log::warn!("Got a packet with an unsupported ICMP type: {}", icmp_type);
+//                 None
+//             }
+//             PacketHandlingError::InterprotoError(
+//                 interproto::error::Error::UnsupportedIcmpv6Type(icmpv6_type),
+//             ) => {
+//                 log::warn!(
+//                     "Got a packet with an unsupported ICMPv6 type: {}",
+//                     icmpv6_type
+//                 );
+//                 None
+//             }
+//             PacketHandlingError::FastNatError(fast_nat::error::Error::Ipv4PoolExhausted) => {
+//                 log::warn!("IPv4 pool exhausted. Dropping packet.");
+//                 None
+//             }
+//             PacketHandlingError::FastNatError(fast_nat::error::Error::InvalidIpv4Address(addr)) => {
+//                 log::warn!("Invalid IPv4 address: {}", addr);
+//                 None
+//             }
+//         },
+//     }
+// }

src/common/profiler.rs

@@ -0,0 +1,20 @@
+use cfg_if::cfg_if;
+
+use crate::args::ProfilerArgs;
+
+cfg_if! {
+    if #[cfg(feature = "profiler")] {
+        pub fn start_puffin_server(args: &ProfilerArgs) -> Option<puffin_http::Server> {
+            if let Some(endpoint) = args.puffin_endpoint {
+                log::info!("Starting puffin server on {}", endpoint);
+                puffin::set_scopes_on(true);
+                Some(puffin_http::Server::new(&endpoint.to_string()).unwrap())
+            } else {
+                None
+            }
+        }
+    } else {
+        #[allow(dead_code)]
+        pub fn start_puffin_server(_args: &ProfilerArgs){}
+    }
+}

src/protomask-clat.rs

@@ -3,7 +3,11 @@
 //! This binary is a Customer-side transLATor (CLAT) that translates all native
 //! IPv4 traffic to IPv6 traffic for transmission over an IPv6-only ISP network.
 
-use crate::common::packet_handler::handle_packet;
+use crate::common::packet_handler::{
+    get_ipv4_src_dst, get_ipv6_src_dst, get_layer_3_proto, handle_translation_error,
+    PacketHandlingError,
+};
+use crate::common::profiler::start_puffin_server;
 use crate::{args::protomask_clat::Args, common::permissions::ensure_root};
 use clap::Parser;
 use common::logging::enable_logger;
@@ -30,6 +34,10 @@ pub async fn main() {
     // We must be root to continue program execution
     ensure_root();
 
+    // Start profiling
+    #[allow(clippy::let_unit_value)]
+    let _server = start_puffin_server(&args.profiler_args);
+
     // Bring up a TUN interface
     let mut tun = Tun::new(&args.interface).unwrap();
 
@@ -81,34 +89,51 @@ pub async fn main() {
     log::info!("Translating packets on {}", tun.name());
     let mut buffer = vec![0u8; 1500];
     loop {
+        // Indicate to the profiler that we are starting a new packet
+        profiling::finish_frame!();
+        profiling::scope!("packet");
+
         // Read a packet
         let len = tun.read(&mut buffer).unwrap();
 
         // Translate it based on the Layer 3 protocol number
-        if let Some(output) = handle_packet(
+        let translation_result: Result<Option<Vec<u8>>, PacketHandlingError> =
-            &buffer[..len],
+            match get_layer_3_proto(&buffer[..len]) {
-            // IPv4 -> IPv6
+                Some(4) => {
-            |packet, source, dest| {
+                    let (source, dest) = get_ipv4_src_dst(&buffer[..len]);
-                Ok(translate_ipv4_to_ipv6(
+                    translate_ipv4_to_ipv6(
-                    packet,
+                        &buffer[..len],
-                    unsafe { embed_ipv4_addr_unchecked(*source, config.embed_prefix) },
+                        unsafe { embed_ipv4_addr_unchecked(source, config.embed_prefix) },
-                    unsafe { embed_ipv4_addr_unchecked(*dest, config.embed_prefix) },
+                        unsafe { embed_ipv4_addr_unchecked(dest, config.embed_prefix) },
-                )
+                    )
-                .map(Some)?)
+                    .map(Some)
-            },
+                    .map_err(PacketHandlingError::from)
-            // IPv6 -> IPv4
+                }
-            |packet, source, dest| {
+                Some(6) => {
-                Ok(translate_ipv6_to_ipv4(
+                    let (source, dest) = get_ipv6_src_dst(&buffer[..len]);
-                    packet,
+                    translate_ipv6_to_ipv4(
-                    unsafe {
+                        &buffer[..len],
-                        extract_ipv4_addr_unchecked(*source, config.embed_prefix.prefix_len())
+                        unsafe {
-                    },
+                            extract_ipv4_addr_unchecked(source, config.embed_prefix.prefix_len())
-                    unsafe { extract_ipv4_addr_unchecked(*dest, config.embed_prefix.prefix_len()) },
+                        },
-                )
+                        unsafe {
-                .map(Some)?)
+                            extract_ipv4_addr_unchecked(dest, config.embed_prefix.prefix_len())
-            },
+                        },
-        ) {
+                    )
-            // Write the packet if we get one back from the handler functions
+                    .map(Some)
+                    .map_err(PacketHandlingError::from)
+                }
+                Some(proto) => {
+                    log::warn!("Unknown Layer 3 protocol: {}", proto);
+                    continue;
+                }
+                None => {
+                    continue;
+                }
+            };
+
+        // Handle any errors and write
+        if let Some(output) = handle_translation_error(translation_result) {
             tun.write_all(&output).unwrap();
         }
     }

src/protomask.rs

@@ -1,4 +1,11 @@
-use crate::common::{packet_handler::handle_packet, permissions::ensure_root};
+use crate::common::{
+    packet_handler::{
+        get_ipv4_src_dst, get_ipv6_src_dst, get_layer_3_proto, handle_translation_error,
+        PacketHandlingError,
+    },
+    permissions::ensure_root,
+    profiler::start_puffin_server,
+};
 use clap::Parser;
 use common::logging::enable_logger;
 use easy_tun::Tun;
@@ -29,6 +36,10 @@ pub async fn main() {
     // We must be root to continue program execution
     ensure_root();
 
+    // Start profiling
+    #[allow(clippy::let_unit_value)]
+    let _server = start_puffin_server(&args.profiler_args);
+
     // Bring up a TUN interface
     log::debug!("Creating new TUN interface");
     let mut tun = Tun::new(&args.interface).unwrap();
@@ -88,38 +99,71 @@ pub async fn main() {
     log::info!("Translating packets on {}", tun.name());
     let mut buffer = vec![0u8; 1500];
     loop {
+        // Indicate to the profiler that we are starting a new packet
+        profiling::finish_frame!();
+        profiling::scope!("packet");
+
         // Read a packet
         let len = tun.read(&mut buffer).unwrap();
 
         // Translate it based on the Layer 3 protocol number
-        if let Some(output) = handle_packet(
+        let translation_result: Result<Option<Vec<u8>>, PacketHandlingError> =
-            &buffer[..len],
+            match get_layer_3_proto(&buffer[..len]) {
-            // IPv4 -> IPv6
+                Some(4) => {
-            |packet, source, dest| match addr_table.borrow().get_ipv6(dest) {
+                    let (source, dest) = get_ipv4_src_dst(&buffer[..len]);
-                Some(new_destination) => Ok(translate_ipv4_to_ipv6(
+                    match addr_table.borrow().get_ipv6(&dest) {
-                    packet,
+                        Some(new_destination) => translate_ipv4_to_ipv6(
-                    unsafe { embed_ipv4_addr_unchecked(*source, config.translation_prefix) },
+                            &buffer[..len],
-                    new_destination,
+                            unsafe { embed_ipv4_addr_unchecked(source, config.translation_prefix) },
-                )
+                            new_destination,
-                .map(Some)?),
+                        )
-                None => {
+                        .map(Some)
-                    protomask_metrics::metric!(PACKET_COUNTER, PROTOCOL_IPV4, STATUS_DROPPED);
+                        .map_err(PacketHandlingError::from),
-                    Ok(None)
+                        None => {
+                            protomask_metrics::metric!(
+                                PACKET_COUNTER,
+                                PROTOCOL_IPV4,
+                                STATUS_DROPPED
+                            );
+                            Ok(None)
+                        }
+                    }
                 }
-            },
+                Some(6) => {
-            // IPv6 -> IPv4
+                    let (source, dest) = get_ipv6_src_dst(&buffer[..len]);
-            |packet, source, dest| {
+                    match addr_table.borrow_mut().get_or_create_ipv4(&source) {
-                Ok(translate_ipv6_to_ipv4(
+                        Ok(new_source) => {
-                    packet,
+                            translate_ipv6_to_ipv4(&buffer[..len], new_source, unsafe {
-                    addr_table.borrow_mut().get_or_create_ipv4(source)?,
+                                extract_ipv4_addr_unchecked(
-                    unsafe {
+                                    dest,
-                        extract_ipv4_addr_unchecked(*dest, config.translation_prefix.prefix_len())
+                                    config.translation_prefix.prefix_len(),
-                    },
+                                )
-                )
+                            })
-                .map(Some)?)
+                            .map(Some)
-            },
+                            .map_err(PacketHandlingError::from)
-        ) {
+                        }
-            // Write the packet if we get one back from the handler functions
+                        Err(error) => {
+                            log::error!("Error getting IPv4 address: {}", error);
+                            protomask_metrics::metric!(
+                                PACKET_COUNTER,
+                                PROTOCOL_IPV6,
+                                STATUS_DROPPED
+                            );
+                            Ok(None)
+                        }
+                    }
+                }
+                Some(proto) => {
+                    log::warn!("Unknown Layer 3 protocol: {}", proto);
+                    continue;
+                }
+                None => {
+                    continue;
+                }
+            };
+
+        // Handle any errors and write
+        if let Some(output) = handle_translation_error(translation_result) {
             tun.write_all(&output).unwrap();
         }
     }