fix(sigwinch): block when waiting for signal

4 files changed, 109 insertions, 64 deletions

diff --git a/src/display/ui.rs b/src/display/ui.rs
index b31f8ea..b263582 100644
--- a/src/display/ui.rs
+++ b/src/display/ui.rs
@@ -8,7 +8,6 @@ use crate::display::UIState;
 use crate::network::{Connection, Utilization};
 
 use ::std::net::Ipv4Addr;
-use ::std::sync::atomic::{AtomicBool, Ordering};
 
 pub struct Ui<B>
 where
@@ -17,7 +16,6 @@ where
     terminal: Terminal<B>,
     state: UIState,
     ip_to_host: HashMap<Ipv4Addr, String>,
-    drawing: AtomicBool,
 }
 
 impl<B> Ui<B>
@@ -32,30 +30,25 @@ where
             terminal,
             state: Default::default(),
             ip_to_host: Default::default(),
-            drawing: AtomicBool::new(false),
         }
     }
     pub fn draw(&mut self) {
-        if !self.drawing.load(Ordering::Acquire) {
-            self.drawing.store(true, Ordering::Release);
-            let state = &self.state;
-            let ip_to_host = &self.ip_to_host;
-            self.terminal
-                .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();
-            self.drawing.store(false, Ordering::Release);
-        }
+        let state = &self.state;
+        let ip_to_host = &self.ip_to_host;
+        self.terminal
+            .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();
     }
     pub fn update_state(
         &mut self,
diff --git a/src/main.rs b/src/main.rs
index 8d21884..33b6bdd 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -28,7 +28,7 @@ use structopt::StructOpt;
 
 #[derive(StructOpt, Debug)]
 #[structopt(name = "what")]
-struct Opt {
+pub struct Opt {
     #[structopt(short, long)]
     interface: String,
 }
@@ -46,12 +46,9 @@ fn try_main() -> Result<(), failure::Error> {
         "Sorry, no implementations for platforms other than linux yet :( - PRs welcome!"
     );
 
-    use os::{
-        get_datalink_channel, get_interface, get_open_sockets, lookup_addr, on_winch,
-        KeyboardEvents,
-    };
-
+    use os::{get_input};
     let opt = Opt::from_args();
+    let os_input = get_input(opt)?;
     let stdout = match io::stdout().into_raw_mode() {
         Ok(stdout) => stdout,
         Err(_) => failure::bail!(
@@ -59,27 +56,6 @@ fn try_main() -> Result<(), failure::Error> {
         ),
     };
     let terminal_backend = TermionBackend::new(stdout);
-
-    let keyboard_events = Box::new(KeyboardEvents);
-    let network_interface = match get_interface(&opt.interface) {
-        Some(interface) => interface,
-        None => {
-            failure::bail!("Cannot find interface {}", opt.interface);
-        }
-    };
-    let network_frames = get_datalink_channel(&network_interface);
-    let lookup_addr = Box::new(lookup_addr);
-    let on_winch = Box::new(on_winch);
-
-    let os_input = OsInput {
-        network_interface,
-        network_frames,
-        get_open_sockets,
-        keyboard_events,
-        lookup_addr,
-        on_winch,
-    };
-
     start(terminal_backend, os_input);
     Ok(())
 }
@@ -91,6 +67,7 @@ pub struct OsInput {
     pub keyboard_events: Box<Iterator<Item = Event> + Send + Sync + 'static>,
     pub lookup_addr: Box<Fn(&IpAddr) -> Option<String> + Send + Sync + 'static>,
     pub on_winch: Box<Fn(Box<Fn() + Send + Sync + 'static>) + Send + Sync + 'static>,
+    pub cleanup: Box<Fn() + Send + Sync + 'static>,
 }
 
 pub fn start<B>(terminal_backend: B, os_input: OsInput)
@@ -103,6 +80,7 @@ where
     let get_open_sockets = os_input.get_open_sockets;
     let lookup_addr = os_input.lookup_addr;
     let on_winch = os_input.on_winch;
+    let cleanup = os_input.cleanup;
 
     let mut sniffer = Sniffer::new(os_input.network_interface, os_input.network_frames);
     let network_utilization = Arc::new(Mutex::new(Utilization::new()));
@@ -170,6 +148,7 @@ where
                 match evt {
                     Event::Key(Key::Ctrl('c')) | Event::Key(Key::Char('q')) => {
                         running.store(false, Ordering::Release);
+                        cleanup();
                         display_handler.unpark();
                         break;
                     }
diff --git a/src/os/linux.rs b/src/os/linux.rs
index fd8905a..2bf979c 100644
--- a/src/os/linux.rs
+++ b/src/os/linux.rs
@@ -9,13 +9,13 @@ use ::std::collections::HashMap;
 use ::std::net::IpAddr;
 use ::std::time;
 
-use ::signal_hook;
-
+use signal_hook::iterator::Signals;
 use ::procfs::FDTarget;
 
 use crate::network::{Connection, Protocol};
+use crate::{Opt, OsInput};
 
-pub struct KeyboardEvents;
+struct KeyboardEvents;
 
 impl Iterator for KeyboardEvents {
     type Item = Event;
@@ -27,7 +27,7 @@ impl Iterator for KeyboardEvents {
     }
 }
 
-pub fn get_datalink_channel(interface: &NetworkInterface) -> Box<DataLinkReceiver> {
+fn get_datalink_channel(interface: &NetworkInterface) -> Box<DataLinkReceiver> {
     let mut config = Config::default();
     config.read_timeout = Some(time::Duration::new(0, 1));
     match datalink::channel(interface, config) {
@@ -40,13 +40,13 @@ pub fn get_datalink_channel(interface: &NetworkInterface) -> Box<DataLinkReceive
     }
 }
 
-pub fn get_interface(interface_name: &str) -> Option<NetworkInterface> {
+fn get_interface(interface_name: &str) -> Option<NetworkInterface> {
     datalink::interfaces()
         .into_iter()
         .find(|iface| iface.name == interface_name)
 }
 
-pub fn get_open_sockets() -> HashMap<Connection, String> {
+fn get_open_sockets() -> HashMap<Connection, String> {
     let mut open_sockets = HashMap::new(); // TODO: better
     let all_procs = procfs::all_processes();
 
@@ -86,15 +86,52 @@ pub fn get_open_sockets() -> HashMap<Connection, String> {
     open_sockets
 }
 
-pub fn lookup_addr(ip: &IpAddr) -> Option<String> {
+fn lookup_addr(ip: &IpAddr) -> Option<String> {
     ::dns_lookup::lookup_addr(ip).ok()
 }
 
-pub fn on_winch<F>(cb: F)
-where
-    F: Fn() + Send + Sync + 'static,
-{
-    unsafe {
-        signal_hook::register(signal_hook::SIGWINCH, cb).unwrap();
-    }
+fn sigwinch () -> (
+    Box<Fn(Box<Fn() + Send + Sync + 'static>) + Send + Sync + 'static>,
+    Box<Fn() + Send + Sync + 'static>
+) {
+    let signals = Signals::new(&[signal_hook::SIGWINCH]).unwrap();
+    let on_winch = {
+        let signals = signals.clone();
+        move |cb: Box<Fn() + Send + Sync + 'static>| {
+            for signal in signals.forever() {
+                match signal {
+                    signal_hook::SIGWINCH => cb(),
+                    _ => unreachable!(),
+                }
+            }
+        }
+    };
+    let cleanup = move || {
+        signals.close();
+    };
+    (Box::new(on_winch), Box::new(cleanup))
+}
+
+pub fn get_input(opt: Opt) -> Result<OsInput, failure::Error> {
+
+    let keyboard_events = Box::new(KeyboardEvents);
+    let network_interface = match get_interface(&opt.interface) {
+        Some(interface) => interface,
+        None => {
+            failure::bail!("Cannot find interface {}", opt.interface);
+        }
+    };
+    let network_frames = get_datalink_channel(&network_interface);
+    let lookup_addr = Box::new(lookup_addr);
+    let (on_winch, cleanup) = sigwinch();
+
+    Ok(OsInput {
+        network_interface,
+        network_frames,
+        get_open_sockets,
+        keyboard_events,
+        lookup_addr,
+        on_winch,
+        cleanup,
+    })
 }
diff --git a/src/tests/mod.rs b/src/tests/mod.rs
index bf1bc1e..e02b4d5 100644
--- a/src/tests/mod.rs
+++ b/src/tests/mod.rs
@@ -73,6 +73,7 @@ fn basic_startup() {
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap::new());
     let on_winch = create_fake_on_winch(false);
+    let cleanup = Box::new(|| {});
 
     let os_input = crate::OsInput {
         network_interface,
@@ -81,6 +82,7 @@ fn basic_startup() {
         keyboard_events,
         lookup_addr,
         on_winch,
+        cleanup,
     };
     crate::start(backend, os_input);
 
@@ -123,6 +125,7 @@ fn one_packet_of_traffic() {
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap::new());
     let on_winch = create_fake_on_winch(false);
+    let cleanup = Box::new(|| {});
 
     let os_input = crate::OsInput {
         network_interface,
@@ -131,6 +134,7 @@ fn one_packet_of_traffic() {
         keyboard_events,
         lookup_addr,
         on_winch,
+        cleanup,
     };
     crate::start(backend, os_input);
 
@@ -185,6 +189,7 @@ fn bi_directional_traffic() {
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap::new());
     let on_winch = create_fake_on_winch(false);
+    let cleanup = Box::new(|| {});
 
     let os_input = crate::OsInput {
         network_interface,
@@ -193,6 +198,7 @@ fn bi_directional_traffic() {
         keyboard_events,
         lookup_addr,
         on_winch,
+        cleanup,
     };
     crate::start(backend, os_input);
 
@@ -245,6 +251,7 @@ fn multiple_packets_of_traffic_from_different_connections() {
         terminal_height,
     );
     let on_winch = create_fake_on_winch(false);
+    let cleanup = Box::new(|| {});
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap::new());
 
@@ -253,6 +260,7 @@ fn multiple_packets_of_traffic_from_different_connections() {
         network_frames,
         get_open_sockets,
         on_winch,
+        cleanup,
         keyboard_events,
         lookup_addr,
     };
@@ -309,6 +317,7 @@ fn multiple_packets_of_traffic_from_single_connection() {
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap::new());
     let on_winch = create_fake_on_winch(false);
+    let cleanup = Box::new(|| {});
 
     let os_input = crate::OsInput {
         network_interface,
@@ -317,6 +326,7 @@ fn multiple_packets_of_traffic_from_single_connection() {
         keyboard_events,
         lookup_addr,
         on_winch,
+        cleanup,
     };
     crate::start(backend, os_input);
 
@@ -371,6 +381,7 @@ fn one_process_with_multiple_connections() {
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap::new());
     let on_winch = create_fake_on_winch(false);
+    let cleanup = Box::new(|| {});
 
     let os_input = crate::OsInput {
         network_interface,
@@ -379,6 +390,7 @@ fn one_process_with_multiple_connections() {
         keyboard_events,
         lookup_addr,
         on_winch,
+        cleanup,
     };
     crate::start(backend, os_input);
 
@@ -447,6 +459,7 @@ fn multiple_processes_with_multiple_connections() {
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap::new());
     let on_winch = create_fake_on_winch(false);
+    let cleanup = Box::new(|| {});
 
     let os_input = crate::OsInput {
         network_interface,
@@ -455,6 +468,7 @@ fn multiple_processes_with_multiple_connections() {
         keyboard_events,
         lookup_addr,
         on_winch,
+        cleanup,
     };
     crate::start(backend, os_input);
 
@@ -509,6 +523,7 @@ fn multiple_connections_from_remote_ip() {
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap::new());
     let on_winch = create_fake_on_winch(false);
+    let cleanup = Box::new(|| {});
 
     let os_input = crate::OsInput {
         network_interface,
@@ -517,6 +532,7 @@ fn multiple_connections_from_remote_ip() {
         keyboard_events,
         lookup_addr,
         on_winch,
+        cleanup,
     };
     crate::start(backend, os_input);
 
@@ -573,6 +589,7 @@ fn sustained_traffic_from_one_process() {
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap::new());
     let on_winch = create_fake_on_winch(false);
+    let cleanup = Box::new(|| {});
 
     let os_input = crate::OsInput {
         network_interface,
@@ -581,6 +598,7 @@ fn sustained_traffic_from_one_process() {
         keyboard_events,
         lookup_addr,
         on_winch,
+        cleanup,
     };
     crate::start(backend, os_input);
 
@@ -651,6 +669,7 @@ fn sustained_traffic_from_multiple_processes() {
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap::new());
     let on_winch = create_fake_on_winch(false);
+    let cleanup = Box::new(|| {});
 
     let os_input = crate::OsInput {
         network_interface,
@@ -659,6 +678,7 @@ fn sustained_traffic_from_multiple_processes() {
         keyboard_events,
         lookup_addr,
         on_winch,
+        cleanup,
     };
     crate::start(backend, os_input);
 
@@ -757,6 +777,7 @@ fn sustained_traffic_from_multiple_processes_bi_directional() {
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap::new());
     let on_winch = create_fake_on_winch(false);
+    let cleanup = Box::new(|| {});
 
     let os_input = crate::OsInput {
         network_interface,
@@ -765,6 +786,7 @@ fn sustained_traffic_from_multiple_processes_bi_directional() {
         keyboard_events,
         lookup_addr,
         on_winch,
+        cleanup,
     };
     crate::start(backend, os_input);
 
@@ -876,6 +898,7 @@ fn traffic_with_host_names() {
     );
     let lookup_addr = create_fake_lookup_addr(ips_to_hostnames);
     let on_winch = create_fake_on_winch(false);
+    let cleanup = Box::new(|| {});
 
     let os_input = crate::OsInput {
         network_interface,
@@ -884,6 +907,7 @@ fn traffic_with_host_names() {
         keyboard_events,
         lookup_addr,
         on_winch,
+        cleanup,
     };
     crate::start(backend, os_input);
 
@@ -929,6 +953,7 @@ fn traffic_with_winch_event() {
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap::new());
     let on_winch = create_fake_on_winch(true);
+    let cleanup = Box::new(|| {});
 
     let os_input = crate::OsInput {
         network_interface,
@@ -937,6 +962,7 @@ fn traffic_with_winch_event() {
         keyboard_events,
         lookup_addr,
         on_winch,
+        cleanup,
     };
     crate::start(backend, os_input);
 
@@ -1006,6 +1032,7 @@ fn layout_full_width_under_30_height() {
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap::new());
     let on_winch = create_fake_on_winch(false);
+    let cleanup = Box::new(|| {});
 
     let os_input = crate::OsInput {
         network_interface,
@@ -1014,6 +1041,7 @@ fn layout_full_width_under_30_height() {
         keyboard_events,
         lookup_addr,
         on_winch,
+        cleanup,
     };
     crate::start(backend, os_input);
 
@@ -1082,6 +1110,7 @@ fn layout_under_150_width_full_height() {
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap::new());
     let on_winch = create_fake_on_winch(false);
+    let cleanup = Box::new(|| {});
 
     let os_input = crate::OsInput {
         network_interface,
@@ -1090,6 +1119,7 @@ fn layout_under_150_width_full_height() {
         keyboard_events,
         lookup_addr,
         on_winch,
+        cleanup,
     };
     crate::start(backend, os_input);
 
@@ -1158,6 +1188,7 @@ fn layout_under_150_width_under_30_height() {
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap::new());
     let on_winch = create_fake_on_winch(false);
+    let cleanup = Box::new(|| {});
 
     let os_input = crate::OsInput {
         network_interface,
@@ -1166,6 +1197,7 @@ fn layout_under_150_width_under_30_height() {
         keyboard_events,
         lookup_addr,
         on_winch,
+        cleanup,
     };
     crate::start(backend, os_input);
 
@@ -1234,6 +1266,7 @@ fn layout_under_120_width_full_height() {
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap::new());
     let on_winch = create_fake_on_winch(false);
+    let cleanup = Box::new(|| {});
 
     let os_input = crate::OsInput {
         network_interface,
@@ -1242,6 +1275,7 @@ fn layout_under_120_width_full_height() {
         keyboard_events,
         lookup_addr,
         on_winch,
+        cleanup,
     };
     crate::start(backend, os_input);
 
@@ -1310,6 +1344,7 @@ fn layout_under_120_width_under_30_height() {
     let network_interface = get_interface();
     let lookup_addr = create_fake_lookup_addr(HashMap::new());
     let on_winch = create_fake_on_winch(false);
+    let cleanup = Box::new(|| {});
 
     let os_input = crate::OsInput {
         network_interface,
@@ -1318,6 +1353,7 @@ fn layout_under_120_width_under_30_height() {
         keyboard_events,
         lookup_addr,
         on_winch,
+        cleanup,
     };
     crate::start(backend, os_input);