From 471fd861d5799b14f819e388246bf8a564ca98a6 Mon Sep 17 00:00:00 2001 From: Kelvin Zhang Date: Tue, 7 Jan 2020 15:08:16 -0500 Subject: [PATCH] Address some comments on PR --- src/main.rs | 14 +++++++++++--- src/network/connection.rs | 4 ++++ src/network/utilization.rs | 4 ++-- src/os/linux.rs | 10 +++++++--- src/os/lsof_utils.rs | 32 +++++++++++--------------------- src/os/macos.rs | 11 +++++++---- src/tests/fakes/fake_input.rs | 10 +++++++--- 7 files changed, 49 insertions(+), 36 deletions(-) diff --git a/src/main.rs b/src/main.rs index b14227c70..9ae11e6dd 100644 --- a/src/main.rs +++ b/src/main.rs @@ -79,10 +79,15 @@ fn try_main() -> Result<(), failure::Error> { Ok(()) } +pub struct OpenSockets { + sockets_to_procs: HashMap, + connections: Vec, +} + pub struct OsInputOutput { pub network_interfaces: Vec, pub network_frames: Vec>, - pub get_open_sockets: fn() -> (HashMap, std::vec::Vec), + pub get_open_sockets: fn() -> OpenSockets, pub keyboard_events: Box + Send>, pub dns_client: Option, pub on_winch: Box, @@ -138,7 +143,10 @@ where while running.load(Ordering::Acquire) { let render_start_time = Instant::now(); let utilization = { network_utilization.lock().unwrap().clone_and_reset() }; - let (connections_to_procs, connections) = get_open_sockets(); + let OpenSockets { + sockets_to_procs, + connections, + } = get_open_sockets(); let mut ip_to_host = IpTable::new(); if let Some(dns_client) = dns_client.as_mut() { ip_to_host = dns_client.cache(); @@ -152,7 +160,7 @@ where } { let mut ui = ui.lock().unwrap(); - ui.update_state(connections_to_procs, utilization, ip_to_host); + ui.update_state(sockets_to_procs, utilization, ip_to_host); if raw_mode { ui.output_text(&mut write_to_stdout); } else { diff --git a/src/network/connection.rs b/src/network/connection.rs index f98d17c17..907e9298f 100644 --- a/src/network/connection.rs +++ b/src/network/connection.rs @@ -11,6 +11,10 @@ pub enum Protocol { } impl Protocol { + // Currently, linux implementation doesn't use this function. + // Without this #[cfg] clippy complains about dead code, and CI refuses + // to pass. + #[cfg(target_os = "macos")] pub fn from_str(string: &str) -> Option { match string { "TCP" => Some(Protocol::Tcp), diff --git a/src/network/utilization.rs b/src/network/utilization.rs index d76ac4445..bdfbb21b5 100644 --- a/src/network/utilization.rs +++ b/src/network/utilization.rs @@ -2,14 +2,14 @@ use crate::network::{Connection, Direction, Segment}; use ::std::collections::HashMap; -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct ConnectionInfo { pub interface_name: String, pub total_bytes_downloaded: u128, pub total_bytes_uploaded: u128, } -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct Utilization { pub connections: HashMap, } diff --git a/src/os/linux.rs b/src/os/linux.rs index 07b8b173a..9ca566cfc 100644 --- a/src/os/linux.rs +++ b/src/os/linux.rs @@ -2,9 +2,10 @@ use ::std::collections::HashMap; use ::procfs::process::FDTarget; -use crate::network::{Connection, LocalSocket, Protocol}; +use crate::network::{Connection, Protocol}; +use crate::OpenSockets; -pub(crate) fn get_open_sockets() -> (HashMap, std::vec::Vec) { +pub(crate) fn get_open_sockets() -> OpenSockets { let mut open_sockets = HashMap::new(); let mut connections = std::vec::Vec::new(); let all_procs = procfs::process::all_processes().unwrap(); @@ -46,5 +47,8 @@ pub(crate) fn get_open_sockets() -> (HashMap, std::vec::Vec connections.push(connection); }; } - (open_sockets, connections) + OpenSockets { + sockets_to_procs: open_sockets, + connections, + } } diff --git a/src/os/lsof_utils.rs b/src/os/lsof_utils.rs index 260c753ea..584acab08 100644 --- a/src/os/lsof_utils.rs +++ b/src/os/lsof_utils.rs @@ -21,16 +21,6 @@ lazy_static! { static ref LISTEN_REGEX: Regex = Regex::new(r"(.*):(.*)").unwrap(); } -fn parse_ip_addr(ip: &str) -> IpAddr { - if let Ok(v4addr) = ip.parse() { - IpAddr::V4(v4addr) - } else if let Ok(v6addr) = ip.parse() { - IpAddr::V6(v6addr) - } else { - panic!("{} is not a valid IP address", ip) - } -} - fn get_null_addr(ip_type: &str) -> &str { if ip_type.contains('4') { "0.0.0.0" @@ -39,16 +29,12 @@ fn get_null_addr(ip_type: &str) -> &str { } } -#[allow(clippy::needless_return)] impl RawConnection { pub fn new(raw_line: &str) -> Option { + // Example row + // com.apple 664 user 198u IPv4 0xeb179a6650592b8d 0t0 TCP 192.168.1.187:58535->1.2.3.4:443 (ESTABLISHED) let columns: Vec<&str> = raw_line.split_ascii_whitespace().collect(); if columns.len() < 9 { - println!( - "lsof's output string has {} columns, different than expected: {:#?}", - columns.len(), - columns - ); return None; } let process_name = columns[0].replace("\\x20", " "); @@ -71,6 +57,10 @@ impl RawConnection { // let connection_state = columns[9]; // If this socket is in a "connected" state if let Some(caps) = CONNECTION_REGEX.captures(connection_str) { + // Example + // 192.168.1.187:64230->0.1.2.3:5228 + // *:* + // *:4567 let local_ip = String::from(caps.get(1).unwrap().as_str()); let local_port = String::from(caps.get(2).unwrap().as_str()); let remote_ip = String::from(caps.get(3).unwrap().as_str()); @@ -113,23 +103,23 @@ impl RawConnection { } pub fn get_protocol(&self) -> Protocol { - return Protocol::from_str(&self.protocol).unwrap(); + Protocol::from_str(&self.protocol).unwrap() } pub fn get_remote_ip(&self) -> IpAddr { - return parse_ip_addr(&self.remote_ip); + self.remote_ip.parse().unwrap() } pub fn get_remote_port(&self) -> u16 { - return self.remote_port.parse::().unwrap(); + self.remote_port.parse::().unwrap() } pub fn get_local_ip(&self) -> IpAddr { - return parse_ip_addr(&self.local_ip); + self.local_ip.parse().unwrap() } pub fn get_local_port(&self) -> u16 { - return self.local_port.parse::().unwrap(); + self.local_port.parse::().unwrap() } } diff --git a/src/os/macos.rs b/src/os/macos.rs index 2594330e1..68b4badd6 100644 --- a/src/os/macos.rs +++ b/src/os/macos.rs @@ -1,6 +1,7 @@ use ::std::collections::HashMap; -use crate::network::{Connection, LocalSocket}; +use crate::network::Connection; +use crate::OpenSockets; use super::lsof_utils; use std::net::SocketAddr; @@ -14,8 +15,7 @@ struct RawConnection { process_name: String, } -#[allow(clippy::needless_return)] -pub(crate) fn get_open_sockets() -> (HashMap, std::vec::Vec) { +pub(crate) fn get_open_sockets() -> OpenSockets { let mut open_sockets = HashMap::new(); let mut connections_vec = std::vec::Vec::new(); @@ -35,5 +35,8 @@ pub(crate) fn get_open_sockets() -> (HashMap, std::vec::Vec connections_vec.push(connection); } - return (open_sockets, connections_vec); + OpenSockets { + sockets_to_procs: open_sockets, + connections: connections_vec, + } } diff --git a/src/tests/fakes/fake_input.rs b/src/tests/fakes/fake_input.rs index 8faaa41c4..38b14cbdc 100644 --- a/src/tests/fakes/fake_input.rs +++ b/src/tests/fakes/fake_input.rs @@ -13,9 +13,10 @@ use ::termion::event::Event; use crate::{ network::{ dns::{self, Lookup}, - Connection, LocalSocket, Protocol, + Connection, Protocol, }, os::OnSigWinch, + OpenSockets, }; pub struct KeyboardEvents { @@ -85,7 +86,7 @@ impl DataLinkReceiver for NetworkFrames { } } -pub fn get_open_sockets() -> (HashMap, std::vec::Vec) { +pub fn get_open_sockets() -> OpenSockets { let mut open_sockets = HashMap::new(); let local_ip = IpAddr::V4(Ipv4Addr::new(10, 0, 0, 2)); open_sockets.insert( @@ -145,7 +146,10 @@ pub fn get_open_sockets() -> (HashMap, std::vec::Vec Vec {