diff options
author | Kelvin Zhang <zhangxp1998@gmail.com> | 2020-01-07 16:48:06 -0500 |
---|---|---|
committer | GitHub <noreply@github.com> | 2020-01-07 16:48:06 -0500 |
commit | 5826f04ca64071f304286e646c1b262154202b74 (patch) | |
tree | 3a9b7f27adfedc52b85051ec1f8cd05fd2e3dbc6 /src | |
parent | 73ac4dce7a399658369e5e4e59bc9dd69e572183 (diff) | |
parent | 471fd861d5799b14f819e388246bf8a564ca98a6 (diff) |
Merge pull request #82 from zhangxp1998/udp
Fix bug that non-connected Udp sockets aren't displayed
Diffstat (limited to 'src')
34 files changed, 342 insertions, 200 deletions
diff --git a/src/display/ui.rs b/src/display/ui.rs index e58c3e9..e61f5e7 100644 --- a/src/display/ui.rs +++ b/src/display/ui.rs @@ -5,7 +5,7 @@ use ::tui::Terminal; use crate::display::components::{Layout, Table, TotalBandwidth}; use crate::display::UIState; -use crate::network::{display_connection_string, display_ip_or_host, Connection, Utilization}; +use crate::network::{display_connection_string, display_ip_or_host, LocalSocket, Utilization}; use ::std::net::Ipv4Addr; @@ -94,7 +94,7 @@ where } pub fn update_state( &mut self, - connections_to_procs: HashMap<Connection, String>, + connections_to_procs: HashMap<LocalSocket, String>, utilization: Utilization, ip_to_host: HashMap<Ipv4Addr, String>, ) { diff --git a/src/display/ui_state.rs b/src/display/ui_state.rs index e5c2217..acbd04f 100644 --- a/src/display/ui_state.rs +++ b/src/display/ui_state.rs @@ -1,7 +1,7 @@ use ::std::collections::{BTreeMap, HashMap}; -use ::std::net::Ipv4Addr; +use ::std::net::{IpAddr, Ipv4Addr, Ipv6Addr}; -use crate::network::{Connection, Utilization}; +use crate::network::{Connection, LocalSocket, Utilization}; pub trait Bandwidth { fn get_total_bytes_downloaded(&self) -> u128; @@ -51,38 +51,61 @@ pub struct UIState { } impl UIState { + fn get_proc_name<'a>( + connections_to_procs: &'a HashMap<LocalSocket, String>, + local_socket: &LocalSocket, + ) -> Option<&'a String> { + if let Some(process_name) = connections_to_procs.get(local_socket) { + Some(process_name) + } else if let Some(process_name) = connections_to_procs.get(&LocalSocket { + ip: IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), + port: local_socket.port, + protocol: local_socket.protocol, + }) { + Some(process_name) + } else { + connections_to_procs.get(&LocalSocket { + ip: IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 0)), + port: local_socket.port, + protocol: local_socket.protocol, + }) + } + } pub fn new( - connections_to_procs: HashMap<Connection, String>, - mut network_utilization: Utilization, + connections_to_procs: HashMap<LocalSocket, String>, + network_utilization: Utilization, ) -> Self { let mut processes: BTreeMap<String, NetworkData> = BTreeMap::new(); let mut remote_addresses: 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_info) = network_utilization.connections.remove(&connection) { - let data_for_remote_address = remote_addresses - .entry(connection.remote_socket.ip) - .or_default(); - let connection_data = connections.entry(connection).or_default(); - let data_for_process = processes.entry(process_name.clone()).or_default(); + for (connection, connection_info) in network_utilization.connections { + let connection_data = connections.entry(connection).or_default(); + if let Some(process_name) = + UIState::get_proc_name(&connections_to_procs, &connection.local_socket) + { + let data_for_process = processes.entry(process_name.clone()).or_default(); data_for_process.total_bytes_downloaded += connection_info.total_bytes_downloaded; data_for_process.total_bytes_uploaded += connection_info.total_bytes_uploaded; data_for_process.connection_count += 1; - connection_data.total_bytes_downloaded += connection_info.total_bytes_downloaded; - connection_data.total_bytes_uploaded += connection_info.total_bytes_uploaded; - connection_data.process_name = process_name; - connection_data.interface_name = connection_info.interface_name; - data_for_remote_address.total_bytes_downloaded += - connection_info.total_bytes_downloaded; - data_for_remote_address.total_bytes_uploaded += - connection_info.total_bytes_uploaded; - data_for_remote_address.connection_count += 1; - total_bytes_downloaded += connection_info.total_bytes_downloaded; - total_bytes_uploaded += connection_info.total_bytes_uploaded; + connection_data.process_name = process_name.clone(); + } else { + connection_data.process_name = String::from("<UNKNOWN>"); } + let data_for_remote_address = remote_addresses + .entry(connection.remote_socket.ip) + .or_default(); + connection_data.total_bytes_downloaded += connection_info.total_bytes_downloaded; + connection_data.total_bytes_uploaded += connection_info.total_bytes_uploaded; + connection_data.interface_name = connection_info.interface_name; + data_for_remote_address.total_bytes_downloaded += + connection_info.total_bytes_downloaded; + data_for_remote_address.total_bytes_uploaded += connection_info.total_bytes_uploaded; + data_for_remote_address.connection_count += 1; + total_bytes_downloaded += connection_info.total_bytes_downloaded; + total_bytes_uploaded += connection_info.total_bytes_uploaded; } UIState { processes, diff --git a/src/main.rs b/src/main.rs index 589e000..ea40de0 100644 --- a/src/main.rs +++ b/src/main.rs @@ -9,7 +9,7 @@ mod tests; use display::{RawTerminalBackend, Ui}; use network::{ dns::{self, IpTable}, - Connection, Sniffer, Utilization, + Connection, LocalSocket, Sniffer, Utilization, }; use os::OnSigWinch; @@ -79,10 +79,15 @@ fn try_main() -> Result<(), failure::Error> { Ok(()) } +pub struct OpenSockets { + sockets_to_procs: HashMap<LocalSocket, String>, + connections: Vec<Connection>, +} + pub struct OsInputOutput { pub network_interfaces: Vec<NetworkInterface>, pub network_frames: Vec<Box<dyn DataLinkReceiver>>, - pub get_open_sockets: fn() -> HashMap<Connection, String>, + pub get_open_sockets: fn() -> OpenSockets, pub keyboard_events: Box<dyn Iterator<Item = Event> + Send>, pub dns_client: Option<dns::Client>, pub on_winch: Box<OnSigWinch>, @@ -138,12 +143,15 @@ 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 = 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(); - let unresolved_ips = connections_to_procs - .keys() + let unresolved_ips = connections + .iter() .filter(|conn| !ip_to_host.contains_key(&conn.remote_socket.ip)) .map(|conn| conn.remote_socket.ip) .collect::<Vec<_>>(); @@ -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 b4e0de2..907e929 100644 --- a/src/network/connection.rs +++ b/src/network/connection.rs @@ -1,16 +1,20 @@ use ::std::collections::HashMap; use ::std::fmt; -use ::std::net::Ipv4Addr; +use ::std::net::{IpAddr, Ipv4Addr}; use ::std::net::SocketAddr; -#[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord, Debug)] +#[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord, Debug, Copy)] pub enum Protocol { Tcp, Udp, } 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<Self> { match string { "TCP" => Some(Protocol::Tcp), @@ -29,17 +33,23 @@ impl fmt::Display for Protocol { } } -#[derive(Clone, Ord, PartialOrd, PartialEq, Eq, Hash)] +#[derive(Clone, Ord, PartialOrd, PartialEq, Eq, Hash, Debug, Copy)] pub struct Socket { pub ip: Ipv4Addr, pub port: u16, } -#[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord)] +#[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord, Debug, Copy)] +pub struct LocalSocket { + pub ip: IpAddr, + pub port: u16, + pub protocol: Protocol, +} + +#[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord, Debug, Copy)] pub struct Connection { pub remote_socket: Socket, - pub protocol: Protocol, - pub local_port: u16, + pub local_socket: LocalSocket, } pub fn display_ip_or_host(ip: Ipv4Addr, ip_to_host: &HashMap<Ipv4Addr, String>) -> String { @@ -57,23 +67,31 @@ pub fn display_connection_string( format!( "<{}>:{} => {}:{} ({})", interface_name, - connection.local_port, + connection.local_socket.port, display_ip_or_host(connection.remote_socket.ip, ip_to_host), connection.remote_socket.port, - connection.protocol, + connection.local_socket.protocol, ) } impl Connection { - pub fn new(remote_socket: SocketAddr, local_port: u16, protocol: Protocol) -> Option<Self> { + pub fn new( + remote_socket: SocketAddr, + local_ip: IpAddr, + 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, + local_socket: LocalSocket { + ip: local_ip, + port: local_port, + protocol, + }, }), _ => None, } diff --git a/src/network/sniffer.rs b/src/network/sniffer.rs index c7c5ba7..e0a5a39 100644 --- a/src/network/sniffer.rs +++ b/src/network/sniffer.rs @@ -88,8 +88,10 @@ impl Sniffer { let to = SocketAddr::new(IpAddr::V4(ip_packet.get_destination()), destination_port); let connection = match direction { - Direction::Download => Connection::new(from, destination_port, protocol)?, - Direction::Upload => Connection::new(to, source_port, protocol)?, + Direction::Download => { + Connection::new(from, to.ip(), destination_port, protocol)? + } + Direction::Upload => Connection::new(to, from.ip(), source_port, protocol)?, }; Some(Segment { interface_name, diff --git a/src/os/linux.rs b/src/os/linux.rs index 4d65edb..9ca566c 100644 --- a/src/os/linux.rs +++ b/src/os/linux.rs @@ -3,9 +3,11 @@ use ::std::collections::HashMap; use ::procfs::process::FDTarget; use crate::network::{Connection, Protocol}; +use crate::OpenSockets; -pub(crate) fn get_open_sockets() -> HashMap<Connection, String> { +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(); let mut inode_to_procname = HashMap::new(); @@ -23,23 +25,30 @@ pub(crate) fn get_open_sockets() -> HashMap<Connection, String> { let tcp = ::procfs::net::tcp().unwrap(); for entry in tcp.into_iter() { let local_port = entry.local_address.port(); + let local_ip = entry.local_address.ip(); if let (Some(connection), Some(procname)) = ( - Connection::new(entry.remote_address, local_port, Protocol::Tcp), + Connection::new(entry.remote_address, local_ip, local_port, Protocol::Tcp), inode_to_procname.get(&entry.inode), ) { - open_sockets.insert(connection, procname.clone()); + open_sockets.insert(connection.local_socket, procname.clone()); + connections.push(connection); }; } let udp = ::procfs::net::udp().unwrap(); for entry in udp.into_iter() { let local_port = entry.local_address.port(); + let local_ip = entry.local_address.ip(); if let (Some(connection), Some(procname)) = ( - Connection::new(entry.remote_address, local_port, Protocol::Udp), + Connection::new(entry.remote_address, local_ip, local_port, Protocol::Udp), inode_to_procname.get(&entry.inode), ) { - open_sockets.insert(connection, procname.clone()); + open_sockets.insert(connection.local_socket, procname.clone()); + connections.push(connection); }; } - open_sockets + OpenSockets { + sockets_to_procs: open_sockets, + connections, + } } diff --git a/src/os/lsof_utils.rs b/src/os/lsof_utils.rs index d1ca51b..584acab 100644 --- a/src/os/lsof_utils.rs +++ b/src/os/lsof_utils.rs @@ -7,7 +7,8 @@ use std::process::Command; #[derive(Debug, Clone)] pub struct RawConnection { - ip: String, + remote_ip: String, + local_ip: String, local_port: String, remote_port: String, protocol: String, @@ -16,48 +17,109 @@ pub struct RawConnection { lazy_static! { static ref CONNECTION_REGEX: Regex = - Regex::new(r"([^\s]+).*(TCP|UDP).*:(.*)->(.*):(\d*)(\s|$)").unwrap(); + Regex::new(r"\[?([^\s\]]*)\]?:(\d+)->\[?([^\s\]]*)\]?:(\d+)").unwrap(); + static ref LISTEN_REGEX: Regex = Regex::new(r"(.*):(.*)").unwrap(); +} + +fn get_null_addr(ip_type: &str) -> &str { + if ip_type.contains('4') { + "0.0.0.0" + } else { + "::0" + } } impl RawConnection { pub fn new(raw_line: &str) -> Option<RawConnection> { - let raw_connection_iter = CONNECTION_REGEX.captures_iter(raw_line).filter_map(|cap| { - let process_name = String::from(cap.get(1).unwrap().as_str()).replace("\\x20", " "); - let protocol = String::from(cap.get(2).unwrap().as_str()); - let local_port = String::from(cap.get(3).unwrap().as_str()); - let ip = String::from(cap.get(4).unwrap().as_str()); - let remote_port = String::from(cap.get(5).unwrap().as_str()); + // 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 { + return None; + } + let process_name = columns[0].replace("\\x20", " "); + // Unneeded + // let pid = columns[1]; + // let username = columns[2]; + // let fd = columns[3]; + + // IPv4 or IPv6 + let ip_type = columns[4]; + // let device = columns[5]; + // let size = columns[6]; + // UDP/TCP + let protocol = columns[7].to_ascii_uppercase(); + if protocol != "TCP" && protocol != "UDP" { + return None; + } + let connection_str = columns[8]; + // "(LISTEN)" or "(ESTABLISHED)", this column may or may not be present + // 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()); + let remote_port = String::from(caps.get(4).unwrap().as_str()); let connection = RawConnection { + local_ip, + local_port, + remote_ip, + remote_port, + protocol, process_name, - ip, + }; + Some(connection) + } else if let Some(caps) = LISTEN_REGEX.captures(connection_str) { + let local_ip = if caps.get(1).unwrap().as_str() == "*" { + get_null_addr(ip_type) + } else { + caps.get(1).unwrap().as_str() + }; + let local_ip = String::from(local_ip); + let local_port = String::from(if caps.get(2).unwrap().as_str() == "*" { + "0" + } else { + caps.get(2).unwrap().as_str() + }); + let remote_ip = String::from(get_null_addr(ip_type)); + let remote_port = String::from("0"); + let connection = RawConnection { + local_ip, local_port, + remote_ip, remote_port, protocol, + process_name, }; Some(connection) - }); - let raw_connection_vec = raw_connection_iter.map(|m| m).collect::<Vec<_>>(); - if raw_connection_vec.is_empty() { - None } else { - Some(raw_connection_vec[0].clone()) + None } } pub fn get_protocol(&self) -> Protocol { - return Protocol::from_str(&self.protocol).unwrap(); + Protocol::from_str(&self.protocol).unwrap() } - pub fn get_ip_address(&self) -> IpAddr { - return IpAddr::V4(self.ip.parse().unwrap()); + pub fn get_remote_ip(&self) -> IpAddr { + self.remote_ip.parse().unwrap() } pub fn get_remote_port(&self) -> u16 { - return self.remote_port.parse::<u16>().unwrap(); + self.remote_port.parse::<u16>().unwrap() + } + + pub fn get_local_ip(&self) -> IpAddr { + self.local_ip.parse().unwrap() } pub fn get_local_port(&self) -> u16 { - return self.local_port.parse::<u16>().unwrap(); + self.local_port.parse::<u16>().unwrap() } } @@ -160,7 +222,7 @@ com.apple 590 etoledom 204u IPv4 0x28ffb9c04111253f 0t0 TCP 192.168.1. fn test_raw_connection_parse_ip_address() { let connection = RawConnection::new(LINE_RAW_OUTPUT).unwrap(); assert_eq!( - connection.get_ip_address().to_string(), + connection.get_remote_ip().to_string(), String::from("198.252.206.25") ); } diff --git a/src/os/macos.rs b/src/os/macos.rs index cf68d04..68b4bad 100644 --- a/src/os/macos.rs +++ b/src/os/macos.rs @@ -1,6 +1,7 @@ use ::std::collections::HashMap; use crate::network::Connection; +use crate::OpenSockets; use super::lsof_utils; use std::net::SocketAddr; @@ -14,22 +15,28 @@ struct RawConnection { process_name: String, } -pub(crate) fn get_open_sockets() -> HashMap<Connection, String> { +pub(crate) fn get_open_sockets() -> OpenSockets { let mut open_sockets = HashMap::new(); + let mut connections_vec = std::vec::Vec::new(); let connections = lsof_utils::get_connections(); for raw_connection in connections { let protocol = raw_connection.get_protocol(); - let ip_address = raw_connection.get_ip_address(); + let remote_ip = raw_connection.get_remote_ip(); + let local_ip = raw_connection.get_local_ip(); let remote_port = raw_connection.get_remote_port(); let local_port = raw_connection.get_local_port(); - let socket_addr = SocketAddr::new(ip_address, remote_port); - let connection = Connection::new(socket_addr, local_port, protocol).unwrap(); + let socket_addr = SocketAddr::new(remote_ip, remote_port); + let connection = Connection::new(socket_addr, local_ip, local_port, protocol).unwrap(); - open_sockets.insert(connection, raw_connection.process_name.clone()); + open_sockets.insert(connection.local_socket, raw_connection.process_name.clone()); + connections_vec.push(connection); } - return open_sockets; + OpenSockets { + sockets_to_procs: open_sockets, + connections: connections_vec, + } } diff --git a/src/tests/cases/raw_mode.rs b/src/tests/cases/raw_mode.rs index 8741a44..5353fd0 100644 --- a/src/tests/cases/raw_mode.rs +++ b/src/tests/cases/raw_mode.rs @@ -95,17 +95,17 @@ fn bi_directional_traffic() { fn multiple_packets_of_traffic_from_different_connections() { let network_frames = vec![NetworkFrames::new(vec![ Some(build_tcp_packet( - "1.1.1.1", + "2.2.2.2", "10.0.0.2", 12345, 443, - b"I have come from 1.1.1.1", + b"I have come from 2.2.2.2", )), Some(build_tcp_packet( "2.2.2.2", "10.0.0.2", 54321, - 443, + 4434, b"I come from 2.2.2.2", )), ]) as Box<dyn DataLinkReceiver>]; @@ -158,11 +158,11 @@ fn one_process_with_multiple_connections() { b"I have come from 1.1.1.1", )), Some(build_tcp_packet( - "3.3.3.3", + "1.1.1.1", "10.0.0.2", - 1337, + 12346, 443, - b"Funny that, I'm from 3.3.3.3", + b"Funny that, I'm from 1.1.1.1", )), ]) as Box<dyn DataLinkReceiver>]; let (_, _, backend) = test_backend_factory(190, 50); @@ -189,21 +189,21 @@ fn multiple_processes_with_multiple_connections() { "3.3.3.3", "10.0.0.2", 1337, - 443, + 4435, b"Awesome, I'm from 3.3.3.3", )), Some(build_tcp_packet( "2.2.2.2", "10.0.0.2", 54321, - 443, + 4434, b"You know, 2.2.2.2 is really nice!", )), Some(build_tcp_packet( "4.4.4.4", "10.0.0.2", 1337, - 443, + 4432, b"I'm partial to 4.4.4.4", )), ]) as Box<dyn DataLinkReceiver>]; @@ -290,7 +290,7 @@ fn sustained_traffic_from_multiple_processes() { "3.3.3.3", "10.0.0.2", 1337, - 443, + 4435, b"I come from 3.3.3.3", )), None, // sleep @@ -305,7 +305,7 @@ fn sustained_traffic_from_multiple_processes() { "3.3.3.3", "10.0.0.2", 1337, - 443, + 4435, b"I come 3.3.3.3 one second later", )), ]) as Box<dyn DataLinkReceiver>]; @@ -326,7 +326,7 @@ fn sustained_traffic_from_multiple_processes_bi_directional() { Some(build_tcp_packet( "10.0.0.2", "3.3.3.3", - 443, + 4435, 1337, b"omw to 3.3.3.3", )), @@ -334,7 +334,7 @@ fn sustained_traffic_from_multiple_processes_bi_directional() { "3.3.3.3", "10.0.0.2", 1337, - 443, + 4435, b"I was just there!", )), Some(build_tcp_packet( @@ -355,7 +355,7 @@ fn sustained_traffic_from_multiple_processes_bi_directional() { Some(build_tcp_packet( "10.0.0.2", "3.3.3.3", - 443, + 4435, 1337, b"Wait for me!", )), @@ -363,7 +363,7 @@ fn sustained_traffic_from_multiple_processes_bi_directional() { "3.3.3.3", "10.0.0.2", 1337, - 443, + 4435, b"They're waiting for you...", )), Some(build_tcp_packet( @@ -398,7 +398,7 @@ fn traffic_with_host_names() { Some(build_tcp_packet( "10.0.0.2", "3.3.3.3", - 443, + 4435, 1337, b"omw to 3.3.3.3", )), @@ -406,7 +406,7 @@ fn traffic_with_host_names() { "3.3.3.3", "10.0.0.2", 1337, |