feat(ui): components, new details and responsive layout

14 files changed, 456 insertions, 206 deletions

diff --git a/src/display/components/display_bandwidth.rs b/src/display/components/display_bandwidth.rs
new file mode 100644
index 0000000..31653e6
--- /dev/null
+++ b/src/display/components/display_bandwidth.rs
@@ -0,0 +1,17 @@
+use ::std::fmt;
+
+pub struct DisplayBandwidth(pub f64);
+
+impl fmt::Display for DisplayBandwidth {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        if self.0 > 999_999_999.0 {
+            write!(f, "{:.2}GBps", self.0 / 1_000_000_000.0)
+        } else if self.0 > 999_999.0 {
+            write!(f, "{:.2}MBps", self.0 / 1_000_000.0)
+        } else if self.0 > 999.0 {
+            write!(f, "{:.2}KBps", self.0 / 1000.0)
+        } else {
+            write!(f, "{}Bps", self.0)
+        }
+    }
+}
diff --git a/src/display/components/layout.rs b/src/display/components/layout.rs
new file mode 100644
index 0000000..27f565a
--- /dev/null
+++ b/src/display/components/layout.rs
@@ -0,0 +1,65 @@
+use ::tui::backend::Backend;
+use ::tui::layout::{Constraint, Direction, Rect};
+use ::tui::terminal::Frame;
+
+use super::Table;
+use super::TotalBandwidth;
+
+const FIRST_HEIGHT_BREAKPOINT: u16 = 30;
+const FIRST_WIDTH_BREAKPOINT: u16 = 120;
+const SECOND_WIDTH_BREAKPOINT: u16 = 150;
+
+fn leave_gap_on_top_of_rect(rect: Rect) -> Rect {
+    let app = ::tui::layout::Layout::default()
+        .direction(Direction::Vertical)
+        .margin(0)
+        .constraints([Constraint::Length(1), Constraint::Length(rect.height - 1)].as_ref())
+        .split(rect);
+    return app[1];
+}
+
+pub struct Layout<'a> {
+    pub header: TotalBandwidth<'a>,
+    pub children: Vec<Table<'a>>,
+}
+
+impl<'a> Layout<'a> {
+    fn split_rect(&self, rect: Rect, splits: Vec<Direction>) -> Vec<Rect> {
+        let mut ret = vec![rect]; // TODO: use fold
+        for direction in splits {
+            let last_split = ret.pop().unwrap();
+            let mut halves = ::tui::layout::Layout::default()
+                .direction(direction)
+                .margin(0)
+                .constraints([Constraint::Percentage(50), Constraint::Percentage(50)].as_ref())
+                .split(last_split);
+            ret.append(&mut halves);
+        }
+        ret
+    }
+    fn get_render_order(&self, rect: &Rect) -> Vec<Rect> {
+        if rect.height < FIRST_HEIGHT_BREAKPOINT && rect.width < FIRST_WIDTH_BREAKPOINT {
+            self.split_rect(*rect, vec![])
+        } else if rect.height < FIRST_HEIGHT_BREAKPOINT {
+            self.split_rect(*rect, vec![Direction::Horizontal])
+        } else if rect.width < FIRST_WIDTH_BREAKPOINT {
+            self.split_rect(*rect, vec![Direction::Vertical])
+        } else if rect.width < SECOND_WIDTH_BREAKPOINT {
+            self.split_rect(*rect, vec![Direction::Vertical, Direction::Horizontal])
+        } else {
+            self.split_rect(*rect, vec![Direction::Horizontal, Direction::Vertical])
+        }
+    }
+    pub fn render(&self, mut frame: &mut Frame<impl Backend>, rect: Rect) {
+        let app = leave_gap_on_top_of_rect(rect);
+        let render_order = self.get_render_order(&app);
+        for i in 0..render_order.len() {
+            if let Some(rect) = render_order.get(i) {
+                if let Some(child) = self.children.get(i) {
+                    child.render(&mut frame, *rect);
+                }
+            }
+        }
+        self.header.render(&mut frame, rect);
+    }
+}
diff --git a/src/display/components/mod.rs b/src/display/components/mod.rs
new file mode 100644
index 0000000..1be5a33
--- /dev/null
+++ b/src/display/components/mod.rs
@@ -0,0 +1,9 @@
+mod table;
+mod layout;
+mod total_bandwidth;
+mod display_bandwidth;
+
+pub use table::*;
+pub use layout::*;
+pub use total_bandwidth::*;
+pub use display_bandwidth::*;
diff --git a/src/display/components/table.rs b/src/display/components/table.rs
new file mode 100644
index 0000000..6cfc2d7
--- /dev/null
+++ b/src/display/components/table.rs
@@ -0,0 +1,185 @@
+use ::std::collections::HashMap;
+
+use ::tui::backend::Backend;
+use ::tui::layout::Rect;
+use ::tui::style::{Color, Style};
+use ::tui::terminal::Frame;
+use ::tui::widgets::{Block, Borders, Row, Widget};
+
+use crate::display::{DisplayBandwidth, Bandwidth, UIState};
+use crate::network::Connection;
+
+use ::std::net::Ipv4Addr;
+use std::iter::FromIterator;
+
+const FIRST_WIDTH_BREAKPOINT: u16 = 50;
+const SECOND_WIDTH_BREAKPOINT: u16 = 71;
+const THIRD_WIDTH_BREAKPOINT: u16 = 95;
+
+const FIRST_COLUMN_WIDTHS: [u16; 4] = [20, 30, 40, 50];
+const SECOND_COLUMN_WIDTHS: [u16; 1] = [20];
+const THIRD_COLUMN_WIDTHS: [u16; 4] = [10, 20, 20, 20];
+
+
+fn display_upload_and_download(bandwidth: &impl Bandwidth) -> String {
+    format!(
+        "{}/{}",
+        DisplayBandwidth(bandwidth.get_total_bytes_uploaded() as f64),
+        DisplayBandwidth(bandwidth.get_total_bytes_downloaded() as f64)
+    )
+}
+
+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 sort_by_bandwidth<'a, T>(
+    list: &'a mut Vec<(T, &impl Bandwidth)>,
+) -> &'a Vec<(T, &'a impl Bandwidth)> {
+    list.sort_by(|(_, a), (_, b)| {
+        let a_highest = if a.get_total_bytes_downloaded() > a.get_total_bytes_uploaded() {
+            a.get_total_bytes_downloaded()
+        } else {
+            a.get_total_bytes_uploaded()
+        };
+        let b_highest = if b.get_total_bytes_downloaded() > b.get_total_bytes_uploaded() {
+            b.get_total_bytes_downloaded()
+        } else {
+            b.get_total_bytes_uploaded()
+        };
+        b_highest.cmp(&a_highest)
+    });
+    list
+}
+
+fn display_connection_string(
+    connection: &Connection,
+    ip_to_host: &HashMap<Ipv4Addr, String>,
+) -> String {
+    format!(
+        ":{} => {}:{} ({})",
+        connection.local_port,
+        display_ip_or_host(connection.remote_socket.ip, ip_to_host),
+        connection.remote_socket.port,
+        connection.protocol,
+    )
+}
+
+pub struct Table<'a> {
+    title: &'a str,
+    column_names: &'a [&'a str],
+    rows: Vec<Vec<String>>,
+}
+
+impl <'a>Table<'a> {
+    pub fn create_connections_table(state: &UIState, ip_to_host: &HashMap<Ipv4Addr, String>) -> Self {
+        let mut connections_list = Vec::from_iter(&state.connections);
+        sort_by_bandwidth(&mut connections_list);
+        let connections_rows = connections_list
+            .iter()
+            .map(|(connection, connection_data)| {
+                vec![
+                    display_connection_string(&connection, &ip_to_host),
+                    connection_data.process_name.to_string(),
+                    display_upload_and_download(*connection_data),
+                ]
+            })
+            .collect();
+        let connections_title = "Utilization by connection";
+        let connections_column_names = &["Connection", "Process", "Rate Up/Down"];
+        Table {
+            title: connections_title,
+            column_names: connections_column_names,
+            rows: connections_rows,
+        }
+    }
+    pub fn create_processes_table(state: &UIState) -> Self {
+        let mut processes_list = Vec::from_iter(&state.processes);
+        sort_by_bandwidth(&mut processes_list);
+        let processes_rows = processes_list
+            .iter()
+            .map(|(process_name, data_for_process)| {
+                vec![
+                    process_name.to_string(),
+                    data_for_process.connection_count.to_string(),
+                    display_upload_and_download(*data_for_process),
+                ]
+            })
+            .collect();
+        let processes_title = "Utilization by process name";
+        let processes_column_names = &["Process", "Connection count", "Rate Up/Down"];
+        Table {
+            title: processes_title,
+            column_names: processes_column_names,
+            rows: processes_rows,
+        }
+    }
+    pub fn create_remote_ips_table(state: &UIState, ip_to_host: &HashMap<Ipv4Addr, String>) -> Self {
+        let mut remote_ips_list = Vec::from_iter(&state.remote_ips);
+        sort_by_bandwidth(&mut remote_ips_list);
+        let remote_ips_rows = remote_ips_list
+            .iter()
+            .map(|(remote_ip, data_for_remote_ip)| {
+                let remote_ip = display_ip_or_host(**remote_ip, &ip_to_host);
+                vec![
+                    remote_ip,
+                    data_for_remote_ip.connection_count.to_string(),
+                    display_upload_and_download(*data_for_remote_ip),
+                ]
+            })
+            .collect();
+        let remote_ips_title = "Utilization by remote ip";
+        let remote_ips_column_names =
+            &["Remote Address", "Connection Count", "Rate Up/Down"];
+        Table {
+            title: remote_ips_title,
+            column_names: remote_ips_column_names,
+            rows: remote_ips_rows,
+        }
+    }
+    pub fn render(&self, frame: &mut Frame<impl Backend>, rect: Rect) {
+        // the second column is only rendered if there is enough room for it
+        // (over third breakpoint)
+        let widths = if rect.width < FIRST_WIDTH_BREAKPOINT {
+            vec![FIRST_COLUMN_WIDTHS[0], THIRD_COLUMN_WIDTHS[0]]
+        } else if rect.width < SECOND_WIDTH_BREAKPOINT {
+            vec![FIRST_COLUMN_WIDTHS[1], THIRD_COLUMN_WIDTHS[1]]
+        } else if rect.width < THIRD_WIDTH_BREAKPOINT {
+            vec![FIRST_COLUMN_WIDTHS[2], THIRD_COLUMN_WIDTHS[2]]
+        } else {
+            vec![FIRST_COLUMN_WIDTHS[3], SECOND_COLUMN_WIDTHS[0], THIRD_COLUMN_WIDTHS[2]]
+        };
+
+        let column_names = if rect.width < THIRD_WIDTH_BREAKPOINT {
+            vec![self.column_names[0], self.column_names[2]]
+        } else {
+            vec![
+                self.column_names[0],
+                self.column_names[1],
+                self.column_names[2],
+            ]
+        };
+
+        let rows = self.rows.iter().map(|row| {
+            if rect.width < THIRD_WIDTH_BREAKPOINT {
+                vec![&row[0], &row[2]]
+            } else {
+                vec![&row[0], &row[1], &row[2]]
+            }
+        });
+
+        let table_rows =
+            rows.map(|row| Row::StyledData(row.into_iter(), Style::default().fg(Color::White)));
+
+        ::tui::widgets::Table::new(column_names.into_iter(), table_rows)
+            .block(Block::default().title(self.title).borders(Borders::ALL))
+            .header_style(Style::default().fg(Color::Yellow))
+            .widths(&widths[..])
+            .style(Style::default().fg(Color::White))
+            .column_spacing(2)
+            .render(frame, rect);
+    }
+}
diff --git a/src/display/components/total_bandwidth.rs b/src/display/components/total_bandwidth.rs
new file mode 100644
index 0000000..823b972
--- /dev/null
+++ b/src/display/components/total_bandwidth.rs
@@ -0,0 +1,28 @@
+use ::tui::backend::Backend;
+use ::tui::layout::{Alignment, Rect};
+use ::tui::style::{Color, Modifier, Style};
+use ::tui::terminal::Frame;
+use ::tui::widgets::{Block, Borders, Paragraph, Text, Widget};
+
+use crate::display::{DisplayBandwidth, UIState};
+
+pub struct TotalBandwidth<'a> {
+    pub state: &'a UIState,
+}
+
+impl<'a>TotalBandwidth<'a> {
+    pub fn render(&self, frame: &mut Frame<impl Backend>, rect: Rect) {
+        let title_text = [Text::styled(
+            format!(
+                " Total Rate Up/Down: {}/{}",
+                DisplayBandwidth(self.state.total_bytes_uploaded as f64),
+                DisplayBandwidth(self.state.total_bytes_downloaded as f64)
+            ),
+            Style::default().fg(Color::Green).modifier(Modifier::BOLD),
+        )];
+        Paragraph::new(title_text.iter())
+            .block(Block::default().borders(Borders::NONE))
+            .alignment(Alignment::Left)
+            .render(frame, rect);
+    }
+}
diff --git a/src/display/mod.rs b/src/display/mod.rs
index c586665..73b3a9c 100644
--- a/src/display/mod.rs
+++ b/src/display/mod.rs
@@ -1,5 +1,7 @@
 mod ui;
 mod ui_state;
+mod components;
 
 pub use ui::*;
 pub use ui_state::*;
+pub use components::*;
diff --git a/src/display/ui.rs b/src/display/ui.rs
index a3f9b58..e603601 100644
--- a/src/display/ui.rs
+++ b/src/display/ui.rs
@@ -1,159 +1,14 @@
 use ::std::collections::HashMap;
-use ::std::fmt;
+
 use ::tui::backend::Backend;
-use ::tui::layout::{Constraint, Direction, Layout, Rect};
-use ::tui::style::{Color, Style};
-use ::tui::terminal::Frame;
-use ::tui::widgets::{Block, Borders, Row, Table, Widget};
 use ::tui::Terminal;
 
-use crate::display::{Bandwidth, UIState};
+use crate::display::UIState;
 use crate::network::{Connection, Utilization};
+use crate::display::components::{Table, Layout, TotalBandwidth};
 
 use ::std::net::Ipv4Addr;
 
-struct DisplayBandwidth(f64);
-
-impl fmt::Display for DisplayBandwidth {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        if self.0 > 999_999_999.0 {
-            write!(f, "{:.2}GBps", self.0 / 1_000_000_000.0)
-        } else if self.0 > 999_999.0 {
-            write!(f, "{:.2}MBps", self.0 / 1_000_000.0)
-        } else if self.0 > 999.0 {
-            write!(f, "{:.2}KBps", self.0 / 1000.0)
-        } else {
-            write!(f, "{}Bps", self.0)
-        }
-    }
-}
-
-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],
-    rows: impl Iterator<Item = Vec<String>> + 'a,
-    widths: &'a [u16],
-) -> impl Widget + 'a {
-    let table_rows =
-        rows.map(|row| Row::StyledData(row.into_iter(), Style::default().fg(Color::White)));
-    Table::new(column_names.iter(), table_rows)
-        .block(Block::default().title(title).borders(Borders::ALL))
-        .header_style(Style::default().fg(Color::Yellow))
-        .widths(widths)
-        .style(Style::default().fg(Color::White))
-        .column_spacing(1)
-}
-
-fn format_row_data(
-    first_cell: String,
-    second_cell: String,
-    bandwidth: &impl Bandwidth,
-) -> Vec<String> {
-    vec![
-        first_cell,
-        second_cell,
-        format!(
-            "{}/{}",
-            DisplayBandwidth(bandwidth.get_total_bytes_uploaded() as f64),
-            DisplayBandwidth(bandwidth.get_total_bytes_downloaded() as f64)
-        ),
-    ]
-}
-
-fn split(direction: Direction, rect: Rect) -> Vec<Rect> {
-    Layout::default()
-        .direction(direction)
-        .margin(0)
-        .constraints([Constraint::Percentage(50), Constraint::Percentage(50)].as_ref())
-        .split(rect)
-}
-
-fn render_process_table(state: &UIState, frame: &mut Frame<impl Backend>, rect: Rect) {
-    let rows = state
-        .processes
-        .iter()
-        .map(|(process_name, data_for_process)| {
-            format_row_data(
-                process_name.to_string(),
-                data_for_process.connection_count.to_string(),
-                data_for_process,
-            )
-        });
-    let mut table = create_table(
-        "Utilization by process name",
-        &["Process", "Connection Count", "Total Bytes"],
-        rows,
-        &[30, 30, 30],
-    );
-    table.render(frame, 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_string,
-                connection_data.process_name.to_string(),
-                connection_data,
-            )
-        });
-    let mut table = create_table(
-        "Utilization by connection",
-        &["Connection", "Processes", "Total Bytes Up/Down"],
-        rows,
-        &[50, 20, 20],
-    );
-    table.render(frame, 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,
-                data_for_remote_ip.connection_count.to_string(),
-                data_for_remote_ip,
-            )
-        });
-    let mut table = create_table(
-        "Utilization by remote ip",
-        &["Remote Address", "Connection Count", "Total Bytes"],
-        rows,
-        &[50, 20, 20],
-    );
-    table.render(frame, rect);
-}
-
 pub struct Ui<B>
 where
     B: Backend,
@@ -181,13 +36,17 @@ where
         let state = &self.state;
         let ip_to_host = &self.ip_to_host;
         self.terminal
-            .draw(|mut f| {
-                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], 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], ip_to_host);
+            .draw(|mut frame| {
+                let size = frame.size();
+                let connections = Table::create_connections_table(&state, &ip_to_host);
+                let processes = Table::create_processes_table(&state);
+                let remote_ips = Table::create_remote_ips_table(&state, &ip_to_host);
+                let total_bandwidth = TotalBandwidth { state: &state };
+                let layout = Layout {
+                    header: total_bandwidth,
+                    children: vec![connections, processes, remote_ips],
+                };
+                layout.render(&mut frame, size);
             })
             .unwrap();
     }
diff --git a/src/display/ui_state.rs b/src/display/ui_state.rs
index 440260d..ac60933 100644
--- a/src/display/ui_state.rs
+++ b/src/display/ui_state.rs
@@ -45,6 +45,8 @@ pub struct UIState {
     pub processes: BTreeMap<String, NetworkData>,
     pub remote_ips: BTreeMap<Ipv4Addr, NetworkData>,
     pub connections: BTreeMap<Connection, ConnectionData>,
+    pub total_bytes_downloaded: u128,
+    pub total_bytes_uploaded: u128,
 }
 
 impl UIState {
@@ -55,6 +57,8 @@ impl UIState {
         let mut processes: BTreeMap<String, NetworkData> = BTreeMap::new();
         let mut remote_ips: BTreeMap<Ipv4Addr, NetworkData> = BTreeMap::new();
         let mut connections: BTreeMap<Connection, ConnectionData> = BTreeMap::new();
+        let mut total_bytes_downloaded: u128 = 0;
+        let mut total_bytes_uploaded: u128 = 0;
         for (connection, process_name) in connections_to_procs {
             if let Some(connection_bandwidth_utilization) =
                 network_utilization.connections.get(&connection)
@@ -78,12 +82,16 @@ impl UIState {
                 data_for_remote_ip.total_bytes_uploaded +=
                     connection_bandwidth_utilization.total_bytes_uploaded;
                 data_for_remote_ip.connection_count += 1;
+                total_bytes_downloaded += connection_bandwidth_utilization.total_bytes_downloaded;
+                total_bytes_uploaded += connection_bandwidth_utilization.total_bytes_uploaded;
             }
         }
         UIState {
             processes,
             remote_ips,
             connections,
+            total_bytes_downloaded,
+            total_bytes_uploaded,
         }
     }
 }
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,
diff --git a/src/os/linux.rs b/src/os/linux.rs
index 39bb5c0..c63bdf6 100644
--- a/src/os/linux.rs
+++ b/src/os/linux.rs
@@ -67,8 +67,9 @@ pub fn get_open_sockets() -> HashMap<Connection, String> {
 
     let tcp = ::procfs::tcp().unwrap();
     for entry in tcp.into_iter() {
+        let local_port = entry.local_address.port();
         if let (Some(connection), Some(procname)) = (
-            Connection::new(entry.local_address, entry.remote_address, Protocol::Tcp),
+            Connection::new(entry.remote_address, local_port, Protocol::Tcp),
             inode_to_procname.get(&entry.inode),
         ) {
             open_sockets.insert(connection, procname.clone());
@@ -77,8 +78,9 @@ pub fn get_open_sockets() -> HashMap<Connection, String> {
 
     let udp = ::procfs::udp().unwrap();
     for entry in udp.into_iter() {
+        let local_port = entry.local_address.port();
         if let (Some(connection), Some(procname)) = (
-            Connection::new(entry.local_address, entry.remote_address, Protocol::Udp),
+            Connection::new(entry.remote_address, local_port, Protocol::Udp),
             inode_to_procname.get(&entry.inode),
         ) {
             open_sockets.insert(connection, procname.clone());
diff --git a/src/tests/fakes/fake_output.rs b/src/tests/fakes/fake_output.rs
index 2112371..1950dec 100644
--- a/src/tests/fakes/fake_output.rs
+++ b/src/tests/fakes/fake_output.rs
@@ -26,16 +26,14 @@ impl TestBackend {
     pub fn new(
         log: Arc<Mutex<Vec<TerminalEvent>>>,
         draw_log: Arc<Mutex<Vec<String>>>,
-        terminal_width: Arc<Mutex<u16>>, // 190
-        terminal_height: Arc<Mutex<u16>>, // 50
+        terminal_width: Arc<Mutex<u16>>,
+        terminal_height: Arc<Mutex<u16>>,
     ) -> TestBackend {
         TestBackend {
             events: log,
             draw_events: draw_log,
             terminal_width,
             terminal_height,
-//            terminal_width: 190,
-//            terminal_height: 50,
         }
     }
 }
diff --git a/src/tests/mod.rs b/src/tests/mod.rs
index bb6da95..846e91a 100644
--- a/src/tests/mod.rs
+++ b/src/tests/mod.rs
@@ -64,7 +64,12 @@ fn basic_startup() {
     let terminal_events = LogWithMirror::new(Vec::new());
     let terminal_draw_events = LogWithMirror::new(Vec::new());
 
-    let backend = TestBackend::new(terminal_events.write, terminal_draw_events.write, terminal_width, terminal_height);
+    let backend = TestBackend::new(
+        terminal_events.write,
+        terminal_draw_events.write,
+        terminal_width,
+        terminal_height,
+    );
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap::new());
     let receive_winch = Box::new(create_fake_receive_winch(false));
@@ -109,7 +114,12 @@ fn one_packet_of_traffic() {
     let terminal_events = LogWithMirror::new(Vec::new());
     let terminal_draw_events = LogWithMirror::new(Vec::new());
 
-    let backend = TestBackend::new(terminal_events.write, terminal_draw_events.write, terminal_width, terminal_height);
+    let backend = TestBackend::new(
+        terminal_events.write,
+        terminal_draw_events.write,
+        terminal_width,
+        terminal_height,
+    );
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap::new());
     let receive_winch = Box::new(create_fake_receive_winch(false));
@@ -166,7 +176,12 @@ fn bi_directional_traffic() {
     let terminal_events = LogWithMirror::new(Vec::new());
     let terminal_draw_events = LogWithMirror::new(Vec::new());
 
-    let backend = TestBackend::new(terminal_events.write, terminal_draw_events.write, terminal_width, terminal_height);
+    let backend = TestBackend::new(
+        terminal_events.write,
+        terminal_draw_events.write,
+        terminal_width,
+        terminal_height,
+    );
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap::new());
     let receive_winch = Box::new(create_fake_receive_winch(false));
@@ -223,7 +238,12 @@ fn multiple_packets_of_traffic_from_different_connections() {
     let terminal_events = LogWithMirror::new(Vec::new());
     let terminal_draw_events = LogWithMirror::new(Vec::new());
 
-    let backend = TestBackend::new(terminal_events.write, terminal_draw_events.write, terminal_width, terminal_height);
+    let backend = TestBackend::new(
+        terminal_events.write,
+        terminal_draw_events.write,
+        terminal_width,
+        terminal_height,
+    );
     let receive_winch = Box::new(create_fake_receive_winch(false));
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap::new());
@@ -280,7 +300,12 @@ fn multiple_packets_of_traffic_from_single_connection() {
     let terminal_events = LogWithMirror::new(Vec::new());
     let terminal_draw_events = LogWithMirror::new(Vec::new());
 
-    let backend = TestBackend::new(terminal_events.write, terminal_draw_events.write, terminal_width, terminal_height);
+    let backend = TestBackend::new(
+        terminal_events.write,
+        terminal_draw_events.write,
+        terminal_width,
+        terminal_height,
+    );
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap::new());
     let receive_winch = Box::new(create_fake_receive_winch(false));
@@ -337,7 +362,12 @@ fn one_process_with_multiple_connections() {
     let terminal_events = LogWithMirror::new(Vec::new());
     let terminal_draw_events = LogWithMirror::new(Vec::new());
 
-    let backend = TestBackend::new(terminal_events.write, terminal_draw_events.write, terminal_width, terminal_height);
+    let backend = TestBackend::new(
+        terminal_events.write,
+        terminal_draw_events.write,
+        terminal_width,
+        terminal_height,
+    );
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap::new());
     let receive_winch = Box::new(create_fake_receive_winch(false));
@@ -408,7 +438,12 @@ fn multiple_processes_with_multiple_connections() {
     let terminal_events = LogWithMirror::new(Vec::new());
     let terminal_draw_events = LogWithMirror::new(Vec::new());
 
-    let backend = TestBackend::new(terminal_events.write, terminal_draw_events.write, terminal_width, terminal_height);
+    let backend = TestBackend::new(
+        terminal_events.write,
+        terminal_draw_events.write,
+        terminal_width,
+        terminal_height,
+    );
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap::new());
     let receive_winch = Box::new(create_fake_receive_winch(false));
@@ -465,7 +500,12 @@ fn multiple_connections_from_remote_ip() {
     let terminal_events = LogWithMirror::new(Vec::new());
     let terminal_draw_events = LogWithMirror::new(Vec::new());
 
-    let backend = TestBackend::new(terminal_events.write, terminal_draw_events.write, terminal_width, terminal_height);
+    let backend = TestBackend::new(
+        terminal_events.write,
+        terminal_draw_events.write,
+        terminal_width,
+        terminal_height,
+    );
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap::new());
     let receive_winch = Box::new(create_fake_receive_winch(false));
@@ -524,7 +564,12 @@ fn sustained_traffic_from_one_process() {
     let terminal_events = LogWithMirror::new(Vec::new());
     let terminal_draw_events = LogWithMirror::new(Vec::new());
 
-    let backend = TestBackend::new(terminal_events.write, terminal_draw_events.write, terminal_width, terminal_height);
+    let backend = TestBackend::new(
+        terminal_events.write,
+        terminal_draw_events.write,
+        terminal_width,
+        terminal_height,
+    );
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap::new());
     let receive_winch = Box::new(create_fake_receive_winch(false));
@@ -597,7 +642,12 @@ fn sustained_traffic_from_multiple_processes() {
     let terminal_events = LogWithMirror::new(Vec::new());
     let terminal_draw_events = LogWithMirror::new(Vec::new());
 
-    let backend = TestBackend::new(terminal_events.write, terminal_draw_events.write, terminal_width, terminal_height);
+    let backend = TestBackend::new(
+        terminal_events.write,
+        terminal_draw_events.write,
+        terminal_width,
+        terminal_height,
+    );
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap: