Merge pull request #82 from zhangxp1998/udp

Fix bug that non-connected Udp sockets aren't displayed

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>,
     ) {

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,

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;
 
@@ -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<Connection, String>,
+    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,12 +143,15 @@ 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 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();
-                        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<_>>();
@@ -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 {

src/network/connection.rs

@@ -1,16 +1,20 @@
 use ::std::collections::HashMap;
 use ::std::fmt;
-use ::std::net::Ipv4Addr;
+use ::std::net::{IpAddr, Ipv4Addr};
 
 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,
 }
 
 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),
@@ -29,17 +33,23 @@ impl fmt::Display for Protocol {
     }
 }
 
-#[derive(Clone, Ord, PartialOrd, PartialEq, Eq, Hash)]
+#[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)]
+#[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, Copy)]
 pub struct Connection {
     pub remote_socket: Socket,
-    pub protocol: Protocol,
-    pub local_port: u16,
+    pub local_socket: LocalSocket,
 }
 
 pub fn display_ip_or_host(ip: Ipv4Addr, ip_to_host: &HashMap<Ipv4Addr, String>) -> String {
@@ -57,23 +67,31 @@ pub fn display_connection_string(
     format!(
         "<{}>:{} => {}:{} ({})",
         interface_name,
-        connection.local_port,
+        connection.local_socket.port,
         display_ip_or_host(connection.remote_socket.ip, ip_to_host),
         connection.remote_socket.port,
-        connection.protocol,
+        connection.local_socket.protocol,
     )
 }
 
 impl Connection {
-    pub fn new(remote_socket: SocketAddr, local_port: u16, protocol: Protocol) -> Option<Self> {
+    pub fn new(
+        remote_socket: SocketAddr,
+        local_ip: IpAddr,
+        local_port: u16,
+        protocol: Protocol,
+    ) -> Option<Self> {
         match remote_socket {
             SocketAddr::V4(remote_socket) => Some(Connection {
                 remote_socket: Socket {
                     ip: *remote_socket.ip(),
                     port: remote_socket.port(),
                 },
-                protocol,
-                local_port,
+                local_socket: LocalSocket {
+                    ip: local_ip,
+                    port: local_port,
+                    protocol,
+                },
             }),
             _ => None,
         }

src/network/sniffer.rs

@@ -88,8 +88,10 @@ impl Sniffer {
                 let to = SocketAddr::new(IpAddr::V4(ip_packet.get_destination()), destination_port);
 
                 let connection = match direction {
-                    Direction::Download => Connection::new(from, destination_port, protocol)?,
-                    Direction::Upload => Connection::new(to, source_port, protocol)?,
+                    Direction::Download => {
+                        Connection::new(from, to.ip(), destination_port, protocol)?
+                    }
+                    Direction::Upload => Connection::new(to, from.ip(), source_port, protocol)?,
                 };
                 Some(Segment {
                     interface_name,

src/os/linux.rs

@@ -3,9 +3,11 @@ use ::std::collections::HashMap;
 use ::procfs::process::FDTarget;
 
 use crate::network::{Connection, Protocol};
+use crate::OpenSockets;
 
-pub(crate) fn get_open_sockets() -> HashMap<Connection, String> {
+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();
 
     let mut inode_to_procname = HashMap::new();
@@ -23,23 +25,30 @@ pub(crate) fn get_open_sockets() -> HashMap<Connection, String> {
     let tcp = ::procfs::net::tcp().unwrap();
     for entry in tcp.into_iter() {
         let local_port = entry.local_address.port();
+        let local_ip = entry.local_address.ip();
         if let (Some(connection), Some(procname)) = (
-            Connection::new(entry.remote_address, local_port, Protocol::Tcp),
+            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);
         };
     }
 
     let udp = ::procfs::net::udp().unwrap();
     for entry in udp.into_iter() {
         let local_port = entry.local_address.port();
+        let local_ip = entry.local_address.ip();
         if let (Some(connection), Some(procname)) = (
-            Connection::new(entry.remote_address, local_port, Protocol::Udp),
+            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
+    OpenSockets {
+        sockets_to_procs: open_sockets,
+        connections,
+    }
 }