feat(ui): react to SIGWINCH

8 files changed, 176 insertions, 74 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 69f5b09..79e5a75 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -22,6 +22,11 @@ dependencies = [
 ]
 
 [[package]]
+name = "arc-swap"
+version = "0.4.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+
+[[package]]
 name = "assert_cmd"
 version = "0.10.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -902,6 +907,24 @@ dependencies = [
 ]
 
 [[package]]
+name = "signal-hook"
+version = "0.1.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+dependencies = [
+ "libc 0.2.60 (registry+https://github.com/rust-lang/crates.io-index)",
+ "signal-hook-registry 1.1.1 (registry+https://github.com/rust-lang/crates.io-index)",
+]
+
+[[package]]
+name = "signal-hook-registry"
+version = "1.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+dependencies = [
+ "arc-swap 0.4.3 (registry+https://github.com/rust-lang/crates.io-index)",
+ "libc 0.2.60 (registry+https://github.com/rust-lang/crates.io-index)",
+]
+
+[[package]]
 name = "smallvec"
 version = "0.6.10"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1189,6 +1212,7 @@ dependencies = [
  "pnet 0.22.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "pnet_base 0.22.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "procfs 0.5.3 (registry+https://github.com/rust-lang/crates.io-index)",
+ "signal-hook 0.1.10 (registry+https://github.com/rust-lang/crates.io-index)",
  "structopt 0.3.1 (registry+https://github.com/rust-lang/crates.io-index)",
  "termion 1.5.3 (registry+https://github.com/rust-lang/crates.io-index)",
  "tui 0.5.1 (registry+https://github.com/rust-lang/crates.io-index)",
@@ -1252,6 +1276,7 @@ dependencies = [
 "checksum adler32 1.0.3 (registry+https://github.com/rust-lang/crates.io-index)" = "7e522997b529f05601e05166c07ed17789691f562762c7f3b987263d2dedee5c"
 "checksum aho-corasick 0.6.10 (registry+https://github.com/rust-lang/crates.io-index)" = "81ce3d38065e618af2d7b77e10c5ad9a069859b4be3c2250f674af3840d9c8a5"
 "checksum ansi_term 0.11.0 (registry+https://github.com/rust-lang/crates.io-index)" = "ee49baf6cb617b853aa8d93bf420db2383fab46d314482ca2803b40d5fde979b"
+"checksum arc-swap 0.4.3 (registry+https://github.com/rust-lang/crates.io-index)" = "f1a1eca3195b729bbd64e292ef2f5fff6b1c28504fed762ce2b1013dde4d8e92"
 "checksum assert_cmd 0.10.2 (registry+https://github.com/rust-lang/crates.io-index)" = "b7ac5c260f75e4e4ba87b7342be6edcecbcb3eb6741a0507fda7ad115845cc65"
 "checksum atty 0.2.13 (registry+https://github.com/rust-lang/crates.io-index)" = "1803c647a3ec87095e7ae7acfca019e98de5ec9a7d01343f611cf3152ed71a90"
 "checksum autocfg 0.1.5 (registry+https://github.com/rust-lang/crates.io-index)" = "22130e92352b948e7e82a49cdb0aa94f2211761117f29e052dd397c1ac33542b"
@@ -1356,6 +1381,8 @@ dependencies = [
 "checksum serde_json 1.0.40 (registry+https://github.com/rust-lang/crates.io-index)" = "051c49229f282f7c6f3813f8286cc1e3323e8051823fce42c7ea80fe13521704"
 "checksum serde_yaml 0.8.9 (registry+https://github.com/rust-lang/crates.io-index)" = "38b08a9a90e5260fe01c6480ec7c811606df6d3a660415808c3c3fa8ed95b582"
 "checksum sha-1 0.8.1 (registry+https://github.com/rust-lang/crates.io-index)" = "23962131a91661d643c98940b20fcaffe62d776a823247be80a48fcb8b6fce68"
+"checksum signal-hook 0.1.10 (registry+https://github.com/rust-lang/crates.io-index)" = "4f61c4d59f3aaa9f61bba6450a9b80ba48362fd7d651689e7a10c453b1f6dc68"
+"checksum signal-hook-registry 1.1.1 (registry+https://github.com/rust-lang/crates.io-index)" = "1797d48f38f91643908bb14e35e79928f9f4b3cefb2420a564dde0991b4358dc"
 "checksum smallvec 0.6.10 (registry+https://github.com/rust-lang/crates.io-index)" = "ab606a9c5e214920bb66c458cd7be8ef094f813f20fe77a54cc7dbfff220d4b7"
 "checksum socket2 0.3.11 (registry+https://github.com/rust-lang/crates.io-index)" = "e8b74de517221a2cb01a53349cf54182acdc31a074727d3079068448c0676d85"
 "checksum strsim 0.8.0 (registry+https://github.com/rust-lang/crates.io-index)" = "8ea5119cdb4c55b55d432abb513a0429384878c15dde60cc77b1c99de1a95a6a"
diff --git a/Cargo.toml b/Cargo.toml
index 43ce2a0..9322aa2 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -12,6 +12,7 @@ tui = "0.5"
 termion = "1.5"
 structopt = "0.3"
 dns-lookup = "*"
+signal-hook = "*"
 
 [dev-dependencies]
 assert_cmd = "0.10"
diff --git a/src/display/ui.rs b/src/display/ui.rs
index 1607e7c..30d36d0 100644
--- a/src/display/ui.rs
+++ b/src/display/ui.rs
@@ -144,21 +144,52 @@ fn render_remote_ip_table(state: &UIState, frame: &mut Frame<impl Backend>, rect
     table.render(frame, rect);
 }
 
-pub fn display_loop(
-    network_utilization: &Utilization,
-    terminal: &mut Terminal<impl Backend>,
-    connections_to_procs: HashMap<Connection, String>,
-    ip_to_host: &HashMap<Ipv4Addr, String>,
-) {
-    let state = UIState::new(connections_to_procs, &network_utilization);
-    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);
-        })
-        .unwrap();
+pub struct Ui<B>
+where B: Backend
+{
+    terminal: Terminal<B>,
+    state: UIState,
+    ip_to_host: HashMap<Ipv4Addr, String>,
+}
+
+impl <B>Ui<B>
+where B: Backend
+{
+    pub fn new (terminal_backend: B) -> Self {
+        let mut terminal = Terminal::new(terminal_backend).unwrap();
+        terminal.clear().unwrap();
+        terminal.hide_cursor().unwrap();
+        Ui {
+            terminal: terminal,
+            state: Default::default(),
+            ip_to_host: Default::default(),
+        }
+    }
+    pub fn draw (&mut self) {
+        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);
+            })
+            .unwrap();
+    }
+    pub fn update_state(
+        &mut self,
+        connections_to_procs: HashMap<Connection, String>,
+        utilization: Utilization,
+        ip_to_host: HashMap<Ipv4Addr, String>,
+    ) {
+        self.state = UIState::new(connections_to_procs, utilization);
+        self.ip_to_host = ip_to_host;
+    }
+    pub fn end(&mut self) { // TODO: destroy?
+        self.terminal.clear().unwrap();
+        self.terminal.show_cursor().unwrap();
+    }
 }
diff --git a/src/display/ui_state.rs b/src/display/ui_state.rs
index 83c1382..3469518 100644
--- a/src/display/ui_state.rs
+++ b/src/display/ui_state.rs
@@ -40,6 +40,7 @@ impl Bandwidth for NetworkData {
     }
 }
 
+#[derive(Default)]
 pub struct UIState {
     pub processes: BTreeMap<String, NetworkData>,
     pub remote_ips: BTreeMap<Ipv4Addr, NetworkData>,
@@ -49,7 +50,7 @@ pub struct UIState {
 impl UIState {
     pub fn new(
         connections_to_procs: HashMap<Connection, String>,
-        network_utilization: &Utilization,
+        network_utilization: Utilization,
     ) -> Self {
         let mut processes: BTreeMap<String, NetworkData> = BTreeMap::new();
         let mut remote_ips: BTreeMap<Ipv4Addr, NetworkData> = BTreeMap::new();
diff --git a/src/main.rs b/src/main.rs
index dbebc46..0d44a61 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -4,19 +4,19 @@ mod os;
 #[cfg(test)]
 mod tests;
 
-use display::display_loop;
+use display::Ui;
 use network::{Connection, DnsQueue, Sniffer, Utilization};
 
-use ::std::net::IpAddr;
+use ::std::net:: IpAddr;
 
 use ::pnet::datalink::{DataLinkReceiver, NetworkInterface};
 use ::std::collections::HashMap;
 use ::std::sync::atomic::{AtomicBool, Ordering};
 use ::std::sync::{Arc, Mutex};
 use ::std::{thread, time};
+use ::std::thread::park_timeout;
 use ::termion::event::{Event, Key};
 use ::tui::backend::Backend;
-use ::tui::Terminal;
 
 use ::std::io;
 use ::termion::raw::IntoRawMode;
@@ -37,7 +37,7 @@ fn main() {
         "Sorry, no implementations for platforms other than linux yet :( - PRs welcome!"
     );
 
-    use os::{get_datalink_channel, get_interface, get_open_sockets, lookup_addr, KeyboardEvents};
+    use os::{get_datalink_channel, get_interface, get_open_sockets, lookup_addr, receive_winch, KeyboardEvents};
 
     let opt = Opt::from_args();
     let stdout = io::stdout().into_raw_mode().unwrap();
@@ -47,6 +47,7 @@ fn main() {
     let network_interface = get_interface(&opt.interface).unwrap();
     let network_frames = get_datalink_channel(&network_interface);
     let lookup_addr = Box::new(lookup_addr);
+    let receive_winch = Box::new(receive_winch);
 
     let os_input = OsInput {
         network_interface,
@@ -54,6 +55,7 @@ fn main() {
         get_open_sockets,
         keyboard_events,
         lookup_addr,
+        receive_winch,
     };
 
     start(terminal_backend, os_input)
@@ -65,33 +67,19 @@ pub struct OsInput {
     pub get_open_sockets: fn() -> HashMap<Connection, String>,
     pub keyboard_events: Box<Iterator<Item = Event> + Send + Sync + 'static>,
     pub lookup_addr: Box<Fn(&IpAddr) -> Option<String> + Send + Sync + 'static>,
+    pub receive_winch: Box<Fn(&Arc<AtomicBool>)>,
 }
 
 pub fn start<B>(terminal_backend: B, os_input: OsInput)
 where
-    B: Backend + Send + 'static,
+    B: Backend + Send + Sync + 'static,
 {
     let running = Arc::new(AtomicBool::new(true));
 
     let keyboard_events = os_input.keyboard_events; // TODO: as methods in os_interface
     let get_open_sockets = os_input.get_open_sockets;
     let lookup_addr = os_input.lookup_addr;
-
-    let stdin_handler = thread::spawn({
-        let running = running.clone();
-        move || {
-            for evt in keyboard_events {
-                match evt {
-                    Event::Key(Key::Ctrl('c')) | Event::Key(Key::Char('q')) => {
-                        // TODO: exit faster
-                        running.store(false, Ordering::Relaxed);
-                        break;
-                    }
-                    _ => (),
-                };
-            }
-        }
-    });
+    let receive_winch = os_input.receive_winch;
 
     let mut sniffer = Sniffer::new(os_input.network_interface, os_input.network_frames);
     let network_utilization = Arc::new(Mutex::new(Utilization::new()));
@@ -111,43 +99,75 @@ where
         }
     });
 
+    let ui = Arc::new(Mutex::new(Ui::new(terminal_backend)));
+    let winch = Arc::new(AtomicBool::new(false));
+    receive_winch(&winch);
+
+    let resize_handler = thread::spawn({
+        let running = running.clone();
+        let winch = winch.clone();
+        let ui = ui.clone();
+        move || {
+            while running.load(Ordering::Acquire) {
+                while winch.load(Ordering::Acquire) {
+                    let mut ui = ui.lock().unwrap();
+                    ui.draw();
+                    winch.store(false, Ordering::Release);
+                }
+            }
+        }
+    });
+
     let display_handler = thread::spawn({
         let running = running.clone();
         let network_utilization = network_utilization.clone();
         let ip_to_host = ip_to_host.clone();
         let dns_queue = dns_queue.clone();
+        let ui = ui.clone();
         move || {
-            let mut terminal = Terminal::new(terminal_backend).unwrap();
-            terminal.clear().unwrap();
-            terminal.hide_cursor().unwrap();
-            while running.load(Ordering::Relaxed) {
-                {
-                    let connections_to_procs = get_open_sockets();
-                    let ip_to_host = ip_to_host.lock().unwrap();
-                    let unresolved_ips = connections_to_procs.keys().fold(vec![], |mut unresolved_ips, connection| {
-                        if !ip_to_host.contains_key(&connection.local_socket.ip) {
-                            unresolved_ips.push(connection.local_socket.ip.clone());
-                        }
-                        if !ip_to_host.contains_key(&connection.remote_socket.ip) {
-                            unresolved_ips.push(connection.remote_socket.ip.clone());
-                        }
-                        unresolved_ips
-                    });
+            while running.load(Ordering::Acquire) {
+                let connections_to_procs = get_open_sockets();
+                let ip_to_host = {
+                    ip_to_host.lock().unwrap().clone()
+                };
+                let utilization = {
                     let mut network_utilization = network_utilization.lock().unwrap();
-                    let utilization = network_utilization.clone_and_reset();
-                    dns_queue.add_ips_to_resolve(unresolved_ips);
-                    display_loop(&utilization, &mut terminal, connections_to_procs, &ip_to_host);
+                    network_utilization.clone_and_reset()
+                };
+                dns_queue.find_ips_to_resolve(&connections_to_procs, &ip_to_host);
+                {
+                    let mut ui = ui.lock().unwrap();
+                    ui.update_state(connections_to_procs, utilization, ip_to_host);
+                    ui.draw();
                 }
-                thread::sleep(time::Duration::from_secs(1));
+                park_timeout(time::Duration::from_secs(1));
             }
-            terminal.clear().unwrap();
-            terminal.show_cursor().unwrap();
+            let mut ui = ui.lock().unwrap();
+            ui.end();
             dns_queue.end();
         }
     });
 
+    let stdin_handler = thread::spawn({
+        let running = running.clone();
+        let display_handler = display_handler.thread().clone(); // TODO: better
+        move || {
+            for evt in keyboard_events {
+                match evt {
+                    Event::Key(Key::Ctrl('c')) | Event::Key(Key::Char('q')) => {
+                        running.store(false, Ordering::Release);
+                        display_handler.unpark();
+                        break;
+                    }
+                    _ => (),
+                };
+            }
+        }
+    });
+
+
     let sniffing_handler = thread::spawn(move || {
-        while running.load(Ordering::Relaxed) {
+        while running.load(Ordering::Acquire) {
             if let Some(segment) = sniffer.next() {
                 network_utilization.lock().unwrap().update(&segment)
             }
@@ -157,4 +177,5 @@ where
     sniffing_handler.join().unwrap();
     stdin_handler.join().unwrap();
     dns_handler.join().unwrap();
+    resize_handler.join().unwrap();
 }
diff --git a/src/network/dns_queue.rs b/src/network/dns_queue.rs
index 9f8d406..f504f81 100644
--- a/src/network/dns_queue.rs
+++ b/src/network/dns_queue.rs
@@ -1,6 +1,8 @@
 use ::std::net::Ipv4Addr;
-
 use ::std::sync::{Condvar, Mutex};
+use ::std::collections::HashMap;
+
+use crate::network::Connection;
 
 pub struct DnsQueue {
     jobs: Mutex<Vec<Option<Ipv4Addr>>>,
@@ -17,11 +19,20 @@ impl DnsQueue {
 }
 
 impl DnsQueue {
-    pub fn add_ips_to_resolve(&self, unresolved_ips: Vec<Ipv4Addr>) {
+    pub fn find_ips_to_resolve(
+        &self,
+        connections_to_procs: &HashMap<Connection, String>,
+        ip_to_host: &HashMap<Ipv4Addr, String>,
+    ) {
         let mut queue = self.jobs.lock().unwrap();
-        for ip in unresolved_ips {
-            queue.push(Some(ip));
-        }
+        for connection in connections_to_procs.keys() {
+            if !ip_to_host.contains_key(&connection.local_socket.ip) {
+                queue.push(Some(connection.local_socket.ip.clone()));
+            }
+            if !ip_to_host.contains_key(&connection.remote_socket.ip) {
+                queue.push(Some(connection.remote_socket.ip.clone()));
+            }
+        };
         self.cvar.notify_all();
     }
     pub fn wait_for_job(&self) -> Option<Ipv4Addr> {
diff --git a/src/network/sniffer.rs b/src/network/sniffer.rs
index c1381e9..93f55a4 100644
--- a/src/network/sniffer.rs
+++ b/src/network/sniffer.rs
@@ -50,11 +50,7 @@ impl Sniffer {
         }
     }
     pub fn next(&mut self) -> Option<Segment> {
-        // TODO: https://github.com/libpnet/libpnet/issues/343
-        // make this non-blocking for faster exits
-        let bytes = self.network_frames.next().unwrap_or_else(|e| {
-            panic!("An error occurred while reading: {}", e);
-        });
+        let bytes = self.network_frames.next().ok()?;
         let packet = EthernetPacket::new(bytes)?;
         match packet.get_ethertype() {
             EtherType(2048) => {
diff --git a/src/os/linux.rs b/src/os/linux.rs
index f73a079..7ac5e9f 100644
--- a/src/os/linux.rs
+++ b/src/os/linux.rs
@@ -1,12 +1,18 @@
 use ::pnet::datalink::Channel::Ethernet;
 use ::pnet::datalink::DataLinkReceiver;
-use ::pnet::datalink::{self, NetworkInterface};
+use ::pnet::datalink::{self, NetworkInterface, Config};
 use ::std::io::stdin;
 use ::termion::event::Event;
 use ::termion::input::TermRead;
 
 use ::std::collections::HashMap;
 use ::std::net::IpAddr;
+use ::std::sync::Arc;
+use ::std::time;
+
+use ::std::sync::atomic::AtomicBool;
+
+use ::signal_hook;
 
 use ::procfs::FDTarget;
 
@@ -25,7 +31,9 @@ impl Iterator for KeyboardEvents {
 }
 
 pub fn get_datalink_channel(interface: &NetworkInterface) -> Box<DataLinkReceiver> {
-    match datalink::channel(interface, Default::default()) {
+    let mut config = Config::default();
+    config.read_timeout = Some(time::Duration::new(0, 0));
+    match datalink::channel(interface, config) {
         Ok(Ethernet(_tx, rx)) => rx,
         Ok(_) => panic!("Unhandled channel type"),
         Err(e) => panic!(
@@ -82,3 +90,9 @@ pub fn get_open_sockets() -> HashMap<Connection, String> {
 pub fn lookup_addr(ip: &IpAddr) -> Option<String> {
     ::dns_lookup::lookup_addr(ip).ok()
 }
+
+pub fn receive_winch(winch: &Arc<AtomicBool>) {
+    ::signal_hook::flag::register(
+        signal_hook::SIGWINCH, winch.clone()
+    ).expect("Failed to register SIGWINCH");
+}