From edd83fad99332329c5b92415a0b82452c0a0af01 Mon Sep 17 00:00:00 2001 From: Evan Pratten Date: Wed, 19 Jul 2023 10:36:53 -0400 Subject: [PATCH] Add experimental profiling support --- Cargo.toml | 10 +++++-- src/cli/main.rs | 11 +++++++ src/nat/error.rs | 17 +++++++++++ src/nat/mod.rs | 48 +++++++++++++++++++------------ src/nat/table.rs | 7 +++++ src/packet/protocols/icmp.rs | 3 +- src/packet/protocols/icmpv6.rs | 3 ++ src/packet/protocols/ipv4.rs | 3 ++ src/packet/protocols/ipv6.rs | 2 ++ src/packet/protocols/tcp.rs | 3 ++ src/packet/protocols/udp.rs | 3 ++ src/packet/xlat/icmp/mod.rs | 4 ++- src/packet/xlat/icmp/type_code.rs | 2 ++ src/packet/xlat/ip.rs | 2 ++ src/packet/xlat/tcp.rs | 2 ++ src/packet/xlat/udp.rs | 2 ++ 16 files changed, 100 insertions(+), 22 deletions(-) create mode 100644 src/nat/error.rs diff --git a/Cargo.toml b/Cargo.toml index 32eaaec..5dcb7db 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/cli/main.rs b/src/cli/main.rs index 4fdb9ec..5d8e158 100644 --- a/src/cli/main.rs +++ b/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, diff --git a/src/nat/error.rs b/src/nat/error.rs new file mode 100644 index 0000000..5a2c5f7 --- /dev/null +++ b/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>), +} diff --git a/src/nat/mod.rs b/src/nat/mod.rs index aa1ff8f..581d42d 100644 --- a/src/nat/mod.rs +++ b/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>), -} - 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> = packet.try_into()?; @@ -113,14 +102,37 @@ impl Nat64 { }); } 6 => { + profiling::scope!("Ipv6 to IPv4"); + // Parse the packet let packet: Ipv6Packet> = 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 = diff --git a/src/nat/table.rs b/src/nat/table.rs index 30226f1..6d40cb2 100644 --- a/src/nat/table.rs +++ b/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 { // 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 { // 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(); diff --git a/src/packet/protocols/icmp.rs b/src/packet/protocols/icmp.rs index dce1f92..a1977b1 100644 --- a/src/packet/protocols/icmp.rs +++ b/src/packet/protocols/icmp.rs @@ -29,6 +29,7 @@ where { type Error = PacketError; + #[profiling::function] fn try_from(bytes: Vec) -> Result { // Parse the packet let packet = @@ -43,11 +44,11 @@ where } } - impl Into> for IcmpPacket where T: Into>, { + #[profiling::function] fn into(self) -> Vec { // Convert the payload into raw bytes let payload: Vec = self.payload.into(); diff --git a/src/packet/protocols/icmpv6.rs b/src/packet/protocols/icmpv6.rs index 75ebff6..022428b 100644 --- a/src/packet/protocols/icmpv6.rs +++ b/src/packet/protocols/icmpv6.rs @@ -42,6 +42,7 @@ where T: From>, { /// 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 { /// 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 Into> for Icmpv6Packet where T: Into>, { + #[profiling::function] fn into(self) -> Vec { // Convert the payload into raw bytes let payload: Vec = self.payload.into(); diff --git a/src/packet/protocols/ipv4.rs b/src/packet/protocols/ipv4.rs index 9728084..fb2c752 100644 --- a/src/packet/protocols/ipv4.rs +++ b/src/packet/protocols/ipv4.rs @@ -53,6 +53,7 @@ impl Ipv4Packet { } } + #[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) -> Result { // Parse the packet let packet = @@ -94,6 +96,7 @@ impl Into> for Ipv4Packet where T: Into> + Clone, { + #[profiling::function] fn into(self) -> Vec { // Convert the payload into raw bytes let payload: Vec = self.payload.clone().into(); diff --git a/src/packet/protocols/ipv6.rs b/src/packet/protocols/ipv6.rs index 8d28375..fc6d41f 100644 --- a/src/packet/protocols/ipv6.rs +++ b/src/packet/protocols/ipv6.rs @@ -44,6 +44,7 @@ where { type Error = PacketError; + #[profiling::function] fn try_from(bytes: Vec) -> Result { // Parse the packet let packet = @@ -66,6 +67,7 @@ impl Into> for Ipv6Packet where T: Into>, { + #[profiling::function] fn into(self) -> Vec { // Convert the payload into raw bytes let payload: Vec = self.payload.into(); diff --git a/src/packet/protocols/tcp.rs b/src/packet/protocols/tcp.rs index 1413c20..b498cc4 100644 --- a/src/packet/protocols/tcp.rs +++ b/src/packet/protocols/tcp.rs @@ -113,6 +113,7 @@ where T: From>, { /// 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 { /// 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 Into> for TcpPacket where T: Into>, { + #[profiling::function] fn into(self) -> Vec { // Get the options length in words let options_length = self.options_length(); diff --git a/src/packet/protocols/udp.rs b/src/packet/protocols/udp.rs index 21eebfa..65716b3 100644 --- a/src/packet/protocols/udp.rs +++ b/src/packet/protocols/udp.rs @@ -84,6 +84,7 @@ where T: From>, { /// 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 { /// 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 Into> for UdpPacket where T: Into>, { + #[profiling::function] fn into(self) -> Vec { // Convert the payload into raw bytes let payload: Vec = self.payload.into(); diff --git a/src/packet/xlat/icmp/mod.rs b/src/packet/xlat/icmp/mod.rs index 967cd68..c271b17 100644 --- a/src/packet/xlat/icmp/mod.rs +++ b/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, 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, new_source: Ipv4Addr, @@ -83,7 +85,7 @@ pub fn translate_icmpv6_to_icmp( buffer.extend_from_slice(&inner_payload); buffer }) - }, + } _ => input.payload, }; diff --git a/src/packet/xlat/icmp/type_code.rs b/src/packet/xlat/icmp/type_code.rs index e7cb6cd..1185761 100644 --- a/src/packet/xlat/icmp/type_code.rs +++ b/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, diff --git a/src/packet/xlat/ip.rs b/src/packet/xlat/ip.rs index 540638c..c56d47d 100644 --- a/src/packet/xlat/ip.rs +++ b/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>, 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>, new_source: Ipv4Addr, diff --git a/src/packet/xlat/tcp.rs b/src/packet/xlat/tcp.rs index 7dffe47..b35dfac 100644 --- a/src/packet/xlat/tcp.rs +++ b/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, 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, new_source_addr: Ipv4Addr, diff --git a/src/packet/xlat/udp.rs b/src/packet/xlat/udp.rs index cd8a34b..f0003f6 100644 --- a/src/packet/xlat/udp.rs +++ b/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, 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, new_source_addr: Ipv4Addr,