1 files changed, 19 insertions, 10 deletions
diff --git a/src/network/connection.rs b/src/network/connection.rs
index b4e0de2..a8a17cb 100644
--- a/src/network/connection.rs
+++ b/src/network/connection.rs
@@ -1,6 +1,6 @@
 use ::std::collections::HashMap;
 use ::std::fmt;
-use ::std::net::Ipv4Addr;
+use ::std::net::{Ipv4Addr, IpAddr};
 
 use ::std::net::SocketAddr;
 
@@ -29,17 +29,23 @@ impl fmt::Display for Protocol {
     }
 }
 
-#[derive(Clone, Ord, PartialOrd, PartialEq, Eq, Hash)]
+#[derive(Clone, Ord, PartialOrd, PartialEq, Eq, Hash, Debug)]
 pub struct Socket {
     pub ip: Ipv4Addr,
     pub port: u16,
 }
 
-#[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord)]
+#[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord, Debug)]
+pub struct LocalSocket {
+    pub ip: IpAddr,
+    pub port: u16,
+    pub protocol: Protocol,
+}
+
+#[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord, Debug)]
 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 +63,26 @@ 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,
         }