Support PROCESS-NAME rule

Author: eycorsican

leaf/Cargo.toml

@@ -95,6 +95,9 @@ openssl-tls = ["openssl", "tokio-openssl", "openssl-probe"]
 config-conf = ["regex"]
 config-json = ["serde", "serde_derive", "serde_json"]
 
+# Router rules
+rule-process-name = ["regex"]
+
 # Outbounds
 outbound-direct = []
 outbound-drop = []
@@ -164,6 +167,7 @@ colored = "2.1"
 maxminddb = { version = "0.24", features = ["mmap"] }
 memmap2 = "0.9"
 cidr = "0.2"
+regex = { version = "1.10", optional = true }
 
 # outbound-select
 directories = { version = "4.0", optional = true }
@@ -178,7 +182,6 @@ serde_derive = { version = "1.0", optional = true }
 serde = { version = "1.0", optional = true }
 
 # config-conf
-regex = { version = "1.10", optional = true }
 
 openssl = { version = "0.10", features = ["vendored"], optional = true }
 

leaf/src/app/dispatcher.rs

@@ -1,5 +1,6 @@
 use std::convert::TryFrom;
 use std::io::{self, ErrorKind};
+use std::path::Path;
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -43,15 +44,43 @@ fn log_request(
     } else {
         (sess.network.to_string(), outbound_tag.to_string())
     };
-    info!(
-        "[{}] [{}] [{}] [{}] [{}] [{}]",
-        sess.forwarded_source.unwrap_or_else(|| sess.source.ip()),
-        network,
-        &sess.inbound_tag,
-        outbound_tag,
-        hs,
-        &sess.destination,
-    );
+
+    #[cfg(feature = "rule-process-name")]
+    {
+        let process_name = sess
+            .process_name
+            .as_ref()
+            .map(|x| {
+                Path::new(x)
+                    .file_name()
+                    .and_then(|name| name.to_str())
+                    .unwrap_or(x)
+            })
+            .unwrap_or("");
+        info!(
+            "[{}] [{}] [{}] [{}] [{}] [{}] [{}]",
+            process_name,
+            sess.forwarded_source.unwrap_or_else(|| sess.source.ip()),
+            network,
+            &sess.inbound_tag,
+            outbound_tag,
+            hs,
+            &sess.destination,
+        );
+    }
+
+    #[cfg(not(feature = "rule-process-name"))]
+    {
+        info!(
+            "[{}] [{}] [{}] [{}] [{}] [{}]",
+            sess.forwarded_source.unwrap_or_else(|| sess.source.ip()),
+            network,
+            &sess.inbound_tag,
+            outbound_tag,
+            hs,
+            &sess.destination,
+        );
+    }
 }
 
 pub struct Dispatcher {

leaf/src/app/nat_manager.rs

@@ -133,13 +133,14 @@ impl NatManager {
             source: dgram_src.address,
             destination: pkt.dst_addr.clone(),
             inbound_tag: inbound_tag.to_string(),
+            process_name: dgram_src.process_name.clone(),
             ..Default::default()
         });
         if sess.inbound_tag.is_empty() {
             sess.inbound_tag = inbound_tag.to_string();
         }
 
-        self.add_session(sess, *dgram_src, client_ch_tx.clone(), &mut guard)
+        self.add_session(sess, dgram_src.clone(), client_ch_tx.clone(), &mut guard)
             .await;
 
         debug!(
@@ -170,36 +171,38 @@ impl NatManager {
             mpsc::channel(*crate::option::UDP_UPLINK_CHANNEL_SIZE);
         let (downlink_abort_tx, downlink_abort_rx) = oneshot::channel();
 
-        guard.insert(raddr, (target_ch_tx, downlink_abort_tx, Instant::now()));
+        guard.insert(raddr.clone(), (target_ch_tx, downlink_abort_tx, Instant::now()));
 
         let dispatcher = self.dispatcher.clone();
         let sessions = self.sessions.clone();
 
         // Spawns a new task for dispatching to avoid blocking the current task,
         // because we have stream type transports for UDP traffic, establishing a
         // TCP stream would block the task.
+        let raddr_cloned = raddr.clone();
         tokio::spawn(async move {
             // new socket to communicate with the target.
             let socket = match dispatcher.dispatch_datagram(sess).await {
                 Ok(s) => s,
                 Err(e) => {
-                    debug!("dispatch {} failed: {}", &raddr, e);
-                    sessions.lock().await.remove(&raddr);
+                    debug!("dispatch {} failed: {}", &raddr_cloned, e);
+                    sessions.lock().await.remove(&raddr_cloned);
                     return;
                 }
             };
 
             let (mut target_sock_recv, mut target_sock_send) = socket.split();
 
             // downlink
+            let raddr_downlink = raddr_cloned.clone();
             let downlink_task = async move {
                 let mut buf = vec![0u8; *crate::option::DATAGRAM_BUFFER_SIZE * 1024];
                 loop {
                     match target_sock_recv.recv_from(&mut buf).await {
                         Err(err) => {
                             debug!(
                                 "Failed to receive downlink packets on session {}: {}",
-                                &raddr, err
+                                &raddr_downlink, err
                             );
                             break;
                         }
@@ -208,20 +211,20 @@ impl NatManager {
                             let pkt = UdpPacket::new(
                                 buf[..n].to_vec(),
                                 addr.clone(),
-                                SocksAddr::from(raddr.address),
+                                SocksAddr::from(raddr_downlink.address),
                             );
                             if let Err(err) = client_ch_tx.send(pkt).await {
                                 debug!(
                                     "Failed to send downlink packets on session {} to {}: {}",
-                                    &raddr, &addr, err
+                                    &raddr_downlink, &addr, err
                                 );
                                 break;
                             }
 
                             // activity update
                             {
                                 let mut sessions = sessions.lock().await;
-                                if let Some(sess) = sessions.get_mut(&raddr) {
+                                if let Some(sess) = sessions.get_mut(&raddr_downlink) {
                                     if addr.port() == 53 {
                                         // If the destination port is 53, we assume it's a
                                         // DNS query and set a negative timeout so it will
@@ -237,7 +240,7 @@ impl NatManager {
                         }
                     }
                 }
-                sessions.lock().await.remove(&raddr);
+                sessions.lock().await.remove(&raddr_downlink);
             };
 
             let (downlink_task, downlink_task_handle) = abortable(downlink_task);
@@ -250,6 +253,7 @@ impl NatManager {
             });
 
             // uplink
+            let raddr_uplink = raddr_cloned.clone();
             tokio::spawn(async move {
                 while let Some(pkt) = target_ch_rx.recv().await {
                     trace!(
@@ -260,13 +264,13 @@ impl NatManager {
                     if let Err(e) = target_sock_send.send_to(&pkt.data, &pkt.dst_addr).await {
                         debug!(
                             "Failed to send uplink packets on session {} to {}: {:?}",
-                            &raddr, &pkt.dst_addr, e
+                            &raddr_uplink, &pkt.dst_addr, e
                         );
                         break;
                     }
                 }
                 if let Err(e) = target_sock_send.close().await {
-                    debug!("Failed to close outbound datagram {}: {}", &raddr, e);
+                    debug!("Failed to close outbound datagram {}: {}", &raddr_uplink, e);
                 }
             });
         });

leaf/src/app/router.rs

@@ -8,6 +8,8 @@ use futures::TryFutureExt;
 use maxminddb::geoip2::Country;
 use maxminddb::Mmap;
 use tracing::{debug, warn};
+#[cfg(feature = "rule-process-name")]
+use regex::Regex;
 
 use crate::app::SyncDnsClient;
 use crate::config;
@@ -354,6 +356,42 @@ impl Condition for DomainMatcher {
     }
 }
 
+#[cfg(feature = "rule-process-name")]
+pub struct ProcessNameMatcher {
+    regexes: Vec<Regex>,
+}
+
+#[cfg(feature = "rule-process-name")]
+impl ProcessNameMatcher {
+    pub fn new(patterns: Vec<String>) -> Self {
+        let mut regexes = Vec::new();
+        for pattern in patterns {
+            if let Ok(regex) = Regex::new(&pattern) {
+                regexes.push(regex);
+            } else {
+                warn!("Invalid regex pattern: {}", pattern);
+            }
+        }
+        Self { regexes }
+    }
+}
+
+#[cfg(feature = "rule-process-name")]
+impl Condition for ProcessNameMatcher {
+    fn apply(&self, sess: &Session) -> bool {
+        if let Some(process_name) = sess.process_name.as_ref() {
+            for regex in &self.regexes {
+                if regex.is_match(process_name) {
+                    debug!("Matched process_name={} with regex", process_name);
+                    return true;
+                }
+            }
+        }
+        false
+    }
+}
+
+
 struct ConditionAnd {
     conditions: Vec<Box<dyn Condition>>,
 }
@@ -467,6 +505,11 @@ impl Router {
                 cond_and.add(Box::new(InboundTagMatcher::new(&mut rr.inbound_tags)));
             }
 
+            #[cfg(feature = "rule-process-name")]
+            if !rr.process_names.is_empty() {
+                cond_and.add(Box::new(ProcessNameMatcher::new(rr.process_names.clone())));
+            }
+
             if cond_and.is_empty() {
                 warn!("empty rule at target {}", rr.target_tag);
                 continue;

leaf/src/config/conf/config.rs

@@ -724,7 +724,7 @@ pub fn from_lines(lines: Vec<io::Result<String>>) -> Result<Config> {
 
         match rule.type_field.as_str() {
             "IP-CIDR" | "DOMAIN" | "DOMAIN-SUFFIX" | "DOMAIN-KEYWORD" | "GEOIP" | "EXTERNAL"
-            | "PORT-RANGE" | "NETWORK" | "INBOUND-TAG" => {
+            | "PORT-RANGE" | "NETWORK" | "INBOUND-TAG" | "PROCESS-NAME" => {
                 rule.filter = Some(params[1].to_string());
             }
             _ => {}
@@ -1493,6 +1493,10 @@ pub fn to_internal(conf: &mut Config) -> Result<internal::Config> {
                 "INBOUND-TAG" => {
                     rule.inbound_tags.push(ext_filter);
                 }
+                #[cfg(feature = "rule-process-name")]
+                "PROCESS-NAME" => {
+                    rule.process_names.push(ext_filter);
+                }
                 _ => {}
             }
             rules.push(rule);

leaf/src/config/internal/config.proto

@@ -264,6 +264,7 @@ message Router {
 		repeated string port_ranges = 5;
 		repeated string networks = 6;
 		repeated string inbound_tags = 7;
+		repeated string process_names = 8;
 	}
 
 	repeated Rule rules = 1;

leaf/src/config/internal/config.rs

@@ -4053,6 +4053,8 @@ pub mod router {
         pub networks: ::std::vec::Vec<::std::string::String>,
         // @@protoc_insertion_point(field:Router.Rule.inbound_tags)
         pub inbound_tags: ::std::vec::Vec<::std::string::String>,
+        // @@protoc_insertion_point(field:Router.Rule.process_names)
+        pub process_names: ::std::vec::Vec<::std::string::String>,
         // special fields
         // @@protoc_insertion_point(special_field:Router.Rule.special_fields)
         pub special_fields: ::protobuf::SpecialFields,
@@ -4101,6 +4103,9 @@ pub mod router {
                     58 => {
                         self.inbound_tags.push(is.read_string()?);
                     },
+                    66 => {
+                        self.process_names.push(is.read_string()?);
+                    },
                     tag => {
                         ::protobuf::rt::read_unknown_or_skip_group(tag, is, self.special_fields.mut_unknown_fields())?;
                     },
@@ -4136,6 +4141,9 @@ pub mod router {
             for value in &self.inbound_tags {
                 my_size += ::protobuf::rt::string_size(7, &value);
             };
+            for value in &self.process_names {
+                my_size += ::protobuf::rt::string_size(8, &value);
+            };
             my_size += ::protobuf::rt::unknown_fields_size(self.special_fields.unknown_fields());
             self.special_fields.cached_size().set(my_size as u32);
             my_size
@@ -4163,6 +4171,9 @@ pub mod router {
             for v in &self.inbound_tags {
                 os.write_string(7, &v)?;
             };
+            for v in &self.process_names {
+                os.write_string(8, &v)?;
+            };
             os.write_unknown_fields(self.special_fields.unknown_fields())?;
             ::std::result::Result::Ok(())
         }
@@ -4187,6 +4198,7 @@ pub mod router {
             self.port_ranges.clear();
             self.networks.clear();
             self.inbound_tags.clear();
+            self.process_names.clear();
             self.special_fields.clear();
         }
 
@@ -4199,6 +4211,7 @@ pub mod router {
                 port_ranges: ::std::vec::Vec::new(),
                 networks: ::std::vec::Vec::new(),
                 inbound_tags: ::std::vec::Vec::new(),
+                process_names: ::std::vec::Vec::new(),
                 special_fields: ::protobuf::SpecialFields::new(),
             };
             &instance

leaf/src/config/json/config.rs

@@ -275,6 +275,8 @@ pub struct Rule {
     pub network: Option<Vec<String>>,
     #[serde(rename = "inboundTag")]
     pub inbound_tag: Option<Vec<String>>,
+    #[serde(rename = "processName")]
+    pub process_name: Option<Vec<String>>,
     pub target: String,
 }
 
@@ -1076,6 +1078,12 @@ pub fn to_internal(json: &mut Config) -> Result<internal::Config> {
                         rule.inbound_tags.push(it);
                     }
                 }
+                #[cfg(feature = "rule-process-name")]
+                if let Some(ext_process_names) = ext_rule.process_name.as_mut() {
+                    for process_name in ext_process_names.drain(0..) {
+                        rule.process_names.push(process_name);
+                    }
+                }
                 rules.push(rule);
             }
         }

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

@@ -78,17 +78,18 @@ impl InboundDatagramRecvHalf for DatagramRecvHalf {
         assert!(buf.len() >= payload_size);
         let real_payload = &recv_buf[header_size..header_size + payload_size];
 
-        let local_addr = if let Some(info) = super::UDP_LOCAL_INFO.lock().unwrap().get(&id) {
-            info.local_address.clone()
+        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.
+        // Override with real source address and process name.
         src_addr.address = local_addr;
+        src_addr.process_name = process_name;
 
         if dst_addr.port() == 53 {
             match self.1.generate_fake_response(real_payload).await {