diff --git a/Cargo.lock b/Cargo.lock index 8f3a08d..dec9c6b 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" @@ -473,6 +482,21 @@ dependencies = [ "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]] name = "core-error" version = "0.0.0" @@ -1506,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", diff --git a/Cargo.toml b/Cargo.toml index bd4a9aa..d2cd17f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,11 +4,14 @@ 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"] } diff --git a/mantis/Cargo.toml b/mantis/Cargo.toml index 939d20d..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"] } diff --git a/mantis/src/core/ebpf/xsk_manager.rs b/mantis/src/core/ebpf/xsk_manager.rs index 7bab57e..7197476 100644 --- a/mantis/src/core/ebpf/xsk_manager.rs +++ b/mantis/src/core/ebpf/xsk_manager.rs @@ -6,6 +6,7 @@ use std::sync::Arc; use std::thread; use std::time::Duration; +use bytes::Bytes; use aya::Ebpf; use aya::maps::{MapData, XskMap}; use crossbeam::channel::{Receiver, Sender, bounded}; @@ -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,7 +364,7 @@ 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) }; @@ -378,10 +379,10 @@ impl XskPair { engine.process_packet(packet_slice, self.direction == Direction::Ingress); } - // Single copy shared between Suricata and forward_tx via Arc - let packet_data = Arc::new(packet_slice.to_vec()); + // 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(Arc::clone(&packet_data)); + se.inject(packet_data.clone()); } if let Err(e) = forward_tx.try_send(packet_data) { match e { @@ -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>> = 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..32858e2 100644 --- a/mantis/src/detection/ml/feature_extractor.rs +++ b/mantis/src/detection/ml/feature_extractor.rs @@ -26,8 +26,8 @@ impl FlowFeatures { Self { features, feature_num, - src_ip: flow.flow_key.src_ip.clone(), - dst_ip: flow.flow_key.dst_ip.clone(), + src_ip: flow.flow_key.src_ip.to_string(), + dst_ip: flow.flow_key.dst_ip.to_string(), timestamp: flow.start_time_us, } } @@ -221,8 +221,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..11c2866 100644 --- a/mantis/src/detection/ml/flow_tracker.rs +++ b/mantis/src/detection/ml/flow_tracker.rs @@ -1,6 +1,7 @@ -use std::collections::HashMap; use std::time; +use ahash::AHashMap; + use common::model::event::Event; use super::server_ports; @@ -268,14 +269,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..88819a8 100644 --- a/mantis/src/detection/ml/inference.rs +++ b/mantis/src/detection/ml/inference.rs @@ -1,8 +1,10 @@ -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 +26,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 +38,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 +105,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 +127,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 1331095..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: Arc>) { + 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 1f2faec..b9b32e8 100644 --- a/mantis/src/utils/packet_parser.rs +++ b/mantis/src/utils/packet_parser.rs @@ -1,6 +1,7 @@ use std::time; use common::model::event::{Event, IPv4Event, IPv6Event, TcpFlags}; +use compact_str::{format_compact, CompactString}; use network_types::ip::IpProto; use zerocopy::byteorder::{BigEndian, U16, U32}; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout, Ref}; @@ -182,30 +183,18 @@ fn parse_ipv6(ip_bytes: &[u8], timestamp_us: u64) -> Option<(Event, usize)> { )) } -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], - bytes[2], - bytes[3], - bytes[4], - bytes[5], - bytes[6], - bytes[7], - bytes[8], - bytes[9], - bytes[10], - bytes[11], - bytes[12], - bytes[13], - bytes[14], - bytes[15], + bytes[0], bytes[1], bytes[2], bytes[3], + bytes[4], bytes[5], bytes[6], bytes[7], + bytes[8], bytes[9], bytes[10], bytes[11], + bytes[12], bytes[13], bytes[14], bytes[15], ) }