diff options
Diffstat (limited to 'src/main.rs')
| -rw-r--r-- | src/main.rs | 109 |
1 files changed, 65 insertions, 44 deletions
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(); } |