From e4a85edab729e2dfe9581173fa755c72d7bbd625 Mon Sep 17 00:00:00 2001 From: Badr Date: Sun, 15 Sep 2024 22:55:37 +0200 Subject: [PATCH] fix high cpu usage (#14) --- Cargo.lock | 5 +- oryx-common/Cargo.toml | 2 +- oryx-ebpf/Cargo.lock | 5 +- oryx-ebpf/Cargo.toml | 2 +- oryx-tui/src/app.rs | 562 ++++++++++++++++++------------ oryx-tui/src/ebpf.rs | 38 +- oryx-tui/src/event.rs | 1 - oryx-tui/src/filters/fuzzy.rs | 12 +- oryx-tui/src/filters/link.rs | 12 +- oryx-tui/src/filters/network.rs | 12 +- oryx-tui/src/filters/transport.rs | 12 +- oryx-tui/src/handler.rs | 90 +++-- oryx-tui/src/main.rs | 8 +- oryx-tui/src/notification.rs | 4 +- oryx-tui/src/stats.rs | 52 +-- 15 files changed, 462 insertions(+), 355 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 1e0cefa..5e078e7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -560,8 +560,9 @@ dependencies = [ [[package]] name = "network-types" -version = "0.0.6" -source = "git+https://github.com/vadorovsky/network-types.git?rev=e0ee8d5#e0ee8d50a446923fdcbafd36a99921a0f3aca7fd" +version = "0.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e82e9f64c09f56aa7c80c3fa087997bd99a913f91d9c74d36cf5fd75dd5773e6" [[package]] name = "object" diff --git a/oryx-common/Cargo.toml b/oryx-common/Cargo.toml index d557905..c30d7dd 100644 --- a/oryx-common/Cargo.toml +++ b/oryx-common/Cargo.toml @@ -13,4 +13,4 @@ edition.workspace = true path = "src/lib.rs" [dependencies] -network-types = { git = "https://github.com/vadorovsky/network-types.git", rev = "e0ee8d5" } +network-types = "0.0.7" diff --git a/oryx-ebpf/Cargo.lock b/oryx-ebpf/Cargo.lock index 24941da..03ef207 100644 --- a/oryx-ebpf/Cargo.lock +++ b/oryx-ebpf/Cargo.lock @@ -43,8 +43,9 @@ dependencies = [ [[package]] name = "network-types" -version = "0.0.6" -source = "git+https://github.com/vadorovsky/network-types.git?rev=e0ee8d5#e0ee8d50a446923fdcbafd36a99921a0f3aca7fd" +version = "0.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e82e9f64c09f56aa7c80c3fa087997bd99a913f91d9c74d36cf5fd75dd5773e6" [[package]] name = "oryx-common" diff --git a/oryx-ebpf/Cargo.toml b/oryx-ebpf/Cargo.toml index 5f129b0..7685f2d 100644 --- a/oryx-ebpf/Cargo.toml +++ b/oryx-ebpf/Cargo.toml @@ -10,7 +10,7 @@ homepage = "https://github.com/pythops/oryx" [dependencies] aya-ebpf = "0.1.0" oryx-common = { path = "../oryx-common" } -network-types = { git = "https://github.com/vadorovsky/network-types.git", rev = "e0ee8d5" } +network-types = "0.0.7" [[bin]] name = "oryx" diff --git a/oryx-tui/src/app.rs b/oryx-tui/src/app.rs index e7c0d64..3bce104 100644 --- a/oryx-tui/src/app.rs +++ b/oryx-tui/src/app.rs @@ -1,5 +1,5 @@ use oryx_common::ip::IpPacket; -use oryx_common::AppPacket; +use oryx_common::{AppPacket, RawPacket}; use ratatui::layout::{Alignment, Constraint, Direction, Flex, Layout, Margin, Rect}; use ratatui::style::{Color, Style, Stylize}; use ratatui::text::{Line, Span}; @@ -11,13 +11,17 @@ use ratatui::{ }, Frame, }; -use std::error; +use std::borrow::Borrow; use std::net::IpAddr; +use std::sync::{Arc, Mutex}; +use std::time::Duration; +use std::{error, thread}; use tui_big_text::{BigText, PixelSize}; +use crate::bandwidth::Bandwidth; use crate::filters::direction::TrafficDirectionFilter; use crate::filters::fuzzy::{self, Fuzzy}; -use crate::filters::link::{LinkFilter, NB_LINK_PROTOCOL}; +use crate::filters::link::{LinkFilter, LinkProtocol, NB_LINK_PROTOCOL}; use crate::filters::network::{NetworkFilter, NetworkProtocol, NB_NETWORK_PROTOCOL}; use crate::filters::transport::{TransportFilter, TransportProtocol, NB_TRANSPORT_PROTOCOL}; use crate::help::Help; @@ -27,6 +31,8 @@ use crate::stats::Stats; pub type AppResult = std::result::Result>; +pub const TICK_RATE: u64 = 30; + #[derive(Debug, Copy, Clone, PartialEq)] pub enum FocusedBlock { Interface, @@ -45,6 +51,12 @@ pub enum Mode { Stats, } +#[derive(Debug)] +pub struct DataEventHandler { + pub sender: kanal::Sender<[u8; RawPacket::LEN]>, + pub handler: thread::JoinHandle<()>, +} + #[derive(Debug)] pub struct App { pub running: bool, @@ -58,16 +70,18 @@ pub struct App { pub link_filter: LinkFilter, pub traffic_direction_filter: TrafficDirectionFilter, pub start_sniffing: bool, - pub packets: Vec, + pub packets: Arc>>, pub packets_table_state: TableState, - pub fuzzy: Fuzzy, + pub fuzzy: Arc>, pub notifications: Vec, pub manuall_scroll: bool, pub mode: Mode, - pub stats: Stats, + pub stats: Arc>, pub packet_end_index: usize, pub packet_window_size: usize, pub update_filters: bool, + pub data_channel_sender: kanal::Sender<[u8; RawPacket::LEN]>, + pub bandwidth: Arc>>, } impl Default for App { @@ -78,27 +92,101 @@ impl Default for App { impl App { pub fn new() -> Self { + let packets = Arc::new(Mutex::new(Vec::with_capacity(AppPacket::LEN * 1024 * 1024))); + let stats = Arc::new(Mutex::new(Stats::default())); + let fuzzy = Arc::new(Mutex::new(Fuzzy::default())); + + let network_filter = NetworkFilter::default(); + let transport_filter = TransportFilter::default(); + let link_filter = LinkFilter::default(); + + let (sender, receiver) = kanal::unbounded(); + + thread::spawn({ + let packets = packets.clone(); + let stats = stats.clone(); + let transport_filters = transport_filter.applied_protocols.clone(); + let network_filters = network_filter.applied_protocols.clone(); + let link_filters = link_filter.applied_protocols.clone(); + + move || loop { + if let Ok(raw_packet) = receiver.recv() { + App::process( + packets.clone(), + stats.clone(), + transport_filters.clone(), + network_filters.clone(), + link_filters.clone(), + AppPacket::from(raw_packet), + ); + } + } + }); + + let bandwidth = Arc::new(Mutex::new(Bandwidth::new().ok())); + + thread::spawn({ + let bandwidth = bandwidth.clone(); + move || loop { + thread::sleep(Duration::from_secs(1)); + { + let mut bandwidth = bandwidth.lock().unwrap(); + if bandwidth.is_some() { + let _ = bandwidth.as_mut().unwrap().refresh(); + } + } + } + }); + + thread::spawn({ + let fuzzy = fuzzy.clone(); + let packets = packets.clone(); + move || { + let mut last_index = 0; + let mut pattern = String::new(); + loop { + thread::sleep(Duration::from_millis(TICK_RATE)); + let packets = packets.lock().unwrap(); + let mut fuzzy = fuzzy.lock().unwrap(); + + if fuzzy.is_enabled() && !fuzzy.filter.value().is_empty() { + let current_pattern = fuzzy.filter.value().to_owned(); + if current_pattern != pattern { + fuzzy.find(packets.as_slice()); + pattern = current_pattern; + last_index = packets.len(); + } else { + fuzzy.append(&packets.as_slice()[last_index..]); + last_index = packets.len(); + } + } + } + } + }); + Self { running: true, help: Help::new(), focused_block: FocusedBlock::Interface, previous_focused_block: FocusedBlock::Interface, interface: Interface::default(), - network_filter: NetworkFilter::default(), - transport_filter: TransportFilter::default(), - link_filter: LinkFilter::default(), + network_filter, + transport_filter, + link_filter, traffic_direction_filter: TrafficDirectionFilter::default(), start_sniffing: false, - packets: Vec::with_capacity(AppPacket::LEN * 1024 * 1024), + packets, packets_table_state: TableState::default(), - fuzzy: Fuzzy::default(), + fuzzy, notifications: Vec::new(), manuall_scroll: false, mode: Mode::Packet, - stats: Stats::default(), + stats, packet_end_index: 0, packet_window_size: 0, update_filters: false, + data_channel_sender: sender, + bandwidth, } } @@ -243,87 +331,93 @@ impl App { // Filters let widths = [Constraint::Length(10), Constraint::Fill(1)]; - let filters = [ - Row::new(vec![ - Line::styled("Transport", Style::new().bold()), - Line::from_iter(TransportFilter::default().selected_protocols.iter().map( - |filter| { - if self.transport_filter.applied_protocols.contains(filter) { - Span::styled( - format!(" {} ", filter), - Style::default().light_green(), - ) - } else { - Span::styled( - format!("❌{} ", filter), - Style::default().light_red(), - ) - } - }, - )), - ]), - Row::new(vec![ - Line::styled("Network", Style::new().bold()), - Line::from_iter(NetworkFilter::default().selected_protocols.iter().map( - |filter| { - if self.network_filter.applied_protocols.contains(filter) { - Span::styled( - format!(" {} ", filter), - Style::default().light_green(), - ) - } else { - Span::styled( - format!("❌{} ", filter), - Style::default().light_red(), - ) - } - }, - )), - ]), - Row::new(vec![ - Line::styled("Link", Style::new().bold()), - Line::from_iter(LinkFilter::default().selected_protocols.iter().map( - |filter| { - if self.link_filter.applied_protocols.contains(filter) { - Span::styled( - format!(" {} ", filter), - Style::default().light_green(), - ) - } else { - Span::styled( - format!("❌{} ", filter), - Style::default().light_red(), - ) - } - }, - )), - ]), - Row::new(vec![ - Line::styled("Direction", Style::new().bold()), - Line::from_iter( - TrafficDirectionFilter::default() - .selected_direction - .iter() - .map(|filter| { - if self - .traffic_direction_filter - .applied_direction - .contains(filter) - { + let filters = { + let applied_transport_filters = + self.transport_filter.applied_protocols.lock().unwrap(); + let applied_network_filters = self.network_filter.applied_protocols.lock().unwrap(); + let applied_link_filters = self.link_filter.applied_protocols.lock().unwrap(); + [ + Row::new(vec![ + Line::styled("Transport", Style::new().bold()), + Line::from_iter(TransportFilter::default().selected_protocols.iter().map( + |filter| { + if applied_transport_filters.contains(filter) { Span::styled( - format!("󰞁 {} ", filter), + format!(" {} ", filter), Style::default().light_green(), ) } else { Span::styled( - format!("󰿝 {} ", filter), + format!("❌{} ", filter), Style::default().light_red(), ) } - }), - ), - ]), - ]; + }, + )), + ]), + Row::new(vec![ + Line::styled("Network", Style::new().bold()), + Line::from_iter(NetworkFilter::default().selected_protocols.iter().map( + |filter| { + if applied_network_filters.contains(filter) { + Span::styled( + format!(" {} ", filter), + Style::default().light_green(), + ) + } else { + Span::styled( + format!("❌{} ", filter), + Style::default().light_red(), + ) + } + }, + )), + ]), + Row::new(vec![ + Line::styled("Link", Style::new().bold()), + Line::from_iter(LinkFilter::default().selected_protocols.iter().map( + |filter| { + if applied_link_filters.contains(filter) { + Span::styled( + format!(" {} ", filter), + Style::default().light_green(), + ) + } else { + Span::styled( + format!("❌{} ", filter), + Style::default().light_red(), + ) + } + }, + )), + ]), + Row::new(vec![ + Line::styled("Direction", Style::new().bold()), + Line::from_iter( + TrafficDirectionFilter::default() + .selected_direction + .iter() + .map(|filter| { + if self + .traffic_direction_filter + .applied_direction + .contains(filter) + { + Span::styled( + format!("󰞁 {} ", filter), + Style::default().light_green(), + ) + } else { + Span::styled( + format!("󰿝 {} ", filter), + Style::default().light_red(), + ) + } + }), + ), + ]), + ] + }; let filter_table = Table::new(filters, widths).column_spacing(3).block( Block::default() @@ -431,8 +525,16 @@ impl App { } pub fn render_packets_mode(&mut self, frame: &mut Frame, packet_mode_block: Rect) { + let app_packets = self.packets.lock().unwrap(); + let mut fuzzy = self.fuzzy.lock().unwrap(); + let fuzzy_packets = fuzzy.clone().packets.clone(); + + //TODO: ugly + let pattern = fuzzy.clone(); + let pattern = pattern.filter.value(); + let (packet_block, fuzzy_block) = { - if self.fuzzy.is_enabled() { + if fuzzy.is_enabled() { let chunks = Layout::default() .direction(Direction::Vertical) .constraints([Constraint::Fill(1), Constraint::Length(3)]) @@ -467,89 +569,74 @@ impl App { self.packet_end_index = window_size; } - if self.fuzzy.packet_end_index < window_size { - self.fuzzy.packet_end_index = window_size; + if fuzzy.packet_end_index < window_size { + fuzzy.packet_end_index = window_size; } let packets_to_display = match self.manuall_scroll { true => { - if self.fuzzy.is_enabled() & !self.fuzzy.filter.value().is_empty() { - if self.fuzzy.packets.len() > window_size { - &self.fuzzy.packets.as_slice()[self - .fuzzy - .packet_end_index - .saturating_sub(window_size) - ..self.fuzzy.packet_end_index] + if fuzzy.is_enabled() & !fuzzy.filter.value().is_empty() { + if fuzzy_packets.len() > window_size { + &fuzzy_packets[fuzzy.packet_end_index.saturating_sub(window_size) + ..fuzzy.packet_end_index] } else { - &self.fuzzy.packets + &fuzzy_packets } - } else if self.packets.len() > window_size { - &self.packets + } else if app_packets.len() > window_size { + &app_packets [self.packet_end_index.saturating_sub(window_size)..self.packet_end_index] } else { - &self.packets + &app_packets } } false => { - if self.fuzzy.is_enabled() & !self.fuzzy.filter.value().is_empty() { - if self.fuzzy.packets.len() > window_size { - &self.fuzzy.packets[self.fuzzy.packets.len().saturating_sub(window_size)..] + if fuzzy.is_enabled() & !fuzzy.filter.value().is_empty() { + if fuzzy_packets.len() > window_size { + &fuzzy_packets[fuzzy_packets.len().saturating_sub(window_size)..] } else { - &self.fuzzy.packets + &fuzzy_packets } - } else if self.packets.len() > window_size { - &self.packets[self.packets.len().saturating_sub(window_size)..] + } else if app_packets.len() > window_size { + &app_packets[app_packets.len().saturating_sub(window_size)..] } else { - &self.packets + &app_packets } } }; // Style the packets - let packets: Vec = if self.fuzzy.is_enabled() & !self.fuzzy.filter.value().is_empty() { + let packets: Vec = if fuzzy.is_enabled() & !fuzzy.borrow().filter.value().is_empty() { packets_to_display .iter() .map(|app_packet| match app_packet { AppPacket::Arp(packet) => Row::new(vec![ - fuzzy::highlight(self.fuzzy.filter.value(), packet.src_mac.to_string()) - .blue(), + fuzzy::highlight(pattern, packet.src_mac.to_string()).blue(), Cell::from(Line::from("-").centered()).yellow(), - fuzzy::highlight(self.fuzzy.filter.value(), packet.dst_mac.to_string()) - .blue(), + fuzzy::highlight(pattern, packet.dst_mac.to_string()).blue(), Cell::from(Line::from("-").centered()).yellow(), - fuzzy::highlight(self.fuzzy.filter.value(), "ARP".to_string()).cyan(), + fuzzy::highlight(pattern, "ARP".to_string()).cyan(), ]), AppPacket::Ip(packet) => match packet { IpPacket::Tcp(p) => Row::new(vec![ - fuzzy::highlight(self.fuzzy.filter.value(), p.src_ip.to_string()) - .blue(), - fuzzy::highlight(self.fuzzy.filter.value(), p.src_port.to_string()) - .yellow(), - fuzzy::highlight(self.fuzzy.filter.value(), p.dst_ip.to_string()) - .blue(), - fuzzy::highlight(self.fuzzy.filter.value(), p.dst_port.to_string()) - .yellow(), - fuzzy::highlight(self.fuzzy.filter.value(), "TCP".to_string()).cyan(), + fuzzy::highlight(pattern, p.src_ip.to_string()).blue(), + fuzzy::highlight(pattern, p.src_port.to_string()).yellow(), + fuzzy::highlight(pattern, p.dst_ip.to_string()).blue(), + fuzzy::highlight(pattern, p.dst_port.to_string()).yellow(), + fuzzy::highlight(pattern, "TCP".to_string()).cyan(), ]), IpPacket::Udp(p) => Row::new(vec![ - fuzzy::highlight(self.fuzzy.filter.value(), p.src_ip.to_string()) - .blue(), - fuzzy::highlight(self.fuzzy.filter.value(), p.src_port.to_string()) - .yellow(), - fuzzy::highlight(self.fuzzy.filter.value(), p.dst_ip.to_string()) - .blue(), - fuzzy::highlight(self.fuzzy.filter.value(), p.dst_port.to_string()) - .yellow(), - fuzzy::highlight(self.fuzzy.filter.value(), "UDP".to_string()).cyan(), + fuzzy::highlight(pattern, p.src_ip.to_string()).blue(), + fuzzy::highlight(pattern, p.src_port.to_string()).yellow(), + fuzzy::highlight(pattern, p.dst_ip.to_string()).blue(), + fuzzy::highlight(pattern, p.dst_port.to_string()).yellow(), + fuzzy::highlight(pattern, "UDP".to_string()).cyan(), ]), IpPacket::Icmp(p) => Row::new(vec![ - fuzzy::highlight(self.fuzzy.filter.value(), p.src_ip.to_string()) - .blue(), + fuzzy::highlight(pattern, p.src_ip.to_string()).blue(), Cell::from(Line::from("-").centered()).yellow(), - fuzzy::highlight(self.fuzzy.filter.value(), p.dst_ip.to_string()) - .blue(), + fuzzy::highlight(pattern, p.dst_ip.to_string()).blue(), Cell::from(Line::from("-").centered()).yellow(), - fuzzy::highlight(self.fuzzy.filter.value(), "ICMP".to_string()).cyan(), + fuzzy::highlight(pattern, "ICMP".to_string()).cyan(), ]), }, }) @@ -606,10 +693,8 @@ impl App { // Always select the last packet if !self.manuall_scroll { - if self.fuzzy.is_enabled() { - self.fuzzy - .scroll_state - .select(Some(packets_to_display.len())); + if fuzzy.is_enabled() { + fuzzy.scroll_state.select(Some(packets_to_display.len())); } else { self.packets_table_state .select(Some(packets_to_display.len())); @@ -658,8 +743,8 @@ impl App { .border_style(Style::default().green()), ); - if self.fuzzy.is_enabled() { - frame.render_stateful_widget(table, packet_block, &mut self.fuzzy.scroll_state); + if fuzzy.borrow().is_enabled() { + frame.render_stateful_widget(table, packet_block, &mut fuzzy.scroll_state); } else { frame.render_stateful_widget(table, packet_block, &mut self.packets_table_state); } @@ -670,34 +755,33 @@ impl App { .begin_symbol(Some("↑")) .end_symbol(Some("↓")); - let mut scrollbar_state = - if self.fuzzy.is_enabled() && self.fuzzy.packets.len() > window_size { - ScrollbarState::new(self.fuzzy.packets.len()).position({ - if self.manuall_scroll { - if self.fuzzy.packet_end_index == window_size { - 0 - } else { - self.fuzzy.packet_end_index - } + let mut scrollbar_state = if fuzzy.is_enabled() && fuzzy_packets.len() > window_size { + ScrollbarState::new(fuzzy_packets.len()).position({ + if self.manuall_scroll { + if fuzzy.borrow().packet_end_index == window_size { + 0 } else { - self.fuzzy.packets.len() + fuzzy.borrow().packet_end_index } - }) - } else if !self.fuzzy.is_enabled() && self.packets.len() > window_size { - ScrollbarState::new(self.packets.len()).position({ - if self.manuall_scroll { - if self.packet_end_index == window_size { - 0 - } else { - self.packet_end_index - } + } else { + fuzzy.borrow().packets.len() + } + }) + } else if !fuzzy.is_enabled() && app_packets.len() > window_size { + ScrollbarState::new(app_packets.len()).position({ + if self.manuall_scroll { + if self.packet_end_index == window_size { + 0 } else { - self.packets.len() + self.packet_end_index } - }) - } else { - ScrollbarState::default() - }; + } else { + app_packets.len() + } + }) + } else { + ScrollbarState::default() + }; frame.render_stateful_widget( scrollbar, @@ -708,8 +792,8 @@ impl App { &mut scrollbar_state, ); - if self.fuzzy.is_enabled() { - let fuzzy = Paragraph::new(format!("> {}", self.fuzzy.filter.value())) + if fuzzy.borrow().is_enabled() { + let fuzzy = Paragraph::new(format!("> {}", fuzzy.filter.value())) .alignment(Alignment::Left) .style(Style::default().white()) .block( @@ -717,14 +801,14 @@ impl App { .borders(Borders::all()) .title(" Search  ") .title_style({ - if self.fuzzy.is_paused() { + if fuzzy.is_paused() { Style::default().bold().green() } else { Style::default().bold().yellow() } }) .border_style({ - if self.fuzzy.is_paused() { + if fuzzy.is_paused() { Style::default().green() } else { Style::default().yellow() @@ -736,109 +820,125 @@ impl App { } } - pub fn render_stats_mode(&mut self, frame: &mut Frame, stats_block: Rect) { - self.stats - .render(frame, stats_block, &self.interface.selected_interface.name); + pub fn render_stats_mode(&mut self, frame: &mut Frame, block: Rect) { + let stats = self.stats.lock().unwrap(); + let mut bandwidth = self.bandwidth.lock().unwrap(); + + let (bandwidth_block, stats_block) = { + let chunks = Layout::default() + .direction(Direction::Vertical) + .constraints([Constraint::Percentage(50), Constraint::Percentage(50)].as_ref()) + .margin(2) + .split(block); + (chunks[0], chunks[1]) + }; + + frame.render_widget( + Block::default() + .title({ + Line::from(vec![ + Span::from(" Packet ").fg(Color::DarkGray), + Span::styled( + " Stats ", + Style::default().bg(Color::Green).fg(Color::White).bold(), + ), + ]) + }) + .title_alignment(Alignment::Left) + .padding(Padding::top(1)) + .borders(Borders::ALL) + .style(Style::default()) + .border_type(BorderType::default()) + .border_style(Style::default().green()), + block.inner(Margin { + horizontal: 1, + vertical: 0, + }), + ); + + stats.render(frame, stats_block); + + if bandwidth.is_some() { + bandwidth.as_mut().unwrap().render( + frame, + bandwidth_block, + &self.interface.selected_interface.name.clone(), + ); + } } - pub fn process(&mut self, app_packet: AppPacket) { - if self.packets.len() == self.packets.capacity() { - self.packets.reserve(1024 * 1024); + pub fn process( + packets: Arc>>, + stats: Arc>, + transport_filters: Arc>>, + network_filters: Arc>>, + link_filters: Arc>>, + app_packet: AppPacket, + ) { + let mut packets = packets.lock().unwrap(); + + let applied_transport_protocols = transport_filters.lock().unwrap(); + let applied_network_protocols = network_filters.lock().unwrap(); + let applied_link_protocols = link_filters.lock().unwrap(); + + if packets.len() == packets.capacity() { + packets.reserve(1024 * 1024); } match app_packet { AppPacket::Arp(_) => { - if self - .link_filter - .applied_protocols - .contains(&crate::filters::link::LinkProtocol::Arp) - { - self.packets.push(app_packet); + if applied_link_protocols.contains(&LinkProtocol::Arp) { + packets.push(app_packet); } } AppPacket::Ip(packet) => match packet { IpPacket::Tcp(p) => { - if self - .transport_filter - .applied_protocols - .contains(&TransportProtocol::TCP) - { + if applied_transport_protocols.contains(&TransportProtocol::TCP) { match p.src_ip { core::net::IpAddr::V6(_) => { - if self - .network_filter - .applied_protocols - .contains(&NetworkProtocol::Ipv6) - { - self.packets.push(app_packet); + if applied_network_protocols.contains(&NetworkProtocol::Ipv6) { + packets.push(app_packet); } } core::net::IpAddr::V4(_) => { - if self - .network_filter - .applied_protocols - .contains(&NetworkProtocol::Ipv4) - { - self.packets.push(app_packet); + if applied_network_protocols.contains(&NetworkProtocol::Ipv4) { + packets.push(app_packet); } } } } } IpPacket::Udp(p) => { - if self - .transport_filter - .applied_protocols - .contains(&TransportProtocol::UDP) - { + if applied_transport_protocols.contains(&TransportProtocol::UDP) { match p.src_ip { core::net::IpAddr::V6(_) => { - if self - .network_filter - .applied_protocols - .contains(&NetworkProtocol::Ipv6) - { - self.packets.push(app_packet); + if applied_network_protocols.contains(&NetworkProtocol::Ipv6) { + packets.push(app_packet); } } core::net::IpAddr::V4(_) => { - if self - .network_filter - .applied_protocols - .contains(&NetworkProtocol::Ipv4) - { - self.packets.push(app_packet); + if applied_network_protocols.contains(&NetworkProtocol::Ipv4) { + packets.push(app_packet); } } } } } IpPacket::Icmp(_) => { - if self - .network_filter - .applied_protocols - .contains(&NetworkProtocol::Icmp) - { - self.packets.push(app_packet); + if applied_network_protocols.contains(&NetworkProtocol::Icmp) { + packets.push(app_packet); } } }, } - self.stats.refresh(&app_packet); + let mut stats = stats.lock().unwrap(); + + stats.refresh(&app_packet); } pub fn tick(&mut self) { self.notifications.iter_mut().for_each(|n| n.ttl -= 1); self.notifications.retain(|n| n.ttl > 0); - - // Even when not updating the pattern, new packets are still comming - if self.fuzzy.is_enabled() { - self.fuzzy.find(&self.packets); - } - - if let Some(bandwidth) = &mut self.stats.bandwidth { - bandwidth.refresh().ok(); - } } pub fn quit(&mut self) { diff --git a/oryx-tui/src/ebpf.rs b/oryx-tui/src/ebpf.rs index 08cc3ff..67d5d8d 100644 --- a/oryx-tui/src/ebpf.rs +++ b/oryx-tui/src/ebpf.rs @@ -59,10 +59,15 @@ impl Source for RingBuffer<'_> { } impl Ebpf { - pub fn load_ingress(iface: String, sender: kanal::Sender, terminate: Arc) { + pub fn load_ingress( + iface: String, + notification_sender: kanal::Sender, + data_sender: kanal::Sender<[u8; RawPacket::LEN]>, + terminate: Arc, + ) { thread::spawn({ let iface = iface.to_owned(); - let sender = sender.clone(); + let notification_sender = notification_sender.clone(); move || { let rlim = libc::rlimit { @@ -81,7 +86,7 @@ impl Ebpf { Notification::send( format!("Failed to load the ingress eBPF bytecode\n {}", e), NotificationLevel::Error, - sender, + notification_sender, ) .unwrap(); return; @@ -97,7 +102,7 @@ impl Ebpf { Notification::send( format!("Failed to load the ingress eBPF bytecode\n {}", e), NotificationLevel::Error, - sender, + notification_sender, ) .unwrap(); return; @@ -116,7 +121,7 @@ impl Ebpf { e ), NotificationLevel::Error, - sender, + notification_sender, ) .unwrap(); return; @@ -129,7 +134,7 @@ impl Ebpf { e ), NotificationLevel::Error, - sender, + notification_sender, ) .unwrap(); return; @@ -168,7 +173,7 @@ impl Ebpf { } let packet: [u8; RawPacket::LEN] = item.to_owned().try_into().unwrap(); - sender.send(Event::Packet(packet)).ok(); + data_sender.send(packet).ok(); } } } @@ -181,10 +186,15 @@ impl Ebpf { }); } - pub fn load_egress(iface: String, sender: kanal::Sender, terminate: Arc) { + pub fn load_egress( + iface: String, + notification_sender: kanal::Sender, + data_sender: kanal::Sender<[u8; RawPacket::LEN]>, + terminate: Arc, + ) { thread::spawn({ let iface = iface.to_owned(); - let sender = sender.clone(); + let notification_sender = notification_sender.clone(); move || { let rlim = libc::rlimit { @@ -203,7 +213,7 @@ impl Ebpf { Notification::send( format!("Fail to load the egress eBPF bytecode\n {}", e), NotificationLevel::Error, - sender, + notification_sender, ) .unwrap(); return; @@ -219,7 +229,7 @@ impl Ebpf { Notification::send( format!("Failed to load the egress eBPF bytecode\n {}", e), NotificationLevel::Error, - sender, + notification_sender, ) .unwrap(); return; @@ -234,7 +244,7 @@ impl Ebpf { Notification::send( format!("Fail to load the egress eBPF program to the kernel\n{}", e), NotificationLevel::Error, - sender, + notification_sender, ) .unwrap(); return; @@ -247,7 +257,7 @@ impl Ebpf { e ), NotificationLevel::Error, - sender, + notification_sender, ) .unwrap(); return; @@ -286,7 +296,7 @@ impl Ebpf { } let packet: [u8; RawPacket::LEN] = item.to_owned().try_into().unwrap(); - sender.send(Event::Packet(packet)).ok(); + data_sender.send(packet).ok(); } } } diff --git a/oryx-tui/src/event.rs b/oryx-tui/src/event.rs index c07a210..b03c89c 100644 --- a/oryx-tui/src/event.rs +++ b/oryx-tui/src/event.rs @@ -12,7 +12,6 @@ pub enum Event { Key(KeyEvent), Mouse(MouseEvent), Resize(u16, u16), - Packet([u8; 72]), Notification(Notification), Reset, } diff --git a/oryx-tui/src/filters/fuzzy.rs b/oryx-tui/src/filters/fuzzy.rs index ed380b5..6d315fa 100644 --- a/oryx-tui/src/filters/fuzzy.rs +++ b/oryx-tui/src/filters/fuzzy.rs @@ -7,7 +7,7 @@ use tui_input::Input; use oryx_common::AppPacket; -#[derive(Debug, Default)] +#[derive(Debug, Clone, Default)] pub struct Fuzzy { enabled: bool, paused: bool, @@ -26,6 +26,16 @@ impl Fuzzy { .collect::>(); } + pub fn append(&mut self, packets: &[AppPacket]) { + self.packets.append( + &mut packets + .iter() + .copied() + .filter(|p| p.to_string().contains(self.filter.value())) + .collect::>(), + ); + } + pub fn enable(&mut self) { self.enabled = true; } diff --git a/oryx-tui/src/filters/link.rs b/oryx-tui/src/filters/link.rs index 0750200..7985af6 100644 --- a/oryx-tui/src/filters/link.rs +++ b/oryx-tui/src/filters/link.rs @@ -1,4 +1,7 @@ -use std::fmt::Display; +use std::{ + fmt::Display, + sync::{Arc, Mutex}, +}; use ratatui::{ layout::{Alignment, Constraint, Direction, Flex, Layout, Rect}, @@ -15,7 +18,7 @@ pub const NB_LINK_PROTOCOL: u16 = 1; pub struct LinkFilter { pub state: TableState, pub selected_protocols: Vec, - pub applied_protocols: Vec, + pub applied_protocols: Arc>>, } #[derive(Debug, Copy, Clone, PartialEq)] @@ -34,14 +37,15 @@ impl Default for LinkFilter { Self { state: TableState::default(), selected_protocols: vec![LinkProtocol::Arp], - applied_protocols: Vec::new(), + applied_protocols: Arc::new(Mutex::new(Vec::new())), } } } impl LinkFilter { pub fn apply(&mut self) { - self.applied_protocols = self.selected_protocols.clone(); + let mut applied_protocols = self.applied_protocols.lock().unwrap(); + *applied_protocols = self.selected_protocols.clone(); self.selected_protocols.clear(); } diff --git a/oryx-tui/src/filters/network.rs b/oryx-tui/src/filters/network.rs index 8f80916..0876afb 100644 --- a/oryx-tui/src/filters/network.rs +++ b/oryx-tui/src/filters/network.rs @@ -1,4 +1,7 @@ -use std::fmt::Display; +use std::{ + fmt::Display, + sync::{Arc, Mutex}, +}; use ratatui::{ layout::{Alignment, Constraint, Direction, Flex, Layout, Rect}, @@ -15,7 +18,7 @@ pub const NB_NETWORK_PROTOCOL: u16 = 3; pub struct NetworkFilter { pub state: TableState, pub selected_protocols: Vec, - pub applied_protocols: Vec, + pub applied_protocols: Arc>>, } impl Default for NetworkFilter { @@ -27,7 +30,7 @@ impl Default for NetworkFilter { NetworkProtocol::Ipv6, NetworkProtocol::Icmp, ], - applied_protocols: Vec::new(), + applied_protocols: Arc::new(Mutex::new(Vec::new())), } } } @@ -51,7 +54,8 @@ impl Display for NetworkProtocol { impl NetworkFilter { pub fn apply(&mut self) { - self.applied_protocols = self.selected_protocols.clone(); + let mut applied_protocols = self.applied_protocols.lock().unwrap(); + *applied_protocols = self.selected_protocols.clone(); self.selected_protocols.clear(); } pub fn render(&mut self, frame: &mut Frame, block: Rect, focused_block: &FocusedBlock) { diff --git a/oryx-tui/src/filters/transport.rs b/oryx-tui/src/filters/transport.rs index 1c08694..7a6d8cb 100644 --- a/oryx-tui/src/filters/transport.rs +++ b/oryx-tui/src/filters/transport.rs @@ -1,4 +1,7 @@ -use std::fmt::Display; +use std::{ + fmt::Display, + sync::{Arc, Mutex}, +}; use ratatui::{ layout::{Alignment, Constraint, Direction, Flex, Layout, Rect}, @@ -15,7 +18,7 @@ pub const NB_TRANSPORT_PROTOCOL: u16 = 2; pub struct TransportFilter { pub state: TableState, pub selected_protocols: Vec, - pub applied_protocols: Vec, + pub applied_protocols: Arc>>, } #[derive(Debug, Copy, Clone, PartialEq)] @@ -38,14 +41,15 @@ impl Default for TransportFilter { Self { state: TableState::default(), selected_protocols: vec![TransportProtocol::TCP, TransportProtocol::UDP], - applied_protocols: Vec::new(), + applied_protocols: Arc::new(Mutex::new(Vec::new())), } } } impl TransportFilter { pub fn apply(&mut self) { - self.applied_protocols = self.selected_protocols.clone(); + let mut applied_protocols = self.applied_protocols.lock().unwrap(); + *applied_protocols = self.selected_protocols.clone(); self.selected_protocols.clear(); } diff --git a/oryx-tui/src/handler.rs b/oryx-tui/src/handler.rs index 9af8709..d760a90 100644 --- a/oryx-tui/src/handler.rs +++ b/oryx-tui/src/handler.rs @@ -27,7 +27,10 @@ pub fn handle_key_events( app: &mut App, sender: kanal::Sender, ) -> AppResult<()> { - if app.fuzzy.is_enabled() { + let fuzzy = app.fuzzy.clone(); + let mut fuzzy = fuzzy.lock().unwrap(); + + if fuzzy.is_enabled() { match key_event.code { KeyCode::Esc => { if app.focused_block == FocusedBlock::Help { @@ -40,14 +43,14 @@ pub fn handle_key_events( return Ok(()); } - if app.fuzzy.is_paused() { + if fuzzy.is_paused() { if app.manuall_scroll { app.manuall_scroll = false; } else { - app.fuzzy.disable(); + fuzzy.disable(); } } else { - app.fuzzy.pause(); + fuzzy.pause(); } } @@ -141,15 +144,15 @@ pub fn handle_key_events( if app.focused_block == FocusedBlock::Help { return Ok(()); } - if !app.fuzzy.is_paused() && !app.update_filters { - app.fuzzy + if !fuzzy.is_paused() && !app.update_filters { + fuzzy .filter .handle_event(&crossterm::event::Event::Key(key_event)); } else { match key_event.code { KeyCode::Char('/') => { if !app.update_filters { - app.fuzzy.unpause(); + fuzzy.unpause(); } } @@ -162,14 +165,19 @@ pub fn handle_key_events( app.update_filters = true; app.focused_block = FocusedBlock::TransportFilter; + let applied_network_protocols = + app.network_filter.applied_protocols.lock().unwrap(); app.network_filter.selected_protocols = - app.network_filter.applied_protocols.clone(); + applied_network_protocols.clone(); + let applied_transport_protocols = + app.transport_filter.applied_protocols.lock().unwrap(); app.transport_filter.selected_protocols = - app.transport_filter.applied_protocols.clone(); + applied_transport_protocols.clone(); - app.link_filter.selected_protocols = - app.link_filter.applied_protocols.clone(); + let applied_link_protocols = + app.link_filter.applied_protocols.lock().unwrap(); + app.link_filter.selected_protocols = applied_link_protocols.clone(); app.traffic_direction_filter.selected_direction = app.traffic_direction_filter.applied_direction.clone(); @@ -183,26 +191,26 @@ pub fn handle_key_events( if !app.manuall_scroll { app.manuall_scroll = true; // Record the last position. Usefull for selecting the packets to display - app.fuzzy.packet_end_index = app.fuzzy.packets.len(); + fuzzy.packet_end_index = fuzzy.packets.len(); } - let i = match app.fuzzy.scroll_state.selected() { + let i = match fuzzy.scroll_state.selected() { Some(i) => { if i < app.packet_window_size - 1 { i + 1 } else if i == app.packet_window_size - 1 - && app.fuzzy.packets.len() > app.fuzzy.packet_end_index + && fuzzy.packets.len() > fuzzy.packet_end_index { // shit the window by one - app.fuzzy.packet_end_index += 1; + fuzzy.packet_end_index += 1; i + 1 } else { i } } - None => app.fuzzy.packets.len(), + None => fuzzy.packets.len(), }; - app.fuzzy.scroll_state.select(Some(i)); + fuzzy.scroll_state.select(Some(i)); } else { match &app.focused_block { FocusedBlock::NetworkFilter => { @@ -263,26 +271,26 @@ pub fn handle_key_events( if !app.manuall_scroll { app.manuall_scroll = true; // Record the last position. Usefull for selecting the packets to display - app.fuzzy.packet_end_index = app.fuzzy.packets.len(); + fuzzy.packet_end_index = fuzzy.packets.len(); } - let i = match app.fuzzy.scroll_state.selected() { + let i = match fuzzy.scroll_state.selected() { Some(i) => { if i > 1 { i - 1 } else if i == 0 - && app.fuzzy.packet_end_index > app.packet_window_size + && fuzzy.packet_end_index > app.packet_window_size { // shit the window by one - app.fuzzy.packet_end_index -= 1; + fuzzy.packet_end_index -= 1; 0 } else { 0 } } - None => app.fuzzy.packets.len(), + None => fuzzy.packets.len(), }; - app.fuzzy.scroll_state.select(Some(i)); + fuzzy.scroll_state.select(Some(i)); } else { match &app.focused_block { FocusedBlock::NetworkFilter => { @@ -402,13 +410,16 @@ pub fn handle_key_events( app.focused_block = FocusedBlock::TransportFilter; - app.network_filter.selected_protocols = - app.network_filter.applied_protocols.clone(); + let applied_network_protocols = + app.network_filter.applied_protocols.lock().unwrap(); + app.network_filter.selected_protocols = applied_network_protocols.clone(); - app.transport_filter.selected_protocols = - app.transport_filter.applied_protocols.clone(); + let applied_transport_protocols = + app.transport_filter.applied_protocols.lock().unwrap(); + app.transport_filter.selected_protocols = applied_transport_protocols.clone(); - app.link_filter.selected_protocols = app.link_filter.applied_protocols.clone(); + let applied_link_protocols = app.link_filter.applied_protocols.lock().unwrap(); + app.link_filter.selected_protocols = applied_link_protocols.clone(); app.traffic_direction_filter.selected_direction = app.traffic_direction_filter.applied_direction.clone(); @@ -423,14 +434,15 @@ pub fn handle_key_events( } if app.start_sniffing { - if app.packets.is_empty() { + let app_packets = app.packets.lock().unwrap(); + if app_packets.is_empty() { Notification::send( "There is no packets".to_string(), NotificationLevel::Info, sender, )?; } else { - match export(&app.packets) { + match export(&app_packets) { Ok(_) => { Notification::send( "Packets exported to ~/oryx/capture file".to_string(), @@ -455,8 +467,8 @@ pub fn handle_key_events( return Ok(()); } if app.start_sniffing { - app.fuzzy.enable(); - app.fuzzy.unpause(); + fuzzy.enable(); + fuzzy.unpause(); } } @@ -493,6 +505,7 @@ pub fn handle_key_events( Ebpf::load_ingress( iface.clone(), sender.clone(), + app.data_channel_sender.clone(), app.traffic_direction_filter.terminate_ingress.clone(), ); } @@ -505,6 +518,7 @@ pub fn handle_key_events( Ebpf::load_egress( iface, sender.clone(), + app.data_channel_sender.clone(), app.traffic_direction_filter.terminate_egress.clone(), ); } @@ -541,6 +555,7 @@ pub fn handle_key_events( Ebpf::load_egress( iface, sender.clone(), + app.data_channel_sender.clone(), app.traffic_direction_filter.terminate_egress.clone(), ); } @@ -574,6 +589,7 @@ pub fn handle_key_events( Ebpf::load_ingress( iface, sender.clone(), + app.data_channel_sender.clone(), app.traffic_direction_filter.terminate_ingress.clone(), ); } @@ -860,19 +876,20 @@ pub fn handle_key_events( if let FocusedBlock::Help = app.focused_block { return Ok(()); } + let app_packets = app.packets.lock().unwrap(); // Sniff mode if app.start_sniffing && !app.update_filters { if !app.manuall_scroll { app.manuall_scroll = true; // Record the last position. Usefull for selecting the packets to display - app.packet_end_index = app.packets.len(); + app.packet_end_index = app_packets.len(); } let i = match app.packets_table_state.selected() { Some(i) => { if i < app.packet_window_size - 1 { i + 1 } else if i == app.packet_window_size - 1 - && app.packets.len() > app.packet_end_index + && app_packets.len() > app.packet_end_index { // shit the window by one app.packet_end_index += 1; @@ -881,7 +898,7 @@ pub fn handle_key_events( i } } - None => app.packets.len(), + None => app_packets.len(), }; app.packets_table_state.select(Some(i)); @@ -960,6 +977,7 @@ pub fn handle_key_events( } KeyCode::Char('k') | KeyCode::Up => { + let app_packets = app.packets.lock().unwrap(); if let FocusedBlock::Help = app.focused_block { return Ok(()); } @@ -967,7 +985,7 @@ pub fn handle_key_events( if !app.manuall_scroll { app.manuall_scroll = true; // Record the last position. Usefull for selecting the packets to display - app.packet_end_index = app.packets.len(); + app.packet_end_index = app_packets.len(); } let i = match app.packets_table_state.selected() { Some(i) => { diff --git a/oryx-tui/src/main.rs b/oryx-tui/src/main.rs index 7206c8c..623240a 100644 --- a/oryx-tui/src/main.rs +++ b/oryx-tui/src/main.rs @@ -4,8 +4,7 @@ static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; use std::io; use clap::{crate_description, crate_version, Command}; -use oryx_common::AppPacket; -use oryx_tui::app::{App, AppResult}; +use oryx_tui::app::{App, AppResult, TICK_RATE}; use oryx_tui::event::{Event, EventHandler}; use oryx_tui::handler::handle_key_events; use oryx_tui::tui::Tui; @@ -27,7 +26,7 @@ fn main() -> AppResult<()> { let backend = CrosstermBackend::new(io::stdout()); let terminal = Terminal::new(backend)?; - let events = EventHandler::new(1_000); + let events = EventHandler::new(TICK_RATE); let mut tui = Tui::new(terminal, events); tui.init()?; @@ -38,9 +37,6 @@ fn main() -> AppResult<()> { Event::Key(key_event) => { handle_key_events(key_event, &mut app, tui.events.sender.clone())? } - Event::Packet(packet) => { - app.process(AppPacket::from(packet)); - } Event::Notification(notification) => { app.notifications.push(notification); } diff --git a/oryx-tui/src/notification.rs b/oryx-tui/src/notification.rs index 7d77bec..034d999 100644 --- a/oryx-tui/src/notification.rs +++ b/oryx-tui/src/notification.rs @@ -12,7 +12,7 @@ use crate::{app::AppResult, event::Event}; pub struct Notification { pub message: String, pub level: NotificationLevel, - pub ttl: u8, + pub ttl: u16, } #[derive(Debug, Clone)] @@ -69,7 +69,7 @@ impl Notification { let notif = Notification { message: message.to_string(), level, - ttl: 3, + ttl: 500, }; sender.send(Event::Notification(notif))?; diff --git a/oryx-tui/src/stats.rs b/oryx-tui/src/stats.rs index f1768e7..02e9252 100644 --- a/oryx-tui/src/stats.rs +++ b/oryx-tui/src/stats.rs @@ -5,16 +5,14 @@ use std::net::{IpAddr, Ipv4Addr}; use oryx_common::ip::IpPacket; -use ratatui::layout::{Alignment, Constraint, Direction, Flex, Layout, Margin, Rect}; -use ratatui::style::{Color, Style, Stylize}; -use ratatui::text::{Line, Span}; +use ratatui::layout::{Alignment, Constraint, Direction, Flex, Layout, Rect}; +use ratatui::style::{Color, Style}; +use ratatui::text::Line; use ratatui::{ - widgets::{Bar, BarChart, BarGroup, Block, BorderType, Borders, Padding}, + widgets::{Bar, BarChart, BarGroup, Block, Padding}, Frame, }; -use crate::bandwidth::Bandwidth; - #[derive(Debug)] pub struct Stats { pub total: usize, @@ -23,7 +21,6 @@ pub struct Stats { pub transport: TransportStats, pub link: LinkStats, pub addresses: HashMap, usize)>, - pub bandwidth: Option, } impl Default for Stats { @@ -41,7 +38,6 @@ impl Stats { transport: TransportStats::default(), link: LinkStats::default(), addresses: HashMap::with_capacity(1024), - bandwidth: Bandwidth::new().ok(), } } pub fn get_top_10(&self) -> Vec<(&Ipv4Addr, &(Option, usize))> { @@ -159,16 +155,7 @@ impl Stats { self.total += 1; } - pub fn render(&self, frame: &mut Frame, stats_block: Rect, network_interface: &str) { - let (graph_block, barchart_block) = { - let chunks = Layout::default() - .direction(Direction::Vertical) - .constraints([Constraint::Percentage(50), Constraint::Percentage(50)].as_ref()) - .margin(2) - .split(stats_block); - (chunks[0], chunks[1]) - }; - + pub fn render(&self, frame: &mut Frame, stats_block: Rect) { let (address_block, network_block, transport_block, link_block) = { let chunks = Layout::default() .direction(Direction::Horizontal) @@ -183,33 +170,10 @@ impl Stats { ) .margin(1) .flex(Flex::SpaceBetween) - .split(barchart_block); + .split(stats_block); (chunks[0], chunks[1], chunks[2], chunks[3]) }; - frame.render_widget( - Block::default() - .title({ - Line::from(vec![ - Span::from(" Packet ").fg(Color::DarkGray), - Span::styled( - " Stats ", - Style::default().bg(Color::Green).fg(Color::White).bold(), - ), - ]) - }) - .title_alignment(Alignment::Left) - .padding(Padding::top(1)) - .borders(Borders::ALL) - .style(Style::default()) - .border_type(BorderType::default()) - .border_style(Style::default().green()), - stats_block.inner(Margin { - horizontal: 1, - vertical: 0, - }), - ); - let link_chart = BarChart::default() .bar_width(3) .bar_gap(1) @@ -354,10 +318,6 @@ impl Stats { frame.render_widget(transport_chart, transport_block); frame.render_widget(network_chart, network_block); frame.render_widget(link_chart, link_block); - - if let Some(bandwidth) = &self.bandwidth { - bandwidth.render(frame, graph_block, network_interface) - } } }