diff --git a/mantis/src/core/ebpf/xsk_manager.rs b/mantis/src/core/ebpf/xsk_manager.rs index f4bb5c9..a28d99c 100644 --- a/mantis/src/core/ebpf/xsk_manager.rs +++ b/mantis/src/core/ebpf/xsk_manager.rs @@ -6,12 +6,11 @@ use std::sync::Arc; use std::thread; use std::time::Duration; -use libc; - use aya::Ebpf; use aya::maps::{MapData, XskMap}; use crossbeam::channel::{Receiver, Sender, bounded}; use crossbeam::queue::SegQueue; +use libc; use macros::log; use parking_lot::Mutex; use tokio::sync::oneshot; @@ -177,7 +176,10 @@ impl XskPair { XskBindMode::Copy => "copy", XskBindMode::Zero => "zero-copy", }; - log!(EbpfLog::XskBindModeSet { mode: mode_str.to_string(), queue_id }); + log!(EbpfLog::XskBindModeSet { + mode: mode_str.to_string(), + queue_id + }); let socket_config = SocketConfig::builder() .tx_queue_size(tx_queue_size) .rx_queue_size(rx_queue_size) @@ -218,7 +220,10 @@ impl XskPair { ); log!(EbpfLog::BusyPollEnabled { queue_id }); } else { - log!(EbpfLog::BusyPollUnavailable { queue_id, errno: *libc::__errno_location() }); + log!(EbpfLog::BusyPollUnavailable { + queue_id, + errno: *libc::__errno_location() + }); } } @@ -396,7 +401,8 @@ impl XskPair { // them out once space is available. let produced = unsafe { let (fill, rx) = (&mut self.fill_queue, &mut self.rx); - fill.produce_and_wakeup(&rx_descs[..rx_count], rx.fd_mut(), 0).unwrap_or(0) + fill.produce_and_wakeup(&rx_descs[..rx_count], rx.fd_mut(), 0) + .unwrap_or(0) }; if produced < rx_count { let mut pool = self.frame_pool.lock(); @@ -432,7 +438,8 @@ impl XskPair { let end = (offset + BATCH).min(all_frames.len()); let added = unsafe { let (fill, rx) = (&mut self.fill_queue, &mut self.rx); - fill.produce_and_wakeup(&all_frames[offset..end], rx.fd_mut(), 0).unwrap_or(0) + fill.produce_and_wakeup(&all_frames[offset..end], rx.fd_mut(), 0) + .unwrap_or(0) }; if added == 0 { break; // fill queue full @@ -508,4 +515,4 @@ impl XskPair { Ok(nb_submitted) } -} \ No newline at end of file +} diff --git a/mantis/src/detection/ml/engine.rs b/mantis/src/detection/ml/engine.rs index 2d05b18..4a02169 100644 --- a/mantis/src/detection/ml/engine.rs +++ b/mantis/src/detection/ml/engine.rs @@ -89,7 +89,7 @@ impl Engine { continue; }; let total_flows = t.flow_count(); - let flows = t.drain_ready_flows(self.flow_timeout_us); + let flows = t.drain_ready_flows(self.flow_timeout_us, self.flow_timeout_us / 2); 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. diff --git a/mantis/src/detection/ml/flow_tracker.rs b/mantis/src/detection/ml/flow_tracker.rs index 339c463..34ee07f 100644 --- a/mantis/src/detection/ml/flow_tracker.rs +++ b/mantis/src/detection/ml/flow_tracker.rs @@ -231,6 +231,40 @@ impl FlowData { pub fn packet_count(&self) -> usize { self.fwd_packets.len() + self.bwd_packets.len() } + + // Reset per-window counters so this entry can be reused for the next + // export window without removing it from the flow table. + // Flow key, direction, and last_time_us are preserved. + // init_win_bytes are cleared — they are only meaningful at connection start. + pub fn reset(&mut self, now_us: u64) { + self.fwd_packets.clear(); + self.bwd_packets.clear(); + self.fwd_total_bytes = 0; + self.bwd_total_bytes = 0; + self.fwd_header_bytes = 0; + self.bwd_header_bytes = 0; + self.fin_count = 0; + self.syn_count = 0; + self.rst_count = 0; + self.psh_count = 0; + self.ack_count = 0; + self.urg_count = 0; + self.cwe_count = 0; + self.ece_count = 0; + self.init_win_bytes_fwd = 0; + self.init_win_bytes_bwd = 0; + self.active_periods.clear(); + self.idle_periods.clear(); + self.last_packet_time = now_us; + self.last_time_us = now_us; + self.fwd_bulk_state = BulkState::default(); + self.bwd_bulk_state = BulkState::default(); + self.act_data_pkt_fwd = 0; + self.start_time_us = now_us; + self.active_start_us = now_us; + self.fwd_fin_seen = false; + self.bwd_fin_seen = false; + } } pub struct FlowTracker { @@ -312,23 +346,44 @@ impl FlowTracker { } } - pub fn drain_ready_flows(&mut self, timeout_us: u64) -> Vec { + // Export flows that are finished or idle. + // Active flows alive >= force_interval_us are cloned and exported, + // then reset in place so the next window starts accumulating immediately. + // This ensures active attack flows reach ML/logger within force_interval_us + // and keeps training and inference feature distributions identical. + pub fn drain_ready_flows(&mut self, timeout_us: u64, force_interval_us: u64) -> Vec { let now = time::SystemTime::now() .duration_since(time::UNIX_EPOCH) .map(|d| d.as_micros() as u64) .unwrap_or(0); - let ready_keys: Vec = self + let terminal_keys: Vec = self .flows .iter() .filter(|(_, flow)| flow.is_finished() || now.saturating_sub(flow.last_time_us) >= timeout_us) .map(|(key, _)| key.clone()) .collect(); - ready_keys + let mut result: Vec = terminal_keys .into_iter() .filter_map(|key| self.flows.remove(&key)) - .collect() + .filter(|flow| flow.packet_count() > 0) + .collect(); + + for flow in self.flows.values_mut() { + if now.saturating_sub(flow.start_time_us) >= force_interval_us { + if flow.packet_count() > 0 { + result.push(flow.clone()); + flow.reset(now); + } else { + // No packets in this window; advance start so we don't re-check + // every tick until the next packet arrives. + flow.start_time_us = now; + } + } + } + + result } pub fn active_src_ips(&self) -> impl Iterator { diff --git a/mantis/src/model/config.rs b/mantis/src/model/config.rs index c698384..242a8d0 100644 --- a/mantis/src/model/config.rs +++ b/mantis/src/model/config.rs @@ -114,4 +114,4 @@ fn default_fusion_window_secs() -> u64 { fn default_ae_threshold_method() -> String { "95".to_string() -} \ No newline at end of file +} diff --git a/mantis/src/model/log/ebpf.rs b/mantis/src/model/log/ebpf.rs index 9336b27..de87a71 100644 --- a/mantis/src/model/log/ebpf.rs +++ b/mantis/src/model/log/ebpf.rs @@ -63,4 +63,4 @@ loggable! { #[error("SO_PREFER_BUSY_POLL not supported on this kernel/driver (queue {queue_id}, errno {errno})")] BusyPollUnavailable { queue_id: u32, errno: i32 } => tracing::Level::WARN, } -} \ No newline at end of file +}