From 1c996c38f4df2996b61080a7457f1dc1cf5afbfc Mon Sep 17 00:00:00 2001
From: cyqsimon <28627918+cyqsimon@users.noreply.github.com>
Date: Wed, 11 Oct 2023 19:45:48 +0800
Subject: Ignore connections that fail parsing instead of panicking on BSD
 (#288)

* Ignore connections that fail parsing instead of panicking on BSD

- Tentative fix for #217

* Log when a connection fails parsing

- I actually love unreadable code

* Fix clippy complaint
---
 src/os/lsof.rs       | 25 +++++--------------------
 src/os/lsof_utils.rs | 38 +++++++++++++++++++++++++++++---------
 2 files changed, 34 insertions(+), 29 deletions(-)

(limited to 'src')

diff --git a/src/os/lsof.rs b/src/os/lsof.rs
index f6d1e3f..088d9cf 100644
--- a/src/os/lsof.rs
+++ b/src/os/lsof.rs
@@ -1,24 +1,9 @@
-use std::collections::HashMap;
-
-use crate::{network::LocalSocket, os::lsof_utils::get_connections, OpenSockets};
+use crate::{os::lsof_utils::get_connections, OpenSockets};
 
 pub(crate) fn get_open_sockets() -> OpenSockets {
-    let mut open_sockets = HashMap::new();
-
-    let connections = get_connections();
-
-    for raw_connection in connections {
-        open_sockets.insert(
-            LocalSocket {
-                ip: raw_connection.get_local_ip(),
-                port: raw_connection.get_local_port(),
-                protocol: raw_connection.get_protocol(),
-            },
-            raw_connection.process_name.clone(),
-        );
-    }
+    let sockets_to_procs = get_connections()
+        .filter_map(|raw| raw.as_local_socket().map(|s| (s, raw.process_name)))
+        .collect();
 
-    OpenSockets {
-        sockets_to_procs: open_sockets,
-    }
+    OpenSockets { sockets_to_procs }
 }
diff --git a/src/os/lsof_utils.rs b/src/os/lsof_utils.rs
index 21735f1..5e55dad 100644
--- a/src/os/lsof_utils.rs
+++ b/src/os/lsof_utils.rs
@@ -1,8 +1,9 @@
 use std::{ffi::OsStr, net::IpAddr, process::Command, sync::OnceLock};
 
+use log::warn;
 use regex::Regex;
 
-use crate::network::Protocol;
+use crate::network::{LocalSocket, Protocol};
 
 #[allow(dead_code)]
 #[derive(Debug, Clone)]
@@ -103,16 +104,35 @@ impl RawConnection {
         }
     }
 
-    pub fn get_protocol(&self) -> Protocol {
-        Protocol::from_str(&self.protocol).unwrap()
+    pub fn get_protocol(&self) -> Option<Protocol> {
+        Protocol::from_str(&self.protocol)
     }
 
-    pub fn get_local_ip(&self) -> IpAddr {
-        self.local_ip.parse().unwrap()
+    pub fn get_local_ip(&self) -> Option<IpAddr> {
+        self.local_ip.parse().ok()
     }
 
-    pub fn get_local_port(&self) -> u16 {
-        self.local_port.parse::<u16>().unwrap()
+    pub fn get_local_port(&self) -> Option<u16> {
+        self.local_port.parse::<u16>().ok()
+    }
+
+    pub fn as_local_socket(&self) -> Option<LocalSocket> {
+        let process = &self.process_name;
+
+        let Some(ip) = self.get_local_ip() else {
+            warn!(r#"Failed to get the local IP of a connection belonging to "{process}"."#);
+            return None;
+        };
+        let Some(port) = self.get_local_port() else {
+            warn!(r#"Failed to get the local port of a connection belonging to "{process}"."#);
+            return None;
+        };
+        let Some(protocol) = self.get_protocol() else {
+            warn!(r#"Failed to get the protocol of a connection belonging to "{process}"."#);
+            return None;
+        };
+
+        Some(LocalSocket { ip, port, protocol })
     }
 }
 
@@ -203,7 +223,7 @@ com.apple   590 etoledom  204u  IPv4 0x28ffb9c04111253f      0t0  TCP 192.168.1.
     }
     fn test_raw_connection_parse_local_port(raw_output: &str) {
         let connection = RawConnection::new(raw_output).unwrap();
-        assert_eq!(connection.get_local_port(), 1111);
+        assert_eq!(connection.get_local_port(), Some(1111));
     }
 
     #[test]
@@ -216,7 +236,7 @@ com.apple   590 etoledom  204u  IPv4 0x28ffb9c04111253f      0t0  TCP 192.168.1.
     }
     fn test_raw_connection_parse_protocol(raw_line: &str) {
         let connection = RawConnection::new(raw_line).unwrap();
-        assert_eq!(connection.get_protocol(), Protocol::Udp);
+        assert_eq!(connection.get_protocol(), Some(Protocol::Udp));
     }
 
     #[test]
-- 
cgit v1.2.3