From 2d5551ea03a585f84ae45508ea62460359e49078 Mon Sep 17 00:00:00 2001
From: Evan Pratten <ewpratten@gmail.com>
Date: Mon, 17 Jul 2023 19:18:43 -0400
Subject: [PATCH] Allow remote profiling with build flag

---
 Cargo.toml       | 12 +++++++++++-
 Dockerfile       |  1 +
 src/cli.rs       |  6 +++++-
 src/main.rs      | 10 ++++++++++
 src/nat/mod.rs   |  2 ++
 src/nat/table.rs |  4 ++++
 6 files changed, 33 insertions(+), 2 deletions(-)
 create mode 100644 Dockerfile

diff --git a/Cargo.toml b/Cargo.toml
index 46d0c17..4936115 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -12,12 +12,22 @@ license = "GPL-3.0"
 keywords = []
 categories = []
 
+[features]
+default = []
+enable-profiling = ["profiling/profile-with-puffin", "puffin_http"]
+
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 [dependencies]
-tokio = { version = "1.29.1", features = ["macros", "rt-multi-thread", "process"] }
+tokio = { version = "1.29.1", features = [
+    "macros",
+    "rt-multi-thread",
+    "process"
+] }
 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 }
+profiling = "1.0.8"
 toml = "0.7.6"
 log = "0.4.19"
 fern = "0.6.2"
diff --git a/Dockerfile b/Dockerfile
new file mode 100644
index 0000000..0a443d8
--- /dev/null
+++ b/Dockerfile
@@ -0,0 +1 @@
+# FROM 
\ No newline at end of file
diff --git a/src/cli.rs b/src/cli.rs
index 227f9ba..8637b3b 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -5,11 +5,15 @@ use clap::Parser;
 #[derive(Parser, Debug)]
 #[command(author, version, about, long_about = None)]
 pub struct Args {
-
     /// Path to the config file
     pub config_file: PathBuf,
 
     /// Enable verbose logging
     #[clap(short, long)]
     pub verbose: bool,
+
+    /// Enable the puffin profiling server for debugging
+    #[cfg(feature = "enable-profiling")]
+    #[clap(long)]
+    pub enable_profiling: bool,
 }
diff --git a/src/main.rs b/src/main.rs
index 74a0e64..e6ec994 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -49,6 +49,16 @@ pub async fn main() {
         log::debug!("Verbose logging enabled");
     }
 
+    // If the binary was built with profiling support, enable it
+    #[cfg(feature = "enable-profiling")]
+    let _puffin_server: puffin_http::Server;
+    #[cfg(feature = "enable-profiling")]
+    if args.enable_profiling {
+        _puffin_server =
+            puffin_http::Server::new(&format!("0.0.0.0:{}", puffin_http::DEFAULT_PORT)).unwrap();
+        log::info!("Puffin profiling server started");
+    }
+
     // Parse the config file
     let config = Config::load(args.config_file).unwrap();
 
diff --git a/src/nat/mod.rs b/src/nat/mod.rs
index f1c925a..4dce403 100644
--- a/src/nat/mod.rs
+++ b/src/nat/mod.rs
@@ -119,6 +119,8 @@ impl Nat64 {
 }
 
 impl Nat64 {
+
+    #[profiling::function]
     async fn process_packet<'a>(
         &mut self,
         packet: IpPacket<'a>,
diff --git a/src/nat/table.rs b/src/nat/table.rs
index 1f357f5..a254457 100644
--- a/src/nat/table.rs
+++ b/src/nat/table.rs
@@ -67,6 +67,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();
@@ -100,6 +101,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();
@@ -124,6 +126,7 @@ impl Nat64Table {
     }
 
     /// Calculate the translated version of any address
+    #[profiling::function]
     pub fn calculate_xlat_addr(
         &mut self,
         input: &IpAddr,
@@ -176,6 +179,7 @@ impl Nat64Table {
 
 impl Nat64Table {
     /// Prune old reservations
+    #[profiling::function]
     pub fn prune(&mut self) {
         let now = Instant::now();