Support for TCP and UDP

.github/workflows/ci.yml

@@ -7,6 +7,7 @@ on:
   pull_request:
     branches:
       - '*'
+  workflow_dispatch:
 
 env:
   CARGO_TERM_COLOR: always

src/lib.rs

@@ -15,6 +15,7 @@ pub struct Listener {
     pub process: Process,
     /// The TCP socket this listener is listening on.
     pub socket: SocketAddr,
+    pub protocol: Protocol,
 }
 
 /// A process, characterized by its PID and name.
@@ -26,6 +27,11 @@ pub struct Process {
     pub name: String,
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum Protocol {
+    TCP,
+    UDP,
+}
 /// Returns all the [Listener]s.
 ///
 /// # Errors
@@ -144,9 +150,13 @@ pub fn get_ports_by_process_name(name: &str) -> Result<HashSet<u16>> {
 }
 
 impl Listener {
-    fn new(pid: u32, name: String, socket: SocketAddr) -> Self {
+    fn new(pid: u32, name: String, socket: SocketAddr, protocol: Protocol) -> Self {
         let process = Process::new(pid, name);
-        Self { process, socket }
+        Self {
+            process,
+            socket,
+            protocol,
+        }
     }
 }
 
@@ -158,9 +168,13 @@ impl Process {
 
 impl Display for Listener {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        let Listener { process, socket } = self;
+        let Listener {
+            process,
+            socket,
+            protocol,
+        } = self;
         let process = process.to_string();
-        write!(f, "{process:<55} Socket: {socket}",)
+        write!(f, "{process:<55} Socket: {socket:<30} Protocol: {protocol}",)
     }
 }
 
@@ -171,22 +185,32 @@ impl Display for Process {
     }
 }
 
+impl Display for Protocol {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match *self {
+            Protocol::TCP => write!(f, "TCP"),
+            Protocol::UDP => write!(f, "UDP"),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
 
-    use crate::{Listener, Process};
+    use crate::{Listener, Process, Protocol};
 
     #[test]
     fn test_v4_listener_to_string() {
         let listener = Listener::new(
             455,
             "rapportd".to_string(),
             SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 51189),
+            Protocol::TCP,
         );
         assert_eq!(
             listener.to_string(),
-            "PID: 455        Process name: rapportd                  Socket: 0.0.0.0:51189"
+            "PID: 455        Process name: rapportd                  Socket: 0.0.0.0:51189                  Protocol: TCP"
         );
     }
 
@@ -196,10 +220,11 @@ mod tests {
             160,
             "mysqld".to_string(),
             SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 3306),
+            Protocol::TCP,
         );
         assert_eq!(
             listener.to_string(),
-            "PID: 160        Process name: mysqld                    Socket: [::]:3306"
+            "PID: 160        Process name: mysqld                    Socket: [::]:3306                      Protocol: TCP"
         );
     }
 

src/platform/linux/mod.rs

@@ -19,7 +19,12 @@ pub(crate) fn get_all() -> crate::Result<HashSet<Listener>> {
 
     for tcp_listener in TcpListener::get_all()? {
         if let Some(p) = inode_proc_map.get(&tcp_listener.inode()) {
-            let listener = Listener::new(p.pid(), p.name(), tcp_listener.local_addr());
+            let listener = Listener::new(
+                p.pid(),
+                p.name(),
+                tcp_listener.local_addr(),
+                tcp_listener.protocol(),
+            );
             listeners.insert(listener);
         }
     }

src/platform/linux/tcp_listener.rs

@@ -1,3 +1,4 @@
+use crate::Protocol;
 use std::fs::File;
 use std::io::{BufRead, BufReader};
 use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
@@ -7,10 +8,11 @@ use std::str::FromStr;
 pub(super) struct TcpListener {
     local_addr: SocketAddr,
     inode: u64,
+    protocol: Protocol,
 }
 
 impl TcpListener {
-    const LISTEN_STATE: &'static str = "0A";
+    // const LISTEN_STATE: &'static str = "0A";
 
     pub(super) fn local_addr(&self) -> SocketAddr {
         self.local_addr
@@ -20,30 +22,50 @@ impl TcpListener {
         self.inode
     }
 
+    pub(super) fn protocol(&self) -> Protocol {
+        self.protocol
+    }
+
     pub(super) fn get_all() -> crate::Result<Vec<TcpListener>> {
         let mut table = Vec::new();
         let tcp_table = File::open("/proc/net/tcp")?;
         for line in BufReader::new(tcp_table).lines().map_while(Result::ok) {
-            if let Ok(l) = TcpListener::from_tcp_table_entry(&line) {
+            if let Ok(l) = TcpListener::from_protocol_table_entry(&line, Protocol::TCP) {
                 table.push(l);
             }
         }
+
         let tcp6_table = File::open("/proc/net/tcp6")?;
         for line in BufReader::new(tcp6_table).lines().map_while(Result::ok) {
-            if let Ok(l) = TcpListener::from_tcp6_table_entry(&line) {
+            if let Ok(l) = TcpListener::from_protocolv6_table_entry(&line, Protocol::TCP) {
                 table.push(l);
             }
         }
+
+        let udp_table = File::open("/proc/net/udp")?;
+        for line in BufReader::new(udp_table).lines().map_while(Result::ok) {
+            // the lines/fields for tcp and udp are identical as far as Listeners is concerend
+            if let Ok(l) = TcpListener::from_protocol_table_entry(&line, Protocol::UDP) {
+                table.push(l)
+            }
+        }
+
+        let udp_table = File::open("/proc/net/udp6")?;
+        for line in BufReader::new(udp_table).lines().map_while(Result::ok) {
+            // the lines/fields for tcp and udp are identical as far as Listeners is concerend
+            if let Ok(l) = TcpListener::from_protocolv6_table_entry(&line, Protocol::UDP) {
+                table.push(l)
+            }
+        }
         Ok(table)
     }
 
-    fn from_tcp_table_entry(line: &str) -> crate::Result<Self> {
+    fn from_protocol_table_entry(line: &str, protocol: Protocol) -> crate::Result<Self> {
         let mut s = line.split_whitespace();
 
         let local_addr_hex = s.nth(1).ok_or("Failed to get local address")?;
-        let Some(Self::LISTEN_STATE) = s.nth(1) else {
-            return Err("Not a listening socket".into());
-        };
+        // consider all states
+        let _ = s.nth(1).ok_or("Failed to get state")?;
 
         let local_ip_port = local_addr_hex
             .split(':')
@@ -59,10 +81,14 @@ impl TcpListener {
         let inode_n = s.nth(5).ok_or("Failed to get inode")?;
         let inode = u64::from_str(inode_n)?;
 
-        Ok(Self { local_addr, inode })
+        Ok(Self {
+            local_addr,
+            inode,
+            protocol,
+        })
     }
 
-    fn from_tcp6_table_entry(line: &str) -> crate::Result<Self> {
+    fn from_protocolv6_table_entry(line: &str, protocol: Protocol) -> crate::Result<Self> {
         #[cfg(target_endian = "little")]
         let read_endian = u32::from_le_bytes;
         #[cfg(target_endian = "big")]
@@ -71,9 +97,8 @@ impl TcpListener {
         let mut s = line.split_whitespace();
 
         let local_addr_hex = s.nth(1).ok_or("Failed to get local address")?;
-        let Some(Self::LISTEN_STATE) = s.nth(1) else {
-            return Err("Not a listening socket".into());
-        };
+        // consider all states
+        let _ = s.nth(1).ok_or("Failed to get state")?;
 
         let mut local_ip_port = local_addr_hex.split(':');
 
@@ -108,6 +133,10 @@ impl TcpListener {
         let inode_n = s.nth(5).ok_or("Failed to get inode")?;
         let inode = u64::from_str(inode_n)?;
 
-        Ok(Self { local_addr, inode })
+        Ok(Self {
+            local_addr,
+            inode,
+            protocol,
+        })
     }
 }

src/platform/macos/c_socket_fd_info.rs

@@ -3,8 +3,10 @@ use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
 
 use byteorder::{ByteOrder, NetworkEndian};
 
-use crate::platform::macos::statics::SOCKET_STATE_LISTEN;
 use crate::platform::macos::tcp_listener::TcpListener;
+use crate::Protocol;
+
+use super::statics::{IPPROTO_TCP, IPPROTO_UDP};
 
 #[repr(C)]
 pub(super) struct CSocketFdInfo {
@@ -16,18 +18,31 @@ impl CSocketFdInfo {
     pub(super) fn to_tcp_listener(&self) -> crate::Result<TcpListener> {
         let sock_info = self.psi;
         let family = sock_info.soi_family;
+        let transport_protocol = sock_info.soi_protocol;
 
-        let tcp_in = unsafe { sock_info.soi_proto.pri_tcp };
+        let general_sock_info = unsafe {
+            match transport_protocol {
+                IPPROTO_TCP => sock_info.soi_proto.pri_tcp.tcpsi_ini,
+                IPPROTO_UDP => sock_info.soi_proto.pri_in,
+                _ => return Err("Unsupported protocol".into()),
+            }
+        };
 
-        if tcp_in.tcpsi_state != SOCKET_STATE_LISTEN {
-            return Err("Socket is not in listen state".into());
-        }
+        // if tcp, do not filter on state (get em all)
+        // if tcp_in.tcpsi_state != SOCKET_STATE_LISTEN && ip_protocol == IPPROT_TCP {
+        //     return Err("Socket is not in listening state".into());
+        // }
 
-        let tcp_sockaddr_in = tcp_in.tcpsi_ini;
-        let lport_bytes: [u8; 4] = i32::to_le_bytes(tcp_sockaddr_in.insi_lport);
-        let local_address = Self::get_local_addr(family, tcp_sockaddr_in)?;
+        // let tcp_sockaddr_in = tcp_in.tcpsi_ini;
+        let lport_bytes: [u8; 4] = i32::to_le_bytes(general_sock_info.insi_lport);
+        let local_address = Self::get_local_addr(family, general_sock_info)?;
+        let protocol = Self::get_protocol(family, transport_protocol)?;
 
-        let socket_info = TcpListener::new(local_address, NetworkEndian::read_u16(&lport_bytes));
+        let socket_info = TcpListener::new(
+            local_address,
+            NetworkEndian::read_u16(&lport_bytes),
+            protocol,
+        );
 
         Ok(socket_info)
     }
@@ -49,6 +64,14 @@ impl CSocketFdInfo {
             _ => Err("Unsupported socket family".into()),
         }
     }
+
+    fn get_protocol(family: c_int, ip_protocol: c_int) -> crate::Result<Protocol> {
+        match (family, ip_protocol) {
+            (2 | 30, IPPROTO_TCP) => Ok(Protocol::TCP),
+            (2 | 30, IPPROTO_UDP) => Ok(Protocol::UDP),
+            (_, _) => Err("unsupported protocol".into()),
+        }
+    }
 }
 
 #[repr(C)]

src/platform/macos/mod.rs

@@ -23,7 +23,12 @@ pub(crate) fn get_all() -> crate::Result<HashSet<Listener>> {
         for fd in SocketFd::get_all_of_pid(pid).iter().flatten() {
             if let Ok(tcp_listener) = TcpListener::from_pid_fd(pid, fd) {
                 if let Ok(ProcName(name)) = ProcName::from_pid(pid) {
-                    let listener = Listener::new(pid.as_u_32()?, name, tcp_listener.socket_addr());
+                    let listener = Listener::new(
+                        pid.as_u_32()?,
+                        name,
+                        tcp_listener.socket_addr(),
+                        tcp_listener.protocol(),
+                    );
                     listeners.insert(listener);
                 }
             }

src/platform/macos/statics.rs

@@ -4,5 +4,7 @@ pub(super) const PROC_ALL_PIDS: u32 = 1;
 pub(super) const PROC_PID_LIST_FDS: c_int = 1;
 pub(super) const PROC_PID_FD_SOCKET_INFO: c_int = 3;
 pub(super) const FD_TYPE_SOCKET: u32 = 2;
-pub(super) const SOCKET_STATE_LISTEN: c_int = 1;
+// pub(super) const SOCKET_STATE_LISTEN: c_int = 1;
 pub(super) const PROC_PID_PATH_INFO_MAXSIZE: usize = 4096;
+pub(super) const IPPROTO_TCP: c_int = 6;
+pub(super) const IPPROTO_UDP: c_int = 17;

src/platform/macos/tcp_listener.rs

@@ -9,16 +9,28 @@ use crate::platform::macos::proc_pid::ProcPid;
 use crate::platform::macos::socket_fd::SocketFd;
 use crate::platform::macos::statics::PROC_PID_FD_SOCKET_INFO;
 
+use crate::Protocol;
+
 #[derive(Debug)]
-pub(super) struct TcpListener(SocketAddr);
+pub(super) struct TcpListener {
+    local_addr: SocketAddr,
+    protocol: Protocol,
+}
 
 impl TcpListener {
-    pub(super) fn new(addr: IpAddr, port: u16) -> Self {
-        TcpListener(SocketAddr::new(addr, port))
+    pub(super) fn new(addr: IpAddr, port: u16, protocol: Protocol) -> Self {
+        TcpListener {
+            local_addr: SocketAddr::new(addr, port),
+            protocol,
+        }
     }
 
     pub(super) fn socket_addr(&self) -> SocketAddr {
-        self.0
+        self.local_addr
+    }
+
+    pub(super) fn protocol(&self) -> Protocol {
+        self.protocol
     }
 
     pub(super) fn from_pid_fd(pid: ProcPid, fd: &SocketFd) -> crate::Result<Self> {

src/platform/windows/c_iphlpapi.rs

@@ -12,3 +12,16 @@ extern "system" {
         Reserved: c_ulong,
     ) -> c_ulong;
 }
+
+#[allow(non_snake_case)]
+#[link(name = "iphlpapi")]
+extern "system" {
+    pub(super) fn GetExtendedUdpTable(
+        pUdpTable: *mut c_void,
+        pdwSize: *mut c_ulong,
+        bOrder: c_int,
+        ulAf: c_ulong,
+        TableClass: c_ulong,
+        Reserved: c_ulong,
+    ) -> c_ulong;
+}

src/platform/windows/mod.rs

@@ -10,6 +10,8 @@ mod statics;
 mod tcp6_table;
 mod tcp_listener;
 mod tcp_table;
+mod udp6_table;
+mod udp_table;
 
 pub(crate) fn get_all() -> crate::Result<HashSet<Listener>> {
     let mut listeners = HashSet::new();