diff --git a/Cargo.lock b/Cargo.lock index 2352848..f29d984 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -407,6 +407,15 @@ dependencies = [ "thiserror 2.0.16", ] +[[package]] +name = "castaway" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a" +dependencies = [ + "rustversion", +] + [[package]] name = "cc" version = "1.2.47" @@ -470,6 +479,22 @@ dependencies = [ "aya-ebpf", "network-types", "serde", + "zerocopy 0.8.33", +] + +[[package]] +name = "compact_str" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7fd622ebbb56a5b2ccb651b32b911cdeb2a9b4b11776b2473bf26a26a286244e" +dependencies = [ + "castaway", + "cfg-if", + "itoa", + "rustversion", + "ryu", + "serde", + "static_assertions", ] [[package]] @@ -1505,14 +1530,17 @@ dependencies = [ name = "mantis" version = "0.1.0" dependencies = [ + "ahash", "argon2", "axum", "aya", "aya-log", + "bytes", "cargo_metadata", "cc", "chrono", "common", + "compact_str", "crossbeam", "dotenvy", "futures", @@ -1545,6 +1573,7 @@ dependencies = [ "url", "uuid", "xsk-rs", + "zerocopy 0.8.33", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 6330603..d2cd17f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,14 +4,18 @@ members = ["mantis", "common", "macros", "ingress-ebpf", "egress-ebpf"] default-members = ["mantis", "common"] [workspace.dependencies] +ahash = { version = "0.8", default-features = false, features = ["std"] } aya = { version = "0.13.1", default-features = false } aya-ebpf = { version = "0.1.1", default-features = false } aya-log = { version = "0.2.1", default-features = false } aya-log-ebpf = { version = "0.1.0", default-features = false } +bytes = { version = "1" } cargo_metadata = { version = "0.23.1", default-features = false } +compact_str = { version = "0.8", features = ["serde"] } libc = { version = "0.2.159", default-features = false } network-types = "0.1.0" serde = { version = "1.0.215", features = ["derive"] } +zerocopy = { version = "0.8", default-features = false, features = ["derive"] } xsk-rs = { version = "0.8.0", default-features = false } [profile.dev] diff --git a/common/Cargo.toml b/common/Cargo.toml index 1cfaa89..2a36f92 100644 --- a/common/Cargo.toml +++ b/common/Cargo.toml @@ -5,13 +5,14 @@ edition = "2024" [features] default = [] -user = ["aya", "serde"] +user = ["aya", "serde", "dep:zerocopy"] kernel = ["aya-ebpf"] [dependencies] aya = { workspace = true, optional = true } aya-ebpf = { workspace = true, optional = true } serde = { workspace = true, optional = true } +zerocopy = { workspace = true, optional = true } network-types = { workspace = true } [lib] diff --git a/common/src/model/flow_stats.rs b/common/src/model/flow_stats.rs index d9ccb5c..31949e2 100644 --- a/common/src/model/flow_stats.rs +++ b/common/src/model/flow_stats.rs @@ -2,10 +2,15 @@ use aya::Pod; #[cfg(feature = "user")] use serde::Serialize; +#[cfg(feature = "user")] +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; #[repr(C, align(8))] #[derive(Clone, Copy)] -#[cfg_attr(feature = "user", derive(Serialize, Debug))] +#[cfg_attr( + feature = "user", + derive(Serialize, Debug, FromBytes, IntoBytes, KnownLayout, Immutable) +)] pub struct FlowStats { pub bytes: u64, pub packets: u64, diff --git a/common/src/model/ip_address.rs b/common/src/model/ip_address.rs index 84a0e99..bc09241 100644 --- a/common/src/model/ip_address.rs +++ b/common/src/model/ip_address.rs @@ -1,5 +1,7 @@ #[cfg(feature = "user")] use aya::Pod; +#[cfg(feature = "user")] +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; pub type IPv4 = u32; pub type IPv6 = u128; @@ -7,6 +9,7 @@ pub type Port = u16; #[repr(transparent)] #[derive(Debug, Copy, Clone)] +#[cfg_attr(feature = "user", derive(FromBytes, IntoBytes, KnownLayout, Immutable))] pub struct AddrPortV4([u8; 8]); impl AddrPortV4 { @@ -43,6 +46,7 @@ unsafe impl Pod for AddrPortV4 {} #[repr(transparent)] #[derive(Debug, Copy, Clone)] +#[cfg_attr(feature = "user", derive(FromBytes, IntoBytes, KnownLayout, Immutable))] pub struct AddrPortV6([u8; 32]); impl AddrPortV6 { diff --git a/mantis/Cargo.toml b/mantis/Cargo.toml index a1cdae1..fc48c71 100644 --- a/mantis/Cargo.toml +++ b/mantis/Cargo.toml @@ -6,6 +6,9 @@ edition = "2024" [dependencies] common = { path = "../common", features = ["user"] } macros = { path = "../macros" } +ahash = { workspace = true } +bytes = { workspace = true } +compact_str = { workspace = true } axum = { version = "0.8", features = ["ws", "macros"] } tower = { version = "0.5", features = ["util"] } @@ -21,6 +24,7 @@ parking_lot = "0.12.5" rust-embed = "8.7.2" serde = { workspace = true } serde_json = "1.0.143" +zerocopy = { workspace = true } sysinfo = "0.39.2" thiserror = "2.0.3" tokio = { version = "1.40.0", features = ["full", "macros"] } diff --git a/mantis/build.rs b/mantis/build.rs index 2187468..73a0c1a 100644 --- a/mantis/build.rs +++ b/mantis/build.rs @@ -37,6 +37,12 @@ fn main() { println!("cargo:rustc-env=DB_PATH={}", db_dir.display()); println!("cargo:rustc-env=RULE_EVE_PATH={}", suricata_eve_socket.display()); + // Watch output artifacts: if they are missing Cargo treats the path as non-existent + // and unconditionally re-runs this script, which is exactly what we want after a + // git clean, fresh clone, or manual deletion. + println!("cargo:rerun-if-changed={}", ingress_edpf_dir.display()); + println!("cargo:rerun-if-changed={}", egress_edpf_dir.display()); + for item in &[ "src", "public", diff --git a/mantis/src/core/ebpf/xsk_manager.rs b/mantis/src/core/ebpf/xsk_manager.rs index 187aed0..a496d93 100644 --- a/mantis/src/core/ebpf/xsk_manager.rs +++ b/mantis/src/core/ebpf/xsk_manager.rs @@ -8,6 +8,7 @@ use std::time::Duration; use aya::Ebpf; use aya::maps::{MapData, XskMap}; +use bytes::Bytes; use crossbeam::channel::{Receiver, Sender, bounded}; use crossbeam::queue::SegQueue; use libc; @@ -61,8 +62,8 @@ impl XskManager { let combined_queue_count = config.combined_queue_count; for queue_id in 0..combined_queue_count { - let (ingress_to_egress_tx, ingress_to_egress_rx) = bounded(config.channel_size); - let (egress_to_ingress_tx, egress_to_ingress_rx) = bounded(config.channel_size); + let (ingress_to_egress_tx, ingress_to_egress_rx) = bounded::(config.channel_size); + let (egress_to_ingress_tx, egress_to_ingress_rx) = bounded::(config.channel_size); let ingress_xsk = XskPair::new( config.clone(), @@ -258,8 +259,8 @@ impl XskPair { pub fn run( mut self, - forward_tx: Sender>, - forward_rx: Receiver>, + forward_tx: Sender, + forward_rx: Receiver, ) -> Result, EbpfError> { let (shutdown_tx, shutdown_rx) = oneshot::channel(); @@ -363,26 +364,26 @@ impl XskPair { Ok(nb_completed) } - fn process_rx_queue(&mut self, forward_tx: &Sender>) -> Result { + fn process_rx_queue(&mut self, forward_tx: &Sender) -> Result { let mut rx_descs = vec![FrameDesc::default(); 256]; let rx_count = unsafe { self.rx.consume(&mut rx_descs) }; if rx_count > 0 { for rx_desc in rx_descs.iter().take(rx_count) { - let lengths = rx_desc.lengths(); - let packet_len = lengths.data() as usize; - + let packet_len = rx_desc.lengths().data() as usize; let data = unsafe { self.umem.data(rx_desc) }; - let packet_data = data.contents()[..packet_len].to_vec(); + let packet_slice = &data.contents()[..packet_len]; + // ML engine reads directly from UMEM — no copy for this consumer + if let Some(ref engine) = self.engine { + engine.process_packet(packet_slice, self.direction == Direction::Ingress); + } + + // Bytes::clone is a refcount increment — no copy for Suricata vs forward_tx + let packet_data = Bytes::copy_from_slice(packet_slice); if let Some(ref se) = self.suricata_engine { se.inject(packet_data.clone()); } - - if let Some(ref engine) = self.engine { - engine.process_packet(&packet_data, self.direction == Direction::Ingress); - } - if let Err(e) = forward_tx.try_send(packet_data) { match e { crossbeam::channel::TrySendError::Full(_) => { @@ -455,7 +456,7 @@ impl XskPair { } } - fn process_tx_queue(&mut self, forward_rx: &Receiver>) -> Result { + fn process_tx_queue(&mut self, forward_rx: &Receiver) -> Result { // Drain completed TX frames first to maximise pool availability. let _ = self.process_comp_queue(); @@ -474,7 +475,7 @@ impl XskPair { // Consume at most min(pool_size, 64) packets so we never over-commit. let max_to_send = pool_size.min(64); - let mut packets_to_send = Vec::with_capacity(max_to_send); + let mut packets_to_send: Vec = Vec::with_capacity(max_to_send); while let Ok(packet) = forward_rx.try_recv() { packets_to_send.push(packet); if packets_to_send.len() >= max_to_send { @@ -508,7 +509,7 @@ impl XskPair { self.umem .data_mut(frame) .cursor() - .write_all(packet) + .write_all(packet.as_ref()) .map_err(EbpfError::AfXdpSetFailed)?; } } diff --git a/mantis/src/detection/ml/aggregator.rs b/mantis/src/detection/ml/aggregator.rs index cadccbe..2f7e418 100644 --- a/mantis/src/detection/ml/aggregator.rs +++ b/mantis/src/detection/ml/aggregator.rs @@ -1,6 +1,9 @@ -use std::collections::{HashMap, HashSet}; +use std::collections::HashSet; use std::time::{Duration, Instant}; +use ahash::AHashMap; +use compact_str::CompactString; + use crate::model::ml_detection::FlowKey; // L2 thresholds: anomalous flows per src_ip within the aggregation window @@ -13,8 +16,8 @@ struct SrcIpState { } pub struct AttackAggregator { - detections: HashMap>, - src_ip_states: HashMap, + detections: AHashMap>, + src_ip_states: AHashMap, window_duration: Duration, min_detections: usize, alert_threshold_multiplier: f32, @@ -23,8 +26,8 @@ pub struct AttackAggregator { impl AttackAggregator { pub fn new(window_secs: u64, min_detections: usize) -> Self { Self { - detections: HashMap::new(), - src_ip_states: HashMap::new(), + detections: AHashMap::new(), + src_ip_states: AHashMap::new(), window_duration: Duration::from_secs(window_secs), min_detections, alert_threshold_multiplier: 1.2, @@ -53,7 +56,7 @@ impl AttackAggregator { let state = self .src_ip_states - .entry(src_ip.to_string()) + .entry(CompactString::from(src_ip)) .or_insert_with(|| SrcIpState { events: Vec::new(), last_alert: None, diff --git a/mantis/src/detection/ml/engine.rs b/mantis/src/detection/ml/engine.rs index 4a02169..d395e46 100644 --- a/mantis/src/detection/ml/engine.rs +++ b/mantis/src/detection/ml/engine.rs @@ -1,6 +1,8 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; +use ahash::AHashSet; +use compact_str::CompactString; use macros::log; use tokio::sync::oneshot; use tokio::time::interval; @@ -93,10 +95,10 @@ impl Engine { t.cleanup_old_flows(self.flow_timeout_us); // active_ips covers both the just-drained flows and flows still in the // tracker (ongoing connections), so their LSTM buffers are preserved. - let active_ips: std::collections::HashSet = flows + let active_ips: AHashSet = flows .iter() .map(|f| f.flow_key.src_ip.clone()) - .chain(t.active_src_ips().map(str::to_owned)) + .chain(t.active_src_ips().map(CompactString::from)) .collect(); (total_flows, flows, active_ips) }; @@ -168,7 +170,7 @@ impl Engine { )); if let Ok(mut aggregator) = self.aggregator.lock() { - let mut alerted_src_ips = std::collections::HashSet::new(); + let mut alerted_src_ips: AHashSet = AHashSet::new(); for result in &results { if result.is_attack { // L1: full 5-tuple aggregation for persistent same-port attacks diff --git a/mantis/src/detection/ml/feature_extractor.rs b/mantis/src/detection/ml/feature_extractor.rs index a5e3d63..b56aae9 100644 --- a/mantis/src/detection/ml/feature_extractor.rs +++ b/mantis/src/detection/ml/feature_extractor.rs @@ -9,10 +9,6 @@ use crate::model::ml_detection::{ClipParams, PacketData}; pub struct FlowFeatures { pub features: Vec, pub feature_num: usize, - - pub src_ip: String, - pub dst_ip: String, - pub timestamp: u64, } impl FlowFeatures { @@ -23,13 +19,7 @@ impl FlowFeatures { for name in feature_names { features.push(Self::get_feature_by_name(flow, name.trim())); } - Self { - features, - feature_num, - src_ip: flow.flow_key.src_ip.clone(), - dst_ip: flow.flow_key.dst_ip.clone(), - timestamp: flow.start_time_us, - } + Self { features, feature_num } } fn get_feature_by_name(flow: &FlowData, feature_name: &str) -> f64 { @@ -221,8 +211,8 @@ impl FlowFeatures { pub fn get_csv_column(flow: &FlowData, column: &str) -> String { match column { - "Source IP" => flow.flow_key.src_ip.clone(), - "Destination IP" => flow.flow_key.dst_ip.clone(), + "Source IP" => flow.flow_key.src_ip.to_string(), + "Destination IP" => flow.flow_key.dst_ip.to_string(), "Timestamp" => { let ts_ms = (flow.start_time_us / 1000) as i64; match Utc.timestamp_millis_opt(ts_ms) { diff --git a/mantis/src/detection/ml/flow_tracker.rs b/mantis/src/detection/ml/flow_tracker.rs index e9888c4..b870d90 100644 --- a/mantis/src/detection/ml/flow_tracker.rs +++ b/mantis/src/detection/ml/flow_tracker.rs @@ -1,6 +1,6 @@ -use std::collections::HashMap; use std::time; +use ahash::AHashMap; use common::model::event::Event; use super::server_ports; @@ -268,14 +268,14 @@ impl FlowData { } pub struct FlowTracker { - flows: HashMap, + flows: AHashMap, max_flows: usize, } impl FlowTracker { pub fn new(max_flows: usize) -> Self { Self { - flows: HashMap::new(), + flows: AHashMap::new(), max_flows, } } diff --git a/mantis/src/detection/ml/inference.rs b/mantis/src/detection/ml/inference.rs index 04520fd..f783b62 100644 --- a/mantis/src/detection/ml/inference.rs +++ b/mantis/src/detection/ml/inference.rs @@ -1,8 +1,9 @@ -use std::collections::HashMap; use std::collections::VecDeque; use std::sync::{Arc, Mutex}; use std::time::Instant; +use ahash::{AHashMap, AHashSet}; +use compact_str::CompactString; use macros::log; use ndarray::Array3; use ort::{inputs, value::TensorRef}; @@ -24,7 +25,7 @@ pub struct Inference { // inflated MSE regardless of traffic type, causing false positives. min_window_fill: usize, // per-src_ip sliding window buffer: src_ip -> deque of feature vectors - flow_buffers: Mutex>>>, + flow_buffers: Mutex>>>, } impl Inference { @@ -36,11 +37,11 @@ impl Inference { config, threshold, min_window_fill, - flow_buffers: Mutex::new(HashMap::new()), + flow_buffers: Mutex::new(AHashMap::new()), } } - pub fn cleanup_buffers(&self, active_src_ips: &std::collections::HashSet) { + pub fn cleanup_buffers(&self, active_src_ips: &AHashSet) { let Ok(mut buffers) = self.flow_buffers.lock() else { log!(MLError::InferenceLockPoisoned); return; @@ -103,7 +104,7 @@ impl Inference { let t4 = Instant::now(); log!(MLLog::InferenceTiming( - flow.flow_key.src_ip.clone(), + flow.flow_key.src_ip.to_string(), t1.duration_since(t0).as_millis() as u64, t2.duration_since(t1).as_millis() as u64, t3.duration_since(t2).as_millis() as u64, @@ -125,7 +126,7 @@ impl Inference { .collect::>() .join("\n"); log!(MLLog::WindowDebug( - flow.flow_key.src_ip.clone(), + flow.flow_key.src_ip.to_string(), pad, window_size, ae_score, diff --git a/mantis/src/detection/suricata/engine.rs b/mantis/src/detection/suricata/engine.rs index 2490fb8..0d06b68 100644 --- a/mantis/src/detection/suricata/engine.rs +++ b/mantis/src/detection/suricata/engine.rs @@ -6,6 +6,7 @@ use std::sync::Arc; use std::thread; use std::time::Duration; +use bytes::Bytes; use crossbeam::channel::{Sender, bounded}; use macros::log; @@ -22,7 +23,7 @@ const SURICATA_LOG: &str = "/tmp/suricata.log"; const CHANNEL_CAP: usize = 4096; pub struct SuricataEngine { - tx: Sender>, + tx: Sender, child: std::sync::Mutex, } @@ -133,7 +134,7 @@ impl SuricataEngine { }) .map_err(|e| SuricataError::ProcessSpawnFailed { reason: e.to_string() })?; - let (tx, rx) = bounded::>(CHANNEL_CAP); + let (tx, rx) = bounded::(CHANNEL_CAP); thread::Builder::new() .name("suricata-mirror".into()) @@ -182,7 +183,7 @@ impl SuricataEngine { } /* Non-blocking: drops silently when the channel is full under load. */ - pub fn inject(&self, data: Vec) { + pub fn inject(&self, data: Bytes) { match self.tx.try_send(data) { Ok(()) => {} Err(crossbeam::channel::TrySendError::Full(_)) => { diff --git a/mantis/src/model/ml_detection.rs b/mantis/src/model/ml_detection.rs index c383ffb..5772ccd 100644 --- a/mantis/src/model/ml_detection.rs +++ b/mantis/src/model/ml_detection.rs @@ -1,4 +1,5 @@ use common::model::event::{Event, TcpFlags}; +use compact_str::CompactString; use serde::{Deserialize, Serialize}; use crate::model::direction::Direction; @@ -12,8 +13,8 @@ pub struct ClipParams { } #[derive(Debug, Clone, Hash, Eq, PartialEq, Serialize, Deserialize)] pub struct FlowKey { - pub src_ip: String, - pub dst_ip: String, + pub src_ip: CompactString, + pub dst_ip: CompactString, pub src_port: u16, pub dst_port: u16, pub protocol: u8, @@ -178,8 +179,8 @@ impl UnifiedAlert { Self { timestamp: now_secs(), flow_key: result.flow_key.clone(), - src_ip: result.flow_key_raw.src_ip.clone(), - dst_ip: result.flow_key_raw.dst_ip.clone(), + src_ip: result.flow_key_raw.src_ip.to_string(), + dst_ip: result.flow_key_raw.dst_ip.to_string(), src_port: result.flow_key_raw.src_port, dst_port: result.flow_key_raw.dst_port, protocol: result.flow_key_raw.protocol, @@ -221,8 +222,8 @@ impl UnifiedAlert { Self { timestamp: now_secs(), flow_key: result.flow_key.clone(), - src_ip: result.flow_key_raw.src_ip.clone(), - dst_ip: result.flow_key_raw.dst_ip.clone(), + src_ip: result.flow_key_raw.src_ip.to_string(), + dst_ip: result.flow_key_raw.dst_ip.to_string(), src_port: result.flow_key_raw.src_port, dst_port: result.flow_key_raw.dst_port, protocol: result.flow_key_raw.protocol, diff --git a/mantis/src/utils/packet_parser.rs b/mantis/src/utils/packet_parser.rs index be4f2a9..e688ff1 100644 --- a/mantis/src/utils/packet_parser.rs +++ b/mantis/src/utils/packet_parser.rs @@ -1,171 +1,220 @@ -use std::mem; use std::time; use common::model::event::{Event, IPv4Event, IPv6Event, TcpFlags}; +use compact_str::{CompactString, format_compact}; use network_types::ip::IpProto; +use zerocopy::byteorder::{BigEndian, U16, U32}; +use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout, Ref}; + +#[repr(C, packed)] +#[derive(FromBytes, IntoBytes, KnownLayout, Immutable, Clone, Copy)] +struct EthHdr { + _dst_mac: [u8; 6], + _src_mac: [u8; 6], + ether_type: U16, +} + +#[repr(C, packed)] +#[derive(FromBytes, IntoBytes, KnownLayout, Immutable, Clone, Copy)] +struct Ipv4Hdr { + version_ihl: u8, + _dscp_ecn: u8, + total_len: U16, + _ident: U16, + _flags_frag: U16, + _ttl: u8, + protocol: u8, + _checksum: U16, + src_addr: U32, + dst_addr: U32, +} + +#[repr(C, packed)] +#[derive(FromBytes, IntoBytes, KnownLayout, Immutable, Clone, Copy)] +struct Ipv6Hdr { + _version_tc_fl: U32, + payload_len: U16, + next_hdr: u8, + _hop_limit: u8, + src_addr: [u8; 16], + dst_addr: [u8; 16], +} + +#[repr(C, packed)] +#[derive(FromBytes, IntoBytes, KnownLayout, Immutable, Clone, Copy)] +struct TcpHdr { + src_port: U16, + dst_port: U16, + _seq_num: U32, + _ack_num: U32, + data_off_flags: U16, + window: U16, + _checksum: U16, + _urgent_ptr: U16, +} + +#[repr(C, packed)] +#[derive(FromBytes, IntoBytes, KnownLayout, Immutable, Clone, Copy)] +struct UdpHdr { + src_port: U16, + dst_port: U16, + _length: U16, + _checksum: U16, +} + +const ETH_HDR_LEN: usize = 14; +const IPV6_HDR_LEN: usize = 40; pub fn parse_packet(packet_data: &[u8]) -> Option<(Event, usize)> { - if packet_data.len() < 14 { - return None; - } - - let eth_type = u16::from_be_bytes([packet_data[12], packet_data[13]]); + let (eth_hdr, ip_rest) = Ref::<&[u8], EthHdr>::from_prefix(packet_data).ok()?; let timestamp_us = time::SystemTime::now() .duration_since(time::UNIX_EPOCH) .ok()? .as_micros() as u64; - match eth_type { - 0x0800 => parse_ipv4(packet_data, timestamp_us), - 0x86DD => parse_ipv6(packet_data, timestamp_us), + match eth_hdr.ether_type.get() { + 0x0800 => parse_ipv4(ip_rest, timestamp_us), + 0x86DD => parse_ipv6(ip_rest, timestamp_us), _ => None, } } -fn parse_ipv4(packet_data: &[u8], timestamp_us: u64) -> Option<(Event, usize)> { - if packet_data.len() < 34 { - return None; - } - - let ip_header = &packet_data[14..]; - - let protocol_byte = ip_header[9]; - let protocol = unsafe { mem::transmute::(protocol_byte) }; - - let src_ip = u32::from_be_bytes([ip_header[12], ip_header[13], ip_header[14], ip_header[15]]); - let dst_ip = u32::from_be_bytes([ip_header[16], ip_header[17], ip_header[18], ip_header[19]]); +fn parse_ipv4(ip_bytes: &[u8], timestamp_us: u64) -> Option<(Event, usize)> { + let (ipv4_hdr, _) = Ref::<&[u8], Ipv4Hdr>::from_prefix(ip_bytes).ok()?; + let protocol_byte = ipv4_hdr.protocol; if protocol_byte != 6 && protocol_byte != 17 { return None; } + let protocol = unsafe { std::mem::transmute::(protocol_byte) }; - let ihl = (ip_header[0] & 0x0F) as usize * 4; - let total_len = u16::from_be_bytes([ip_header[2], ip_header[3]]) as u32; + let src_ip = ipv4_hdr.src_addr.get(); + let dst_ip = ipv4_hdr.dst_addr.get(); + let total_len = ipv4_hdr.total_len.get() as u32; + let ihl = (ipv4_hdr.version_ihl & 0x0F) as usize * 4; - if packet_data.len() < 14 + ihl + 4 { + if ip_bytes.len() < ihl + 4 { return None; } + let transport_bytes = &ip_bytes[ihl..]; - let transport_header = &ip_header[ihl..]; - let src_port = u16::from_be_bytes([transport_header[0], transport_header[1]]); - let dst_port = u16::from_be_bytes([transport_header[2], transport_header[3]]); - - let (tcp_flags, tcp_window_size, header_length) = if protocol_byte == 6 { - if packet_data.len() < 14 + ihl + 20 { - return None; - } - - let data_offset = (transport_header[12] >> 4) as u16 * 4; - let flags = TcpFlags::from_byte(transport_header[13]); - let window = u16::from_be_bytes([transport_header[14], transport_header[15]]); - - (flags, window, data_offset) - } else if protocol_byte == 17 { - (TcpFlags::default(), 0, 8) + let (src_port, dst_port, tcp_flags, tcp_window_size, header_length) = if protocol_byte == 6 { + let (tcp_hdr, _) = Ref::<&[u8], TcpHdr>::from_prefix(transport_bytes).ok()?; + let dof = tcp_hdr.data_off_flags.get(); + let data_offset = (dof >> 12) as u16 * 4; + let flags = TcpFlags::from_byte((dof & 0xFF) as u8); + ( + tcp_hdr.src_port.get(), + tcp_hdr.dst_port.get(), + flags, + tcp_hdr.window.get(), + data_offset, + ) } else { - (TcpFlags::default(), 0, 0) + let (udp_hdr, _) = Ref::<&[u8], UdpHdr>::from_prefix(transport_bytes).ok()?; + ( + udp_hdr.src_port.get(), + udp_hdr.dst_port.get(), + TcpFlags::default(), + 0u16, + 8u16, + ) }; let payload_length = total_len.saturating_sub(ihl as u32 + header_length as u32); - let payload_start = 14 + ihl + header_length as usize; + let payload_start = ETH_HDR_LEN + ihl + header_length as usize; - let event = IPv4Event { - protocol, - src_ip, - dst_ip, - src_port, - dst_port, - packet_length: total_len, - payload_length, - header_length, - timestamp_us, - tcp_flags, - tcp_window_size, - is_forward: false, - }; - - Some((Event::IPv4(event), payload_start)) + Some(( + Event::IPv4(IPv4Event { + protocol, + src_ip, + dst_ip, + src_port, + dst_port, + packet_length: total_len, + payload_length, + header_length, + timestamp_us, + tcp_flags, + tcp_window_size, + is_forward: false, + }), + payload_start, + )) } -fn parse_ipv6(packet_data: &[u8], timestamp_us: u64) -> Option<(Event, usize)> { - if packet_data.len() < 54 { - return None; - } - - let ip_header = &packet_data[14..]; - - let protocol_byte = ip_header[6]; - let protocol = unsafe { mem::transmute::(protocol_byte) }; - - let mut source_ip_bytes = [0u8; 16]; - source_ip_bytes.copy_from_slice(&ip_header[8..24]); - let src_ip = u128::from_be_bytes(source_ip_bytes); - - let mut dest_ip_bytes = [0u8; 16]; - dest_ip_bytes.copy_from_slice(&ip_header[24..40]); - let dst_ip = u128::from_be_bytes(dest_ip_bytes); +fn parse_ipv6(ip_bytes: &[u8], timestamp_us: u64) -> Option<(Event, usize)> { + let (ipv6_hdr, transport_bytes) = Ref::<&[u8], Ipv6Hdr>::from_prefix(ip_bytes).ok()?; + let protocol_byte = ipv6_hdr.next_hdr; if protocol_byte != 6 && protocol_byte != 17 { return None; } + let protocol = unsafe { std::mem::transmute::(protocol_byte) }; - let payload_len = u16::from_be_bytes([ip_header[4], ip_header[5]]) as u32; + let src_ip = u128::from_be_bytes(ipv6_hdr.src_addr); + let dst_ip = u128::from_be_bytes(ipv6_hdr.dst_addr); + let payload_len = ipv6_hdr.payload_len.get() as u32; let total_len = payload_len + 40; - if packet_data.len() < 54 + 4 { + if transport_bytes.len() < 4 { return None; } - let transport_header = &ip_header[40..]; - let src_port = u16::from_be_bytes([transport_header[0], transport_header[1]]); - let dst_port = u16::from_be_bytes([transport_header[2], transport_header[3]]); - - let (tcp_flags, tcp_window_size, header_length) = if protocol_byte == 6 { - if packet_data.len() < 54 + 20 { - return None; - } - - let data_offset = (transport_header[12] >> 4) as u16 * 4; - let flags = TcpFlags::from_byte(transport_header[13]); - let window = u16::from_be_bytes([transport_header[14], transport_header[15]]); - - (flags, window, data_offset) - } else if protocol_byte == 17 { - (TcpFlags::default(), 0, 8) + let (src_port, dst_port, tcp_flags, tcp_window_size, header_length) = if protocol_byte == 6 { + let (tcp_hdr, _) = Ref::<&[u8], TcpHdr>::from_prefix(transport_bytes).ok()?; + let dof = tcp_hdr.data_off_flags.get(); + let data_offset = (dof >> 12) as u16 * 4; + let flags = TcpFlags::from_byte((dof & 0xFF) as u8); + ( + tcp_hdr.src_port.get(), + tcp_hdr.dst_port.get(), + flags, + tcp_hdr.window.get(), + data_offset, + ) } else { - (TcpFlags::default(), 0, 0) + let (udp_hdr, _) = Ref::<&[u8], UdpHdr>::from_prefix(transport_bytes).ok()?; + ( + udp_hdr.src_port.get(), + udp_hdr.dst_port.get(), + TcpFlags::default(), + 0u16, + 8u16, + ) }; let payload_length = total_len.saturating_sub(40 + header_length as u32); - let payload_start = 14 + 40 + header_length as usize; + let payload_start = ETH_HDR_LEN + IPV6_HDR_LEN + header_length as usize; - let event = IPv6Event { - protocol, - src_ip, - dst_ip, - src_port, - dst_port, - packet_length: total_len, - payload_length, - header_length, - timestamp_us, - tcp_flags, - tcp_window_size, - is_forward: false, - }; - - Some((Event::IPv6(event), payload_start)) + Some(( + Event::IPv6(IPv6Event { + protocol, + src_ip, + dst_ip, + src_port, + dst_port, + packet_length: total_len, + payload_length, + header_length, + timestamp_us, + tcp_flags, + tcp_window_size, + is_forward: false, + }), + payload_start, + )) } -pub fn format_ipv4(addr: u32) -> String { +pub fn format_ipv4(addr: u32) -> CompactString { let bytes = addr.to_be_bytes(); - format!("{}.{}.{}.{}", bytes[0], bytes[1], bytes[2], bytes[3],) + format_compact!("{}.{}.{}.{}", bytes[0], bytes[1], bytes[2], bytes[3]) } -pub fn format_ipv6(addr: u128) -> String { +pub fn format_ipv6(addr: u128) -> CompactString { let bytes = addr.to_be_bytes(); - format!( + format_compact!( "{:02x}{:02x}:{:02x}{:02x}:{:02x}{:02x}:{:02x}{:02x}:{:02x}{:02x}:{:02x}{:02x}:{:02x}{:02x}:{:02x}{:02x}", bytes[0], bytes[1], @@ -182,6 +231,6 @@ pub fn format_ipv6(addr: u128) -> String { bytes[12], bytes[13], bytes[14], - bytes[15] + bytes[15], ) }