more refactoring and attempt to simplify get proc name on windows

GyulyVGC · Mar 13, 2024 · c5cdade · c5cdade
1 parent 418fb16
commit c5cdade
Show file tree

Hide file tree

Showing 6 changed files with 25 additions and 164 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -5,10 +5,7 @@ edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
-[dependencies]
-
-#[target.'cfg(target_os = "windows")'.dependencies]
-netstat2 = "0.9.1"
+[target.'cfg(target_os = "windows")'.dependencies]
 windows = { version = "0.54.0", features = ["Win32_Foundation", "Win32_System_Diagnostics_ToolHelp"] }
 
 [target.'cfg(target_os = "macos")'.dependencies]

diff --git a/src/lib.rs b/src/lib.rs
@@ -3,8 +3,6 @@ use std::net::SocketAddr;
 
 pub use platform::get_all;
 
-// use netstat2::{get_sockets_info, AddressFamilyFlags, ProtocolFlags, ProtocolSocketInfo, TcpState};
-
 mod platform;
 
 type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;
@@ -13,32 +11,17 @@ type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;
 #[derive(Eq, PartialEq, Hash, Debug)]
 pub struct Listener {
     /// The process ID of the listener process
-    pid: u32,
+    pub pid: u32,
     /// The name of the listener process
-    pname: String,
+    pub pname: String,
     /// The local socket this listener is listening on
-    socket: SocketAddr,
+    pub socket: SocketAddr,
 }
 
 impl Listener {
     fn new(pid: u32, pname: String, socket: SocketAddr) -> Self {
         Self { pid, pname, socket }
     }
-
-    #[must_use]
-    pub fn pid(&self) -> u32 {
-        self.pid
-    }
-
-    #[must_use]
-    pub fn pname(&self) -> &str {
-        &self.pname
-    }
-
-    #[must_use]
-    pub fn socket(&self) -> &SocketAddr {
-        &self.socket
-    }
 }
 
 impl Display for Listener {
@@ -50,66 +33,3 @@ impl Display for Listener {
         )
     }
 }
-
-// #[must_use]
-// pub fn get_all() -> Vec<Listener> {
-//     get_with_filters(
-//         AddressFamilyFlags::IPV4 | AddressFamilyFlags::IPV6,
-//         ProtocolFlags::TCP | ProtocolFlags::UDP,
-//         None,
-//         None,
-//     )
-// }
-//
-// #[must_use]
-// pub fn get_for_nullnet(ip_addr: IpAddr) -> HashSet<String> {
-//     get_with_filters(
-//         AddressFamilyFlags::IPV4,
-//         ProtocolFlags::TCP,
-//         Some(TcpState::Listen),
-//         Some(ip_addr),
-//     )
-//     .iter()
-//     .map(|l| l.pname.clone())
-//     .collect()
-// }
-//
-// fn get_with_filters(
-//     af_flags: AddressFamilyFlags,
-//     proto_flags: ProtocolFlags,
-//     tcp_state: Option<TcpState>,
-//     ip_addr: Option<IpAddr>,
-// ) -> Vec<Listener> {
-//     let mut listeners = Vec::new();
-//
-//     let sockets_info = get_sockets_info(af_flags, proto_flags).unwrap_or_default();
-//
-//     let mut add_listeners = |pids: Vec<u32>, ip: IpAddr, port: u16| {
-//         for pid in pids {
-//             if let Some(pname) = get_name_from_pid(pid) {
-//                 if ip.is_unspecified() || ip_addr.is_none() || ip_addr.unwrap() == ip {
-//                     listeners.push(Listener {
-//                         pid,
-//                         pname,
-//                         socket: SocketAddr::new(ip, port),
-//                     });
-//                 }
-//             }
-//         }
-//     };
-//
-//     for si in sockets_info {
-//         match si.protocol_socket_info {
-//             ProtocolSocketInfo::Tcp(tcp_si) => {
-//                 if tcp_state.is_none() || tcp_si.state == tcp_state.unwrap() {
-//                     add_listeners(si.associated_pids, tcp_si.local_addr, tcp_si.local_port);
-//                 }
-//             }
-//             ProtocolSocketInfo::Udp(udp_si) => {
-//                 add_listeners(si.associated_pids, udp_si.local_addr, udp_si.local_port);
-//             }
-//         }
-//     }
-//
-//     listeners
-// }
diff --git a/src/main.rs b/src/main.rs
@@ -1,6 +1,5 @@
 fn main() {
-    let listeners = listeners::get_all().unwrap();
-    for listener in listeners {
+    for listener in listeners::get_all().iter().flatten() {
         println!("{listener}");
     }
 }
diff --git a/src/platform/macos/mod.rs b/src/platform/macos/mod.rs
@@ -22,9 +22,10 @@ pub fn get_all() -> crate::Result<HashSet<Listener>> {
     for pid in ProcPid::get_all().into_iter().flatten() {
         for fd in SocketFd::get_all_of_pid(pid).iter().flatten() {
             if let Ok(tcp_listener) = TcpListener::from_pid_fd(pid, fd) {
-                let ProcName(name) = ProcName::from_pid(pid).unwrap_or_default();
-                let listener = Listener::new(pid.as_u_32()?, name, tcp_listener.socket_addr());
-                listeners.insert(listener);
+                if let Ok(ProcName(name)) = ProcName::from_pid(pid) {
+                    let listener = Listener::new(pid.as_u_32()?, name, tcp_listener.socket_addr());
+                    listeners.insert(listener);
+                }
             }
         }
     }

diff --git a/src/platform/windows/tcp_listener.rs b/src/platform/windows/tcp_listener.rs
@@ -2,7 +2,13 @@ use crate::platform::windows::socket_table::SocketTable;
 use crate::platform::windows::tcp6_table::Tcp6Table;
 use crate::platform::windows::tcp_table::TcpTable;
 use crate::Listener;
+use std::mem::size_of;
+use std::mem::zeroed;
 use std::net::{IpAddr, SocketAddr};
+use windows::Win32::Foundation::CloseHandle;
+use windows::Win32::System::Diagnostics::ToolHelp::{
+    CreateToolhelp32Snapshot, Process32First, Process32Next, PROCESSENTRY32, TH32CS_SNAPPROCESS,
+};
 
 #[derive(Debug)]
 pub(super) struct TcpListener {
@@ -40,17 +46,9 @@ impl TcpListener {
     }
 
     fn pname(&self) -> Option<String> {
-        use std::mem::size_of;
-        use std::mem::zeroed;
-        use windows::Win32::Foundation::CloseHandle;
-        use windows::Win32::System::Diagnostics::ToolHelp::{
-            CreateToolhelp32Snapshot, Process32First, Process32Next, PROCESSENTRY32,
-            TH32CS_SNAPPROCESS,
-        };
-
         let pid = self.pid;
 
-        let h = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0).unwrap() };
+        let h = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0)? };
 
         let mut process = unsafe { zeroed::<PROCESSENTRY32>() };
         process.dwSize = size_of::<PROCESSENTRY32>() as u32;
@@ -68,16 +66,14 @@ impl TcpListener {
             }
         }
 
-        unsafe { CloseHandle(h).unwrap() };
+        unsafe { CloseHandle(h)? };
 
         let name = process.szExeFile;
-        let mut temp: Vec<u8> = vec![];
-        let len = name.iter().position(|&x| x == 0).unwrap();
+        let len = name.iter().position(|&x| x == 0)?;
 
-        for i in name.iter() {
-            temp.push(*i as u8);
-        }
-        Some(String::from_utf8(temp[0..len].to_vec()).unwrap_or_default())
+        Some(String::from_utf8(
+            name[0..len].iter().map(|e| *e as u8).collect(),
+        )?)
     }
 
     pub(super) fn to_listener(&self) -> Option<Listener> {