diff --git a/Cargo.toml b/Cargo.toml index 770445c..68466a8 100644 --- a/Cargo.toml +++ b/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" diff --git a/libs/easy-tun/Cargo.toml b/libs/easy-tun/Cargo.toml index 2d56d75..35bc5a0 100644 --- a/libs/easy-tun/Cargo.toml +++ b/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" diff --git a/libs/easy-tun/src/tun.rs b/libs/easy-tun/src/tun.rs index 4ff5bdf..74efd79 100644 --- a/libs/easy-tun/src/tun.rs +++ b/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 { diff --git a/libs/fast-nat/Cargo.toml b/libs/fast-nat/Cargo.toml index 416670a..570ee9e 100644 --- a/libs/fast-nat/Cargo.toml +++ b/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" diff --git a/libs/interproto/Cargo.toml b/libs/interproto/Cargo.toml index bfa7b2b..3ae9d31 100644 --- a/libs/interproto/Cargo.toml +++ b/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 } diff --git a/src/args/mod.rs b/src/args/mod.rs index 3268015..4f4fa42 100644 --- a/src/args/mod.rs +++ b/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, + } + } else { + #[derive(Debug, clap::Args)] + pub struct ProfilerArgs; + } +} diff --git a/src/args/protomask.rs b/src/args/protomask.rs index 90a120f..92265b4 100644 --- a/src/args/protomask.rs +++ b/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 for (Ipv4Addr, Ipv6Addr) { (val.ipv4, val.ipv6) } } + diff --git a/src/args/protomask_clat.rs b/src/args/protomask_clat.rs index 8fa352b..e3aaba7 100644 --- a/src/args/protomask_clat.rs +++ b/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, diff --git a/src/common/mod.rs b/src/common/mod.rs index e00f3b9..b10937a 100644 --- a/src/common/mod.rs +++ b/src/common/mod.rs @@ -4,3 +4,4 @@ pub mod logging; pub mod packet_handler; pub mod permissions; pub mod rfc6052; +pub mod profiler; \ No newline at end of file diff --git a/src/common/packet_handler.rs b/src/common/packet_handler.rs index 24366c7..621f3a5 100644 --- a/src/common/packet_handler.rs +++ b/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( - packet: &[u8], - mut ipv4_handler: Ipv4Handler, - mut ipv6_handler: Ipv6Handler, -) -> Option> -where - Ipv4Handler: FnMut(&[u8], &Ipv4Addr, &Ipv4Addr) -> Result>, PacketHandlingError>, - Ipv6Handler: FnMut(&[u8], &Ipv6Addr, &Ipv6Addr) -> Result>, PacketHandlingError>, -{ + +/// Get the layer 3 protocol of a packet +pub fn get_layer_3_proto(packet: &[u8]) -> Option { // 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 { - // 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())); + Some(layer_3_proto) +} - // Call the handler - ipv4_handler(packet, &source_addr, &destination_addr) - } +/// Get the source and destination addresses of an IPv4 packet +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 - 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())); + (source_addr, destination_addr) +} - // Call the handler - ipv6_handler(packet, &source_addr, &destination_addr) - } +/// Get the source and destination addresses of an IPv6 packet +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 - proto => { - log::warn!("Unknown Layer 3 protocol: {}", proto); - return None; - } - }; + (source_addr, destination_addr) +} - // The response from the handler may or may not be a warn-able error - match handler_response { +/// Appropriately handle a translation error +pub fn handle_translation_error( + result: Result>, PacketHandlingError>, +) -> Option> { + // 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( +// packet: &[u8], +// mut ipv4_handler: Ipv4Handler, +// mut ipv6_handler: Ipv6Handler, +// ) -> Option> +// where +// Ipv4Handler: FnMut(&[u8], &Ipv4Addr, &Ipv4Addr) -> Result>, PacketHandlingError>, +// Ipv6Handler: FnMut(&[u8], &Ipv6Addr, &Ipv6Addr) -> Result>, 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 +// } +// }, +// } +// } diff --git a/src/common/profiler.rs b/src/common/profiler.rs new file mode 100644 index 0000000..8f9fefd --- /dev/null +++ b/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 { + 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){} + } +} diff --git a/src/protomask-clat.rs b/src/protomask-clat.rs index 8983cb7..42f785c 100644 --- a/src/protomask-clat.rs +++ b/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( - &buffer[..len], - // IPv4 -> IPv6 - |packet, source, dest| { - Ok(translate_ipv4_to_ipv6( - packet, - unsafe { embed_ipv4_addr_unchecked(*source, config.embed_prefix) }, - unsafe { embed_ipv4_addr_unchecked(*dest, config.embed_prefix) }, - ) - .map(Some)?) - }, - // IPv6 -> IPv4 - |packet, source, dest| { - Ok(translate_ipv6_to_ipv4( - packet, - unsafe { - extract_ipv4_addr_unchecked(*source, config.embed_prefix.prefix_len()) - }, - unsafe { extract_ipv4_addr_unchecked(*dest, config.embed_prefix.prefix_len()) }, - ) - .map(Some)?) - }, - ) { - // Write the packet if we get one back from the handler functions + let translation_result: Result>, PacketHandlingError> = + match get_layer_3_proto(&buffer[..len]) { + Some(4) => { + let (source, dest) = get_ipv4_src_dst(&buffer[..len]); + translate_ipv4_to_ipv6( + &buffer[..len], + unsafe { embed_ipv4_addr_unchecked(source, config.embed_prefix) }, + unsafe { embed_ipv4_addr_unchecked(dest, config.embed_prefix) }, + ) + .map(Some) + .map_err(PacketHandlingError::from) + } + Some(6) => { + let (source, dest) = get_ipv6_src_dst(&buffer[..len]); + translate_ipv6_to_ipv4( + &buffer[..len], + unsafe { + extract_ipv4_addr_unchecked(source, config.embed_prefix.prefix_len()) + }, + unsafe { + extract_ipv4_addr_unchecked(dest, config.embed_prefix.prefix_len()) + }, + ) + .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(); } } diff --git a/src/protomask.rs b/src/protomask.rs index 35ec8f7..dd1493c 100644 --- a/src/protomask.rs +++ b/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( - &buffer[..len], - // IPv4 -> IPv6 - |packet, source, dest| match addr_table.borrow().get_ipv6(dest) { - Some(new_destination) => Ok(translate_ipv4_to_ipv6( - packet, - unsafe { embed_ipv4_addr_unchecked(*source, config.translation_prefix) }, - new_destination, - ) - .map(Some)?), - None => { - protomask_metrics::metric!(PACKET_COUNTER, PROTOCOL_IPV4, STATUS_DROPPED); - Ok(None) + let translation_result: Result>, PacketHandlingError> = + match get_layer_3_proto(&buffer[..len]) { + Some(4) => { + let (source, dest) = get_ipv4_src_dst(&buffer[..len]); + match addr_table.borrow().get_ipv6(&dest) { + Some(new_destination) => translate_ipv4_to_ipv6( + &buffer[..len], + unsafe { embed_ipv4_addr_unchecked(source, config.translation_prefix) }, + new_destination, + ) + .map(Some) + .map_err(PacketHandlingError::from), + None => { + protomask_metrics::metric!( + PACKET_COUNTER, + PROTOCOL_IPV4, + STATUS_DROPPED + ); + Ok(None) + } + } } - }, - // IPv6 -> IPv4 - |packet, source, dest| { - Ok(translate_ipv6_to_ipv4( - packet, - addr_table.borrow_mut().get_or_create_ipv4(source)?, - unsafe { - extract_ipv4_addr_unchecked(*dest, config.translation_prefix.prefix_len()) - }, - ) - .map(Some)?) - }, - ) { - // Write the packet if we get one back from the handler functions + Some(6) => { + let (source, dest) = get_ipv6_src_dst(&buffer[..len]); + match addr_table.borrow_mut().get_or_create_ipv4(&source) { + Ok(new_source) => { + translate_ipv6_to_ipv4(&buffer[..len], new_source, unsafe { + extract_ipv4_addr_unchecked( + dest, + config.translation_prefix.prefix_len(), + ) + }) + .map(Some) + .map_err(PacketHandlingError::from) + } + 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(); } }