optimize socket info cache

leaf/src/app/process_finder.rs

@@ -2,25 +2,15 @@ use std::net::IpAddr;
 use std::sync::{Arc, Mutex};
 
 use lazy_static::lazy_static;
-use log::debug;
-use netstat2::{
-    get_sockets_info, AddressFamilyFlags, ProtocolFlags, ProtocolSocketInfo, SocketInfo,
-};
-use sysinfo::{Pid, PidExt, ProcessExt, ProcessRefreshKind, System, SystemExt};
+use netstat2::{get_sockets_info, AddressFamilyFlags, ProtocolFlags, SocketInfo};
+use sysinfo::{Pid, ProcessExt, ProcessRefreshKind, System, SystemExt};
 
 lazy_static! {
-    pub static ref SYSTEM: Arc<Mutex<System>> = {
+    static ref SYSTEM: Arc<Mutex<System>> = {
         let system = System::new_all();
         Arc::new(Mutex::new(system))
     };
-}
-
-#[derive(Debug)]
-pub struct PortInfo {
-    pub address: IpAddr,
-    pub port: u16,
-    pub protocol: String,
-    pub process_info: Option<ProcessInfo>,
+    static ref CACHE: Mutex<Option<Vec<SocketInfo>>> = Mutex::new(None);
 }
 
 #[derive(Debug)]
@@ -30,72 +20,59 @@ pub struct ProcessInfo {
     pub process_path: String,
 }
 
-impl From<SocketInfo> for PortInfo {
-    fn from(socket_info: SocketInfo) -> Self {
-        let protocol = match socket_info.protocol_socket_info {
-            ProtocolSocketInfo::Tcp(_) => "TCP",
-            ProtocolSocketInfo::Udp(_) => "UDP",
-        };
-        let pid = socket_info.associated_pids.first().unwrap();
-        let mut system = SYSTEM.lock().unwrap();
+impl ProcessInfo {
+    fn from_socket_info(socket_info: SocketInfo) -> Option<Self> {
+        let pid = socket_info.associated_pids.first()?.to_owned();
+        let mut system = SYSTEM.lock().ok()?;
         system.refresh_processes_specifics(ProcessRefreshKind::default());
-        let process_info =
-            system
-                .process(Pid::from(pid.to_owned() as usize))
-                .map(|p| ProcessInfo {
-                    name: p.name().to_owned(),
-                    pid: p.pid().as_u32(),
-                    process_path: p.exe().to_string_lossy().to_string(),
-                });
-        Self {
-            address: socket_info.local_addr(),
-            port: socket_info.local_port(),
-            protocol: protocol.to_string(),
-            process_info,
-        }
+        let process = system.process(Pid::from(pid.to_owned() as usize))?;
+        let name = process.name().to_owned();
+        let process_path = process.exe().to_string_lossy().to_string();
+        Some(ProcessInfo {
+            name,
+            pid,
+            process_path,
+        })
     }
 }
 
 fn get_socket_info(protocol: &str, ip: &IpAddr, port: u16) -> Option<SocketInfo> {
-    let mut af_flags: AddressFamilyFlags = AddressFamilyFlags::from_bits(0).unwrap();
-    if ip.is_ipv6() {
-        af_flags |= AddressFamilyFlags::IPV6;
-    }
-    if ip.is_ipv4() {
-        af_flags |= AddressFamilyFlags::IPV4;
-    }
-    let mut proto_flags: ProtocolFlags = ProtocolFlags::from_bits(0).unwrap();
-    if protocol == "udp" {
-        proto_flags |= ProtocolFlags::UDP;
-    }
-    if protocol == "tcp" {
-        proto_flags |= ProtocolFlags::TCP;
-    }
-    let sockets = get_sockets_info(af_flags, proto_flags).unwrap_or_default();
-    let socket_info = sockets
+    let af_flags = match ip {
+        IpAddr::V4(_) => AddressFamilyFlags::IPV4,
+        IpAddr::V6(_) => AddressFamilyFlags::IPV6,
+    };
+
+    let proto_flags = match protocol {
+        "udp" => ProtocolFlags::UDP,
+        "tcp" => ProtocolFlags::TCP,
+        _ => ProtocolFlags::from_bits(0).unwrap(),
+    };
+
+    let mut cache = CACHE.lock().unwrap();
+    let mut sockets = cache.clone().unwrap_or_else(|| {
+        let new_sockets = get_sockets_info(af_flags, proto_flags).unwrap_or_default();
+        *cache = Some(new_sockets.clone());
+        new_sockets
+    });
+
+    sockets
         .into_iter()
-        .find(|p| p.local_addr() == ip.to_owned() && p.local_port() == port);
-    return socket_info;
+        .find(|p| p.local_addr() == ip.to_owned() && p.local_port() == port)
+        .or_else(|| {
+            let new_sockets = get_sockets_info(af_flags, proto_flags).unwrap_or_default();
+            *cache = Some(new_sockets.clone());
+            new_sockets.into_iter().find(|p| p.local_addr() == ip.to_owned() && p.local_port() == port)
+        })
 }
 
-pub fn find_process_id(protocol: &str, ip: &IpAddr, port: u16) -> Option<u32> {
-    let start_time = tokio::time::Instant::now();
+pub fn find_pid(protocol: &str, ip: &IpAddr, port: u16) -> Option<u32> {
     let socket_info = get_socket_info(protocol, ip, port);
-    let pid = socket_info.map(|s| s.associated_pids.first().unwrap().to_owned());
-    if let Some(ref pid) = pid {
-        let elapsed = tokio::time::Instant::now().duration_since(start_time);
-        debug!("found process id [{}ms] {:?}", elapsed.as_millis(), pid);
-    }
-    pid
+    socket_info.map(|s| s.associated_pids.first().unwrap().to_owned())
 }
 
-pub fn find_process(protocol: &str, ip: &IpAddr, port: u16) -> Option<PortInfo> {
-    let start_time = tokio::time::Instant::now();
+pub fn find_process(protocol: &str, ip: &IpAddr, port: u16) -> Option<ProcessInfo> {
     let socket_info = get_socket_info(protocol, ip, port);
-    let port_info = socket_info.map(|socket_info| PortInfo::from(socket_info));
-    if let Some(ref p) = port_info {
-        let elapsed = tokio::time::Instant::now().duration_since(start_time);
-        debug!("found process [{}ms] {:?}", elapsed.as_millis(), p);
-    }
-    return port_info;
+    socket_info
+        .map(|s| ProcessInfo::from_socket_info(s))
+        .flatten()
 }

leaf/src/app/router.rs

@@ -371,7 +371,7 @@ impl ProcessPidMatcher {
 #[cfg(any(target_os = "windows", target_os = "macos", target_os = "linux"))]
 impl Condition for ProcessPidMatcher {
     fn apply(&self, sess: &Session) -> bool {
-        let process_id = process_finder::find_process_id(
+        let process_id = process_finder::find_pid(
             &sess.network.to_string(),
             &sess.source.ip(),
             sess.source.port(),
@@ -401,24 +401,22 @@ impl ProcessNameMatcher {
 #[cfg(any(target_os = "windows", target_os = "macos", target_os = "linux"))]
 impl Condition for ProcessNameMatcher {
     fn apply(&self, sess: &Session) -> bool {
-        let port_info = process_finder::find_process(
+        let process_info = process_finder::find_process(
             &sess.network.to_string(),
             &sess.source.ip(),
             sess.source.port(),
         );
-        if let Some(port) = port_info {
-            if let Some(info) = port.process_info {
-                if info
-                    .process_path
-                    .to_lowercase()
-                    .contains(&self.value.to_lowercase())
-                {
-                    debug!(
-                        "[{}] matches process name [{}]",
-                        info.process_path, &self.value
-                    );
-                    return true;
-                }
+        if let Some(process) = process_info {
+            if process
+                .process_path
+                .to_lowercase()
+                .contains(&self.value.to_lowercase())
+            {
+                debug!(
+                    "[{}] matches process name [{}]",
+                    process.process_path, &self.value
+                );
+                return true;
             }
         }
         false

leaf/src/config/internal/config.rs

@@ -4282,8 +4282,8 @@ pub mod router {
             }
         }
 
-        #[derive(PartialEq,Clone,Default,Debug)]
         // @@protoc_insertion_point(message:Router.Rule.Process)
+        #[derive(PartialEq,Clone,Default,Debug)]
         pub struct Process {
             // message fields
             // @@protoc_insertion_point(field:Router.Rule.Process.type)
@@ -4411,6 +4411,14 @@ pub mod router {
                     }
                 }
 
+                fn from_str(str: &str) -> ::std::option::Option<Type> {
+                    match str {
+                        "PID" => ::std::option::Option::Some(Type::PID),
+                        "NAME" => ::std::option::Option::Some(Type::NAME),
+                        _ => ::std::option::Option::None
+                    }
+                }
+
                 const VALUES: &'static [Type] = &[
                     Type::PID,
                     Type::NAME,