Use <local_ip, local_port, protocol> to indentify processes

9 files changed, 84 insertions, 45 deletions

diff --git a/src/display/ui.rs b/src/display/ui.rs
index e58c3e9..e61f5e7 100644
--- a/src/display/ui.rs
+++ b/src/display/ui.rs
@@ -5,7 +5,7 @@ use ::tui::Terminal;
 
 use crate::display::components::{Layout, Table, TotalBandwidth};
 use crate::display::UIState;
-use crate::network::{display_connection_string, display_ip_or_host, Connection, Utilization};
+use crate::network::{display_connection_string, display_ip_or_host, LocalSocket, Utilization};
 
 use ::std::net::Ipv4Addr;
 
@@ -94,7 +94,7 @@ where
     }
     pub fn update_state(
         &mut self,
-        connections_to_procs: HashMap<Connection, String>,
+        connections_to_procs: HashMap<LocalSocket, String>,
         utilization: Utilization,
         ip_to_host: HashMap<Ipv4Addr, String>,
     ) {
diff --git a/src/display/ui_state.rs b/src/display/ui_state.rs
index e5c2217..acbd04f 100644
--- a/src/display/ui_state.rs
+++ b/src/display/ui_state.rs
@@ -1,7 +1,7 @@
 use ::std::collections::{BTreeMap, HashMap};
-use ::std::net::Ipv4Addr;
+use ::std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
 
-use crate::network::{Connection, Utilization};
+use crate::network::{Connection, LocalSocket, Utilization};
 
 pub trait Bandwidth {
     fn get_total_bytes_downloaded(&self) -> u128;
@@ -51,38 +51,61 @@ pub struct UIState {
 }
 
 impl UIState {
+    fn get_proc_name<'a>(
+        connections_to_procs: &'a HashMap<LocalSocket, String>,
+        local_socket: &LocalSocket,
+    ) -> Option<&'a String> {
+        if let Some(process_name) = connections_to_procs.get(local_socket) {
+            Some(process_name)
+        } else if let Some(process_name) = connections_to_procs.get(&LocalSocket {
+            ip: IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)),
+            port: local_socket.port,
+            protocol: local_socket.protocol,
+        }) {
+            Some(process_name)
+        } else {
+            connections_to_procs.get(&LocalSocket {
+                ip: IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 0)),
+                port: local_socket.port,
+                protocol: local_socket.protocol,
+            })
+        }
+    }
     pub fn new(
-        connections_to_procs: HashMap<Connection, String>,
-        mut network_utilization: Utilization,
+        connections_to_procs: HashMap<LocalSocket, String>,
+        network_utilization: Utilization,
     ) -> Self {
         let mut processes: BTreeMap<String, NetworkData> = BTreeMap::new();
         let mut remote_addresses: BTreeMap<Ipv4Addr, NetworkData> = BTreeMap::new();
         let mut connections: BTreeMap<Connection, ConnectionData> = BTreeMap::new();
         let mut total_bytes_downloaded: u128 = 0;
         let mut total_bytes_uploaded: u128 = 0;
-        for (connection, process_name) in connections_to_procs {
-            if let Some(connection_info) = network_utilization.connections.remove(&connection) {
-                let data_for_remote_address = remote_addresses
-                    .entry(connection.remote_socket.ip)
-                    .or_default();
-                let connection_data = connections.entry(connection).or_default();
-                let data_for_process = processes.entry(process_name.clone()).or_default();
+        for (connection, connection_info) in network_utilization.connections {
+            let connection_data = connections.entry(connection).or_default();
 
+            if let Some(process_name) =
+                UIState::get_proc_name(&connections_to_procs, &connection.local_socket)
+            {
+                let data_for_process = processes.entry(process_name.clone()).or_default();
                 data_for_process.total_bytes_downloaded += connection_info.total_bytes_downloaded;
                 data_for_process.total_bytes_uploaded += connection_info.total_bytes_uploaded;
                 data_for_process.connection_count += 1;
-                connection_data.total_bytes_downloaded += connection_info.total_bytes_downloaded;
-                connection_data.total_bytes_uploaded += connection_info.total_bytes_uploaded;
-                connection_data.process_name = process_name;
-                connection_data.interface_name = connection_info.interface_name;
-                data_for_remote_address.total_bytes_downloaded +=
-                    connection_info.total_bytes_downloaded;
-                data_for_remote_address.total_bytes_uploaded +=
-                    connection_info.total_bytes_uploaded;
-                data_for_remote_address.connection_count += 1;
-                total_bytes_downloaded += connection_info.total_bytes_downloaded;
-                total_bytes_uploaded += connection_info.total_bytes_uploaded;
+                connection_data.process_name = process_name.clone();
+            } else {
+                connection_data.process_name = String::from("<UNKNOWN>");
             }
+            let data_for_remote_address = remote_addresses
+                .entry(connection.remote_socket.ip)
+                .or_default();
+            connection_data.total_bytes_downloaded += connection_info.total_bytes_downloaded;
+            connection_data.total_bytes_uploaded += connection_info.total_bytes_uploaded;
+            connection_data.interface_name = connection_info.interface_name;
+            data_for_remote_address.total_bytes_downloaded +=
+                connection_info.total_bytes_downloaded;
+            data_for_remote_address.total_bytes_uploaded += connection_info.total_bytes_uploaded;
+            data_for_remote_address.connection_count += 1;
+            total_bytes_downloaded += connection_info.total_bytes_downloaded;
+            total_bytes_uploaded += connection_info.total_bytes_uploaded;
         }
         UIState {
             processes,
diff --git a/src/main.rs b/src/main.rs
index 77424c2..b14227c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -9,7 +9,7 @@ mod tests;
 use display::{RawTerminalBackend, Ui};
 use network::{
     dns::{self, IpTable},
-    Connection, Sniffer, Utilization,
+    Connection, LocalSocket, Sniffer, Utilization,
 };
 use os::OnSigWinch;
 
@@ -82,7 +82,7 @@ fn try_main() -> Result<(), failure::Error> {
 pub struct OsInputOutput {
     pub network_interfaces: Vec<NetworkInterface>,
     pub network_frames: Vec<Box<dyn DataLinkReceiver>>,
-    pub get_open_sockets: fn() -> HashMap<Connection, String>,
+    pub get_open_sockets: fn() -> (HashMap<LocalSocket, String>, std::vec::Vec<Connection>),
     pub keyboard_events: Box<dyn Iterator<Item = Event> + Send>,
     pub dns_client: Option<dns::Client>,
     pub on_winch: Box<OnSigWinch>,
@@ -138,12 +138,12 @@ 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 = get_open_sockets();
+                    let (connections_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();
-                        let unresolved_ips = connections_to_procs
-                            .keys()
+                        let unresolved_ips = connections
+                            .iter()
                             .filter(|conn| !ip_to_host.contains_key(&conn.remote_socket.ip))
                             .map(|conn| conn.remote_socket.ip)
                             .collect::<Vec<_>>();
diff --git a/src/network/connection.rs b/src/network/connection.rs
index a8a17cb..c12311e 100644
--- a/src/network/connection.rs
+++ b/src/network/connection.rs
@@ -4,7 +4,7 @@ use ::std::net::{Ipv4Addr, IpAddr};
 
 use ::std::net::SocketAddr;
 
-#[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord, Debug)]
+#[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord, Debug, Copy)]
 pub enum Protocol {
     Tcp,
     Udp,
@@ -29,20 +29,20 @@ impl fmt::Display for Protocol {
     }
 }
 
-#[derive(Clone, Ord, PartialOrd, PartialEq, Eq, Hash, Debug)]
+#[derive(Clone, Ord, PartialOrd, PartialEq, Eq, Hash, Debug, Copy)]
 pub struct Socket {
     pub ip: Ipv4Addr,
     pub port: u16,
 }
 
-#[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord, Debug)]
+#[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord, Debug, Copy)]
 pub struct LocalSocket {
     pub ip: IpAddr,
     pub port: u16,
     pub protocol: Protocol,
 }
 
-#[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord, Debug)]
+#[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord, Debug, Copy)]
 pub struct Connection {
     pub remote_socket: Socket,
     pub local_socket: LocalSocket,
diff --git a/src/network/utilization.rs b/src/network/utilization.rs
index bdfbb21..d76ac44 100644
--- 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)]
+#[derive(Clone, Debug)]
 pub struct ConnectionInfo {
     pub interface_name: String,
     pub total_bytes_downloaded: u128,
     pub total_bytes_uploaded: u128,
 }
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct Utilization {
     pub connections: HashMap<Connection, ConnectionInfo>,
 }
diff --git a/src/os/linux.rs b/src/os/linux.rs
index a9b1ca5..07b8b17 100644
--- a/src/os/linux.rs
+++ b/src/os/linux.rs
@@ -6,6 +6,7 @@ use crate::network::{Connection, LocalSocket, Protocol};
 
 pub(crate) fn get_open_sockets() -> (HashMap<LocalSocket, String>, std::vec::Vec<Connection>) {
     let mut open_sockets = HashMap::new();
+    let mut connections = std::vec::Vec::new();
     let all_procs = procfs::process::all_processes().unwrap();
 
     let mut inode_to_procname = HashMap::new();
@@ -28,7 +29,8 @@ pub(crate) fn get_open_sockets() -> (HashMap<LocalSocket, String>, std::vec::Vec
             Connection::new(entry.remote_address, local_ip, local_port, Protocol::Tcp),
             inode_to_procname.get(&entry.inode),
         ) {
-            open_sockets.insert(connection, procname.clone());
+            open_sockets.insert(connection.local_socket, procname.clone());
+            connections.push(connection);
         };
     }
 
@@ -40,7 +42,8 @@ pub(crate) fn get_open_sockets() -> (HashMap<LocalSocket, String>, std::vec::Vec
             Connection::new(entry.remote_address, local_ip, local_port, Protocol::Udp),
             inode_to_procname.get(&entry.inode),
         ) {
-            open_sockets.insert(connection, procname.clone());
+            open_sockets.insert(connection.local_socket, procname.clone());
+            connections.push(connection);
         };
     }
     (open_sockets, connections)
diff --git a/src/os/lsof_utils.rs b/src/os/lsof_utils.rs
index 891e6c6..4464abc 100644
--- a/src/os/lsof_utils.rs
+++ b/src/os/lsof_utils.rs
@@ -59,7 +59,10 @@ impl RawConnection {
         // let device = columns[5];
         // let size = columns[6];
         // UDP/TCP
-        let protocol = String::from(columns[7]);
+        let protocol = columns[7].to_ascii_uppercase();
+        if protocol != "TCP" && protocol != "UDP" {
+            return None;
+        }
         let connection_str = columns[8];
         // "(LISTEN)" or "(ESTABLISHED)",  this column may or may not be present
         // let connection_state = columns[9];
diff --git a/src/os/macos.rs b/src/os/macos.rs
index 1ed6ecf..35f0d8f 100644
--- a/src/os/macos.rs
+++ b/src/os/macos.rs
@@ -1,6 +1,6 @@
 use ::std::collections::HashMap;
 
-use crate::network::Connection;
+use crate::network::{LocalSocket, Connection};
 
 use super::lsof_utils;
 use std::net::SocketAddr;
@@ -14,8 +14,10 @@ struct RawConnection {
     process_name: String,
 }
 
-pub(crate) fn get_open_sockets() -> HashMap<Connection, String> {
+#[allow(clippy::needless_return)]
+pub(crate) fn get_open_sockets() -> (HashMap<LocalSocket, String>, std::vec::Vec<Connection>) {
     let mut open_sockets = HashMap::new();
+    let mut connections_vec = std::vec::Vec::new();
 
     let connections = lsof_utils::get_connections();
 
@@ -29,8 +31,9 @@ pub(crate) fn get_open_sockets() -> HashMap<Connection, String> {
         let socket_addr = SocketAddr::new(remote_ip, remote_port);
         let connection = Connection::new(socket_addr, local_ip, local_port, protocol).unwrap();
 
-        open_sockets.insert(connection, raw_connection.process_name.clone());
+        open_sockets.insert(connection.local_socket, raw_connection.process_name.clone());
+        connections_vec.push(connection);
     }
 
-    return open_sockets;
+    return (open_sockets, connections_vec);
 }
diff --git a/src/tests/fakes/fake_input.rs b/src/tests/fakes/fake_input.rs
index 7f384f6..1294e4c 100644
--- a/src/tests/fakes/fake_input.rs
+++ b/src/tests/fakes/fake_input.rs
@@ -13,7 +13,7 @@ use ::termion::event::Event;
 use crate::{
     network::{
         dns::{self, Lookup},
-        Connection, Protocol,
+        Connection, Protocol, LocalSocket
     },
     os::OnSigWinch,
 };
@@ -85,7 +85,7 @@ impl DataLinkReceiver for NetworkFrames {
     }
 }
 
-pub fn get_open_sockets() -> HashMap<Connection, String> {
+pub fn get_open_sockets() -> (HashMap<LocalSocket, String>, std::vec::Vec<Connection>) {
     let mut open_sockets = HashMap::new();
     let local_ip = IpAddr::V4(Ipv4Addr::new(10, 0, 0, 2));
     open_sockets.insert(
@@ -138,7 +138,14 @@ pub fn get_open_sockets() -> HashMap<Connection, String> {
         .unwrap(),
         String::from("3"),
     );
-    open_sockets
+    let mut local_socket_to_procs = HashMap::new();
+    let mut connections = std::vec::Vec::new();
+    for (connection, process_name) in open_sockets {
+        local_socket_to_procs.insert(connection.local_socket, process_name);
+        connections.push(connection);
+    }
+
+    (local_socket_to_procs, connections)
 }
 
 pub fn get_interfaces() -> Vec<NetworkInterface> {