diff --git a/crates/rustnet-core/src/network/types.rs b/crates/rustnet-core/src/network/types.rs index 49d9f3f..acbacb6 100644 --- a/crates/rustnet-core/src/network/types.rs +++ b/crates/rustnet-core/src/network/types.rs @@ -1613,34 +1613,38 @@ pub struct AppProtocolDistribution { } impl AppProtocolDistribution { + /// Add one connection to the distribution. + pub fn record_connection(&mut self, conn: &Connection) { + if let Some(dpi_info) = &conn.dpi_info { + match &dpi_info.application { + ApplicationProtocol::Https(_) => self.https_count += 1, + ApplicationProtocol::Http(_) => self.http_count += 1, + ApplicationProtocol::Quic(_) => self.quic_count += 1, + ApplicationProtocol::Dns(_) => self.dns_count += 1, + ApplicationProtocol::Ssh(_) => self.ssh_count += 1, + ApplicationProtocol::Ntp(_) + | ApplicationProtocol::Mdns(_) + | ApplicationProtocol::Llmnr(_) + | ApplicationProtocol::Dhcp(_) + | ApplicationProtocol::Snmp(_) + | ApplicationProtocol::Ssdp(_) + | ApplicationProtocol::NetBios(_) + | ApplicationProtocol::BitTorrent(_) + | ApplicationProtocol::Stun(_) + | ApplicationProtocol::Mqtt(_) + | ApplicationProtocol::Ftp(_) => self.other_count += 1, + } + } else { + self.other_count += 1; + } + } + /// Calculate distribution from a list of connections pub fn from_connections(connections: &[Connection]) -> Self { let mut dist = Self::default(); for conn in connections { - if let Some(dpi_info) = &conn.dpi_info { - match &dpi_info.application { - ApplicationProtocol::Https(_) => dist.https_count += 1, - ApplicationProtocol::Http(_) => dist.http_count += 1, - ApplicationProtocol::Quic(_) => dist.quic_count += 1, - ApplicationProtocol::Dns(_) => dist.dns_count += 1, - ApplicationProtocol::Ssh(_) => dist.ssh_count += 1, - // New protocols counted as "other" for now - ApplicationProtocol::Ntp(_) - | ApplicationProtocol::Mdns(_) - | ApplicationProtocol::Llmnr(_) - | ApplicationProtocol::Dhcp(_) - | ApplicationProtocol::Snmp(_) - | ApplicationProtocol::Ssdp(_) - | ApplicationProtocol::NetBios(_) - | ApplicationProtocol::BitTorrent(_) - | ApplicationProtocol::Stun(_) - | ApplicationProtocol::Mqtt(_) - | ApplicationProtocol::Ftp(_) => dist.other_count += 1, - } - } else { - dist.other_count += 1; - } + dist.record_connection(conn); } dist diff --git a/src/main.rs b/src/main.rs index ec3d6a6..3818d8a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -795,9 +795,8 @@ where // Ensure we have a valid selection (handles connection removals) if ui_state.grouping_enabled { - ui_state.ensure_valid_grouped_selection(&grouped_rows); let selected_idx = ui_state - .get_selected_grouped_index(&grouped_rows) + .ensure_valid_grouped_selection(&grouped_rows) .unwrap_or(0); ui_state.grouped_scroll_offset = ui::compute_scroll_offset( selected_idx, @@ -806,8 +805,7 @@ where grouped_rows.len(), ); } else { - ui_state.ensure_valid_selection(&connections); - let selected_idx = ui_state.get_selected_index(&connections).unwrap_or(0); + let selected_idx = ui_state.ensure_valid_selection(&connections).unwrap_or(0); ui_state.scroll_offset = ui::compute_scroll_offset( selected_idx, ui_state.scroll_offset, diff --git a/src/ui/state.rs b/src/ui/state.rs index 315afe7..50232ae 100644 --- a/src/ui/state.rs +++ b/src/ui/state.rs @@ -212,6 +212,13 @@ impl ClickableRegions { pub struct UIState { pub selected_tab: usize, pub selected_connection_key: Option, + /// Cached positions for the selected key. Each lookup validates the hint + /// against the current row before using it, so sorting and filtering keep + /// key-based selection semantics without repeatedly scanning the full list. + #[doc(hidden)] + pub selected_connection_index_hint: Cell>, + #[doc(hidden)] + pub selected_grouped_index_hint: Cell>, pub show_help: bool, pub quit_confirmation: bool, pub clear_confirmation: bool, @@ -258,6 +265,8 @@ impl Default for UIState { Self { selected_tab: 0, selected_connection_key: None, + selected_connection_index_hint: Cell::new(None), + selected_grouped_index_hint: Cell::new(None), show_help: false, quit_confirmation: false, clear_confirmation: false, @@ -318,6 +327,8 @@ impl UIState { pub fn set_connection_key(&mut self, key: Option) { if self.selected_connection_key != key { self.details_scroll.reset(); + self.selected_connection_index_hint.set(None); + self.selected_grouped_index_hint.set(None); } self.selected_connection_key = key; } @@ -325,9 +336,19 @@ impl UIState { /// Get the current selected connection index, if any pub fn get_selected_index(&self, connections: &[Connection]) -> Option { if let Some(ref selected_key) = self.selected_connection_key { - connections + if let Some(index) = self.selected_connection_index_hint.get() + && connections + .get(index) + .is_some_and(|conn| conn.key() == *selected_key) + { + return Some(index); + } + + let index = connections .iter() - .position(|conn| conn.key() == *selected_key) + .position(|conn| conn.key() == *selected_key); + self.selected_connection_index_hint.set(index); + index } else if !connections.is_empty() { Some(0) // Default to first connection } else { @@ -339,6 +360,7 @@ impl UIState { pub fn set_selected_by_index(&mut self, connections: &[Connection], index: usize) { if let Some(conn) = connections.get(index) { self.set_connection_key(Some(conn.key())); + self.selected_connection_index_hint.set(Some(index)); } } @@ -463,11 +485,11 @@ impl UIState { } /// Ensure we have a valid selection when connections list changes - pub fn ensure_valid_selection(&mut self, connections: &[Connection]) { + pub fn ensure_valid_selection(&mut self, connections: &[Connection]) -> Option { if connections.is_empty() { log::debug!("ensure_valid_selection: connections list is empty, clearing selection"); self.set_connection_key(None); - return; + return None; } let current_index = self.get_selected_index(connections); @@ -481,6 +503,9 @@ impl UIState { if self.selected_connection_key.is_none() || current_index.is_none() { log::debug!("ensure_valid_selection: selecting first connection (index 0)"); self.set_selected_by_index(connections, 0); + Some(0) + } else { + current_index } } @@ -627,15 +652,32 @@ impl UIState { /// Get the current selected index in the grouped rows pub fn get_selected_grouped_index(&self, grouped_rows: &[GroupedRow]) -> Option { if grouped_rows.is_empty() { + self.selected_grouped_index_hint.set(None); return None; } + if let Some(index) = self.selected_grouped_index_hint.get() + && let Some(row) = grouped_rows.get(index) + { + let matches = if let Some(ref selected_key) = self.selected_connection_key { + matches!(row, GroupedRow::Connection { connection, .. } if connection.key() == *selected_key) + } else if let Some(ref selected_group) = self.selected_group { + matches!(row, GroupedRow::Group { process_name, .. } if process_name == selected_group) + } else { + false + }; + if matches { + return Some(index); + } + } + // First check if we have a selected connection that's visible if let Some(ref selected_key) = self.selected_connection_key { for (idx, row) in grouped_rows.iter().enumerate() { if let GroupedRow::Connection { connection, .. } = row && connection.key() == *selected_key { + self.selected_grouped_index_hint.set(Some(idx)); return Some(idx); } } @@ -647,6 +689,7 @@ impl UIState { if let GroupedRow::Group { process_name, .. } = row && process_name == selected_group { + self.selected_grouped_index_hint.set(Some(idx)); return Some(idx); } } @@ -663,6 +706,7 @@ impl UIState { GroupedRow::Group { process_name, .. } => { self.selected_group = Some(process_name.clone()); self.set_connection_key(None); + self.selected_grouped_index_hint.set(Some(index)); } GroupedRow::Connection { process_name, @@ -671,6 +715,7 @@ impl UIState { } => { self.set_connection_key(Some(connection.key())); self.selected_group = Some(process_name.clone()); + self.selected_grouped_index_hint.set(Some(index)); } } } @@ -737,20 +782,23 @@ impl UIState { } /// Ensure valid selection in grouped view - pub fn ensure_valid_grouped_selection(&mut self, grouped_rows: &[GroupedRow]) { + pub fn ensure_valid_grouped_selection(&mut self, grouped_rows: &[GroupedRow]) -> Option { if grouped_rows.is_empty() { self.selected_group = None; self.set_connection_key(None); - return; + return None; } // If no group is selected, or current selection is not visible, reset to first row // This handles the case when grouping is first enabled - let needs_init = self.selected_group.is_none() - || self.get_selected_grouped_index(grouped_rows).is_none(); + let current_index = self.get_selected_grouped_index(grouped_rows); + let needs_init = self.selected_group.is_none() || current_index.is_none(); if needs_init { self.set_selected_grouped_by_index(grouped_rows, 0); + Some(0) + } else { + current_index } } @@ -846,6 +894,18 @@ pub fn compute_grouped_rows<'a>( mod tests { use super::*; use crate::ui::{HELP_TAB_INDEX, TAB_COUNT}; + use std::net::{IpAddr, Ipv4Addr, SocketAddr}; + + fn test_connection(port: u16, process: &str) -> Connection { + let mut connection = Connection::new( + Protocol::Tcp, + SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(192, 0, 2, 1)), 443), + crate::network::types::ProtocolState::Tcp(crate::network::types::TcpState::Established), + ); + connection.process_name = Some(process.to_string()); + connection + } #[test] fn jump_to_tab_sets_selected_and_help_flag() { @@ -953,4 +1013,38 @@ mod tests { ui.set_connection_key(Some("b".to_string())); assert_eq!(ui.details_scroll.clamp_for_render(20), 0); } + + #[test] + fn selection_hint_recovers_after_reordering() { + let mut connections = vec![ + test_connection(1000, "first"), + test_connection(1001, "second"), + ]; + let mut ui = UIState::default(); + ui.set_selected_by_index(&connections, 1); + assert_eq!(ui.selected_connection_index_hint.get(), Some(1)); + + connections.swap(0, 1); + + assert_eq!(ui.get_selected_index(&connections), Some(0)); + assert_eq!(ui.selected_connection_index_hint.get(), Some(0)); + } + + #[test] + fn grouped_selection_hint_recovers_after_rows_shift() { + let connections = vec![ + test_connection(1000, "alpha"), + test_connection(1001, "beta"), + ]; + let expanded = HashSet::from(["alpha".to_string(), "beta".to_string()]); + let rows = compute_grouped_rows(&connections, &expanded); + let mut ui = UIState::default(); + ui.set_selected_grouped_by_index(&rows, 3); + assert_eq!(ui.selected_grouped_index_hint.get(), Some(3)); + + let rows = compute_grouped_rows(&connections, &HashSet::from(["beta".to_string()])); + + assert_eq!(ui.get_selected_grouped_index(&rows), Some(2)); + assert_eq!(ui.selected_grouped_index_hint.get(), Some(2)); + } } diff --git a/src/ui/tabs/graph.rs b/src/ui/tabs/graph.rs index dad6054..ea7c139 100644 --- a/src/ui/tabs/graph.rs +++ b/src/ui/tabs/graph.rs @@ -2,6 +2,8 @@ //! health, TCP counters, TCP state distribution, application //! protocol distribution, and top processes by bandwidth. +use std::collections::HashMap; + use anyhow::Result; use ratatui::{ Frame, @@ -20,6 +22,70 @@ use crate::ui::{ widgets::braille_graph, }; +const TCP_STATE_NAMES: [&str; 11] = [ + "ESTAB", + "SYN_SENT", + "SYN_RECV", + "FIN_WAIT1", + "FIN_WAIT2", + "TIME_WAIT", + "CLOSE_WAIT", + "LAST_ACK", + "CLOSING", + "CLOSED", + "UNKNOWN", +]; + +/// Data used by the connection-derived graph panels. Building it once keeps +/// the render cost to one pass over active connections. +struct GraphAnalytics<'a> { + app_distribution: AppProtocolDistribution, + process_traffic: HashMap<&'a str, f64>, + tcp_state_counts: [usize; TCP_STATE_NAMES.len()], +} + +impl<'a> GraphAnalytics<'a> { + fn from_connections(connections: &'a [Connection]) -> Self { + let mut analytics = Self { + app_distribution: AppProtocolDistribution::default(), + process_traffic: HashMap::new(), + tcp_state_counts: [0; TCP_STATE_NAMES.len()], + }; + + for conn in connections.iter().filter(|conn| !conn.is_historic) { + analytics.app_distribution.record_connection(conn); + + let name = conn.process_name.as_deref().unwrap_or("Unknown"); + let traffic = conn.current_incoming_rate_bps + conn.current_outgoing_rate_bps; + *analytics.process_traffic.entry(name).or_insert(0.0) += traffic; + + if conn.protocol == Protocol::Tcp + && let ProtocolState::Tcp(state) = &conn.protocol_state + { + analytics.tcp_state_counts[tcp_state_index(state)] += 1; + } + } + + analytics + } +} + +fn tcp_state_index(state: &TcpState) -> usize { + match state { + TcpState::Established => 0, + TcpState::SynSent => 1, + TcpState::SynReceived => 2, + TcpState::FinWait1 => 3, + TcpState::FinWait2 => 4, + TcpState::TimeWait => 5, + TcpState::CloseWait => 6, + TcpState::LastAck => 7, + TcpState::Closing => 8, + TcpState::Closed => 9, + TcpState::Unknown => 10, + } +} + /// Bold default-foreground title span for a graph section header. fn graph_title(text: &str) -> Span<'_> { Span::styled(text, Style::default().add_modifier(Modifier::BOLD)) @@ -47,14 +113,7 @@ pub(in crate::ui) fn draw_graph_tab( connections: &[Connection], area: Rect, ) -> Result<()> { - // Filter out historic connections — graph should only show alive connections - let active_connections: Vec = connections - .iter() - .filter(|c| !c.is_historic) - .cloned() - .collect(); - let connections = &active_connections; - + let analytics = GraphAnalytics::from_connections(connections); let traffic_history = app.get_traffic_history(); // Each panel is a borderless section_header region; layout spacing @@ -98,9 +157,9 @@ pub(in crate::ui) fn draw_graph_tab( draw_connections_sparkline(f, &traffic_history, top_chunks[1]); draw_health_chart(f, &traffic_history, health_chunks[0]); draw_tcp_counters(f, app, health_chunks[1]); - draw_tcp_states(f, connections, health_chunks[2]); - draw_app_distribution(f, connections, bottom_chunks[0]); - draw_top_processes(f, connections, bottom_chunks[1]); + draw_tcp_states(f, &analytics.tcp_state_counts, health_chunks[2]); + draw_app_distribution(f, &analytics.app_distribution, bottom_chunks[0]); + draw_top_processes(f, &analytics.process_traffic, bottom_chunks[1]); Ok(()) } @@ -215,10 +274,9 @@ fn draw_connections_sparkline(f: &mut Frame, history: &TrafficHistory, area: Rec } /// Draw application protocol distribution -fn draw_app_distribution(f: &mut Frame, connections: &[Connection], area: Rect) { +fn draw_app_distribution(f: &mut Frame, dist: &AppProtocolDistribution, area: Rect) { let inner = section_header(f, area, graph_title(" Application Distribution")); - let dist = AppProtocolDistribution::from_connections(connections); let percentages = dist.as_percentages(); // Filter out zero-count protocols and create bars. @@ -273,34 +331,18 @@ fn draw_app_distribution(f: &mut Frame, connections: &[Connection], area: Rect) } /// Draw top processes by bandwidth -fn draw_top_processes<'a>(f: &mut Frame, connections: &'a [Connection], area: Rect) { +fn draw_top_processes(f: &mut Frame, process_traffic: &HashMap<&str, f64>, area: Rect) { use std::borrow::Cow; - use std::collections::HashMap; let inner = section_header(f, area, graph_title(" Top Processes")); - // Aggregate traffic by process; borrow process names from the connections - // slice to avoid cloning one String per connection per frame. - let mut process_traffic: HashMap<&'a str, f64> = HashMap::new(); - for conn in connections { - let name = conn.process_name.as_deref().unwrap_or("Unknown"); - let traffic = conn.current_incoming_rate_bps + conn.current_outgoing_rate_bps; - *process_traffic.entry(name).or_insert(0.0) += traffic; - } - - // Sort by traffic descending, filter out processes with no traffic - let mut sorted: Vec<_> = process_traffic - .into_iter() - .filter(|(_, rate)| *rate > 0.0) - .collect(); - sorted.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); + let top_processes = select_top_processes(process_traffic, 5); // Create rows for top 5 processes. Process name absorbs whatever width is // left after the fixed-width Rate column, and Rate is right-aligned so the // numbers form a clean right edge. - let rows: Vec = sorted + let rows: Vec = top_processes .into_iter() - .take(5) .map(|(name, rate)| { let display_name: Cow = if name.len() > 20 { Cow::Owned(format!("{}...", &name[..17])) @@ -332,6 +374,24 @@ fn draw_top_processes<'a>(f: &mut Frame, connections: &'a [Connection], area: Re f.render_widget(table, inner); } +fn select_top_processes<'a>( + process_traffic: &HashMap<&'a str, f64>, + limit: usize, +) -> Vec<(&'a str, f64)> { + let mut processes: Vec<_> = process_traffic + .iter() + .filter_map(|(&name, &rate)| (rate > 0.0).then_some((name, rate))) + .collect(); + let compare = |a: &(&str, f64), b: &(&str, f64)| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(b.0)); + + if processes.len() > limit { + processes.select_nth_unstable_by(limit, compare); + processes.truncate(limit); + } + processes.sort_unstable_by(compare); + processes +} + /// Draw the network health gauges with RTT and packet loss bars fn draw_health_chart(f: &mut Frame, history: &TrafficHistory, area: Rect) { let inner = section_header(f, area, graph_title(" Network Health")); @@ -506,52 +566,12 @@ fn draw_tcp_counters(f: &mut Frame, app: &App, area: Rect) { } /// Draw TCP connection states breakdown -fn draw_tcp_states(f: &mut Frame, connections: &[Connection], area: Rect) { - use std::collections::HashMap; - - // Count TCP states - let mut state_counts: HashMap<&str, usize> = HashMap::new(); - for conn in connections { - if conn.protocol == Protocol::Tcp - && let ProtocolState::Tcp(tcp_state) = &conn.protocol_state - { - let state_name = match tcp_state { - TcpState::Established => "ESTAB", - TcpState::SynSent => "SYN_SENT", - TcpState::SynReceived => "SYN_RECV", - TcpState::FinWait1 => "FIN_WAIT1", - TcpState::FinWait2 => "FIN_WAIT2", - TcpState::TimeWait => "TIME_WAIT", - TcpState::CloseWait => "CLOSE_WAIT", - TcpState::LastAck => "LAST_ACK", - TcpState::Closing => "CLOSING", - TcpState::Closed => "CLOSED", - TcpState::Unknown => "UNKNOWN", - }; - *state_counts.entry(state_name).or_insert(0) += 1; - } - } - - // Fixed order based on connection lifecycle (most important first) - const STATE_ORDER: &[&str] = &[ - "ESTAB", - "SYN_SENT", - "SYN_RECV", - "FIN_WAIT1", - "FIN_WAIT2", - "TIME_WAIT", - "CLOSE_WAIT", - "LAST_ACK", - "CLOSING", - "CLOSED", - "LISTEN", - "UNKNOWN", - ]; - +fn draw_tcp_states(f: &mut Frame, state_counts: &[usize; TCP_STATE_NAMES.len()], area: Rect) { // Build ordered list with only non-zero counts - let states: Vec<_> = STATE_ORDER + let states: Vec<_> = TCP_STATE_NAMES .iter() - .filter_map(|&name| state_counts.get(name).map(|&count| (name, count))) + .zip(state_counts) + .filter_map(|(&name, &count)| (count > 0).then_some((name, count))) .collect(); let inner = section_header(f, area, graph_title(" TCP States")); @@ -603,3 +623,76 @@ fn draw_tcp_states(f: &mut Frame, connections: &[Connection], area: Rect) { let paragraph = Paragraph::new(lines); f.render_widget(paragraph, inner); } + +#[cfg(test)] +mod tests { + use super::*; + use std::net::{IpAddr, Ipv4Addr, SocketAddr}; + + fn test_connection( + port: u16, + process: &str, + rate: f64, + state: TcpState, + historic: bool, + ) -> Connection { + let mut connection = Connection::new( + Protocol::Tcp, + SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port), + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(192, 0, 2, 1)), 443), + ProtocolState::Tcp(state), + ); + connection.process_name = Some(process.to_string()); + connection.current_incoming_rate_bps = rate; + connection.is_historic = historic; + connection + } + + #[test] + fn analytics_aggregate_active_connections_once() { + let connections = vec![ + test_connection(1000, "alpha", 10.0, TcpState::Established, false), + test_connection(1001, "alpha", 20.0, TcpState::TimeWait, false), + test_connection(1002, "beta", 5.0, TcpState::Established, false), + test_connection(1003, "old", 100.0, TcpState::Closed, true), + ]; + + let analytics = GraphAnalytics::from_connections(&connections); + + assert_eq!(analytics.app_distribution.total(), 3); + assert_eq!(analytics.process_traffic.get("alpha"), Some(&30.0)); + assert_eq!(analytics.process_traffic.get("beta"), Some(&5.0)); + assert!(!analytics.process_traffic.contains_key("old")); + assert_eq!( + analytics.tcp_state_counts[tcp_state_index(&TcpState::Established)], + 2 + ); + assert_eq!( + analytics.tcp_state_counts[tcp_state_index(&TcpState::TimeWait)], + 1 + ); + assert_eq!( + analytics.tcp_state_counts[tcp_state_index(&TcpState::Closed)], + 0 + ); + } + + #[test] + fn top_process_selection_only_sorts_the_requested_prefix() { + let traffic = HashMap::from([ + ("one", 1.0), + ("two", 2.0), + ("three", 3.0), + ("four", 4.0), + ("five", 5.0), + ("six", 6.0), + ("seven", 7.0), + ("eight", 8.0), + ]); + + let top = select_top_processes(&traffic, 3); + + assert_eq!(top, vec![("eight", 8.0), ("seven", 7.0), ("six", 6.0)]); + assert!(select_top_processes(&traffic, 0).is_empty()); + } +}