feat(dns): lazily lookup hostnames for ips

1 files changed, 76 insertions, 16 deletions

diff --git a/src/network/connection.rs b/src/network/connection.rs
index 8676ec1..34194c1 100644
--- a/src/network/connection.rs
+++ b/src/network/connection.rs
@@ -4,7 +4,10 @@ use ::std::net::Ipv4Addr;
 use ::std::mem::swap;
 use ::std::net::SocketAddr;
 
-#[derive(PartialEq, Hash, Eq, Debug, Clone, PartialOrd, Ord)]
+use std::cmp::Ordering;
+use std::hash::{Hash, Hasher};
+
+#[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord)]
 pub enum Protocol {
     Tcp,
     Udp,
@@ -19,19 +22,67 @@ impl fmt::Display for Protocol {
     }
 }
 
-#[derive(PartialEq, Hash, Eq, Debug, Clone, PartialOrd, Ord)]
+#[derive(Clone)]
 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 {
-        write!(f, "{}:{}", self.ip, self.port)
+        match &self.host_addr {
+            Some(host_addr) => {
+                write!(f, "{}:{}", host_addr, self.port)
+            },
+            None => {
+                write!(f, "{}:{}", self.ip, self.port)
+            }
+        }
     }
 }
 
-#[derive(PartialEq, Hash, Eq, Debug, Clone, PartialOrd, Ord)]
+#[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord)]
 pub struct Connection {
     pub local_socket: Socket,
     pub remote_socket: Socket,
@@ -55,22 +106,31 @@ impl Connection {
         protocol: Protocol,
     ) -> Option<Self> {
         match (local_socket, remote_socket) {
-            (SocketAddr::V4(local_socket), SocketAddr::V4(remote_socket)) => Some(Connection {
-                // we use our own Socket here because SocketAddr is not sorable
-                local_socket: Socket {
-                    ip: *local_socket.ip(),
-                    port: local_socket.port(),
-                },
-                remote_socket: Socket {
-                    ip: *remote_socket.ip(),
-                    port: remote_socket.port(),
-                },
-                protocol,
-            }),
+            (SocketAddr::V4(local_socket), SocketAddr::V4(remote_socket)) => {
+                Some(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,
+                })
+            },
             (_, _) => None,
         }
     }
     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));
+    }
 }