Skip to content

Commit

Permalink
fix(rule): find process by both ip & port
Browse files Browse the repository at this point in the history
  • Loading branch information
@lemos1235
lemos1235 committed Apr 6, 2023
1 parent b0a8657 commit 202c882
Show file tree
Hide file tree
Showing 2 changed files with 49 additions and 78 deletions.
36 changes: 19 additions & 17 deletions leaf/src/app/process_finder.rs
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
use log::debug;
use netstat2::{get_sockets_info, AddressFamilyFlags, ProtocolFlags, ProtocolSocketInfo, TcpState};
use std::net;
use std::net::IpAddr;
use log::debug;
use netstat2::{AddressFamilyFlags, get_sockets_info, ProtocolFlags, ProtocolSocketInfo, TcpState};
use sysinfo::{ProcessExt, System, SystemExt, Pid, PidExt};
use sysinfo::{Pid, PidExt, ProcessExt, System, SystemExt};

#[derive(Debug)]
pub struct PortInfo {
Expand All @@ -26,15 +26,15 @@ impl From<netstat2::SocketInfo> for PortInfo {
ProtocolSocketInfo::Udp(_) => "UDP",
};
let system = System::new_all();
// system.refresh_system();
let pid = socket_info.associated_pids.first().unwrap();
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(),
});
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(),
Expand All @@ -45,32 +45,34 @@ impl From<netstat2::SocketInfo> for PortInfo {
}

pub fn find_process(protocol: &str, ip: IpAddr, port: u16) -> Option<PortInfo> {
debug!("looking up process for {}:{}:{}", protocol, ip, port);
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 start_time = tokio::time::Instant::now();
let sockets = get_sockets_info(af_flags, proto_flags).unwrap_or_default();
let mut ports = sockets
let mut port_infos = sockets
.into_iter()
.filter(|socket_info| match &socket_info.protocol_socket_info {
ProtocolSocketInfo::Tcp(tcp) => tcp.state != TcpState::Closed,
ProtocolSocketInfo::Udp(_) => true,
})
.map(|socket_info| PortInfo::from(socket_info));
let port = ports.find(|p| p.port == port);
if let Some(ref p) = port {
debug!("find process port {:?}", p);
let port_info = port_infos.find(|p| p.address == ip && p.port == port);
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;
return port_info;
}
91 changes: 30 additions & 61 deletions leaf/src/app/router.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,11 +13,7 @@ use crate::app::SyncDnsClient;
use crate::config;
use crate::session::{Network, Session, SocksAddr};

#[cfg(any(
target_os = "windows",
target_os = "macos",
target_os = "linux"
))]
#[cfg(any(target_os = "windows", target_os = "macos", target_os = "linux"))]
use crate::app::process_finder;

pub trait Condition: Send + Sync + Unpin {
Expand Down Expand Up @@ -361,35 +357,26 @@ impl Condition for DomainMatcher {
}
}

#[cfg(any(
target_os = "windows",
target_os = "macos",
target_os = "linux"
))]
#[cfg(any(target_os = "windows", target_os = "macos", target_os = "linux"))]
struct ProcessPidMatcher {
value: String,
}

#[cfg(any(
target_os = "windows",
target_os = "macos",
target_os = "linux"
))]
#[cfg(any(target_os = "windows", target_os = "macos", target_os = "linux"))]
impl ProcessPidMatcher {
fn new(value: String) -> Self {
ProcessPidMatcher { value }
}
}

#[cfg(any(
target_os = "windows",
target_os = "macos",
target_os = "linux"
))]
#[cfg(any(target_os = "windows", target_os = "macos", target_os = "linux"))]
impl Condition for ProcessPidMatcher {
fn apply(&self, sess: &Session) -> bool {
let port_info = process_finder::find_process(&sess.network.to_string(),
sess.source.ip(), sess.source.port());
let port_info = process_finder::find_process(
&sess.network.to_string(),
sess.source.ip(),
sess.source.port(),
);
if let Some(port) = port_info {
if let Some(process_info) = port.process_info {
if process_info.pid.to_string() == self.value {
Expand All @@ -402,39 +389,37 @@ impl Condition for ProcessPidMatcher {
}
}

#[cfg(any(
target_os = "windows",
target_os = "macos",
target_os = "linux"
))]
#[cfg(any(target_os = "windows", target_os = "macos", target_os = "linux"))]
struct ProcessNameMatcher {
value: String,
}

#[cfg(any(
target_os = "windows",
target_os = "macos",
target_os = "linux"
))]
#[cfg(any(target_os = "windows", target_os = "macos", target_os = "linux"))]
impl ProcessNameMatcher {
fn new(value: String) -> Self {
ProcessNameMatcher { value }
}
}

#[cfg(any(
target_os = "windows",
target_os = "macos",
target_os = "linux"
))]
#[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(&sess.network.to_string(),
sess.source.ip(), sess.source.port());
let port_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);
if info
.process_path
.to_lowercase()
.contains(&self.value.to_lowercase())
{
debug!(
"[{}] matches process name [{}]",
info.process_path, &self.value
);
return true;
}
}
Expand All @@ -443,20 +428,12 @@ impl Condition for ProcessNameMatcher {
}
}

#[cfg(any(
target_os = "windows",
target_os = "macos",
target_os = "linux"
))]
#[cfg(any(target_os = "windows", target_os = "macos", target_os = "linux"))]
struct ProcessMatcher {
condition: Box<dyn Condition>,
}

#[cfg(any(
target_os = "windows",
target_os = "macos",
target_os = "linux"
))]
#[cfg(any(target_os = "windows", target_os = "macos", target_os = "linux"))]
impl ProcessMatcher {
fn new(processes: &mut Vec<config::router::rule::Process>) -> Self {
let mut cond_or = ConditionOr::new();
Expand All @@ -477,11 +454,7 @@ impl ProcessMatcher {
}
}

#[cfg(any(
target_os = "windows",
target_os = "macos",
target_os = "linux"
))]
#[cfg(any(target_os = "windows", target_os = "macos", target_os = "linux"))]
impl Condition for ProcessMatcher {
fn apply(&self, sess: &Session) -> bool {
self.condition.apply(sess)
Expand Down Expand Up @@ -601,11 +574,7 @@ impl Router {
cond_and.add(Box::new(InboundTagMatcher::new(&mut rr.inbound_tags)));
}

#[cfg(any(
target_os = "windows",
target_os = "macos",
target_os = "linux"
))]
#[cfg(any(target_os = "windows", target_os = "macos", target_os = "linux"))]
if rr.processes.len() > 0 {
cond_and.add(Box::new(ProcessMatcher::new(&mut rr.processes)));
}
Expand Down

0 comments on commit 202c882

Please sign in to comment.