Free nf on shutdown

leaf-cli/Cargo.toml

@@ -32,5 +32,6 @@ features = [
   "outbound-quic",
   "ctrlc",
   "auto-reload",
-  "inbound-nf"
+  "inbound-nf",
+  "rule-process-name"
 ]

leaf/src/app/inbound/manager.rs

@@ -86,8 +86,6 @@ impl InboundManager {
                 "nf" => {
                     let settings: crate::config::NfInboundSettings =
                         protobuf::Message::parse_from_bytes(&inbound.settings)?;
-                    nf::inbound::init(settings.driver_name.clone(), &settings.nfapi)?;
-
                     use crate::app::fake_dns::{FakeDns, FakeDnsMode};
                     let fake_dns_exclude = settings.fake_dns_exclude.clone();
                     let fake_dns_include = settings.fake_dns_include.clone();
@@ -97,11 +95,15 @@ impl InboundManager {
                         (FakeDnsMode::Exclude, fake_dns_exclude)
                     };
                     let fake_dns = Arc::new(FakeDns::new(mode, filters));
-
+                    let manager = Arc::new(nf::inbound::NfManager::new(
+                        settings.driver_name.clone(),
+                        settings.nfapi.clone(),
+                        fake_dns,
+                    )?);
                     let stream = Arc::new(nf::inbound::StreamHandler {
-                        fake_dns: fake_dns.clone(),
+                        manager: manager.clone(),
                     });
-                    let datagram = Arc::new(nf::inbound::DatagramHandler { fake_dns });
+                    let datagram = Arc::new(nf::inbound::DatagramHandler { manager });
                     let handler = Arc::new(crate::proxy::inbound::Handler::new(
                         tag.clone(),
                         Some(stream),

leaf/src/lib.rs

@@ -41,7 +41,6 @@ pub mod mobile;
 #[cfg(all(feature = "inbound-tun", any(target_os = "macos", target_os = "linux")))]
 mod sys;
 
-
 #[cfg(all(feature = "inbound-tun", target_os = "windows"))]
 mod winsys;
 

leaf/src/proxy/nf/inbound/datagram.rs

@@ -13,9 +13,10 @@ use crate::{
 };
 
 use super::packed::{SOCKADDR_IN, SOCKADDR_IN6};
+use super::NfManager;
 
 pub struct Handler {
-    pub fake_dns: Arc<FakeDns>,
+    pub manager: Arc<NfManager>,
 }
 
 #[async_trait]
@@ -24,7 +25,7 @@ impl InboundDatagramHandler for Handler {
         Ok(InboundTransport::Datagram(
             Box::new(Datagram {
                 socket,
-                fake_dns: self.fake_dns.clone(),
+                fake_dns: self.manager.fake_dns.clone(),
             }),
             None,
         ))
@@ -78,14 +79,15 @@ impl InboundDatagramRecvHalf for DatagramRecvHalf {
         assert!(buf.len() >= payload_size);
         let real_payload = &recv_buf[header_size..header_size + payload_size];
 
-        let (local_addr, process_name) = if let Some(info) = super::UDP_LOCAL_INFO.lock().unwrap().get(&id) {
-            (info.local_address.clone(), info.process_name.clone())
-        } else {
-            return Err(ProxyError::DatagramWarn(anyhow!(format!(
-                "local socket not found id={}",
-                id
-            ))));
-        };
+        let (local_addr, process_name) =
+            if let Some(info) = super::UDP_LOCAL_INFO.lock().unwrap().get(&id) {
+                (info.local_address.clone(), info.process_name.clone())
+            } else {
+                return Err(ProxyError::DatagramWarn(anyhow!(format!(
+                    "local socket not found id={}",
+                    id
+                ))));
+            };
 
         // Override with real source address and process name.
         src_addr.address = local_addr;

leaf/src/proxy/nf/inbound/mod.rs

@@ -29,6 +29,8 @@ use tracing::{debug, instrument, trace, warn};
 
 use packed::{SOCKADDR, SOCKADDR_IN, SOCKADDR_IN6};
 
+use crate::app::fake_dns::FakeDns;
+
 const MAX_PATH: usize = 260;
 const IPPROTO_TCP: i32 = 6;
 
@@ -908,24 +910,21 @@ fn init_if_needed<P: AsRef<OsStr>>(driver_name: String, nfapi: P) -> Result<()>
 
 static IS_NF_INITIALIZED: AtomicBool = AtomicBool::new(false);
 
-// TODO Guard initializing.
-pub fn init<P: AsRef<OsStr>>(driver_name: String, nfapi: P) -> Result<()> {
-    if !IS_NF_INITIALIZED.load(Ordering::Relaxed) {
+fn init<P: AsRef<OsStr>>(driver_name: String, nfapi: P) -> Result<()> {
+    if !IS_NF_INITIALIZED.swap(true, Ordering::Relaxed) {
         init_if_needed(driver_name, nfapi)?;
-        IS_NF_INITIALIZED.store(true, Ordering::Relaxed);
     }
     Ok(())
 }
 
 unsafe fn uninit_nf() {
-    if IS_NF_INITIALIZED.load(Ordering::Relaxed) {
+    if IS_NF_INITIALIZED.swap(false, Ordering::Relaxed) {
         NF_FREE.unwrap()();
         if let Some(nfapi) = NFAPI.write().take() {
             if let Err(e) = nfapi.close() {
                 debug!("close nf failed: {}", e);
             }
         }
-        IS_NF_INITIALIZED.store(false, Ordering::Relaxed);
     }
 }
 
@@ -955,3 +954,20 @@ pub unsafe fn get_process_name(pid: u32) -> Result<String> {
     // Return the full path instead of just the filename
     Ok(process_name.to_string_lossy().to_string())
 }
+
+pub struct NfManager {
+    pub fake_dns: Arc<FakeDns>,
+}
+
+impl NfManager {
+    pub fn new(driver_name: String, nfapi: String, fake_dns: Arc<FakeDns>) -> Result<Self> {
+        init(driver_name, nfapi)?;
+        Ok(Self { fake_dns })
+    }
+}
+
+impl Drop for NfManager {
+    fn drop(&mut self) {
+        uninit();
+    }
+}

leaf/src/proxy/nf/inbound/stream.rs

@@ -3,8 +3,10 @@ use async_trait::async_trait;
 use crate::app::fake_dns::FakeDns;
 use crate::{proxy::*, session::Session};
 
+use super::NfManager;
+
 pub struct Handler {
-    pub fake_dns: Arc<FakeDns>,
+    pub manager: Arc<NfManager>,
 }
 
 #[async_trait]
@@ -28,8 +30,8 @@ impl InboundStreamHandler for Handler {
         sess.process_name = process_name;
 
         let remote_ip = remote_addr.ip();
-        if self.fake_dns.is_fake_ip(&remote_ip).await {
-            if let Some(domain) = self.fake_dns.query_domain(&remote_ip).await {
+        if self.manager.fake_dns.is_fake_ip(&remote_ip).await {
+            if let Some(domain) = self.manager.fake_dns.query_domain(&remote_ip).await {
                 sess.destination = SocksAddr::Domain(domain, remote_addr.port());
             } else {
                 if remote_addr.port() != 443 && remote_addr.port() != 80 {