Address some comments on PR

imsnif · Jan 7, 2020 · 471fd86 · 471fd86
1 parent cf772ac
commit 471fd86
Show file tree

Hide file tree

Showing 7 changed files with 49 additions and 36 deletions.
diff --git a/src/main.rs b/src/main.rs
@@ -79,10 +79,15 @@ fn try_main() -> Result<(), failure::Error> {
     Ok(())
 }
 
+pub struct OpenSockets {
+    sockets_to_procs: HashMap<LocalSocket, String>,
+    connections: Vec<Connection>,
+}
+
 pub struct OsInputOutput {
     pub network_interfaces: Vec<NetworkInterface>,
     pub network_frames: Vec<Box<dyn DataLinkReceiver>>,
-    pub get_open_sockets: fn() -> (HashMap<LocalSocket, String>, std::vec::Vec<Connection>),
+    pub get_open_sockets: fn() -> OpenSockets,
     pub keyboard_events: Box<dyn Iterator<Item = Event> + Send>,
     pub dns_client: Option<dns::Client>,
     pub on_winch: Box<OnSigWinch>,
@@ -138,7 +143,10 @@ where
                 while running.load(Ordering::Acquire) {
                     let render_start_time = Instant::now();
                     let utilization = { network_utilization.lock().unwrap().clone_and_reset() };
-                    let (connections_to_procs, connections) = get_open_sockets();
+                    let OpenSockets {
+                        sockets_to_procs,
+                        connections,
+                    } = get_open_sockets();
                     let mut ip_to_host = IpTable::new();
                     if let Some(dns_client) = dns_client.as_mut() {
                         ip_to_host = dns_client.cache();
@@ -152,7 +160,7 @@ where
                     }
                     {
                         let mut ui = ui.lock().unwrap();
-                        ui.update_state(connections_to_procs, utilization, ip_to_host);
+                        ui.update_state(sockets_to_procs, utilization, ip_to_host);
                         if raw_mode {
                             ui.output_text(&mut write_to_stdout);
                         } else {

diff --git a/src/network/connection.rs b/src/network/connection.rs
@@ -11,6 +11,10 @@ pub enum Protocol {
 }
 
 impl Protocol {
+    // Currently, linux implementation doesn't use this function.
+    // Without this #[cfg] clippy complains about dead code, and CI refuses
+    // to pass.
+    #[cfg(target_os = "macos")]
     pub fn from_str(string: &str) -> Option<Self> {
         match string {
             "TCP" => Some(Protocol::Tcp),

diff --git a/src/network/utilization.rs b/src/network/utilization.rs
@@ -2,14 +2,14 @@ use crate::network::{Connection, Direction, Segment};
 
 use ::std::collections::HashMap;
 
-#[derive(Clone, Debug)]
+#[derive(Clone)]
 pub struct ConnectionInfo {
     pub interface_name: String,
     pub total_bytes_downloaded: u128,
     pub total_bytes_uploaded: u128,
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone)]
 pub struct Utilization {
     pub connections: HashMap<Connection, ConnectionInfo>,
 }

diff --git a/src/os/linux.rs b/src/os/linux.rs
@@ -2,9 +2,10 @@ use ::std::collections::HashMap;
 
 use ::procfs::process::FDTarget;
 
-use crate::network::{Connection, LocalSocket, Protocol};
+use crate::network::{Connection, Protocol};
+use crate::OpenSockets;
 
-pub(crate) fn get_open_sockets() -> (HashMap<LocalSocket, String>, std::vec::Vec<Connection>) {
+pub(crate) fn get_open_sockets() -> OpenSockets {
     let mut open_sockets = HashMap::new();
     let mut connections = std::vec::Vec::new();
     let all_procs = procfs::process::all_processes().unwrap();
@@ -46,5 +47,8 @@ pub(crate) fn get_open_sockets() -> (HashMap<LocalSocket, String>, std::vec::Vec
             connections.push(connection);
         };
     }
-    (open_sockets, connections)
+    OpenSockets {
+        sockets_to_procs: open_sockets,
+        connections,
+    }
 }
diff --git a/src/os/lsof_utils.rs b/src/os/lsof_utils.rs
@@ -21,16 +21,6 @@ lazy_static! {
     static ref LISTEN_REGEX: Regex = Regex::new(r"(.*):(.*)").unwrap();
 }
 
-fn parse_ip_addr(ip: &str) -> IpAddr {
-    if let Ok(v4addr) = ip.parse() {
-        IpAddr::V4(v4addr)
-    } else if let Ok(v6addr) = ip.parse() {
-        IpAddr::V6(v6addr)
-    } else {
-        panic!("{} is not a valid IP address", ip)
-    }
-}
-
 fn get_null_addr(ip_type: &str) -> &str {
     if ip_type.contains('4') {
         "0.0.0.0"
@@ -39,16 +29,12 @@ fn get_null_addr(ip_type: &str) -> &str {
     }
 }
 
-#[allow(clippy::needless_return)]
 impl RawConnection {
     pub fn new(raw_line: &str) -> Option<RawConnection> {
+        // Example row
+        // com.apple   664     user  198u  IPv4 0xeb179a6650592b8d      0t0    TCP 192.168.1.187:58535->1.2.3.4:443 (ESTABLISHED)
         let columns: Vec<&str> = raw_line.split_ascii_whitespace().collect();
         if columns.len() < 9 {
-            println!(
-                "lsof's output string has {} columns, different than expected: {:#?}",
-                columns.len(),
-                columns
-            );
             return None;
         }
         let process_name = columns[0].replace("\\x20", " ");
@@ -71,6 +57,10 @@ impl RawConnection {
         // let connection_state = columns[9];
         // If this socket is in a "connected" state
         if let Some(caps) = CONNECTION_REGEX.captures(connection_str) {
+            // Example
+            // 192.168.1.187:64230->0.1.2.3:5228
+            // *:*
+            // *:4567
             let local_ip = String::from(caps.get(1).unwrap().as_str());
             let local_port = String::from(caps.get(2).unwrap().as_str());
             let remote_ip = String::from(caps.get(3).unwrap().as_str());
@@ -113,23 +103,23 @@ impl RawConnection {
     }
 
     pub fn get_protocol(&self) -> Protocol {
-        return Protocol::from_str(&self.protocol).unwrap();
+        Protocol::from_str(&self.protocol).unwrap()
     }
 
     pub fn get_remote_ip(&self) -> IpAddr {
-        return parse_ip_addr(&self.remote_ip);
+        self.remote_ip.parse().unwrap()
     }
 
     pub fn get_remote_port(&self) -> u16 {
-        return self.remote_port.parse::<u16>().unwrap();
+        self.remote_port.parse::<u16>().unwrap()
     }
 
     pub fn get_local_ip(&self) -> IpAddr {
-        return parse_ip_addr(&self.local_ip);
+        self.local_ip.parse().unwrap()
     }
 
     pub fn get_local_port(&self) -> u16 {
-        return self.local_port.parse::<u16>().unwrap();
+        self.local_port.parse::<u16>().unwrap()
     }
 }
 

diff --git a/src/os/macos.rs b/src/os/macos.rs
@@ -1,6 +1,7 @@
 use ::std::collections::HashMap;
 
-use crate::network::{Connection, LocalSocket};
+use crate::network::Connection;
+use crate::OpenSockets;
 
 use super::lsof_utils;
 use std::net::SocketAddr;
@@ -14,8 +15,7 @@ struct RawConnection {
     process_name: String,
 }
 
-#[allow(clippy::needless_return)]
-pub(crate) fn get_open_sockets() -> (HashMap<LocalSocket, String>, std::vec::Vec<Connection>) {
+pub(crate) fn get_open_sockets() -> OpenSockets {
     let mut open_sockets = HashMap::new();
     let mut connections_vec = std::vec::Vec::new();
 
@@ -35,5 +35,8 @@ pub(crate) fn get_open_sockets() -> (HashMap<LocalSocket, String>, std::vec::Vec
         connections_vec.push(connection);
     }
 
-    return (open_sockets, connections_vec);
+    OpenSockets {
+        sockets_to_procs: open_sockets,
+        connections: connections_vec,
+    }
 }
diff --git a/src/tests/fakes/fake_input.rs b/src/tests/fakes/fake_input.rs
@@ -13,9 +13,10 @@ use ::termion::event::Event;
 use crate::{
     network::{
         dns::{self, Lookup},
-        Connection, LocalSocket, Protocol,
+        Connection, Protocol,
     },
     os::OnSigWinch,
+    OpenSockets,
 };
 
 pub struct KeyboardEvents {
@@ -85,7 +86,7 @@ impl DataLinkReceiver for NetworkFrames {
     }
 }
 
-pub fn get_open_sockets() -> (HashMap<LocalSocket, String>, std::vec::Vec<Connection>) {
+pub fn get_open_sockets() -> OpenSockets {
     let mut open_sockets = HashMap::new();
     let local_ip = IpAddr::V4(Ipv4Addr::new(10, 0, 0, 2));
     open_sockets.insert(
@@ -145,7 +146,10 @@ pub fn get_open_sockets() -> (HashMap<LocalSocket, String>, std::vec::Vec<Connec
         connections.push(connection);
     }
 
-    (local_socket_to_procs, connections)
+    OpenSockets {
+        sockets_to_procs: local_socket_to_procs,
+        connections,
+    }
 }
 
 pub fn get_interfaces() -> Vec<NetworkInterface> {