ParrotXray/Mantis
1
0
Fork 0
mirror of https://github.com/ParrotXray/Mantis.git synced 2026-08-24 19:00:27 +09:00
Code Issues Packages Projects Releases Wiki Activity

Add bytes/ahash/compact_str for packet path and ML pipeline

Browse Source 
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0138PxtKH73hqxv7h1oaoSdS
This commit is contained in:
Claude 2026-06-16 22:31:36 +00:00
parent 88cab716e3
commit fdf086b960
No known key found for this signature in database
12 changed files with 95 additions and 62 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
Cargo.lock generated
View File

@ -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",

3
Cargo.toml
View File

@ -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"] }

3
mantis/Cargo.toml
View File

@ -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"] }

23
mantis/src/core/ebpf/xsk_manager.rs
View File

@ -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::<Arc<Vec<u8>>>(config.channel_size);
let (egress_to_ingress_tx, egress_to_ingress_rx) = bounded::<Arc<Vec<u8>>>(config.channel_size);
let (ingress_to_egress_tx, ingress_to_egress_rx) = bounded::<Bytes>(config.channel_size);
let (egress_to_ingress_tx, egress_to_ingress_rx) = bounded::<Bytes>(config.channel_size);
let ingress_xsk = XskPair::new(
config.clone(),
@ -258,8 +259,8 @@ impl XskPair {
pub fn run(
mut self,
forward_tx: Sender<Arc<Vec<u8>>>,
forward_rx: Receiver<Arc<Vec<u8>>>,
forward_tx: Sender<Bytes>,
forward_rx: Receiver<Bytes>,
) -> Result<oneshot::Sender<()>, 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<Arc<Vec<u8>>>) -> Result<usize, EbpfError> {
fn process_rx_queue(&mut self, forward_tx: &Sender<Bytes>) -> Result<usize, EbpfError> {
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<Arc<Vec<u8>>>) -> Result<usize, EbpfError> {
fn process_tx_queue(&mut self, forward_rx: &Receiver<Bytes>) -> Result<usize, EbpfError> {
// 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<Arc<Vec<u8>>> = Vec::with_capacity(max_to_send);
let mut packets_to_send: Vec<Bytes> = 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)?;
}
}

15
mantis/src/detection/ml/aggregator.rs
View File

@ -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<FlowKey, Vec<(Instant, f32)>>,
src_ip_states: HashMap<String, SrcIpState>,
detections: AHashMap<FlowKey, Vec<(Instant, f32)>>,
src_ip_states: AHashMap<CompactString, SrcIpState>,
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,

8
mantis/src/detection/ml/engine.rs
View File

@ -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<String> = flows
let active_ips: AHashSet<CompactString> = 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<CompactString> = AHashSet::new();
for result in &results {
if result.is_attack {
// L1: full 5-tuple aggregation for persistent same-port attacks

8
mantis/src/detection/ml/feature_extractor.rs
View File

@ -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) {

7
mantis/src/detection/ml/flow_tracker.rs
View File

@ -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<FlowKey, FlowData>,
flows: AHashMap<FlowKey, FlowData>,
max_flows: usize,
}
impl FlowTracker {
pub fn new(max_flows: usize) -> Self {
Self {
flows: HashMap::new(),
flows: AHashMap::new(),
max_flows,
}
}

14
mantis/src/detection/ml/inference.rs
View File

@ -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<HashMap<String, VecDeque<Vec<f32>>>>,
flow_buffers: Mutex<AHashMap<CompactString, VecDeque<Vec<f32>>>>,
}
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<String>) {
pub fn cleanup_buffers(&self, active_src_ips: &AHashSet<CompactString>) {
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::<Vec<_>>()
.join("\n");
log!(MLLog::WindowDebug(
flow.flow_key.src_ip.clone(),
flow.flow_key.src_ip.to_string(),
pad,
window_size,
ae_score,

7
mantis/src/detection/suricata/engine.rs
View File

@ -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<Arc<Vec<u8>>>,
tx: Sender<Bytes>,
child: std::sync::Mutex<Child>,
}
@ -133,7 +134,7 @@ impl SuricataEngine {
})
.map_err(|e| SuricataError::ProcessSpawnFailed { reason: e.to_string() })?;
let (tx, rx) = bounded::<Arc<Vec<u8>>>(CHANNEL_CAP);
let (tx, rx) = bounded::<Bytes>(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<Vec<u8>>) {
pub fn inject(&self, data: Bytes) {
match self.tx.try_send(data) {
Ok(()) => {}
Err(crossbeam::channel::TrySendError::Full(_)) => {

13
mantis/src/model/ml_detection.rs
View File

@ -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,

29
mantis/src/utils/packet_parser.rs
View File

@ -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],
)
}