3 files changed, 11 insertions, 24 deletions
diff --git a/src/network/connection.rs b/src/network/connection.rs
index e778f88..9a47fb0 100644
--- a/src/network/connection.rs
+++ b/src/network/connection.rs
@@ -1,7 +1,6 @@
 use ::std::fmt;
 use ::std::net::Ipv4Addr;
 
-use ::std::mem::swap;
 use ::std::net::SocketAddr;
 
 #[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord)]
@@ -27,33 +26,23 @@ pub struct Socket {
 
 #[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord)]
 pub struct Connection {
-    pub local_socket: Socket,
     pub remote_socket: Socket,
     pub protocol: Protocol,
+    pub local_port: u16,
 }
 
 impl Connection {
-    pub fn new(
-        local_socket: SocketAddr,
-        remote_socket: SocketAddr,
-        protocol: Protocol,
-    ) -> Option<Self> {
-        match (local_socket, remote_socket) {
-            (SocketAddr::V4(local_socket), SocketAddr::V4(remote_socket)) => Some(Connection {
-                local_socket: Socket {
-                    ip: *local_socket.ip(),
-                    port: local_socket.port(),
-                },
+    pub fn new(remote_socket: SocketAddr, 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,
             }),
-            (_, _) => None,
+            _ => None,
         }
     }
-    pub fn swap_direction(&mut self) {
-        swap(&mut self.local_socket, &mut self.remote_socket);
-    }
 }
diff --git a/src/network/dns_queue.rs b/src/network/dns_queue.rs
index ddb2d9a..dea30ca 100644
--- a/src/network/dns_queue.rs
+++ b/src/network/dns_queue.rs
@@ -26,9 +26,6 @@ impl DnsQueue {
     ) {
         let mut queue = self.jobs.lock().unwrap();
         for connection in connections_to_procs.keys() {
-            if !ip_to_host.contains_key(&connection.local_socket.ip) {
-                queue.push(Some(connection.local_socket.ip));
-            }
             if !ip_to_host.contains_key(&connection.remote_socket.ip) {
                 queue.push(Some(connection.remote_socket.ip));
             }
diff --git a/src/network/sniffer.rs b/src/network/sniffer.rs
index 4c1c3e2..4957681 100644
--- a/src/network/sniffer.rs
+++ b/src/network/sniffer.rs
@@ -81,10 +81,11 @@ impl Sniffer {
                 let direction = Direction::new(&self.network_interface.ips, &ip_packet);
                 let from = SocketAddr::new(IpAddr::V4(ip_packet.get_source()), source_port);
                 let to = SocketAddr::new(IpAddr::V4(ip_packet.get_destination()), destination_port);
-                let mut connection = Connection::new(from, to, protocol)?;
-                if let Direction::Download = direction {
-                    connection.swap_direction();
-                }
+
+                let connection = match direction {
+                    Direction::Download => Connection::new(from, destination_port, protocol)?,
+                    Direction::Upload => Connection::new(to, source_port, protocol)?,
+                };
                 Some(Segment {
                     connection,
                     data_length,