From 5f5b82567675b157faac19bbb9c96f4faf11ade3 Mon Sep 17 00:00:00 2001 From: Evan Pratten Date: Mon, 17 Jul 2023 19:51:59 -0400 Subject: [PATCH] Speed optimizations --- Makefile | 4 +++ src/lib.rs | 1 + src/logging.rs | 58 ++++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 39 +++-------------------------- src/nat/mod.rs | 14 +++++------ src/nat/xlat/icmp.rs | 10 ++++---- 6 files changed, 78 insertions(+), 48 deletions(-) create mode 100644 src/logging.rs diff --git a/Makefile b/Makefile index c585828..a032b13 100644 --- a/Makefile +++ b/Makefile @@ -2,4 +2,8 @@ SRC=$(wildcard src/*.rs) $(wildcard src/**/*.rs) $(wildcard src/**/**/*.rs) Carg target/debug/protomask: $(SRC) cross build --target x86_64-unknown-linux-musl + sudo setcap cap_net_admin=eip $@ + +target/release/protomask: $(SRC) + cross build --target x86_64-unknown-linux-musl --release sudo setcap cap_net_admin=eip $@ \ No newline at end of file diff --git a/src/lib.rs b/src/lib.rs index a04114b..ca84c3a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,3 +1,4 @@ #![deny(unsafe_code)] pub mod nat; +mod logging; diff --git a/src/logging.rs b/src/logging.rs new file mode 100644 index 0000000..9950daa --- /dev/null +++ b/src/logging.rs @@ -0,0 +1,58 @@ +use std::sync::OnceLock; + +use colored::Colorize; + +/// A global variable that is used to early-kill attempts to write debug logs if debug logging is disabled +pub static DEBUG_ENABLED: OnceLock = OnceLock::new(); + +/// A macro that can completely skip the debug step if debug logging is disabled +#[macro_export] +macro_rules! debug { + ($($arg:tt)*) => { + if *$crate::logging::DEBUG_ENABLED.get().unwrap_or(&false) { + log::debug!($($arg)*); + } + }; +} + +/// Enable the logger +#[allow(dead_code)] +pub fn enable_logger(verbose: bool) { + fern::Dispatch::new() + .format(move |out, message, record| { + out.finish(format_args!( + "{}: {}", + format!( + "{}{}", + // Level messages are padded to keep the output looking somewhat sane + match record.level() { + log::Level::Error => "ERROR".red().bold().to_string(), + log::Level::Warn => "WARN ".yellow().bold().to_string(), + log::Level::Info => "INFO ".green().bold().to_string(), + log::Level::Debug => "DEBUG".bright_blue().bold().to_string(), + log::Level::Trace => "TRACE".bright_white().bold().to_string(), + }, + // Only show the outer package name if verbose logging is enabled (otherwise nothing) + match verbose { + true => format!(" [{}]", record.target().split("::").nth(0).unwrap()), + false => String::new(), + } + .bright_black() + ), + message + )) + }) + .level(match verbose { + true => log::LevelFilter::Debug, + false => log::LevelFilter::Info, + }) + .chain(std::io::stdout()) + .apply() + .unwrap(); + if verbose { + log::debug!("Verbose logging enabled"); + } + + // Set the global debug enabled variable + DEBUG_ENABLED.set(verbose).unwrap(); +} diff --git a/src/main.rs b/src/main.rs index e6ec994..bca525e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,11 +1,12 @@ use clap::Parser; -use colored::Colorize; use config::Config; +use logging::enable_logger; use nat::Nat64; mod cli; mod config; mod nat; +mod logging; #[tokio::main] pub async fn main() { @@ -13,41 +14,7 @@ pub async fn main() { let args = cli::Args::parse(); // Set up logging - let log_verbose = args.verbose; - fern::Dispatch::new() - .format(move |out, message, record| { - out.finish(format_args!( - "{}: {}", - format!( - "{}{}", - // Level messages are padded to keep the output looking somewhat sane - match record.level() { - log::Level::Error => "ERROR".red().bold().to_string(), - log::Level::Warn => "WARN ".yellow().bold().to_string(), - log::Level::Info => "INFO ".green().bold().to_string(), - log::Level::Debug => "DEBUG".bright_blue().bold().to_string(), - log::Level::Trace => "TRACE".bright_white().bold().to_string(), - }, - // Only show the outer package name if verbose logging is enabled (otherwise nothing) - match log_verbose { - true => format!(" [{}]", record.target().split("::").nth(0).unwrap()), - false => String::new(), - } - .bright_black() - ), - message - )) - }) - .level(match args.verbose { - true => log::LevelFilter::Debug, - false => log::LevelFilter::Info, - }) - .chain(std::io::stdout()) - .apply() - .unwrap(); - if args.verbose { - log::debug!("Verbose logging enabled"); - } + enable_logger(args.verbose); // If the binary was built with profiling support, enable it #[cfg(feature = "enable-profiling")] diff --git a/src/nat/mod.rs b/src/nat/mod.rs index 4dce403..1685763 100644 --- a/src/nat/mod.rs +++ b/src/nat/mod.rs @@ -76,7 +76,7 @@ impl Nat64 { match self.interface.recv(&mut buffer) { Ok(packet_len) => { // Parse in to a more friendly format - log::debug!("--- NEW PACKET ---"); + crate::debug!("--- NEW PACKET ---"); match IpPacket::new(&buffer[..packet_len]) { // Try to process the packet Ok(inbound_packet) => match self.process_packet(inbound_packet).await { @@ -84,11 +84,11 @@ impl Nat64 { // If data is returned, send it back out the interface Some(outbound_packet) => { let packet_bytes = outbound_packet.to_bytes(); - log::debug!( + crate::debug!( "Outbound packet next header: {}", outbound_packet.get_next_header().0 ); - log::debug!("Sending packet: {:?}", packet_bytes); + crate::debug!("Sending packet: {:?}", packet_bytes); self.interface.send(&packet_bytes).unwrap(); } // Otherwise, we can assume that the packet was dealt with, and can move on @@ -98,7 +98,7 @@ impl Nat64 { // Some errors are non-critical as far as this loop is concerned Err(error) => match error { Nat64Error::TableError(TableError::NoIpv6Mapping(address)) => { - log::debug!("No IPv6 mapping for {}", address); + crate::debug!("No IPv6 mapping for {}", address); } error => { return Err(error); @@ -130,7 +130,7 @@ impl Nat64 { IpAddr::V4(ipv4_addr) => !self.table.is_address_within_pool(&ipv4_addr), IpAddr::V6(ipv6_addr) => !self.ipv6_nat_prefix.contains(&ipv6_addr), } { - log::debug!( + crate::debug!( "Packet destination {} is not within the NAT64 prefix or IPv4 pool", packet.get_destination(), ); @@ -148,12 +148,12 @@ impl Nat64 { .calculate_xlat_addr(&destination, &self.ipv6_nat_prefix)?; // Log information about the packet - log::debug!( + crate::debug!( "Received packet traveling from {} to {}", source, destination ); - log::debug!( + crate::debug!( "New path shall become: {} -> {}", new_source, new_destination diff --git a/src/nat/xlat/icmp.rs b/src/nat/xlat/icmp.rs index b683432..10d34f7 100644 --- a/src/nat/xlat/icmp.rs +++ b/src/nat/xlat/icmp.rs @@ -144,7 +144,7 @@ pub fn translate_icmp_4_to_6( // if the original payload's next header is ICMP, we need to translated the inner payload's ICMP type if original_payload.get_next_level_protocol() == IpNextHeaderProtocols::Icmp { - log::debug!("Time Exceeded packet contains another ICMP packet.. Translating"); + crate::debug!("Time Exceeded packet contains another ICMP packet.. Translating"); if let Some((icmpv6_type, icmpv6_code)) = translate_type_and_code_4_to_6( IcmpType(original_payload_inner[0]), IcmpCode(original_payload_inner[1]), @@ -157,7 +157,7 @@ pub fn translate_icmp_4_to_6( &original_payload_inner[4..] ); original_payload_inner = inner_icmpv6.packet().to_vec(); - log::debug!( + crate::debug!( "Translated inner ICMPv6 packet: {:?}", original_payload_inner ); @@ -194,7 +194,7 @@ pub fn translate_icmp_4_to_6( output.set_icmpv6_code(icmpv6_code); // Set the payload - log::debug!("Setting ICMPv6 payload: {:?}", output_payload); + crate::debug!("Setting ICMPv6 payload: {:?}", output_payload); output.set_payload(&output_payload); // Calculate the checksum @@ -256,7 +256,7 @@ pub fn translate_icmp_6_to_4( // if the original payload's next header is ICMPv6, we need to translated the inner payload's ICMPv6 type if original_payload.get_next_header() == IpNextHeaderProtocols::Icmpv6 { - log::debug!("Time Exceeded packet contains another ICMPv6 packet.. Translating"); + crate::debug!("Time Exceeded packet contains another ICMPv6 packet.. Translating"); if let Some((icmp_type, icmp_code)) = translate_type_and_code_6_to_4( Icmpv6Type(original_payload_inner[0]), Icmpv6Code(original_payload_inner[1]), @@ -264,7 +264,7 @@ pub fn translate_icmp_6_to_4( let inner_icmp = icmp_packet!(icmp_type, icmp_code, &original_payload_inner[8..]); original_payload_inner = inner_icmp.packet().to_vec(); - log::debug!("Translated inner ICMP packet: {:?}", original_payload_inner); + crate::debug!("Translated inner ICMP packet: {:?}", original_payload_inner); } }