refactor(traffic): use own socket struct so it sorts

3 files changed, 57 insertions, 50 deletions

diff --git a/src/display/ui_state.rs b/src/display/ui_state.rs
index b8cbb1d..2fff47a 100644
--- a/src/display/ui_state.rs
+++ b/src/display/ui_state.rs
@@ -59,8 +59,9 @@ impl UIState {
             if let Some(connection_bandwidth_utilization) =
                 network_utilization.connections.get(&connection)
             {
-                let data_for_remote_ip =
-                    remote_ips.entry(connection.remote_ip.clone()).or_default();
+                let data_for_remote_ip = remote_ips
+                    .entry(connection.remote_socket.ip.clone())
+                    .or_default();
                 let connection_data = connections.entry(connection).or_default();
                 for process in &associated_processes {
                     let data_for_process = processes.entry(process.to_string()).or_default();
@@ -70,9 +71,7 @@ impl UIState {
                         &connection_bandwidth_utilization.total_bytes_uploaded;
                     data_for_process.connection_count += 1;
                 }
-                connection_data
-                    .processes
-                    .append(&mut associated_processes);
+                connection_data.processes.append(&mut associated_processes);
                 connection_data.total_bytes_downloaded +=
                     &connection_bandwidth_utilization.total_bytes_downloaded;
                 connection_data.total_bytes_uploaded +=
diff --git a/src/store/current_connections.rs b/src/store/current_connections.rs
index 43ed829..fc7b5fc 100644
--- a/src/store/current_connections.rs
+++ b/src/store/current_connections.rs
@@ -1,4 +1,4 @@
-use crate::traffic::{Connection, Protocol};
+use crate::traffic::{Connection, Protocol, Socket};
 
 use ::netstat::{ProtocolSocketInfo, SocketInfo};
 use ::std::collections::HashMap;
@@ -20,10 +20,14 @@ fn build_ipv4_connection(
 ) -> Option<Connection> {
     match (local_ip, remote_ip) {
         (Some(local_ip), Some(remote_ip)) => Some(Connection {
-            local_ip,
-            remote_ip,
-            local_port,
-            remote_port,
+            local_socket: Socket {
+                ip: local_ip,
+                port: local_port,
+            },
+            remote_socket: Socket {
+                ip: remote_ip,
+                port: remote_port,
+            },
             protocol,
         }),
         (_, _) => None,
diff --git a/src/traffic.rs b/src/traffic.rs
index 72cff7c..7b475b3 100644
--- a/src/traffic.rs
+++ b/src/traffic.rs
@@ -1,6 +1,6 @@
 use ::std::boxed::Box;
 use ::std::fmt;
-use ::std::net::{Ipv4Addr, SocketAddrV4};
+use ::std::net::Ipv4Addr;
 
 use ::pnet::datalink::{DataLinkReceiver, NetworkInterface};
 use ::pnet::packet::ethernet::{EtherType, EthernetPacket};
@@ -12,6 +12,18 @@ use ::pnet::packet::Packet;
 
 use ::ipnetwork::IpNetwork;
 
+#[derive(PartialEq, Hash, Eq, Debug, Clone, PartialOrd, Ord)]
+pub struct Socket {
+    pub ip: Ipv4Addr,
+    pub port: u16,
+}
+
+impl fmt::Display for Socket {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.ip, self.port)
+    }
+}
+
 pub struct Sniffer {
     network_interface: NetworkInterface,
     network_frames: Box<DataLinkReceiver>,
@@ -19,10 +31,8 @@ pub struct Sniffer {
 
 #[derive(PartialEq, Hash, Eq, Debug, Clone, PartialOrd, Ord)]
 pub struct Connection {
-    pub local_ip: Ipv4Addr,
-    pub remote_ip: Ipv4Addr,
-    pub local_port: u16,
-    pub remote_port: u16,
+    pub local_socket: Socket,
+    pub remote_socket: Socket,
     pub protocol: Protocol,
 }
 
@@ -30,8 +40,8 @@ impl fmt::Display for Connection {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(
             f,
-            "{}:{} => {}:{} ({})",
-            self.local_ip, self.local_port, self.remote_ip, self.remote_port, self.protocol
+            "{} => {} ({})",
+            self.local_socket, self.remote_socket, self.protocol
         )
     }
 }
@@ -48,12 +58,6 @@ pub enum Protocol {
     Udp,
 }
 
-#[derive(PartialEq, Hash, Eq, Debug, Clone, PartialOrd)]
-pub enum Direction {
-    Download,
-    Upload,
-}
-
 impl fmt::Display for Protocol {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match *self {
@@ -63,42 +67,39 @@ impl fmt::Display for Protocol {
     }
 }
 
-fn find_direction(network_interface_ips: &Vec<IpNetwork>, ip_packet: &Ipv4Packet) -> Direction {
-    match network_interface_ips
-        .iter()
-        .any(|ip_network| ip_network.ip() == ip_packet.get_source())
-    {
-        true => Direction::Upload,
-        false => Direction::Download,
-    }
+#[derive(PartialEq, Hash, Eq, Debug, Clone, PartialOrd)]
+pub enum Direction {
+    Download,
+    Upload,
 }
 
 impl Direction {
-    pub fn make_connection(
-        &self,
-        from: SocketAddrV4,
-        to: SocketAddrV4,
-        protocol: Protocol,
-    ) -> Connection {
+    pub fn make_connection(&self, from: Socket, to: Socket, protocol: Protocol) -> Connection {
         match self {
             Direction::Upload => Connection {
-                local_ip: *from.ip(),
-                remote_ip: *to.ip(),
-                local_port: from.port(),
-                remote_port: to.port(),
+                local_socket: from,
+                remote_socket: to,
                 protocol,
             },
             Direction::Download => Connection {
-                local_ip: *to.ip(),
-                remote_ip: *from.ip(),
-                local_port: to.port(),
-                remote_port: from.port(),
+                local_socket: to,
+                remote_socket: from,
                 protocol,
             },
         }
     }
 }
 
+fn find_direction(network_interface_ips: &Vec<IpNetwork>, ip_packet: &Ipv4Packet) -> Direction {
+    match network_interface_ips
+        .iter()
+        .any(|ip_network| ip_network.ip() == ip_packet.get_source())
+    {
+        true => Direction::Upload,
+        false => Direction::Download,
+    }
+}
+
 impl Sniffer {
     pub fn new(network_interface: NetworkInterface, network_frames: Box<DataLinkReceiver>) -> Self {
         Sniffer {
@@ -114,13 +115,11 @@ impl Sniffer {
         });
         let packet = EthernetPacket::new(bytes)?;
         match packet.get_ethertype() {
-            // TODO: better way through the module?
             EtherType(2048) => {
                 let ip_packet = Ipv4Packet::new(packet.payload())?;
                 let (protocol, source_port, destination_port) =
                     match ip_packet.get_next_level_protocol() {
                         IpNextHeaderProtocol(6) => {
-                            // tcp
                             let message = TcpPacket::new(ip_packet.payload())?;
                             (
                                 Protocol::Tcp,
@@ -129,7 +128,6 @@ impl Sniffer {
                             )
                         }
                         IpNextHeaderProtocol(17) => {
-                            // udp
                             let datagram = UdpPacket::new(ip_packet.payload())?;
                             (
                                 Protocol::Udp,
@@ -140,8 +138,14 @@ impl Sniffer {
                         _ => return None,
                     };
                 let direction = find_direction(&self.network_interface.ips, &ip_packet);
-                let from = SocketAddrV4::new(ip_packet.get_source(), source_port);
-                let to = SocketAddrV4::new(ip_packet.get_destination(), destination_port);
+                let from = Socket {
+                    ip: ip_packet.get_source(),
+                    port: source_port,
+                };
+                let to = Socket {
+                    ip: ip_packet.get_destination(),
+                    port: destination_port,
+                };
                 let connection = direction.make_connection(from, to, protocol);
                 let ip_length = ip_packet.get_total_length() as u128;
                 Some(Segment {