refactor(dns): handle ip to host conversion in the display layer

6 files changed, 44 insertions, 119 deletions

diff --git a/src/display/ui.rs b/src/display/ui.rs
index ced02ec..1607e7c 100644
--- a/src/display/ui.rs
+++ b/src/display/ui.rs
@@ -10,6 +10,8 @@ use ::tui::Terminal;
 use crate::display::{Bandwidth, UIState};
 use crate::network::{Connection, Utilization};
 
+use ::std::net::Ipv4Addr;
+
 struct DisplayBandwidth(f64);
 
 impl fmt::Display for DisplayBandwidth {
@@ -26,6 +28,13 @@ impl fmt::Display for DisplayBandwidth {
     }
 }
 
+fn display_ip_or_host (ip: &Ipv4Addr, ip_to_host: &HashMap<Ipv4Addr, String>) -> String {
+    match ip_to_host.get(ip) {
+        Some(host) => host.clone(),
+        None => ip.to_string()
+    }
+}
+
 fn create_table<'a>(
     title: &'a str,
     column_names: &'a [&'a str],
@@ -86,13 +95,21 @@ fn render_process_table(state: &UIState, frame: &mut Frame<impl Backend>, rect:
     table.render(frame, rect);
 }
 
-fn render_connections_table(state: &UIState, frame: &mut Frame<impl Backend>, rect: Rect) {
+fn render_connections_table(state: &UIState, frame: &mut Frame<impl Backend>, rect: Rect, ip_to_host: &HashMap<Ipv4Addr, String>) {
     let rows = state
         .connections
         .iter()
         .map(|(connection, connection_data)| {
+            let connection_string = format!(
+                "{}:{} => {}:{} ({})",
+                display_ip_or_host(&connection.local_socket.ip, ip_to_host),
+                connection.local_socket.port,
+                display_ip_or_host(&connection.remote_socket.ip, ip_to_host),
+                connection.remote_socket.port,
+                connection.protocol,
+            );
             format_row_data(
-                connection.to_string(),
+                connection_string,
                 connection_data.process_name.to_string(),
                 connection_data,
             )
@@ -106,13 +123,14 @@ fn render_connections_table(state: &UIState, frame: &mut Frame<impl Backend>, re
     table.render(frame, rect);
 }
 
-fn render_remote_ip_table(state: &UIState, frame: &mut Frame<impl Backend>, rect: Rect) {
+fn render_remote_ip_table(state: &UIState, frame: &mut Frame<impl Backend>, rect: Rect, ip_to_host: &HashMap<Ipv4Addr, String>) {
     let rows = state
         .remote_ips
         .iter()
         .map(|(remote_ip, data_for_remote_ip)| {
+            let remote_ip = display_ip_or_host(remote_ip, &ip_to_host);
             format_row_data(
-                remote_ip.to_string(),
+                remote_ip,
                 data_for_remote_ip.connection_count.to_string(),
                 data_for_remote_ip,
             )
@@ -130,6 +148,7 @@ pub fn display_loop(
     network_utilization: &Utilization,
     terminal: &mut Terminal<impl Backend>,
     connections_to_procs: HashMap<Connection, String>,
+    ip_to_host: &HashMap<Ipv4Addr, String>,
 ) {
     let state = UIState::new(connections_to_procs, &network_utilization);
     terminal
@@ -137,9 +156,9 @@ pub fn display_loop(
             let screen_horizontal_halves = split(Direction::Horizontal, f.size());
             let right_side_vertical_halves =
                 split(Direction::Vertical, screen_horizontal_halves[1]);
-            render_connections_table(&state, &mut f, screen_horizontal_halves[0]);
+            render_connections_table(&state, &mut f, screen_horizontal_halves[0], ip_to_host);
             render_process_table(&state, &mut f, right_side_vertical_halves[0]);
-            render_remote_ip_table(&state, &mut f, right_side_vertical_halves[1]);
+            render_remote_ip_table(&state, &mut f, right_side_vertical_halves[1], ip_to_host);
         })
         .unwrap();
 }
diff --git a/src/display/ui_state.rs b/src/display/ui_state.rs
index 1d3233b..83c1382 100644
--- a/src/display/ui_state.rs
+++ b/src/display/ui_state.rs
@@ -1,4 +1,5 @@
 use ::std::collections::{BTreeMap, HashMap};
+use ::std::net::Ipv4Addr;
 
 use crate::network::{Connection, Utilization};
 
@@ -41,7 +42,7 @@ impl Bandwidth for NetworkData {
 
 pub struct UIState {
     pub processes: BTreeMap<String, NetworkData>,
-    pub remote_ips: BTreeMap<String, NetworkData>,
+    pub remote_ips: BTreeMap<Ipv4Addr, NetworkData>,
     pub connections: BTreeMap<Connection, ConnectionData>,
 }
 
@@ -51,14 +52,14 @@ impl UIState {
         network_utilization: &Utilization,
     ) -> Self {
         let mut processes: BTreeMap<String, NetworkData> = BTreeMap::new();
-        let mut remote_ips: BTreeMap<String, NetworkData> = BTreeMap::new();
+        let mut remote_ips: BTreeMap<Ipv4Addr, NetworkData> = BTreeMap::new();
         let mut connections: BTreeMap<Connection, ConnectionData> = BTreeMap::new();
         for (connection, process_name) in connections_to_procs {
             if let Some(connection_bandwidth_utilization) =
                 network_utilization.connections.get(&connection)
             {
                 let data_for_remote_ip = remote_ips
-                    .entry(connection.remote_socket.clone_host_or_ip())
+                    .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();
diff --git a/src/main.rs b/src/main.rs
index 730fabc..dbebc46 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -5,7 +5,7 @@ mod os;
 mod tests;
 
 use display::display_loop;
-use network::{resolve_connections, Connection, DnsQueue, Sniffer, Utilization};
+use network::{Connection, DnsQueue, Sniffer, Utilization};
 
 use ::std::net::IpAddr;
 
@@ -121,18 +121,22 @@ where
             terminal.clear().unwrap();
             terminal.hide_cursor().unwrap();
             while running.load(Ordering::Relaxed) {
-                let connections_to_procs = {
-                    let open_sockets = get_open_sockets();
-                    let ip_to_host = ip_to_host.lock().unwrap();
-                    let (unresolved_ips, connections_to_procs) =
-                        resolve_connections(open_sockets, &ip_to_host);
-                    dns_queue.add_ips_to_resolve(unresolved_ips);
-                    connections_to_procs
-                };
                 {
+                    let connections_to_procs = get_open_sockets();
+                    let ip_to_host = ip_to_host.lock().unwrap();
+                    let unresolved_ips = connections_to_procs.keys().fold(vec![], |mut unresolved_ips, connection| {
+                        if !ip_to_host.contains_key(&connection.local_socket.ip) {
+                            unresolved_ips.push(connection.local_socket.ip.clone());
+                        }
+                        if !ip_to_host.contains_key(&connection.remote_socket.ip) {
+                            unresolved_ips.push(connection.remote_socket.ip.clone());
+                        }
+                        unresolved_ips
+                    });
                     let mut network_utilization = network_utilization.lock().unwrap();
                     let utilization = network_utilization.clone_and_reset();
-                    display_loop(&utilization, &mut terminal, connections_to_procs);
+                    dns_queue.add_ips_to_resolve(unresolved_ips);
+                    display_loop(&utilization, &mut terminal, connections_to_procs, &ip_to_host);
                 }
                 thread::sleep(time::Duration::from_secs(1));
             }
diff --git a/src/network/connection.rs b/src/network/connection.rs
index 5af3c32..e778f88 100644
--- a/src/network/connection.rs
+++ b/src/network/connection.rs
@@ -4,9 +4,6 @@ use ::std::net::Ipv4Addr;
 use ::std::mem::swap;
 use ::std::net::SocketAddr;
 
-use std::cmp::Ordering;
-use std::hash::{Hash, Hasher};
-
 #[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord)]
 pub enum Protocol {
     Tcp,
@@ -22,60 +19,10 @@ impl fmt::Display for Protocol {
     }
 }
 
-#[derive(Clone)]
+#[derive(Clone, Ord, PartialOrd, PartialEq, Eq, Hash)]
 pub struct Socket {
     pub ip: Ipv4Addr,
     pub port: u16,
-    host_addr: Option<String>,
-}
-
-impl Socket {
-    pub fn clone_host_or_ip(&self) -> String {
-        match &self.host_addr {
-            Some(host_addr) => host_addr.clone(),
-            None => self.ip.to_string(),
-        }
-    }
-}
-
-impl Ord for Socket {
-    fn cmp(&self, other: &Self) -> Ordering {
-        let ip_eq = self.ip.cmp(&other.ip); // TODO: also port
-        match ip_eq {
-            Ordering::Equal => self.port.cmp(&other.port),
-            _ => ip_eq,
-        }
-    }
-}
-
-impl PartialOrd for Socket {
-    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
-        Some(self.cmp(other))
-    }
-}
-
-impl PartialEq for Socket {
-    fn eq(&self, other: &Self) -> bool {
-        self.ip == other.ip && self.port == other.port
-    }
-}
-
-impl Eq for Socket {}
-
-impl Hash for Socket {
-    fn hash<H: Hasher>(&self, state: &mut H) {
-        self.ip.hash(state);
-        self.port.hash(state);
-    }
-}
-
-impl fmt::Display for Socket {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        match &self.host_addr {
-            Some(host_addr) => write!(f, "{}:{}", host_addr, self.port),
-            None => write!(f, "{}:{}", self.ip, self.port),
-        }
-    }
 }
 
 #[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord)]
@@ -85,16 +32,6 @@ pub struct Connection {
     pub protocol: Protocol,
 }
 
-impl fmt::Display for Connection {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(
-            f,
-            "{} => {} ({})",
-            self.local_socket, self.remote_socket, self.protocol
-        )
-    }
-}
-
 impl Connection {
     pub fn new(
         local_socket: SocketAddr,
@@ -106,12 +43,10 @@ impl Connection {
                 local_socket: Socket {
                     ip: *local_socket.ip(),
                     port: local_socket.port(),
-                    host_addr: None,
                 },
                 remote_socket: Socket {
                     ip: *remote_socket.ip(),
                     port: remote_socket.port(),
-                    host_addr: None,
                 },
                 protocol,
             }),
@@ -121,10 +56,4 @@ impl Connection {
     pub fn swap_direction(&mut self) {
         swap(&mut self.local_socket, &mut self.remote_socket);
     }
-    pub fn set_local_host_addr(&mut self, addr: &str) {
-        self.local_socket.host_addr = Some(String::from(addr));
-    }
-    pub fn set_remote_host_addr(&mut self, addr: &str) {
-        self.remote_socket.host_addr = Some(String::from(addr));
-    }
 }
diff --git a/src/network/mod.rs b/src/network/mod.rs
index 7590875..033414b 100644
--- a/src/network/mod.rs
+++ b/src/network/mod.rs
@@ -1,11 +1,9 @@
 mod connection;
 mod dns_queue;
-mod resolve_connections;
 mod sniffer;
 mod utilization;
 
 pub use connection::*;
 pub use dns_queue::*;
-pub use resolve_connections::*;
 pub use sniffer::*;
 pub use utilization::*;
diff --git a/src/network/resolve_connections.rs b/src/network/resolve_connections.rs
deleted file mode 100644
index be09d2d..0000000
--- a/src/network/resolve_connections.rs
+++ /dev/null
@@ -1,26 +0,0 @@
-use crate::Connection;
-
-use ::std::collections::HashMap;
-use ::std::net::Ipv4Addr;
-
-pub fn resolve_connections(
-    open_sockets: HashMap<Connection, String>,
-    ip_to_host: &HashMap<Ipv4Addr, String>,
-) -> (Vec<Ipv4Addr>, HashMap<Connection, String>) {
-    let mut unresolved_ips = vec![];
-    let mut resolved_connections_to_procs: HashMap<Connection, String> = HashMap::new();
-    for connection in open_sockets.keys() {
-        let mut connection = connection.clone();
-        match ip_to_host.get(&connection.local_socket.ip) {
-            Some(local_host_addr) => connection.set_local_host_addr(local_host_addr),
-            None => unresolved_ips.push(connection.local_socket.ip.clone()),
-        }
-        match ip_to_host.get(&connection.remote_socket.ip) {
-            Some(remote_host_addr) => connection.set_remote_host_addr(remote_host_addr),
-            None => unresolved_ips.push(connection.remote_socket.ip.clone()),
-        }
-        let connection_value = open_sockets.get(&connection).unwrap().clone();
-        &resolved_connections_to_procs.insert(connection, connection_value);
-    }
-    (unresolved_ips, resolved_connections_to_procs)
-}