Skip to content

Commit

Permalink
optimize socket info cache
Browse files Browse the repository at this point in the history
  • Loading branch information
@lemos1235
lemos1235 committed Mar 23, 2024
1 parent d274bdb commit 9a8d15f
Show file tree
Hide file tree
Showing 3 changed files with 70 additions and 89 deletions.
121 changes: 48 additions & 73 deletions leaf/src/app/process_finder.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand All @@ -30,72 +20,57 @@ 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
.into_iter()
.find(|p| p.local_addr() == ip.to_owned() && p.local_port() == port);
return socket_info;
pub fn find_pid(protocol: &str, ip: &IpAddr, port: u16) -> Option<u32> {
let socket_info = get_socket_info(protocol, ip, port);
socket_info.map(|s| s.associated_pids.first().unwrap().to_owned())
}

pub fn find_process_id(protocol: &str, ip: &IpAddr, port: u16) -> Option<u32> {
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 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| ProcessInfo::from_socket_info(s))
.flatten()
}

pub fn find_process(protocol: &str, ip: &IpAddr, port: u16) -> Option<PortInfo> {
let start_time = tokio::time::Instant::now();
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;
fn get_socket_info(protocol: &str, ip: &IpAddr, port: u16) -> Option<SocketInfo> {
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 sockets = if let Some(sockets) = cache.as_ref() {
sockets.clone()
} else {
let sockets = get_sockets_info(af_flags, proto_flags).unwrap_or_default();
*cache = Some(sockets.clone());
sockets
};

let socket_info = sockets
.into_iter()
.find(|p| p.local_addr() == ip.to_owned() && p.local_port() == port);
return socket_info;
}
28 changes: 13 additions & 15 deletions leaf/src/app/router.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down Expand Up @@ -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
Expand Down
10 changes: 9 additions & 1 deletion leaf/src/config/internal/config.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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,
Expand Down

0 comments on commit 9a8d15f

Please sign in to comment.