From 19eb710eef47def476c2402072f77bd1b2e40c5f Mon Sep 17 00:00:00 2001 From: Badr Date: Tue, 17 Sep 2024 23:00:55 +0200 Subject: [PATCH] move filters to ebpf --- Cargo.lock | 5 +- oryx-common/src/lib.rs | 1 + oryx-common/src/protocols.rs | 89 ++++++++++++++++++++++ oryx-ebpf/src/main.rs | 67 ++++++++++++++++- oryx-tui/Cargo.toml | 1 + oryx-tui/src/app.rs | 118 ++++++++++-------------------- oryx-tui/src/ebpf.rs | 58 ++++++++++++++- oryx-tui/src/event.rs | 4 +- oryx-tui/src/filters/link.rs | 33 ++------- oryx-tui/src/filters/network.rs | 36 ++------- oryx-tui/src/filters/transport.rs | 37 ++-------- oryx-tui/src/handler.rs | 101 +++++++++++++++++++------ oryx-tui/src/tui.rs | 4 +- 13 files changed, 351 insertions(+), 203 deletions(-) create mode 100644 oryx-common/src/protocols.rs diff --git a/Cargo.lock b/Cargo.lock index 78affc2..a4f7dc1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -147,9 +147,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.1.19" +version = "1.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2d74707dde2ba56f86ae90effb3b43ddd369504387e718014de010cec7959800" +checksum = "45bcde016d64c21da4be18b655631e5ab6d3107607e71a73a9f53eb48aae23fb" dependencies = [ "shlex", ] @@ -598,6 +598,7 @@ version = "0.2.0" dependencies = [ "aya", "clap", + "crossterm", "dirs", "dns-lookup", "itertools", diff --git a/oryx-common/src/lib.rs b/oryx-common/src/lib.rs index 05ec034..fda3a56 100644 --- a/oryx-common/src/lib.rs +++ b/oryx-common/src/lib.rs @@ -12,6 +12,7 @@ use network_types::{arp::ArpHdr, ip::IpHdr}; pub mod arp; pub mod ip; +pub mod protocols; #[repr(C)] pub enum RawPacket { diff --git a/oryx-common/src/protocols.rs b/oryx-common/src/protocols.rs new file mode 100644 index 0000000..037e330 --- /dev/null +++ b/oryx-common/src/protocols.rs @@ -0,0 +1,89 @@ +use core::fmt::Display; + +#[derive(Debug, Copy, Clone)] +#[repr(C, u8)] +pub enum Protocol { + Transport(TransportProtocol) = 0, + Network(NetworkProtocol) = 1, + Link(LinkProtocol) = 2, +} + +// Transport Protocols + +pub const NB_TRANSPORT_PROTOCOL: u16 = 2; + +#[derive(Debug, Copy, Clone, PartialEq)] +#[repr(C)] +pub enum TransportProtocol { + TCP = 0, + UDP = 1, +} + +impl TransportProtocol { + pub fn all() -> [TransportProtocol; 2] { + [TransportProtocol::TCP, TransportProtocol::UDP] + } +} + +impl Display for TransportProtocol { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + match self { + TransportProtocol::TCP => write!(f, "Tcp"), + TransportProtocol::UDP => write!(f, "Udp"), + } + } +} + +// Network Protocols + +pub const NB_NETWORK_PROTOCOL: u16 = 3; + +#[derive(Debug, Copy, Clone, PartialEq)] +#[repr(C)] +pub enum NetworkProtocol { + Ipv4 = 0, + Ipv6 = 1, + Icmp = 2, +} + +impl NetworkProtocol { + pub fn all() -> [NetworkProtocol; 3] { + [ + NetworkProtocol::Ipv4, + NetworkProtocol::Ipv6, + NetworkProtocol::Icmp, + ] + } +} + +impl Display for NetworkProtocol { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + match self { + NetworkProtocol::Ipv4 => write!(f, "Ipv4"), + NetworkProtocol::Ipv6 => write!(f, "Ipv6"), + NetworkProtocol::Icmp => write!(f, "Icmp"), + } + } +} + +// Link Protocols + +pub const NB_LINK_PROTOCOL: u16 = 1; + +#[derive(Debug, Copy, Clone, PartialEq)] +#[repr(C)] +pub enum LinkProtocol { + Arp = 0, +} + +impl LinkProtocol { + pub fn all() -> [LinkProtocol; 1] { + [LinkProtocol::Arp] + } +} + +impl Display for LinkProtocol { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + write!(f, "Arp") + } +} diff --git a/oryx-ebpf/src/main.rs b/oryx-ebpf/src/main.rs index bbafe9a..7cea341 100644 --- a/oryx-ebpf/src/main.rs +++ b/oryx-ebpf/src/main.rs @@ -4,7 +4,7 @@ use aya_ebpf::{ bindings::TC_ACT_PIPE, macros::{classifier, map}, - maps::RingBuf, + maps::{Array, RingBuf}, programs::TcContext, }; use core::mem; @@ -16,11 +16,24 @@ use network_types::{ tcp::TcpHdr, udp::UdpHdr, }; -use oryx_common::{ip::ProtoHdr, RawPacket}; +use oryx_common::{ + ip::ProtoHdr, + protocols::{LinkProtocol, NetworkProtocol, Protocol, TransportProtocol}, + RawPacket, +}; #[map] static DATA: RingBuf = RingBuf::with_byte_size(4096 * RawPacket::LEN as u32, 0); +#[map] +static NETWORK_FILTERS: Array = Array::with_max_entries(8, 0); + +#[map] +static TRANSPORT_FILTERS: Array = Array::with_max_entries(8, 0); + +#[map] +static LINK_FILTERS: Array = Array::with_max_entries(8, 0); + #[classifier] pub fn oryx(ctx: TcContext) -> i32 { match process(ctx) { @@ -48,15 +61,44 @@ fn ptr_at(ctx: &TcContext, offset: usize) -> Result<*const T, ()> { Ok((start + offset) as *const T) } + +#[inline] +fn filter_packet(protocol: Protocol) -> bool { + match protocol { + Protocol::Network(p) => { + if let Some(v) = NETWORK_FILTERS.get(p as u32) { + return *v == 1; + } + } + Protocol::Transport(p) => { + if let Some(v) = TRANSPORT_FILTERS.get(p as u32) { + return *v == 1; + } + } + Protocol::Link(p) => { + if let Some(v) = LINK_FILTERS.get(p as u32) { + return *v == 1; + } + } + } + false +} + #[inline] fn process(ctx: TcContext) -> Result { let ethhdr: EthHdr = ctx.load(0).map_err(|_| ())?; match ethhdr.ether_type { EtherType::Ipv4 => { + if filter_packet(Protocol::Network(NetworkProtocol::Ipv4)) { + return Ok(TC_ACT_PIPE); + } let header: Ipv4Hdr = ctx.load(EthHdr::LEN).map_err(|_| ())?; match header.proto { IpProto::Tcp => { + if filter_packet(Protocol::Transport(TransportProtocol::TCP)) { + return Ok(TC_ACT_PIPE); + } let tcphdr: *const TcpHdr = ptr_at(&ctx, EthHdr::LEN + Ipv4Hdr::LEN)?; submit(RawPacket::Ip( IpHdr::V4(header), @@ -64,6 +106,9 @@ fn process(ctx: TcContext) -> Result { )); } IpProto::Udp => { + if filter_packet(Protocol::Transport(TransportProtocol::UDP)) { + return Ok(TC_ACT_PIPE); + } let udphdr: *const UdpHdr = ptr_at(&ctx, EthHdr::LEN + Ipv4Hdr::LEN)?; submit(RawPacket::Ip( IpHdr::V4(header), @@ -71,6 +116,9 @@ fn process(ctx: TcContext) -> Result { )); } IpProto::Icmp => { + if filter_packet(Protocol::Network(NetworkProtocol::Icmp)) { + return Ok(TC_ACT_PIPE); + } let icmphdr: *const IcmpHdr = ptr_at(&ctx, EthHdr::LEN + Ipv4Hdr::LEN)?; submit(RawPacket::Ip( IpHdr::V4(header), @@ -81,9 +129,15 @@ fn process(ctx: TcContext) -> Result { } } EtherType::Ipv6 => { + if filter_packet(Protocol::Network(NetworkProtocol::Ipv6)) { + return Ok(TC_ACT_PIPE); + } let header: Ipv6Hdr = ctx.load(EthHdr::LEN).map_err(|_| ())?; match header.next_hdr { IpProto::Tcp => { + if filter_packet(Protocol::Transport(TransportProtocol::TCP)) { + return Ok(TC_ACT_PIPE); + } let tcphdr: *const TcpHdr = ptr_at(&ctx, EthHdr::LEN + Ipv6Hdr::LEN)?; submit(RawPacket::Ip( IpHdr::V6(header), @@ -91,6 +145,9 @@ fn process(ctx: TcContext) -> Result { )); } IpProto::Udp => { + if filter_packet(Protocol::Transport(TransportProtocol::UDP)) { + return Ok(TC_ACT_PIPE); + } let udphdr: *const UdpHdr = ptr_at(&ctx, EthHdr::LEN + Ipv6Hdr::LEN)?; submit(RawPacket::Ip( IpHdr::V6(header), @@ -98,6 +155,9 @@ fn process(ctx: TcContext) -> Result { )); } IpProto::Icmp => { + if filter_packet(Protocol::Network(NetworkProtocol::Icmp)) { + return Ok(TC_ACT_PIPE); + } let icmphdr: *const IcmpHdr = ptr_at(&ctx, EthHdr::LEN + Ipv6Hdr::LEN)?; submit(RawPacket::Ip( IpHdr::V6(header), @@ -108,6 +168,9 @@ fn process(ctx: TcContext) -> Result { } } EtherType::Arp => { + if filter_packet(Protocol::Link(LinkProtocol::Arp)) { + return Ok(TC_ACT_PIPE); + } let header: ArpHdr = ctx.load(EthHdr::LEN).map_err(|_| ())?; submit(RawPacket::Arp(header)); } diff --git a/oryx-tui/Cargo.toml b/oryx-tui/Cargo.toml index 6871d6d..9b716a5 100644 --- a/oryx-tui/Cargo.toml +++ b/oryx-tui/Cargo.toml @@ -9,6 +9,7 @@ homepage.workspace = true edition.workspace = true [dependencies] +crossterm = { version = "0.28", default-features = false } ratatui = "0.28" tui-big-text = "0.6" tui-input = "0.10" diff --git a/oryx-tui/src/app.rs b/oryx-tui/src/app.rs index 52863bd..982a8e7 100644 --- a/oryx-tui/src/app.rs +++ b/oryx-tui/src/app.rs @@ -1,4 +1,7 @@ use oryx_common::ip::IpPacket; +use oryx_common::protocols::{ + Protocol, NB_LINK_PROTOCOL, NB_NETWORK_PROTOCOL, NB_TRANSPORT_PROTOCOL, +}; use oryx_common::{AppPacket, RawPacket}; use ratatui::layout::{Alignment, Constraint, Direction, Flex, Layout, Margin, Rect}; use ratatui::style::{Color, Style, Stylize}; @@ -21,9 +24,9 @@ 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, LinkProtocol, NB_LINK_PROTOCOL}; -use crate::filters::network::{NetworkFilter, NetworkProtocol, NB_NETWORK_PROTOCOL}; -use crate::filters::transport::{TransportFilter, TransportProtocol, NB_TRANSPORT_PROTOCOL}; +use crate::filters::link::LinkFilter; +use crate::filters::network::NetworkFilter; +use crate::filters::transport::TransportFilter; use crate::help::Help; use crate::interface::Interface; use crate::notification::Notification; @@ -57,6 +60,21 @@ pub struct DataEventHandler { pub handler: thread::JoinHandle<()>, } +// let x = Protocol::NetworkProtocol(Tcp) + +#[derive(Debug, Clone)] +pub struct FilterChannel { + pub sender: kanal::Sender<(Protocol, bool)>, + pub receiver: kanal::Receiver<(Protocol, bool)>, +} + +impl FilterChannel { + pub fn new() -> Self { + let (sender, receiver) = kanal::unbounded(); + Self { sender, receiver } + } +} + #[derive(Debug)] pub struct App { pub running: bool, @@ -69,6 +87,8 @@ pub struct App { pub transport_filter: TransportFilter, pub link_filter: LinkFilter, pub traffic_direction_filter: TrafficDirectionFilter, + pub ingress_filter_channel: FilterChannel, + pub egress_filter_channel: FilterChannel, pub start_sniffing: bool, pub packets: Arc>>, pub packets_table_state: TableState, @@ -96,29 +116,19 @@ impl App { 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 network_filter = NetworkFilter::new(); + let transport_filter = TransportFilter::new(); + let link_filter = LinkFilter::new(); 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), - ); + App::process(packets.clone(), stats.clone(), AppPacket::from(raw_packet)); } } }); @@ -174,6 +184,8 @@ impl App { transport_filter, link_filter, traffic_direction_filter: TrafficDirectionFilter::default(), + ingress_filter_channel: FilterChannel::new(), + egress_filter_channel: FilterChannel::new(), start_sniffing: false, packets, packets_table_state: TableState::default(), @@ -332,16 +344,12 @@ impl App { // Filters let widths = [Constraint::Length(10), Constraint::Fill(1)]; 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( + Line::from_iter(TransportFilter::new().selected_protocols.iter().map( |filter| { - if applied_transport_filters.contains(filter) { + if self.transport_filter.applied_protocols.contains(filter) { Span::styled( format!(" {} ", filter), Style::default().light_green(), @@ -357,9 +365,9 @@ impl App { ]), Row::new(vec![ Line::styled("Network", Style::new().bold()), - Line::from_iter(NetworkFilter::default().selected_protocols.iter().map( + Line::from_iter(NetworkFilter::new().selected_protocols.iter().map( |filter| { - if applied_network_filters.contains(filter) { + if self.network_filter.applied_protocols.contains(filter) { Span::styled( format!(" {} ", filter), Style::default().light_green(), @@ -375,9 +383,9 @@ impl App { ]), Row::new(vec![ Line::styled("Link", Style::new().bold()), - Line::from_iter(LinkFilter::default().selected_protocols.iter().map( + Line::from_iter(LinkFilter::new().selected_protocols.iter().map( |filter| { - if applied_link_filters.contains(filter) { + if self.link_filter.applied_protocols.contains(filter) { Span::styled( format!(" {} ", filter), Style::default().light_green(), @@ -870,69 +878,17 @@ impl App { 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 applied_link_protocols.contains(&LinkProtocol::Arp) { - packets.push(app_packet); - } - } - AppPacket::Ip(packet) => match packet { - IpPacket::Tcp(p) => { - if applied_transport_protocols.contains(&TransportProtocol::TCP) { - match p.src_ip { - core::net::IpAddr::V6(_) => { - if applied_network_protocols.contains(&NetworkProtocol::Ipv6) { - packets.push(app_packet); - } - } - core::net::IpAddr::V4(_) => { - if applied_network_protocols.contains(&NetworkProtocol::Ipv4) { - packets.push(app_packet); - } - } - } - } - } - IpPacket::Udp(p) => { - if applied_transport_protocols.contains(&TransportProtocol::UDP) { - match p.src_ip { - core::net::IpAddr::V6(_) => { - if applied_network_protocols.contains(&NetworkProtocol::Ipv6) { - packets.push(app_packet); - } - } - core::net::IpAddr::V4(_) => { - if applied_network_protocols.contains(&NetworkProtocol::Ipv4) { - packets.push(app_packet); - } - } - } - } - } - IpPacket::Icmp(_) => { - if applied_network_protocols.contains(&NetworkProtocol::Icmp) { - packets.push(app_packet); - } - } - }, - } - let mut stats = stats.lock().unwrap(); + packets.push(app_packet); + let mut stats = stats.lock().unwrap(); stats.refresh(&app_packet); } diff --git a/oryx-tui/src/ebpf.rs b/oryx-tui/src/ebpf.rs index 67d5d8d..029945d 100644 --- a/oryx-tui/src/ebpf.rs +++ b/oryx-tui/src/ebpf.rs @@ -1,15 +1,16 @@ -use std::io; use std::os::fd::AsRawFd; use std::sync::atomic::AtomicBool; use std::sync::Arc; -use std::thread; +use std::thread::spawn; use std::time::Duration; +use std::{io, thread}; use aya::maps::ring_buf::RingBufItem; -use aya::maps::{MapData, RingBuf}; +use aya::maps::{Array, MapData, RingBuf}; use aya::programs::{tc, SchedClassifier, TcAttachType}; use aya::{include_bytes_aligned, Bpf}; use mio::{Events, Interest, Poll, Registry, Token}; +use oryx_common::protocols::Protocol; use oryx_common::RawPacket; use crate::event::Event; @@ -63,6 +64,7 @@ impl Ebpf { iface: String, notification_sender: kanal::Sender, data_sender: kanal::Sender<[u8; RawPacket::LEN]>, + filter_channel_receiver: kanal::Receiver<(Protocol, bool)>, terminate: Arc, ) { thread::spawn({ @@ -143,6 +145,31 @@ impl Ebpf { let mut poll = Poll::new().unwrap(); let mut events = Events::with_capacity(128); + let mut transport_filters: Array<_, u32> = + Array::try_from(bpf.take_map("TRANSPORT_FILTERS").unwrap()).unwrap(); + + let mut network_filters: Array<_, u32> = + Array::try_from(bpf.take_map("NETWORK_FILTERS").unwrap()).unwrap(); + + let mut link_filters: Array<_, u32> = + Array::try_from(bpf.take_map("LINK_FILTERS").unwrap()).unwrap(); + + spawn(move || loop { + if let Ok((filter, flag)) = filter_channel_receiver.recv() { + match filter { + Protocol::Transport(p) => { + let _ = transport_filters.set(p as u32, flag as u32, 0); + } + Protocol::Network(p) => { + let _ = network_filters.set(p as u32, flag as u32, 0); + } + Protocol::Link(p) => { + let _ = link_filters.set(p as u32, flag as u32, 0); + } + } + } + }); + let mut ring_buf = RingBuffer::new(&mut bpf); poll.registry() @@ -190,6 +217,7 @@ impl Ebpf { iface: String, notification_sender: kanal::Sender, data_sender: kanal::Sender<[u8; RawPacket::LEN]>, + filter_channel_receiver: kanal::Receiver<(Protocol, bool)>, terminate: Arc, ) { thread::spawn({ @@ -266,6 +294,30 @@ impl Ebpf { let mut poll = Poll::new().unwrap(); let mut events = Events::with_capacity(128); + let mut transport_filters: Array<_, u32> = + Array::try_from(bpf.take_map("TRANSPORT_FILTERS").unwrap()).unwrap(); + + let mut network_filters: Array<_, u32> = + Array::try_from(bpf.take_map("NETWORK_FILTERS").unwrap()).unwrap(); + + let mut link_filters: Array<_, u32> = + Array::try_from(bpf.take_map("LINK_FILTERS").unwrap()).unwrap(); + + spawn(move || loop { + if let Ok((filter, flag)) = filter_channel_receiver.recv() { + match filter { + Protocol::Transport(p) => { + let _ = transport_filters.set(p as u32, flag as u32, 0); + } + Protocol::Network(p) => { + let _ = network_filters.set(p as u32, flag as u32, 0); + } + Protocol::Link(p) => { + let _ = link_filters.set(p as u32, flag as u32, 0); + } + } + } + }); let mut ring_buf = RingBuffer::new(&mut bpf); poll.registry() diff --git a/oryx-tui/src/event.rs b/oryx-tui/src/event.rs index b03c89c..5c71572 100644 --- a/oryx-tui/src/event.rs +++ b/oryx-tui/src/event.rs @@ -1,8 +1,6 @@ use crate::app::AppResult; use crate::notification::Notification; -use ratatui::crossterm::event::{ - self, Event as CrosstermEvent, KeyEvent, KeyEventKind, MouseEvent, -}; +use crossterm::event::{self, Event as CrosstermEvent, KeyEvent, KeyEventKind, MouseEvent}; use std::thread; use std::time::{Duration, Instant}; diff --git a/oryx-tui/src/filters/link.rs b/oryx-tui/src/filters/link.rs index 7985af6..212ee85 100644 --- a/oryx-tui/src/filters/link.rs +++ b/oryx-tui/src/filters/link.rs @@ -1,8 +1,4 @@ -use std::{ - fmt::Display, - sync::{Arc, Mutex}, -}; - +use oryx_common::protocols::LinkProtocol; use ratatui::{ layout::{Alignment, Constraint, Direction, Flex, Layout, Rect}, style::{Color, Style, Stylize}, @@ -12,40 +8,23 @@ use ratatui::{ use crate::app::FocusedBlock; -pub const NB_LINK_PROTOCOL: u16 = 1; - #[derive(Debug)] pub struct LinkFilter { pub state: TableState, pub selected_protocols: Vec, - pub applied_protocols: Arc>>, -} - -#[derive(Debug, Copy, Clone, PartialEq)] -pub enum LinkProtocol { - Arp, -} - -impl Display for LinkProtocol { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "Arp") - } + pub applied_protocols: Vec, } -impl Default for LinkFilter { - fn default() -> Self { +impl LinkFilter { + pub fn new() -> Self { Self { state: TableState::default(), selected_protocols: vec![LinkProtocol::Arp], - applied_protocols: Arc::new(Mutex::new(Vec::new())), + applied_protocols: Vec::new(), } } -} - -impl LinkFilter { pub fn apply(&mut self) { - let mut applied_protocols = self.applied_protocols.lock().unwrap(); - *applied_protocols = self.selected_protocols.clone(); + self.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 0876afb..cb5f729 100644 --- a/oryx-tui/src/filters/network.rs +++ b/oryx-tui/src/filters/network.rs @@ -1,8 +1,4 @@ -use std::{ - fmt::Display, - sync::{Arc, Mutex}, -}; - +use oryx_common::protocols::NetworkProtocol; use ratatui::{ layout::{Alignment, Constraint, Direction, Flex, Layout, Rect}, style::{Color, Style, Stylize}, @@ -12,17 +8,15 @@ use ratatui::{ use crate::app::FocusedBlock; -pub const NB_NETWORK_PROTOCOL: u16 = 3; - #[derive(Debug)] pub struct NetworkFilter { pub state: TableState, pub selected_protocols: Vec, - pub applied_protocols: Arc>>, + pub applied_protocols: Vec, } -impl Default for NetworkFilter { - fn default() -> Self { +impl NetworkFilter { + pub fn new() -> Self { NetworkFilter { state: TableState::default(), selected_protocols: vec![ @@ -30,32 +24,14 @@ impl Default for NetworkFilter { NetworkProtocol::Ipv6, NetworkProtocol::Icmp, ], - applied_protocols: Arc::new(Mutex::new(Vec::new())), - } - } -} - -#[derive(Debug, Copy, Clone, PartialEq)] -pub enum NetworkProtocol { - Ipv4, - Ipv6, - Icmp, -} - -impl Display for NetworkProtocol { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - NetworkProtocol::Ipv4 => write!(f, "Ipv4"), - NetworkProtocol::Ipv6 => write!(f, "Ipv6"), - NetworkProtocol::Icmp => write!(f, "Icmp"), + applied_protocols: Vec::new(), } } } impl NetworkFilter { pub fn apply(&mut self) { - let mut applied_protocols = self.applied_protocols.lock().unwrap(); - *applied_protocols = self.selected_protocols.clone(); + self.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 7a6d8cb..4f10066 100644 --- a/oryx-tui/src/filters/transport.rs +++ b/oryx-tui/src/filters/transport.rs @@ -1,8 +1,4 @@ -use std::{ - fmt::Display, - sync::{Arc, Mutex}, -}; - +use oryx_common::protocols::TransportProtocol; use ratatui::{ layout::{Alignment, Constraint, Direction, Flex, Layout, Rect}, style::{Color, Style, Stylize}, @@ -12,44 +8,25 @@ use ratatui::{ use crate::app::FocusedBlock; -pub const NB_TRANSPORT_PROTOCOL: u16 = 2; - #[derive(Debug)] pub struct TransportFilter { pub state: TableState, pub selected_protocols: Vec, - pub applied_protocols: Arc>>, + pub applied_protocols: Vec, } -#[derive(Debug, Copy, Clone, PartialEq)] -pub enum TransportProtocol { - TCP, - UDP, -} - -impl Display for TransportProtocol { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - TransportProtocol::TCP => write!(f, "Tcp"), - TransportProtocol::UDP => write!(f, "Udp"), - } - } -} +impl TransportFilter {} -impl Default for TransportFilter { - fn default() -> Self { +impl TransportFilter { + pub fn new() -> Self { Self { state: TableState::default(), selected_protocols: vec![TransportProtocol::TCP, TransportProtocol::UDP], - applied_protocols: Arc::new(Mutex::new(Vec::new())), + applied_protocols: Vec::new(), } } -} - -impl TransportFilter { pub fn apply(&mut self) { - let mut applied_protocols = self.applied_protocols.lock().unwrap(); - *applied_protocols = self.selected_protocols.clone(); + self.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 d760a90..6628da6 100644 --- a/oryx-tui/src/handler.rs +++ b/oryx-tui/src/handler.rs @@ -1,3 +1,7 @@ +use oryx_common::protocols::{ + LinkProtocol, NetworkProtocol, Protocol, TransportProtocol, NB_LINK_PROTOCOL, + NB_NETWORK_PROTOCOL, NB_TRANSPORT_PROTOCOL, +}; use std::{thread, time::Duration}; use tui_input::backend::crossterm::EventHandler; @@ -6,12 +10,7 @@ use crate::{ ebpf::Ebpf, event::Event, export::export, - filters::{ - direction::TrafficDirection, - link::{LinkProtocol, NB_LINK_PROTOCOL}, - network::{NetworkProtocol, NB_NETWORK_PROTOCOL}, - transport::{TransportProtocol, NB_TRANSPORT_PROTOCOL}, - }, + filters::direction::TrafficDirection, notification::{Notification, NotificationLevel}, }; use ratatui::{ @@ -165,19 +164,14 @@ 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 = - applied_network_protocols.clone(); + app.network_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.transport_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.link_filter.selected_protocols = + app.link_filter.applied_protocols.clone(); app.traffic_direction_filter.selected_direction = app.traffic_direction_filter.applied_direction.clone(); @@ -410,16 +404,13 @@ pub fn handle_key_events( app.focused_block = FocusedBlock::TransportFilter; - let applied_network_protocols = - app.network_filter.applied_protocols.lock().unwrap(); - app.network_filter.selected_protocols = applied_network_protocols.clone(); + app.network_filter.selected_protocols = + app.network_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.transport_filter.selected_protocols = + app.transport_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.link_filter.selected_protocols = app.link_filter.applied_protocols.clone(); app.traffic_direction_filter.selected_direction = app.traffic_direction_filter.applied_direction.clone(); @@ -506,6 +497,7 @@ pub fn handle_key_events( iface.clone(), sender.clone(), app.data_channel_sender.clone(), + app.ingress_filter_channel.receiver.clone(), app.traffic_direction_filter.terminate_ingress.clone(), ); } @@ -519,6 +511,7 @@ pub fn handle_key_events( iface, sender.clone(), app.data_channel_sender.clone(), + app.egress_filter_channel.receiver.clone(), app.traffic_direction_filter.terminate_egress.clone(), ); } @@ -551,11 +544,15 @@ pub fn handle_key_events( .selected_direction .contains(&TrafficDirection::Egress) { + app.traffic_direction_filter + .terminate_egress + .store(false, std::sync::atomic::Ordering::Relaxed); let iface = app.interface.selected_interface.name.clone(); Ebpf::load_egress( iface, sender.clone(), app.data_channel_sender.clone(), + app.egress_filter_channel.receiver.clone(), app.traffic_direction_filter.terminate_egress.clone(), ); } @@ -586,10 +583,14 @@ pub fn handle_key_events( .contains(&TrafficDirection::Ingress) { let iface = app.interface.selected_interface.name.clone(); + app.traffic_direction_filter + .terminate_ingress + .store(false, std::sync::atomic::Ordering::Relaxed); Ebpf::load_ingress( iface, sender.clone(), app.data_channel_sender.clone(), + app.ingress_filter_channel.receiver.clone(), app.traffic_direction_filter.terminate_ingress.clone(), ); } @@ -608,6 +609,60 @@ pub fn handle_key_events( app.update_filters = false; } + + for protocol in TransportProtocol::all().iter() { + if app.transport_filter.applied_protocols.contains(&protocol) { + app.ingress_filter_channel + .sender + .send((Protocol::Transport(*protocol), false))?; + app.egress_filter_channel + .sender + .send((Protocol::Transport(*protocol), false))?; + } else { + app.ingress_filter_channel + .sender + .send((Protocol::Transport(*protocol), true))?; + app.egress_filter_channel + .sender + .send((Protocol::Transport(*protocol), true))?; + } + } + + for protocol in NetworkProtocol::all().iter() { + if app.network_filter.applied_protocols.contains(&protocol) { + app.ingress_filter_channel + .sender + .send((Protocol::Network(*protocol), false))?; + app.egress_filter_channel + .sender + .send((Protocol::Network(*protocol), false))?; + } else { + app.ingress_filter_channel + .sender + .send((Protocol::Network(*protocol), true))?; + app.egress_filter_channel + .sender + .send((Protocol::Network(*protocol), true))?; + } + } + + for protocol in LinkProtocol::all().iter() { + if app.link_filter.applied_protocols.contains(&protocol) { + app.ingress_filter_channel + .sender + .send((Protocol::Link(*protocol), false))?; + app.egress_filter_channel + .sender + .send((Protocol::Link(*protocol), false))?; + } else { + app.ingress_filter_channel + .sender + .send((Protocol::Link(*protocol), true))?; + app.egress_filter_channel + .sender + .send((Protocol::Link(*protocol), true))?; + } + } } KeyCode::Tab => { diff --git a/oryx-tui/src/tui.rs b/oryx-tui/src/tui.rs index fdbfc03..732dc64 100644 --- a/oryx-tui/src/tui.rs +++ b/oryx-tui/src/tui.rs @@ -1,9 +1,9 @@ use crate::app::{App, AppResult}; use crate::event::EventHandler; use crate::ui; +use crossterm::event::{DisableMouseCapture, EnableMouseCapture}; +use crossterm::terminal::{self, EnterAlternateScreen, LeaveAlternateScreen}; use ratatui::backend::Backend; -use ratatui::crossterm::event::{DisableMouseCapture, EnableMouseCapture}; -use ratatui::crossterm::terminal::{self, EnterAlternateScreen, LeaveAlternateScreen}; use ratatui::Terminal; use std::io; use std::panic;