From 6436874c5a0d4a695dbfe1051a27c4365b76862d Mon Sep 17 00:00:00 2001 From: ThetaSinner Date: Wed, 6 Nov 2024 00:08:40 +0000 Subject: [PATCH 1/2] Add macos support and extend test coverage --- Cargo.toml | 24 ++++ sample/port-binder-v6/main.rs | 7 + sample/udp-port-binder-v6/main.rs | 8 ++ sample/udp-port-binder/main.rs | 8 ++ src/common.rs | 14 +- src/error.rs | 2 +- src/port_query.rs | 231 +++++++++++++++++++++++++++++- tests/lib_test.rs | 118 +++++++++++++-- 8 files changed, 395 insertions(+), 17 deletions(-) create mode 100644 sample/port-binder-v6/main.rs create mode 100644 sample/udp-port-binder-v6/main.rs create mode 100644 sample/udp-port-binder/main.rs diff --git a/Cargo.toml b/Cargo.toml index 6fa4393..8620dde 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,6 +14,14 @@ doc = false doctest = false bench = false +[[bin]] +name = "port-binder-v6" +path = "./sample/port-binder-v6/main.rs" +test = false +doc = false +doctest = false +bench = false + [[bin]] name = "proc-runner" path = "./sample/proc-runner/main.rs" @@ -22,6 +30,22 @@ doc = false doctest = false bench = false +[[bin]] +name = "udp-port-binder" +path = "./sample/udp-port-binder/main.rs" +test = false +doc = false +doctest = false +bench = false + +[[bin]] +name = "udp-port-binder-v6" +path = "./sample/udp-port-binder-v6/main.rs" +test = false +doc = false +doctest = false +bench = false + [[bin]] name = "waiter" path = "./sample/waiter/main.rs" diff --git a/sample/port-binder-v6/main.rs b/sample/port-binder-v6/main.rs new file mode 100644 index 0000000..159db90 --- /dev/null +++ b/sample/port-binder-v6/main.rs @@ -0,0 +1,7 @@ +use std::net::TcpListener; + +fn main() { + let listener = TcpListener::bind("[::1]:0").unwrap(); + listener.accept().unwrap(); + println!("Listener finished"); +} diff --git a/sample/udp-port-binder-v6/main.rs b/sample/udp-port-binder-v6/main.rs new file mode 100644 index 0000000..05ad1f8 --- /dev/null +++ b/sample/udp-port-binder-v6/main.rs @@ -0,0 +1,8 @@ +use std::net::UdpSocket; + +fn main() { + let listener = UdpSocket::bind("[::1]:0").unwrap(); + let mut buf = [0; 10]; + listener.recv(&mut buf).unwrap(); + println!("Done receiving on UDP socket"); +} diff --git a/sample/udp-port-binder/main.rs b/sample/udp-port-binder/main.rs new file mode 100644 index 0000000..188141b --- /dev/null +++ b/sample/udp-port-binder/main.rs @@ -0,0 +1,8 @@ +use std::net::UdpSocket; + +fn main() { + let listener = UdpSocket::bind("127.0.0.1:0").unwrap(); + let mut buf = [0; 10]; + listener.recv(&mut buf).unwrap(); + println!("Done receiving on UDP socket"); +} diff --git a/src/common.rs b/src/common.rs index a3f9c72..53df6f8 100644 --- a/src/common.rs +++ b/src/common.rs @@ -1,9 +1,19 @@ -#[cfg(any(target_os = "linux", target_os = "windows", feature = "proc"))] +#[cfg(any( + target_os = "linux", + target_os = "windows", + target_os = "macos", + feature = "proc" +))] pub(crate) trait MaybeHasPid { fn get_pid(&self) -> Option; } -#[cfg(any(target_os = "linux", target_os = "windows", feature = "proc"))] +#[cfg(any( + target_os = "linux", + target_os = "windows", + target_os = "macos", + feature = "proc" +))] pub(crate) fn resolve_pid(maybe_has_pid: &dyn MaybeHasPid) -> crate::ProcCtlResult { match &maybe_has_pid.get_pid() { Some(pid) => Ok(*pid), diff --git a/src/error.rs b/src/error.rs index fe73780..c534e27 100644 --- a/src/error.rs +++ b/src/error.rs @@ -13,7 +13,7 @@ pub enum ProcCtlError { ProcessError(#[from] procfs::ProcError), /// An error occurred while searching process information - #[cfg(target_os = "windows")] + #[cfg(any(target_os = "windows", target_os = "macos"))] #[error("process error")] ProcessError(String), diff --git a/src/port_query.rs b/src/port_query.rs index fda02e5..056e2e0 100644 --- a/src/port_query.rs +++ b/src/port_query.rs @@ -77,9 +77,9 @@ impl PortQuery { /// Execute the query pub fn execute(&self) -> ProcCtlResult> { - #[cfg(any(target_os = "linux", target_os = "windows"))] + #[cfg(any(target_os = "linux", target_os = "windows", target_os = "macos"))] let ports = list_ports_for_pid(self, crate::common::resolve_pid(self)?)?; - #[cfg(not(any(target_os = "linux", target_os = "windows")))] + #[cfg(not(any(target_os = "linux", target_os = "windows", target_os = "macos")))] let ports = Vec::with_capacity(0); if let Some(num) = &self.min_num_ports { @@ -324,7 +324,232 @@ fn load_udp_table( )) } -#[cfg(any(target_os = "linux", target_os = "windows", feature = "proc"))] +#[cfg(target_os = "macos")] +fn list_ports_for_pid(query: &PortQuery, pid: Pid) -> ProcCtlResult> { + let mut out = Vec::new(); + + if query.ipv4_addresses { + if query.tcp_addresses { + match std::process::Command::new("lsof") + .arg("-a") + .arg("-iTCP") + .arg("-i4") + .arg("-sTCP:LISTEN") + .arg("-nP") + .arg("-F0pn") + .output() + { + Ok(output) => out.extend( + find_ports_v4(output.stdout.clone(), pid) + .into_iter() + .map(ProtocolPort::Tcp), + ), + Err(e) => return Err(ProcCtlError::ProcessError(e.to_string())), + } + } + if query.udp_addresses { + match std::process::Command::new("lsof") + .arg("-a") + .arg("-iUDP") + .arg("-i4") + .arg("-nP") + .arg("-F0pn") + .output() + { + Ok(output) => out.extend( + find_ports_v4(output.stdout.clone(), pid) + .into_iter() + .map(ProtocolPort::Udp), + ), + Err(e) => return Err(ProcCtlError::ProcessError(e.to_string())), + } + } + } + if query.ipv6_addresses { + if query.tcp_addresses { + match std::process::Command::new("lsof") + .arg("-a") + .arg("-iTCP") + .arg("-i6") + .arg("-sTCP:LISTEN") + .arg("-nP") + .arg("-F0pn") + .output() + { + Ok(output) => out.extend( + find_ports_v6(output.stdout.clone(), pid) + .into_iter() + .map(ProtocolPort::Tcp), + ), + Err(e) => return Err(ProcCtlError::ProcessError(e.to_string())), + } + } + if query.udp_addresses { + match std::process::Command::new("lsof") + .arg("-a") + .arg("-iUDP") + .arg("-i6") + .arg("-nP") + .arg("-F0pn") + .output() + { + Ok(output) => out.extend( + find_ports_v6(output.stdout.clone(), pid) + .into_iter() + .map(ProtocolPort::Udp), + ), + Err(e) => return Err(ProcCtlError::ProcessError(e.to_string())), + } + } + } + + Ok(out) +} + +#[cfg(target_os = "macos")] +fn find_ports_v4(output: Vec, find_pid: Pid) -> Vec { + let mut out = Vec::new(); + + let mut index = 0; + let len = output.len(); + while index < len { + if output[index] != b'p' { + break; + } + index += 1; + + let start_pid = index; + while index < len && output[index] != 0 { + index += 1; + } + + let Some(pid) = String::from_utf8_lossy(&output[start_pid..index]) + .parse::() + .ok() + else { + break; + }; + index += 1; // 0 + index += 1; // NL + + loop { + if pid == find_pid && index < len && output[index] == b'n' { + while index < len && output[index] != b':' { + index += 1; + } + index += 1; // : + + let start_port = index; + while index < len && output[index] != 0 { + index += 1; + } + + if index >= len { + break; + } + + if let Ok(port) = String::from_utf8_lossy(&output[start_port..index]).parse::() + { + out.push(port); + }; + index += 1; // 0 + } else { + while index < len && output[index] != 0 { + index += 1; + } + index += 1; // 0 + } + + if index < len && output[index] == 10 { + // NL + index += 1; + } + + if index >= len || output[index] == b'p' { + break; + } + } + } + + out +} + +#[cfg(target_os = "macos")] +fn find_ports_v6(output: Vec, find_pid: Pid) -> Vec { + let mut out = Vec::new(); + + let mut index = 0; + let len = output.len(); + while index < len { + if output[index] != b'p' { + break; + } + index += 1; + + let start_pid = index; + while index < len && output[index] != 0 { + index += 1; + } + + let Ok(pid) = String::from_utf8_lossy(&output[start_pid..index]).parse::() else { + break; + }; + index += 1; // 0 + index += 1; // NL + + loop { + if pid == find_pid && index < len && output[index] == b'n' { + while index < len && output[index] != b']' { + index += 1; + } + index += 1; // ] + + if index < len && output[index] != b':' { + break; + } + index += 1; + + let start_port = index; + while index < len && output[index] != 0 { + index += 1; + } + + if index >= len { + break; + } + + if let Ok(port) = String::from_utf8_lossy(&output[start_port..index]).parse::() + { + out.push(port); + }; + index += 1; // 0 + } else { + while index < len && output[index] != 0 { + index += 1; + } + index += 1; // 0 + } + + if index < len && output[index] == 10 { + // NL + index += 1; + } + + if index >= len || output[index] == b'p' { + break; + } + } + } + + out +} + +#[cfg(any( + target_os = "linux", + target_os = "windows", + target_os = "macos", + feature = "proc" +))] impl crate::common::MaybeHasPid for PortQuery { fn get_pid(&self) -> Option { self.process_id diff --git a/tests/lib_test.rs b/tests/lib_test.rs index b764114..451fb31 100644 --- a/tests/lib_test.rs +++ b/tests/lib_test.rs @@ -1,4 +1,9 @@ -#[cfg(any(feature = "proc", target_os = "linux", target_os = "windows"))] +#[cfg(any( + feature = "proc", + target_os = "linux", + target_os = "windows", + target_os = "macos" +))] fn create_command_for_sample(name: &str) -> std::process::Command { let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR")) .join("target") @@ -17,24 +22,44 @@ fn create_command_for_sample(name: &str) -> std::process::Command { std::process::Command::new(path) } -#[cfg(any(feature = "proc", target_os = "linux", target_os = "windows"))] +#[cfg(any( + feature = "proc", + target_os = "linux", + target_os = "windows", + target_os = "macos" +))] struct DropChild(std::process::Child); -#[cfg(any(feature = "proc", target_os = "linux", target_os = "windows"))] +#[cfg(any( + feature = "proc", + target_os = "linux", + target_os = "windows", + target_os = "macos" +))] impl DropChild { fn spawn(mut cmd: std::process::Command) -> Self { DropChild(cmd.spawn().expect("Failed to spawn child process")) } } -#[cfg(any(feature = "proc", target_os = "linux", target_os = "windows"))] +#[cfg(any( + feature = "proc", + target_os = "linux", + target_os = "windows", + target_os = "macos" +))] impl Drop for DropChild { fn drop(&mut self) { self.0.kill().expect("Failed to kill child process"); } } -#[cfg(any(feature = "proc", target_os = "linux", target_os = "windows"))] +#[cfg(any( + feature = "proc", + target_os = "linux", + target_os = "windows", + target_os = "macos" +))] impl std::ops::Deref for DropChild { type Target = std::process::Child; @@ -43,16 +68,21 @@ impl std::ops::Deref for DropChild { } } -#[cfg(any(feature = "proc", target_os = "linux", target_os = "windows"))] +#[cfg(any( + feature = "proc", + target_os = "linux", + target_os = "windows", + target_os = "macos" +))] impl std::ops::DerefMut for DropChild { fn deref_mut(&mut self) -> &mut Self::Target { &mut self.0 } } -#[cfg(any(target_os = "linux", target_os = "windows"))] +#[cfg(any(target_os = "linux", target_os = "windows", target_os = "macos"))] #[test] -fn port_query() { +fn tcp_port_query_v4() { use retry::delay::Fixed; let binder = create_command_for_sample("port-binder"); @@ -71,7 +101,70 @@ fn port_query() { assert_eq!(1, ports.len()); } -#[cfg(any(target_os = "linux", target_os = "windows"))] +#[cfg(any(target_os = "linux", target_os = "windows", target_os = "macos"))] +#[test] +fn tcp_port_query_v6() { + use retry::delay::Fixed; + + let binder = create_command_for_sample("port-binder-v6"); + let mut handle = DropChild::spawn(binder); + + let query = proc_ctl::PortQuery::new() + .tcp_only() + .ip_v6_only() + .process_id(handle.id()) + .expect_min_num_ports(1); + + let ports = retry::retry(Fixed::from_millis(100).take(10), move || query.execute()).unwrap(); + + handle.kill().unwrap(); + + assert_eq!(1, ports.len()); +} + +#[cfg(any(target_os = "linux", target_os = "windows", target_os = "macos"))] +#[test] +fn udp_port_query_v4() { + use retry::delay::Fixed; + + let binder = create_command_for_sample("udp-port-binder"); + let mut handle = DropChild::spawn(binder); + + let query = proc_ctl::PortQuery::new() + .udp_only() + .ip_v4_only() + .process_id(handle.id()) + .expect_min_num_ports(1); + + let ports = retry::retry(Fixed::from_millis(100).take(10), move || query.execute()).unwrap(); + + handle.kill().unwrap(); + + assert_eq!(1, ports.len()); +} + +#[cfg(any(target_os = "linux", target_os = "windows", target_os = "macos"))] +#[test] +fn udp_port_query_v6() { + use retry::delay::Fixed; + + let binder = create_command_for_sample("udp-port-binder-v6"); + let mut handle = DropChild::spawn(binder); + + let query = proc_ctl::PortQuery::new() + .udp_only() + .ip_v6_only() + .process_id(handle.id()) + .expect_min_num_ports(1); + + let ports = retry::retry(Fixed::from_millis(100).take(10), move || query.execute()).unwrap(); + + handle.kill().unwrap(); + + assert_eq!(1, ports.len()); +} + +#[cfg(any(target_os = "linux", target_os = "windows", target_os = "macos"))] #[test] fn port_query_which_expects_too_many_ports() { use retry::delay::Fixed; @@ -95,7 +188,7 @@ fn port_query_which_expects_too_many_ports() { #[cfg(all( feature = "resilience", - any(target_os = "linux", target_os = "windows") + any(target_os = "linux", target_os = "windows", target_os = "macos") ))] #[test] fn port_query_with_sync_retry() { @@ -119,7 +212,10 @@ fn port_query_with_sync_retry() { assert_eq!(1, ports.len()); } -#[cfg(all(feature = "async", any(target_os = "linux", target_os = "windows")))] +#[cfg(all( + feature = "async", + any(target_os = "linux", target_os = "windows", target_os = "macos") +))] #[tokio::test] async fn port_query_with_async_retry() { use std::time::Duration; From 4ceedd0f6730a8a167768653556bc1f4b77126e6 Mon Sep 17 00:00:00 2001 From: ThetaSinner Date: Wed, 6 Nov 2024 00:34:35 +0000 Subject: [PATCH 2/2] Fix behavior on Linux --- src/port_query.rs | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/src/port_query.rs b/src/port_query.rs index 056e2e0..81671b1 100644 --- a/src/port_query.rs +++ b/src/port_query.rs @@ -146,10 +146,11 @@ fn list_ports_for_pid(query: &PortQuery, pid: Pid) -> ProcCtlResult ProcCtlResult