feat(core): BYO ML adapter state machine + Fusion v1 foundation

Cross-cutting rework that lands the multi-source NIDS+IPS fusion groundwork
alongside a manifest-driven ML adapter. Day 1 with no model loaded is a
first-class state, and the remaining 3 sources (Suricata + CV + Graph) fuse
on a shared vocabulary.

ML pipeline (BYO-ready)
- Replace the hardcoded `MLModels { deep_autoencoder, classifier, batch_size }`
  struct with `MLModelAdapter` enum (AutoencoderOnly / ClassifierOnly / MultiTask),
  wrapped in `ModelSourceState { Dormant | Active { adapter, info } | Error { msg,
  since, last_attempted_path } }` and held inside `ArcSwap` on `Inference`.
  Day 1 without a manifest resolves to `Dormant`; reload failures park in
  `Error` carrying the reason and offending path.
- `MultiTask` arm now reads the c2 label index from the manifest instead of
  matching the hardcoded "C2 Communication" string, and replaces the
  `unwrap_or(0.0)` sentinel with an explicit `if let && let && c2 > c2_class_prob`
  chain.
- Aggregator `should_alert` takes `required_confirmations: usize` and
  `alert_multiplier: f32` directly from manifest `labels[*].confirmations` /
  `thresholds.alert_multiplier`; the per-attack-type string match and
  `ALERT_THRESHOLD_MULTIPLIER` constant are gone.
- Manifest validation rejects `confirmations: 0`, case-insensitive duplicate
  label names, and `alert_multiplier <= 0`.
- New `ModelSourceStatus` wire type (serde-tagged Dormant/Active/Error) for
  HTTP / WebSocket consumers.

Fusion core
- Canonical attack-type dictionary (13 seed types + `Unknown` bucket) with
  per-source `translate(source, raw_label)` and `canonical_from_str`. ML /
  Beaconing / Correlation / Suricata raw labels fold into a shared vocabulary
  so the orchestrator's `(source_ip, canonical_attack_type)` dedup key
  actually collides across sources.
- `fusion_math::fused_confidence(&[f32]) -> f32` implements `1 - prod(1 - c_i)`,
  clamped to `[0.0, 1.0]` and safe for empty input.
- `FusionWindowLengths` pins dynamic per-first-source window lengths
  (Suricata 10s / CV 8s / ML 2s / Graph 5s) with `[1, 30]` second clamp.
- Orchestrator canonicalizes raw labels on ingress, opens a fusion window
  keyed by (src_ip, canonical), and silences duplicates for 30s after emit.
- SOAR `Condition::MultiSourceMin { min }` and
  `Condition::SingleSourceHigh { source, min_confidence }` implement the
  "N sources agreed" gate and solo-high-confidence escape hatch.
  `SoarEngine::reload_cache` canonicalizes `trigger_event` on load and logs
  `SoarLog::NonCanonicalTriggerEvent` for anything outside the vocabulary.

Workspace hygiene
- `MODELS_DIR` / `MANIFEST_FILENAME` / `STAGING_SUBDIR` hoisted to
  `model::config::constants`.
- `DetectionSource` gains `FromStr`; `soar/matcher.rs` string matches on
  source become parser-driven.
- Magic numbers in `soar/matcher.rs` (cooldown, min confidence, TTL) and
  `fusion_math::for_source` (window clamp) replaced with named constants.
- CODE_STYLE inline-path violations (3-segment paths to `BTreeMap`,
  `Duration`, `CreateKind`) hoisted to `use`.

Tests: 198 pass (new: 9 fusion_math, 11 attack_type dictionary, 7 adapter
confirmations lookup, 3 model_source serde roundtrip, manifest validator
coverage). \`cargo clippy --package net-guardia -- -D warnings\` clean.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
DaLaw2 2026-04-18 16:05:11 +08:00
parent 231fb88efd
commit 56395f8027
26 changed files with 2182 additions and 548 deletions

View File

@ -1,7 +1,6 @@
# NetGuardia v10 model manifest — authored 2026-04-16 as part of v12 M1.
# Schema: structural/semantic info lives here; preprocessing arrays
# (scaler mean/std, clip params, feature weights) stay in inference_config.json
# referenced by `preprocessing.scaler_sidecar` below.
# NetGuardia model manifest. Structural/semantic fields live here;
# preprocessing arrays (scaler mean/std, clip params, feature weights) stay
# in the JSON sidecar referenced by `preprocessing.scaler_sidecar`.
name: netguardia-v10
adapter: multi_task
@ -46,9 +45,12 @@ features:
- fwd_bwd_bytes_ratio
- iat_cv
# `confirmations` / `playbook` are parsed today; M2 wires them into the aggregator.
# `confirmations` sets the per-class aggregator firing threshold. Classes
# with single-shot semantics (C2 / Bot / DNS tunneling / exploit) use 1 so
# the aggregator alerts on the first detection; noisier classes can raise
# it (DoS/DDoS: 2). Absent entries fall back to the engine default.
labels:
"0": { name: Bot }
"0": { name: Bot, confirmations: 1 }
"1": { name: Brute Force }
"2": { name: C2 Communication, confirmations: 1 }
"3": { name: DNS Tunneling, confirmations: 1 }
@ -64,6 +66,9 @@ thresholds:
c2: 0.9085615873336792
class_min_confidence: 0.4
ae: 0.23011694848537445
# Average score must exceed `class_min_confidence * alert_multiplier`
# before the aggregator fires. Raising this suppresses borderline hits.
alert_multiplier: 1.2
preprocessing:
scaler_sidecar: inference_config.json

View File

@ -0,0 +1,135 @@
//! Fusion policy math primitive — cross-source confidence aggregation.
//!
//! The policy assumes detection sources are conditionally independent given
//! a true attack. In practice ML and CV can share signal on C2 beaconing,
//! so a future calibration pass may introduce per-pair weights; this
//! module stays the canonical home for whichever formula is in force.
use crate::model::event::DetectionSource;
/// Compute `1 ∏(1 c_i)` over the given per-source confidences.
///
/// - Empty input → `0.0` (no evidence).
/// - Single input → returns that confidence unchanged.
/// - Values are clamped to `[0.0, 1.0]` to keep the result bounded even if
/// an upstream source ships noisy unnormalized scores.
pub fn fused_confidence(per_source: &[f32]) -> f32 {
if per_source.is_empty() {
return 0.0;
}
let mut inverse: f64 = 1.0;
for &c in per_source {
let clamped = (c as f64).clamp(0.0, 1.0);
inverse *= 1.0 - clamped;
}
(1.0 - inverse).clamp(0.0, 1.0) as f32
}
/// Default per-source fusion-window length in seconds. Each value scales
/// the orchestrator's lookahead budget when that source opens a dedup key.
/// Slower sources (Suricata signatures) get longer windows so a follow-up
/// ML hit still lands inside; faster sources (ML ticks) use short windows
/// because they'd otherwise waste latency waiting on downstream signals.
#[derive(Debug, Clone, Copy)]
pub struct FusionWindowLengths {
pub suricata_secs: u64,
pub cv_secs: u64,
pub ml_secs: u64,
pub graph_secs: u64,
}
/// Valid range, in seconds, for a fusion window. Clamps protect against a
/// misconfigured source opening a wedged (too-long) or useless (zero) key.
pub const FUSION_WINDOW_MIN_SECS: u64 = 1;
pub const FUSION_WINDOW_MAX_SECS: u64 = 30;
impl Default for FusionWindowLengths {
fn default() -> Self {
Self {
suricata_secs: 10,
cv_secs: 8,
ml_secs: 2,
graph_secs: 5,
}
}
}
impl FusionWindowLengths {
/// Range-clamped lookup for the window length of the first source to
/// open a fusion key.
pub fn for_source(&self, source: DetectionSource) -> u64 {
let raw = match source {
DetectionSource::Suricata => self.suricata_secs,
DetectionSource::Beaconing => self.cv_secs,
DetectionSource::ML => self.ml_secs,
DetectionSource::Correlation => self.graph_secs,
};
raw.clamp(FUSION_WINDOW_MIN_SECS, FUSION_WINDOW_MAX_SECS)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn empty_fusion_returns_zero() {
assert_eq!(fused_confidence(&[]), 0.0);
}
#[test]
fn single_source_passes_through() {
assert!((fused_confidence(&[0.8]) - 0.8).abs() < 1e-5);
assert_eq!(fused_confidence(&[0.0]), 0.0);
assert_eq!(fused_confidence(&[1.0]), 1.0);
}
#[test]
fn two_sources_boost() {
// 1 (1 0.7)(1 0.9) = 1 0.03 = 0.97
let got = fused_confidence(&[0.7, 0.9]);
assert!((got - 0.97).abs() < 1e-5);
}
#[test]
fn three_sources_monotone_in_count() {
let two = fused_confidence(&[0.6, 0.7]);
let three = fused_confidence(&[0.6, 0.7, 0.5]);
assert!(three >= two);
}
#[test]
fn four_sources_bounded() {
let four = fused_confidence(&[0.9, 0.8, 0.7, 0.6]);
assert!(four < 1.0);
assert!(four > 0.99);
}
#[test]
fn out_of_range_confidences_are_clamped() {
// Negative / above-1 inputs don't break the math.
assert_eq!(fused_confidence(&[-0.5, -0.1]), 0.0);
assert_eq!(fused_confidence(&[2.0, 3.0]), 1.0);
}
#[test]
fn window_lengths_defaults() {
let w = FusionWindowLengths::default();
assert_eq!(w.for_source(DetectionSource::Suricata), 10);
assert_eq!(w.for_source(DetectionSource::Beaconing), 8);
assert_eq!(w.for_source(DetectionSource::ML), 2);
assert_eq!(w.for_source(DetectionSource::Correlation), 5);
}
#[test]
fn window_lengths_clamped_to_range() {
let w = FusionWindowLengths {
suricata_secs: 999,
cv_secs: 0,
ml_secs: 2,
graph_secs: 5,
};
assert_eq!(w.for_source(DetectionSource::Suricata), 30);
assert_eq!(w.for_source(DetectionSource::Beaconing), 1);
}
}

View File

@ -1,2 +1,3 @@
pub mod beaconing;
pub mod fusion_math;
pub mod orchestrator;

View File

@ -7,43 +7,65 @@ use macros::log;
use tokio::sync::mpsc;
use tokio::time::interval;
use super::fusion_math::{FusionWindowLengths, fused_confidence};
use crate::infrastructure::communication_manager::CommunicationManager;
use crate::infrastructure::geoip::GeoIpService;
use crate::model::detection::attack_type::translate;
use crate::model::error::system::SystemError;
use crate::model::event::{DetectionEvent, DetectionSource, ThreatDetectedEvent};
use crate::model::log::detection::DetectionLog;
/// Dedup window: detections for the same (source_ip, attack_type) within this window
/// are suppressed after the first emission.
/// Dedup window: detections for the same `(source_ip, canonical_attack_type)`
/// within this window are suppressed after initial fusion-window expiry.
const DEDUP_WINDOW_SECS: u64 = 30;
/// How often to sweep expired dedup entries.
const CLEANUP_INTERVAL_SECS: u64 = 60;
/// Repeat offender detection: same IP within this duration counts as repeat.
const REPEAT_OFFENDER_WINDOW_SECS: u64 = 2 * 60 * 60; // 2 hours
const REPEAT_OFFENDER_WINDOW_SECS: u64 = 2 * 60 * 60;
/// Maximum dedup entries to prevent unbounded memory growth under sustained attack.
const MAX_DEDUP_ENTRIES: usize = 50_000;
struct DedupEntry {
sources: Vec<DetectionSource>,
emitted_at: Instant,
/// Per-source record within an in-flight dedup entry. Keeps the strongest
/// confidence per source so multi-hit from one source doesn't inflate the
/// fused policy.
#[derive(Debug, Clone)]
struct SourceSample {
source: DetectionSource,
confidence: f32,
}
/// Coordinates detections from multiple sources (ML, future: rules, correlation, threat feeds).
/// Deduplicates, enriches with GeoIP/hit count/repeat offender, and emits ThreatDetectedEvent.
struct DedupEntry {
sources: Vec<SourceSample>,
/// When the orchestrator first emitted for this key.
first_emitted_at: Instant,
/// Most recent emit (reset on each fused re-emit within the fusion window).
emitted_at: Instant,
/// Fusion window length for this key, fixed by the first source to arrive.
/// Later arrivals don't reset it so the lookahead budget stays predictable.
fusion_window: Duration,
}
/// Coordinates detections from ML / Suricata / Beaconing / Correlation.
/// Incoming events are canonicalized into a shared attack-type dictionary so
/// dedup keys collide across sources; events sharing a key inside the fusion
/// window accumulate, and additional sources arriving mid-window trigger a
/// re-emit with the combined confidence `1 ∏(1 c_i)`.
pub struct DetectionOrchestrator {
rx: mpsc::Receiver<DetectionEvent>,
comm: Arc<CommunicationManager>,
geoip: Option<Arc<GeoIpService>>,
// Enrichment state
// SAFETY: NonZero::new on a literal is infallible.
src_ip_counts: lru::LruCache<String, u32>,
repeat_tracker: lru::LruCache<String, Instant>,
// Dedup state — LRU-bounded to prevent unbounded growth under sustained attack
// Dedup state — LRU-bounded to prevent unbounded growth under sustained attack.
dedup: lru::LruCache<(String, String), DedupEntry>,
dedup_window: Duration,
/// Per-source fusion window lengths. Future versions may read overrides
/// from DB; the defaults live in `FusionWindowLengths::default`.
fusion_windows: FusionWindowLengths,
}
impl DetectionOrchestrator {
@ -61,10 +83,10 @@ impl DetectionOrchestrator {
repeat_tracker: LruCache::new(NonZero::new(5_000).unwrap()),
dedup: LruCache::new(NonZero::new(MAX_DEDUP_ENTRIES).unwrap()),
dedup_window: Duration::from_secs(DEDUP_WINDOW_SECS),
fusion_windows: FusionWindowLengths::default(),
}
}
/// Spawn the orchestrator as a background task.
pub fn start(self) {
tokio::spawn(async move { self.run().await });
}
@ -79,7 +101,7 @@ impl DetectionOrchestrator {
event = self.rx.recv() => {
match event {
Some(detection) => self.handle_detection(detection).await,
None => break, // All senders dropped
None => break,
}
}
_ = cleanup_interval.tick() => {
@ -89,47 +111,115 @@ impl DetectionOrchestrator {
}
}
async fn handle_detection(&mut self, event: DetectionEvent) {
async fn handle_detection(&mut self, mut event: DetectionEvent) {
// Canonicalize the raw attack_type so Suricata's "brute-force"
// classtype and ML's "Brute Force" class name land on the same dedup
// key — the precondition for cross-source fusion.
let canonical = translate(event.source, &event.attack_type);
event.attack_type = canonical.as_str().to_string();
let key = (event.source_ip.clone(), event.attack_type.clone());
let now = Instant::now();
// Dedup check
if let Some(entry) = self.dedup.get(&key)
&& now.checked_duration_since(entry.emitted_at).unwrap_or(Duration::ZERO) < self.dedup_window
{
// Within window: add source attribution but don't re-emit
if !entry.sources.contains(&event.source) {
// Re-get as mutable to update sources
if let Some(entry) = self.dedup.get_mut(&key) {
entry.sources.push(event.source.clone());
// Path A: existing dedup entry. Decide re-emit (fusion window still
// open) vs silence (window closed but dedup still active).
if let Some(entry) = self.dedup.get_mut(&key) {
let since_first = now.saturating_duration_since(entry.first_emitted_at);
// Post-dedup-window — treat as a brand-new event (fall through).
if since_first >= self.dedup_window {
// Expired dedup; fall through to Path B by dropping the entry.
self.dedup.pop(&key);
} else if since_first < entry.fusion_window {
// Still inside the fusion window — accumulate.
let is_new_source = !entry.sources.iter().any(|s| s.source == event.source);
if is_new_source {
entry.sources.push(SourceSample {
source: event.source,
confidence: event.confidence,
});
entry.emitted_at = now;
} else {
// Same source firing again inside the window — keep the
// strongest confidence for fusion math.
if let Some(existing) = entry.sources.iter_mut().find(|s| s.source == event.source)
&& existing.confidence < event.confidence
{
existing.confidence = event.confidence;
}
}
// Only RE-EMIT when a new source joined — same-source
// refires are silenced to avoid SOAR cooldown churn.
if is_new_source {
self.emit_fused(&event, &key).await;
} else {
log!(DetectionLog::DetectionDeduplicated(
event.source_ip.clone(),
event.attack_type.clone(),
));
}
return;
} else {
// Past fusion window, still inside dedup silence → drop.
log!(DetectionLog::DetectionDeduplicated(
event.source_ip.clone(),
event.attack_type.clone(),
));
return;
}
log!(DetectionLog::DetectionDeduplicated(event.source_ip, event.attack_type,));
return;
}
// Enrich and emit
let threat_event = self.enrich(&event).await;
let sources = vec![event.source.clone()];
log!(DetectionLog::DetectionEmitted(
event.source_ip.clone(),
event.attack_type.clone(),
event.confidence,
event.ae_score,
event.anomaly_score,
event.c2_score,
sources.len(),
));
// Record dedup entry (LRU-bounded)
// Path B: brand-new key (or expired dedup). Emit single-source,
// open a fusion window sized by this source.
let fusion_window = Duration::from_secs(self.fusion_windows.for_source(event.source));
self.dedup.put(
key,
key.clone(),
DedupEntry {
sources,
sources: vec![SourceSample {
source: event.source,
confidence: event.confidence,
}],
first_emitted_at: now,
emitted_at: now,
fusion_window,
},
);
self.emit_fused(&event, &key).await;
}
/// Build the fused ThreatDetectedEvent from the current dedup entry's
/// per-source samples, apply enrichment (hit count / repeat / geoip),
/// and publish. Called both on first emit (single source) and on
/// within-window re-emit (2..=4 sources).
async fn emit_fused(&mut self, trigger_event: &DetectionEvent, key: &(String, String)) {
let (sources_vec, fused, per_source_confs): (Vec<DetectionSource>, f32, Vec<f32>) = {
let entry = match self.dedup.get(key) {
Some(e) => e,
None => return,
};
let confs: Vec<f32> = entry.sources.iter().map(|s| s.confidence).collect();
let sources: Vec<DetectionSource> = entry.sources.iter().map(|s| s.source).collect();
let fused = fused_confidence(&confs);
(sources, fused, confs)
};
let _ = per_source_confs; // currently only used for the log line below
let mut threat_event = self.enrich(trigger_event).await;
threat_event.sources = sources_vec;
threat_event.active_source_count = threat_event.sources.len();
threat_event.fused_confidence = fused;
threat_event.confidence = fused;
log!(DetectionLog::DetectionEmitted(
trigger_event.source_ip.clone(),
trigger_event.attack_type.clone(),
threat_event.confidence,
trigger_event.ae_score,
trigger_event.anomaly_score,
trigger_event.c2_score,
threat_event.active_source_count,
));
if let Err(e) = self.comm.publish_event(threat_event).await {
log!(SystemError::MlSoarBridgeFailed(e));
@ -139,14 +229,12 @@ impl DetectionOrchestrator {
async fn enrich(&mut self, event: &DetectionEvent) -> ThreatDetectedEvent {
let src_ip = &event.source_ip;
// Compute packet rate
let packet_rate = if event.flow_duration_us > 0 {
event.packet_count as f64 / (event.flow_duration_us as f64 / 1_000_000.0)
} else {
0.0
};
// Update hit count (LRU bounded)
let hit_count = match self.src_ip_counts.get_mut(src_ip) {
Some(c) => {
*c = c.saturating_add(1);
@ -158,7 +246,6 @@ impl DetectionOrchestrator {
}
};
// Check repeat offender (same IP within window)
let repeat_window = Duration::from_secs(REPEAT_OFFENDER_WINDOW_SECS);
let now = Instant::now();
let is_repeat = self
@ -167,7 +254,6 @@ impl DetectionOrchestrator {
.is_some_and(|last| now.checked_duration_since(*last).unwrap_or(Duration::ZERO) < repeat_window);
self.repeat_tracker.put(src_ip.clone(), now);
// GeoIP lookup
let geoip_country = if let Some(ref svc) = self.geoip {
if let Ok(ip) = src_ip.parse() {
svc.lookup(ip).await.ok().flatten().and_then(|loc| loc.country_code)
@ -188,7 +274,12 @@ impl DetectionOrchestrator {
protocol: event.protocol,
geoip_country,
is_repeat_offender: is_repeat,
sources: vec![event.source.clone()],
// These three get overwritten in `emit_fused` with the
// accumulated values; initialize to the single-source defaults
// so a direct caller also gets a consistent shape.
sources: vec![event.source],
active_source_count: 1,
fused_confidence: event.confidence,
ae_score: event.ae_score,
anomaly_score: event.anomaly_score,
c2_score: event.c2_score,
@ -198,9 +289,12 @@ impl DetectionOrchestrator {
fn cleanup_expired(&mut self) {
let now = Instant::now();
let window = self.dedup_window;
// Pop expired entries from the LRU (oldest entries are least recently used)
while let Some((_, entry)) = self.dedup.peek_lru() {
if now.checked_duration_since(entry.emitted_at).unwrap_or(Duration::ZERO) >= window {
if now
.checked_duration_since(entry.first_emitted_at)
.unwrap_or(Duration::ZERO)
>= window
{
self.dedup.pop_lru();
} else {
break;
@ -208,3 +302,11 @@ impl DetectionOrchestrator {
}
}
}
#[cfg(test)]
mod tests {
//! Orchestrator integration tests go here once we have an in-memory
//! CommunicationManager harness. The fusion math is covered by
//! `fusion_math::tests` and canonical translation by
//! `model::detection::attack_type::tests`.
}

View File

@ -0,0 +1,235 @@
//! `MLModelAdapter` — the three ONNX adapter shapes the inference pipeline
//! understands, plus the wrapping `ModelSourceState` machine used inside
//! `Inference::models: ArcSwap<ModelSourceState>`.
//!
//! Day 1 is `Dormant` (no model loaded); the happy path is `Active { adapter,
//! info }`; failed loads park in `Error { msg, since, last_attempted_path }`.
use std::collections::BTreeMap;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::SystemTime;
use super::manifest::LabelSpec;
use crate::model::detection::ml_detection::RunnableModel;
use crate::model::detection::model_source::{ModelInfo, ModelSourceStatus};
/// Compile-time sanity: `RunnableModel` must be `Send + Sync` because we
/// stuff it inside an `ArcSwap`. If a future `tract-onnx` upgrade silently
/// drops the bounds, this line stops compiling and we catch it before it
/// becomes a production data race.
const _: fn() = || {
fn assert_send_sync<T: Send + Sync>() {}
assert_send_sync::<RunnableModel>();
};
/// One of three adapter shapes. Each variant carries the RunnableModel(s)
/// it needs plus the metadata required to dispatch inference:
/// batch size, feature counts, label map (classifier / multi-task only).
#[derive(Clone)]
pub enum MLModelAdapter {
/// Pure anomaly-detection autoencoder. Output is a single MSE-style
/// reconstruction error per flow; classifier-less. v1 product main path.
AutoencoderOnly {
model: Arc<RunnableModel>,
batch_size: usize,
n_features: usize,
},
/// Single classifier, no AE. Output is a per-class softmax tensor.
/// Labels come from the manifest's `labels` map; `normal_idx` is
/// pre-resolved so `infer_batch` doesn't re-scan on every tick.
ClassifierOnly {
model: Arc<RunnableModel>,
batch_size: usize,
n_features: usize,
labels: BTreeMap<String, LabelSpec>,
normal_idx: Option<usize>,
},
/// AE + classifier in one adapter. Classifier consumes (ae_features, ae_score).
/// Three output tensors: anomaly_score / class_probs / c2_score.
MultiTask {
ae: Arc<RunnableModel>,
classifier: Arc<RunnableModel>,
batch_size: usize,
n_ae: usize,
n_cls: usize,
labels: BTreeMap<String, LabelSpec>,
normal_idx: Option<usize>,
c2_idx: Option<usize>,
},
}
impl MLModelAdapter {
/// Per-attack-type confirmations lookup: resolve the manifest's
/// `labels[*].confirmations` value whose label name matches
/// `attack_type_name` (case-insensitive). Callers fall back to their own
/// default when this returns `None` — `AutoencoderOnly` carries no labels,
/// and classifier labels may legitimately omit the override.
pub fn confirmations_for(&self, attack_type_name: &str) -> Option<usize> {
match self {
Self::MultiTask { labels, .. } | Self::ClassifierOnly { labels, .. } => {
confirmations_from_labels(labels, attack_type_name)
}
Self::AutoencoderOnly { .. } => None,
}
}
}
/// Pure-data lookup extracted so unit tests can cover the matching rules
/// without constructing a full `MLModelAdapter` (which requires a real
/// `RunnableModel` — expensive and fragile to mock).
fn confirmations_from_labels(labels: &BTreeMap<String, LabelSpec>, attack_type_name: &str) -> Option<usize> {
labels
.values()
.find(|spec| spec.name.eq_ignore_ascii_case(attack_type_name))
.and_then(|spec| spec.confirmations)
}
/// Internal state machine. Held inside `ArcSwap<ModelSourceState>` so the
/// inference pipeline can check state once per tick without locks and
/// atomically swap on upload / deletion / reload failure.
///
/// Variants are deliberately **not** `Serialize` — the `Active` variant
/// holds an `Arc<RunnableModel>` which isn't serde-friendly. Callers that
/// need a wire representation call `to_status()` to produce the lightweight
/// `ModelSourceStatus` consumed by WebSocket / HTTP.
pub enum ModelSourceState {
/// Day 1 default. `models/` empty, or admin deleted the current model.
/// Drift detector becomes a no-op; fusion math still runs with the
/// remaining 3 sources.
Dormant,
/// Model loaded, inference ticks consume flows.
Active { adapter: MLModelAdapter, info: ModelInfo },
/// Last load attempt failed — schema mismatch, timeout, corrupted
/// ONNX. Inference skips; UI renders the reason. Swap to Dormant /
/// Active via normal reload path.
Error {
msg: String,
since: SystemTime,
last_attempted_path: Option<PathBuf>,
},
}
impl ModelSourceState {
pub fn is_active(&self) -> bool {
matches!(self, Self::Active { .. })
}
/// Wire-format snapshot for UI / HTTP. Never borrows the adapter — the
/// returned value is safe to send across WebSocket boundaries.
pub fn to_status(&self) -> ModelSourceStatus {
match self {
Self::Dormant => ModelSourceStatus::Dormant,
Self::Active { info, .. } => ModelSourceStatus::Active { info: info.clone() },
Self::Error {
msg,
since,
last_attempted_path,
} => ModelSourceStatus::Error {
msg: msg.clone(),
since_secs: since
.duration_since(SystemTime::UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0),
last_attempted_path: last_attempted_path.as_ref().map(|p| p.display().to_string()),
},
}
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use super::*;
#[test]
fn dormant_state_serializes_to_dormant_status() {
let state = ModelSourceState::Dormant;
let status = state.to_status();
assert!(status.is_dormant());
}
fn sample_labels() -> BTreeMap<String, LabelSpec> {
let mut m = BTreeMap::new();
m.insert(
"0".into(),
LabelSpec {
name: "Bot".into(),
confirmations: Some(1),
playbook: None,
},
);
m.insert(
"4".into(),
LabelSpec {
name: "DoS/DDoS".into(),
confirmations: Some(2),
playbook: None,
},
);
m.insert(
"7".into(),
LabelSpec {
name: "Normal".into(),
confirmations: None,
playbook: None,
},
);
m
}
#[test]
fn confirmations_lookup_hits_exact_name() {
let labels = sample_labels();
assert_eq!(confirmations_from_labels(&labels, "Bot"), Some(1));
assert_eq!(confirmations_from_labels(&labels, "DoS/DDoS"), Some(2));
}
#[test]
fn confirmations_lookup_is_case_insensitive() {
let labels = sample_labels();
assert_eq!(confirmations_from_labels(&labels, "bot"), Some(1));
assert_eq!(confirmations_from_labels(&labels, "dos/ddos"), Some(2));
}
#[test]
fn confirmations_lookup_label_without_override_returns_none() {
let labels = sample_labels();
assert_eq!(confirmations_from_labels(&labels, "Normal"), None);
}
#[test]
fn confirmations_lookup_unknown_name_returns_none() {
let labels = sample_labels();
assert_eq!(confirmations_from_labels(&labels, "Phantom"), None);
}
#[test]
fn confirmations_lookup_empty_labels_returns_none() {
let labels = BTreeMap::new();
assert_eq!(confirmations_from_labels(&labels, "Bot"), None);
}
#[test]
fn error_state_preserves_details() {
let state = ModelSourceState::Error {
msg: "shape mismatch".to_string(),
since: SystemTime::UNIX_EPOCH + Duration::from_secs(1_700_000_000),
last_attempted_path: Some(PathBuf::from("models/bad.onnx")),
};
let status = state.to_status();
match status {
ModelSourceStatus::Error {
msg,
since_secs,
last_attempted_path,
} => {
assert_eq!(msg, "shape mismatch");
assert_eq!(since_secs, 1_700_000_000);
assert_eq!(last_attempted_path.as_deref(), Some("models/bad.onnx"));
}
_ => panic!("expected Error"),
}
}
}

View File

@ -1,3 +1,9 @@
//! Per-flow detection aggregator.
//!
//! Records per-key hits within a rolling time window; callers decide when to
//! fire based on how many hits a given attack class needs (manifest-driven)
//! and how far the rolling-average score beats the confidence threshold.
use std::collections::HashMap;
use std::time::{Duration, Instant};
@ -6,40 +12,40 @@ use crate::model::detection::ml_detection::FlowKey;
pub struct AttackAggregator {
detections: HashMap<FlowKey, Vec<(Instant, f32)>>,
window_duration: Duration,
min_detections: usize,
alert_threshold_multiplier: f32,
}
impl AttackAggregator {
pub fn new(window_secs: u64, min_detections: usize) -> Self {
pub fn new(window_secs: u64) -> Self {
Self {
detections: HashMap::new(),
window_duration: Duration::from_secs(window_secs),
min_detections,
alert_threshold_multiplier: 1.2,
}
}
pub fn should_alert(&mut self, flow_key: &FlowKey, score: f32, threshold: f32, attack_type: Option<&str>) -> bool {
/// Record a detection and decide whether the flow should fire an alert.
///
/// - `required_confirmations`: in-window hit count the flow must reach.
/// Resolved by the caller from the active manifest's per-label value;
/// validated at manifest load to be ≥ 1, so no runtime floor is needed.
/// - `alert_multiplier`: scales `threshold` before the average-score
/// comparison, driven by the manifest's `thresholds.alert_multiplier`.
pub fn should_alert(
&mut self,
flow_key: &FlowKey,
score: f32,
threshold: f32,
required_confirmations: usize,
alert_multiplier: f32,
) -> bool {
let now = Instant::now();
let detections = self.detections.entry(flow_key.clone()).or_default();
detections.retain(|(time, _)| now.duration_since(*time) < self.window_duration);
detections.push((now, score));
// Per-attack-type adaptive min_detections:
// DoS/DDoS: high frequency, need more confirmations to avoid alert storms
// C2/Bot/DNS Tunneling/Exploitation: low frequency or single-shot, alert immediately
let effective_min = match attack_type {
Some("DoS/DDoS") => self.min_detections.saturating_mul(2).max(1),
Some("C2 Communication") | Some("Bot") | Some("DNS Tunneling") | Some("Exploitation") => 1,
_ => self.min_detections,
};
if detections.len() >= effective_min {
if detections.len() >= required_confirmations {
let avg_score: f32 = detections.iter().map(|(_, s)| s).sum::<f32>() / detections.len() as f32;
return avg_score > threshold * self.alert_threshold_multiplier;
return avg_score > threshold * alert_multiplier;
}
false
@ -58,6 +64,8 @@ impl AttackAggregator {
mod tests {
use super::*;
const TEST_MULTIPLIER: f32 = 1.2;
fn test_key() -> FlowKey {
FlowKey {
src_ip: [192, 168, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
@ -70,61 +78,66 @@ mod tests {
}
#[test]
fn default_attack_type_uses_base_min_detections() {
let mut agg = AttackAggregator::new(60, 3);
fn required_three_fires_on_third_hit() {
let mut agg = AttackAggregator::new(60);
let key = test_key();
// Need 3 detections for default type
assert!(!agg.should_alert(&key, 5.0, 1.0, Some("Brute Force")));
assert!(!agg.should_alert(&key, 5.0, 1.0, Some("Brute Force")));
assert!(agg.should_alert(&key, 5.0, 1.0, Some("Brute Force")));
assert!(!agg.should_alert(&key, 5.0, 1.0, 3, TEST_MULTIPLIER));
assert!(!agg.should_alert(&key, 5.0, 1.0, 3, TEST_MULTIPLIER));
assert!(agg.should_alert(&key, 5.0, 1.0, 3, TEST_MULTIPLIER));
}
#[test]
fn dos_ddos_requires_double_min_detections() {
let mut agg = AttackAggregator::new(60, 3);
fn required_one_fires_immediately() {
let mut agg = AttackAggregator::new(60);
let key = test_key();
assert!(agg.should_alert(&key, 5.0, 1.0, 1, TEST_MULTIPLIER));
}
#[test]
fn required_six_needs_six_hits() {
let mut agg = AttackAggregator::new(60);
let key = test_key();
// DoS/DDoS needs 6 detections (3 * 2)
for _ in 0..5 {
assert!(!agg.should_alert(&key, 5.0, 1.0, Some("DoS/DDoS")));
assert!(!agg.should_alert(&key, 5.0, 1.0, 6, TEST_MULTIPLIER));
}
assert!(agg.should_alert(&key, 5.0, 1.0, Some("DoS/DDoS")));
assert!(agg.should_alert(&key, 5.0, 1.0, 6, TEST_MULTIPLIER));
}
#[test]
fn c2_alerts_on_first_detection() {
let mut agg = AttackAggregator::new(60, 3);
fn average_score_at_or_below_scaled_threshold_does_not_fire() {
let mut agg = AttackAggregator::new(60);
let key = test_key();
// C2 Communication alerts immediately (min=1)
assert!(agg.should_alert(&key, 5.0, 1.0, Some("C2 Communication")));
// score 1.0, threshold 1.0, multiplier 1.2 → gate is 1.2; 1.0 misses.
assert!(!agg.should_alert(&key, 1.0, 1.0, 1, TEST_MULTIPLIER));
}
#[test]
fn exploitation_alerts_on_first_detection() {
let mut agg = AttackAggregator::new(60, 3);
fn larger_multiplier_raises_the_bar() {
let mut agg = AttackAggregator::new(60);
let key = test_key();
assert!(agg.should_alert(&key, 5.0, 1.0, Some("Exploitation")));
// multiplier 2.5, threshold 1.0 → gate is 2.5; score 2.0 misses.
assert!(!agg.should_alert(&key, 2.0, 1.0, 1, 2.5));
}
#[test]
fn bot_alerts_on_first_detection() {
let mut agg = AttackAggregator::new(60, 3);
fn cleanup_preserves_fresh_entries() {
let mut agg = AttackAggregator::new(60);
let key = test_key();
assert!(agg.should_alert(&key, 5.0, 1.0, Some("Bot")));
agg.should_alert(&key, 5.0, 1.0, 10, TEST_MULTIPLIER);
assert!(agg.detections.contains_key(&key));
agg.cleanup();
assert!(agg.detections.contains_key(&key));
}
#[test]
fn dns_tunneling_alerts_on_first_detection() {
let mut agg = AttackAggregator::new(60, 3);
let key = test_key();
assert!(agg.should_alert(&key, 5.0, 1.0, Some("DNS Tunneling")));
}
#[test]
fn none_attack_type_uses_default() {
let mut agg = AttackAggregator::new(60, 3);
let key = test_key();
assert!(!agg.should_alert(&key, 5.0, 1.0, None));
assert!(!agg.should_alert(&key, 5.0, 1.0, None));
assert!(agg.should_alert(&key, 5.0, 1.0, None));
fn independent_flows_track_separately() {
let mut agg = AttackAggregator::new(60);
let key_a = test_key();
let mut key_b = test_key();
key_b.dst_port = 81;
assert!(!agg.should_alert(&key_a, 5.0, 1.0, 2, TEST_MULTIPLIER));
assert!(!agg.should_alert(&key_b, 5.0, 1.0, 2, TEST_MULTIPLIER));
assert!(agg.should_alert(&key_a, 5.0, 1.0, 2, TEST_MULTIPLIER));
assert!(agg.should_alert(&key_b, 5.0, 1.0, 2, TEST_MULTIPLIER));
}
}

View File

@ -95,12 +95,9 @@ fn reconcile_features(
}
fn apply_manifest_overrides(manifest: &ModelManifest, config: &mut MLInferenceConfig) {
// Labels: the manifest is authoritative for the label-name map. `confirmations`
// and `playbook` are parsed but not consumed here — M2's aggregator rewrite owns that.
if !manifest.labels.is_empty() {
config.attack_labels = manifest_labels_to_map(&manifest.labels);
}
// Thresholds: manifest overrides the sidecar when present.
if let Some(v) = manifest.thresholds.anomaly {
config.anomaly_threshold = v;
}
@ -113,6 +110,9 @@ fn apply_manifest_overrides(manifest: &ModelManifest, config: &mut MLInferenceCo
if let Some(v) = manifest.thresholds.ae {
config.ae_threshold = v;
}
if let Some(v) = manifest.thresholds.alert_multiplier {
config.alert_threshold_multiplier = v;
}
}
fn manifest_labels_to_map(labels: &BTreeMap<String, LabelSpec>) -> HashMap<String, String> {
@ -166,6 +166,7 @@ mod tests {
anomaly_threshold: 0.5,
c2_threshold: 0.5,
class_min_confidence: 0.4,
alert_threshold_multiplier: 1.2,
model_type: "MultiTaskModel".into(),
output_names: vec!["anomaly".into(), "class_probs".into(), "c2_score".into()],
ae_feature_weights: HashMap::new(),

View File

@ -1,3 +1,9 @@
//! ML engine — orchestrates flow tracking, feature extraction, adapter-
//! dispatched inference, and alert broadcast. Each inference tick cleans up
//! stale flows, gathers the latest batch, optionally logs a Flow Trace row,
//! and — when the inference source is Active — updates drift and runs
//! adapter-dispatched inference.
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
@ -12,14 +18,17 @@ use super::alert::MLAlert;
use super::drift_detector::DriftDetector;
use super::flow_tracker::{FlowData, FlowTracker};
use super::inference::Inference;
use super::model_loader::MLModels;
use super::traffic_logger::TrafficLogger;
use crate::interface::port::packet_sink::{PacketSink, PacketSinkFactory};
use crate::model::detection::flow_features::FlowFeatures;
use crate::model::detection::ml_detection::{EngineConfig, FlowKey, InferenceStats};
use crate::model::log::ml::MLLog;
use crate::model::monitoring::user_packet::UserPacket;
use crate::model::system::config::MLInferenceConfig;
/// Divisor applied to "ticks per aggregator window" to derive the fallback
/// confirmations count: 2 means a default-behaviour detection must fire
/// across at least half the window's ticks before alerting.
const DEFAULT_CONFIRMATION_WINDOW_FRACTION: u64 = 2;
/// Per-queue tracker. With symmetric hash in eBPF, both directions of a flow
/// land on the same queue, so per-queue trackers correctly see bidirectional flows.
@ -32,29 +41,29 @@ pub struct Engine {
drift_detector: Arc<Mutex<DriftDetector>>,
ml_alert: Arc<MLAlert>,
min_packets: usize,
/// Confirmations count used when the active manifest's label has no
/// explicit `confirmations` override.
default_confirmations: usize,
batch_size: usize,
inference_interval_secs: u64,
traffic_logger: Option<Arc<TrafficLogger>>,
}
impl Engine {
/// Build an Engine around an already-constructed Inference pipeline.
/// The Inference's state (Dormant / Active / Error) is consulted per tick.
pub fn new(
models: Arc<MLModels>,
config: Arc<MLInferenceConfig>,
inference_pipeline: Arc<Inference>,
ml_alert: Arc<MLAlert>,
drift_detector: Arc<Mutex<DriftDetector>>,
engine_config: EngineConfig,
traffic_logger: Option<Arc<TrafficLogger>>,
num_threads: u32,
) -> Self {
let inference_pipeline = Arc::new(Inference::new(models, config));
let min_detections =
((engine_config.aggregator_window_secs / engine_config.inference_interval_secs) / 2).max(1) as usize;
let aggregator = Mutex::new(AttackAggregator::new(
engine_config.aggregator_window_secs,
min_detections,
));
let interval_secs = engine_config.inference_interval_secs.max(1);
let ticks_per_window = engine_config.aggregator_window_secs / interval_secs;
let default_confirmations = (ticks_per_window / DEFAULT_CONFIRMATION_WINDOW_FRACTION).max(1) as usize;
let aggregator = Mutex::new(AttackAggregator::new(engine_config.aggregator_window_secs));
let max_flows_per_thread = engine_config.max_flows / (num_threads as usize).max(1);
let trackers: Vec<ThreadTracker> = (0..num_threads)
@ -68,16 +77,13 @@ impl Engine {
drift_detector,
ml_alert,
min_packets: engine_config.min_packets,
default_confirmations,
batch_size: engine_config.batch_size,
inference_interval_secs: engine_config.inference_interval_secs,
traffic_logger,
}
}
pub fn inference_pipeline(&self) -> &Arc<Inference> {
&self.inference_pipeline
}
/// xsk_manager calls this per queue_id; with symmetric hash each queue has its own tracker.
pub fn tracker(&self, queue_id: u32) -> &ThreadTracker {
&self.trackers[queue_id as usize % self.trackers.len()]
@ -103,21 +109,15 @@ impl Engine {
1 => 1,
// UDP
17 => match flow_key.dst_port {
// DNS: amplification, DGA, tunneling can be 1-2 packets
53 => 1,
// NTP amplification
123 => 2,
// Known mining pool ports
3333 | 45700 => 2,
_ => global,
},
// TCP
6 => match flow_key.dst_port {
// DNS over TCP
53 => 2,
// Common C2 ports: Metasploit, Cobalt Strike, reverse shells
4444 | 8443 | 8080 | 1337 | 31337 => 2,
// Stratum mining
3333 | 45700 => 2,
_ => global,
},
@ -142,7 +142,6 @@ impl Engine {
_ = ticker.tick() => {}
}
// Move CPU-bound ML inference off the tokio executor
let engine = Arc::clone(&self);
let _ = spawn_blocking(move || {
engine.run_inference_tick();
@ -151,6 +150,10 @@ impl Engine {
}
}
/// Phased tick: clean up stale flows, gather the uninferred batch
/// (min-packet filtered), optionally write a Flow Trace row, and when
/// the inference source is Active update drift then run inference.
/// Dormant / Error states short-circuit before drift + inference.
fn run_inference_tick(&self) {
let mut all_flows = Vec::new();
let mut total_count = 0;
@ -159,13 +162,11 @@ impl Engine {
.map(|d| d.as_micros() as u64)
.unwrap_or(0);
// Phase 0: clean up stale / terminated flows
for tracker in &self.trackers {
let mut t = tracker.lock();
t.cleanup_stale_flows(now_us);
}
// Phase 1: short lock per tracker — clone uninferred flows, mark as inferred
for tracker in &self.trackers {
let mut t = tracker.lock();
total_count += t.flow_count();
@ -174,7 +175,6 @@ impl Engine {
flow.packet_count() >= Self::effective_min_packets(&flow.flow_key, self.min_packets)
}),
);
// lock released here
}
log!(MLLog::FlowStats(
@ -190,9 +190,14 @@ impl Engine {
if let Some(ref logger) = self.traffic_logger {
self.log_traffic(&all_flows, logger);
} else {
self.run_inference(&all_flows);
}
if !self.inference_pipeline.is_active() {
return;
}
self.update_drift(&all_flows);
self.run_inference(&all_flows);
}
fn log_traffic(&self, flows: &[FlowData], logger: &TrafficLogger) {
@ -203,32 +208,36 @@ impl Engine {
}
}
/// Update drift baseline — only called when inference state is Active,
/// guaranteeing the ArcSwap snapshot the next `infer_batch` sees matches
/// the features we just normalized against.
fn update_drift(&self, batch: &[FlowData]) {
let batch = &batch[..batch.len().min(self.batch_size)];
let config = &self.inference_pipeline.config;
let mut dd = self.drift_detector.lock();
for flow in batch {
let features = FlowFeatures::extract(flow, &config.ae_feature_names);
let normalized: Vec<f64> = features
.features
.iter()
.zip(config.ae_scaler_mean.iter().zip(config.ae_scaler_std.iter()))
.map(|(&val, (&mean, &std))| if std.abs() > 1e-12 { (val - mean) / std } else { 0.0 })
.collect();
dd.update(&normalized);
}
}
fn run_inference(&self, flows: &[FlowData]) {
let batch = &flows[..flows.len().min(self.batch_size)];
log!(MLLog::RunningInference(batch.len()));
// Feed normalized features into drift detector for each flow in the batch
{
let config = &self.inference_pipeline.config;
let mut dd = self.drift_detector.lock();
for flow in batch.iter() {
let features = FlowFeatures::extract(flow, &config.ae_feature_names);
let normalized: Vec<f64> = features
.features
.iter()
.zip(config.ae_scaler_mean.iter().zip(config.ae_scaler_std.iter()))
.map(|(&val, (&mean, &std))| if std.abs() > 1e-12 { (val - mean) / std } else { 0.0 })
.collect();
dd.update(&normalized);
}
}
let start = Instant::now();
let results = self.inference_pipeline.infer_batch(batch);
let elapsed_us = start.elapsed().as_micros() as u64;
let stats = InferenceStats::from_results(&results, elapsed_us);
self.inference_pipeline.record_tick_qps(stats.flows_per_second);
if results.len() != batch.len() {
log!(MLLog::InferenceResults(batch.len(), results.len()));
@ -242,39 +251,42 @@ impl Engine {
stats.flows_per_second
));
{
let mut aggregator = self.aggregator.lock();
for result in &results {
if result.is_attack {
let should_alert = aggregator.should_alert(
&result.flow_key_raw,
let mut aggregator = self.aggregator.lock();
let config = &self.inference_pipeline.config;
for result in &results {
if result.is_attack {
let required_confirmations = result
.attack_type
.as_deref()
.and_then(|at| self.inference_pipeline.confirmations_for_attack_type(at))
.unwrap_or(self.default_confirmations);
let should_alert = aggregator.should_alert(
&result.flow_key_raw,
result.confidence,
config.class_min_confidence,
required_confirmations,
config.alert_threshold_multiplier,
);
if should_alert {
log!(MLLog::ThreatDetected(
format!("{:?}", result.direction),
result.flow_key.clone(),
result.attack_type.clone().unwrap_or_else(|| "UNKNOWN".to_string()),
result.confidence,
self.inference_pipeline.config.class_min_confidence,
result.attack_type.as_deref(),
);
result.ae_score,
));
if should_alert {
log!(MLLog::ThreatDetected(
format!("{:?}", result.direction),
result.flow_key.clone(),
result.attack_type.clone().unwrap_or_else(|| "UNKNOWN".to_string()),
result.confidence,
result.ae_score,
));
self.ml_alert.broadcast_alert(result);
}
self.ml_alert.broadcast_alert(result);
}
}
aggregator.cleanup();
}
aggregator.cleanup();
}
}
/// Adapter that exposes one `ThreadTracker` (per AF_XDP queue) as a
/// `PacketSink`. `XskManager` holds `Arc<dyn PacketSink>` per queue and never
/// touches `FlowTracker` concrete types.
/// Adapter that exposes one `ThreadTracker` (per AF_XDP queue) as a `PacketSink`.
struct QueueTrackerSink {
tracker: ThreadTracker,
}

View File

@ -1,4 +1,16 @@
//! Inference pipeline — state-aware, adapter-dispatched.
//!
//! The pipeline holds `ArcSwap<ModelSourceState>` so the engine can take a
//! lock-free snapshot per tick and `infer_batch` dispatches on whichever
//! `MLModelAdapter` variant the active state carries.
//!
//! The MultiTask arm fires on any of: anomaly head above threshold,
//! classifier picking a non-Normal class at ≥ `class_min_confidence`, or the
//! C2 head elevated above its own threshold — so AE reconstruction error
//! can't mask a stealthy attack the classifier does recognize.
use std::cmp::Ordering as CmpOrdering;
use std::collections::BTreeMap;
use std::panic::{self, AssertUnwindSafe};
use std::sync::Arc;
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
@ -8,14 +20,16 @@ use arc_swap::ArcSwap;
use macros::log;
use tract_onnx::prelude::*;
use super::adapter::{MLModelAdapter, ModelSourceState};
use super::flow_tracker::FlowData;
use super::model_loader::MLModels;
use super::manifest::LabelSpec;
use crate::model::detection::flow_features::FlowFeatures;
use crate::model::detection::ml_detection::DetectionResult;
use crate::model::detection::ml_detection::{DetectionResult, RunnableModel};
use crate::model::detection::model_source::ModelSourceStatus;
use crate::model::log::ml::MLLog;
use crate::model::system::config::MLInferenceConfig;
/// (anomaly_scores, per_class_probs, c2_scores)
/// (anomaly_scores, per_class_probs, c2_scores) — MultiTask batch output.
type ClassifierBatchOutput = (Vec<f32>, Vec<Vec<f32>>, Vec<f32>);
/// Consecutive failures to trip the circuit breaker.
@ -26,56 +40,87 @@ const CIRCUIT_BREAKER_WINDOW_SECS: u64 = 60;
const CIRCUIT_BREAKER_COOLDOWN_SECS: u64 = 120;
pub struct Inference {
models: ArcSwap<MLModels>,
state: ArcSwap<ModelSourceState>,
pub config: Arc<MLInferenceConfig>,
c2_class_idx: Option<usize>,
normal_class_idx: Option<usize>,
/// Rolling-window QPS estimate, published in `ModelInfo.qps_recent`.
/// Stored as u32 (integer QPS) for lock-free update; fractional QPS
/// information is not useful at the UI grain we're publishing.
qps_recent: AtomicU32,
failure_count: AtomicU32,
failure_window_start: AtomicU64,
circuit_open_since: AtomicU64,
}
impl Inference {
pub fn new(models: Arc<MLModels>, config: Arc<MLInferenceConfig>) -> Self {
let c2_class_idx = Self::find_label_index(&config, "C2 Communication");
let normal_class_idx = Self::find_label_index(&config, "Normal");
/// Build a new Inference with the given initial state. Use
/// `ModelSourceState::Dormant` when no model is loaded (Day 1 default).
pub fn new(initial_state: ModelSourceState, config: Arc<MLInferenceConfig>) -> Self {
Self {
models: ArcSwap::from(models),
state: ArcSwap::from_pointee(initial_state),
config,
c2_class_idx,
normal_class_idx,
qps_recent: AtomicU32::new(0),
failure_count: AtomicU32::new(0),
failure_window_start: AtomicU64::new(0),
circuit_open_since: AtomicU64::new(0),
}
}
/// Atomically swap to new models. Old models are dropped when the last reader finishes.
pub fn swap_models(&self, new_models: Arc<MLModels>) {
self.models.store(new_models);
// Reset circuit breaker on successful model swap
self.failure_count.store(0, Ordering::Relaxed);
self.failure_window_start.store(0, Ordering::Relaxed);
self.circuit_open_since.store(0, Ordering::Relaxed);
/// Atomically swap in a new state. Any transition INTO `Active` resets
/// the circuit breaker so a freshly loaded model starts with a clean
/// failure record; `Error ↔ Dormant` transitions leave the CB alone.
pub fn swap_state(&self, new_state: ModelSourceState) {
let new_is_active = new_state.is_active();
self.state.store(Arc::new(new_state));
if new_is_active {
self.failure_count.store(0, Ordering::Relaxed);
self.failure_window_start.store(0, Ordering::Relaxed);
self.circuit_open_since.store(0, Ordering::Relaxed);
}
}
fn find_label_index(config: &MLInferenceConfig, label: &str) -> Option<usize> {
config
.attack_labels
.iter()
.find(|(_, v)| v.as_str() == label)
.and_then(|(k, _)| k.parse::<usize>().ok())
/// Current state snapshot for wire broadcast. Merges the in-memory
/// `qps_recent` atomic into the Active info so the UI sees live QPS.
pub fn current_status(&self) -> ModelSourceStatus {
let guard = self.state.load();
let mut status = guard.to_status();
if let ModelSourceStatus::Active { ref mut info } = status {
info.qps_recent = self.qps_recent.load(Ordering::Relaxed) as f32;
}
status
}
/// Batched inference with circuit breaker protection.
pub fn is_active(&self) -> bool {
self.state.load().is_active()
}
/// Per-attack-type confirmations from the active manifest's labels.
/// Returns `None` when the source is Dormant/Error or when the label has
/// no `confirmations` override; callers apply their own fallback.
pub fn confirmations_for_attack_type(&self, attack_type_name: &str) -> Option<usize> {
match self.state.load().as_ref() {
ModelSourceState::Active { adapter, .. } => adapter.confirmations_for(attack_type_name),
ModelSourceState::Dormant | ModelSourceState::Error { .. } => None,
}
}
/// Batched inference with circuit breaker protection + state dispatch.
/// Returns an empty Vec when Dormant / Error / circuit-open.
pub fn infer_batch(&self, flows: &[FlowData]) -> Vec<DetectionResult> {
if self.is_circuit_open() {
return Vec::new();
}
match panic::catch_unwind(AssertUnwindSafe(|| self.infer_batch_inner(flows))) {
// Snapshot the state once so the whole batch sees a consistent adapter.
let guard = self.state.load();
let adapter = match guard.as_ref() {
ModelSourceState::Active { adapter, .. } => adapter,
ModelSourceState::Dormant | ModelSourceState::Error { .. } => {
return Vec::new();
}
};
match panic::catch_unwind(AssertUnwindSafe(|| self.infer_batch_inner(adapter, flows))) {
Ok(results) => {
// Success: reset failure counter
self.failure_count.store(0, Ordering::Relaxed);
self.failure_window_start.store(0, Ordering::Relaxed);
results
@ -91,23 +136,68 @@ impl Inference {
}
}
fn infer_batch_inner(&self, flows: &[FlowData]) -> Vec<DetectionResult> {
let models = self.models.load();
let n = flows.len();
let batch_size = models.batch_size;
let n_ae = self.config.num_ae_features();
let n_cls = self.config.num_classifier_features();
/// Dispatch on adapter variant.
fn infer_batch_inner(&self, adapter: &MLModelAdapter, flows: &[FlowData]) -> Vec<DetectionResult> {
match adapter {
MLModelAdapter::MultiTask {
ae,
classifier,
batch_size,
n_ae,
n_cls,
labels,
normal_idx,
c2_idx,
} => self.infer_multitask(
ae,
classifier,
*batch_size,
*n_ae,
*n_cls,
labels,
*normal_idx,
*c2_idx,
flows,
),
MLModelAdapter::AutoencoderOnly {
model,
batch_size,
n_features,
} => self.infer_autoencoder_only(model, *batch_size, *n_features, flows),
MLModelAdapter::ClassifierOnly {
model,
batch_size,
n_features,
labels,
normal_idx,
} => self.infer_classifier_only(model, *batch_size, *n_features, labels, *normal_idx, flows),
}
}
/// MultiTask path. Runs the AE batch → computes per-flow MSE → feeds the
/// classifier over (ae_features ++ ae_score) → fires on anomaly OR
/// non-Normal classifier agreement OR elevated C2 head.
#[allow(clippy::too_many_arguments)]
fn infer_multitask(
&self,
ae: &RunnableModel,
classifier: &RunnableModel,
batch_size: usize,
n_ae: usize,
n_cls: usize,
labels: &BTreeMap<String, LabelSpec>,
normal_idx: Option<usize>,
c2_idx: Option<usize>,
flows: &[FlowData],
) -> Vec<DetectionResult> {
let n = flows.len();
// Phase 1: preprocess all AE features
let all_ae_features: Vec<Vec<f32>> = flows.iter().map(|f| self.preprocess_ae_features(f)).collect();
// Phase 2: run AE in chunks → compute per-flow MSE
let mut ae_scores = Vec::with_capacity(n);
for chunk_start in (0..n).step_by(batch_size) {
let chunk_end = (chunk_start + batch_size).min(n);
let actual = chunk_end - chunk_start;
// Build (batch_size, n_ae) tensor, zero-padded
let ae_input = tract_ndarray::Array2::<f32>::from_shape_fn((batch_size, n_ae), |(i, j)| {
if i < actual {
all_ae_features[chunk_start + i][j]
@ -116,7 +206,7 @@ impl Inference {
}
});
match Self::run_ae_batch(&models, &ae_input, actual, n_ae) {
match run_ae_batch(ae, &ae_input, actual, n_ae) {
Ok(scores) => ae_scores.extend_from_slice(&scores),
Err(e) => {
log!(MLLog::InferenceFailed("DeepAutoEncoder".to_string(), e.to_string()));
@ -126,7 +216,6 @@ impl Inference {
}
}
// Phase 3: build classifier input (ae_features ++ ae_score) and run in chunks
let mut all_anomaly = Vec::with_capacity(n);
let mut all_class_probs = Vec::with_capacity(n);
let mut all_c2_scores = Vec::with_capacity(n);
@ -135,7 +224,6 @@ impl Inference {
let chunk_end = (chunk_start + batch_size).min(n);
let actual = chunk_end - chunk_start;
// Build (batch_size, n_cls) tensor
let cls_input = tract_ndarray::Array2::<f32>::from_shape_fn((batch_size, n_cls), |(i, j)| {
if i < actual {
if j < n_ae {
@ -148,7 +236,7 @@ impl Inference {
}
});
match Self::run_classifier_batch(&models, &cls_input, actual) {
match run_classifier_batch(classifier, &cls_input, actual) {
Ok((anomaly, class_probs, c2)) => {
all_anomaly.extend_from_slice(&anomaly);
all_class_probs.extend(class_probs);
@ -162,159 +250,192 @@ impl Inference {
}
}
// Phase 4: build DetectionResults
let mut results = Vec::with_capacity(n);
let class_min_conf = self.config.class_min_confidence;
let anomaly_thr = self.config.anomaly_threshold;
let c2_thr = self.config.c2_threshold;
for i in 0..n {
let flow = &flows[i];
let flow_key = format!(
"{}:{} -> {}:{} (proto {}) [{}]",
flow.flow_key.src_ip_string(),
flow.flow_key.src_port,
flow.flow_key.dst_ip_string(),
flow.flow_key.dst_port,
flow.flow_key.protocol,
flow.direction
);
let class_probs = &all_class_probs[i];
let anomaly = all_anomaly[i];
let c2 = all_c2_scores[i];
let result = self.build_detection_result(
flow,
flow_key,
ae_scores[i],
all_anomaly[i],
&all_class_probs[i],
all_c2_scores[i],
);
results.push(result);
let (predicted_class, class_confidence) = argmax(class_probs);
let attack_type_name = labels
.get(&predicted_class.to_string())
.map(|l| l.name.clone())
.unwrap_or_else(|| "UNKNOWN".to_string());
let classifier_fires =
normal_idx.is_none_or(|ni| predicted_class != ni) && class_confidence >= class_min_conf;
let c2_fires = c2 > c2_thr;
let mut is_attack = anomaly > anomaly_thr || classifier_fires || c2_fires;
// "Normal" with no C2 elevation stays benign regardless of AE noise.
if normal_idx == Some(predicted_class) && !c2_fires {
is_attack = false;
}
// When the manifest declares a C2 class and its head score beats
// the classifier's probability for that same class, relabel the
// event with the manifest's C2 label and use the head score as
// the outgoing confidence. Manifests without a C2 class keep the
// argmax label and class-confidence untouched.
let mut attack_type = attack_type_name;
let mut confidence = class_confidence;
if c2_fires
&& let Some(idx) = c2_idx
&& let Some(c2_class_prob) = class_probs.get(idx).copied()
&& c2 > c2_class_prob
{
if let Some(spec) = labels.get(&idx.to_string()) {
attack_type = spec.name.clone();
}
confidence = c2;
}
results.push(DetectionResult {
flow_key: build_flow_key_label(flow),
flow_key_raw: flow.flow_key.clone(),
direction: flow.direction,
is_attack,
attack_type: if is_attack { Some(attack_type) } else { None },
confidence,
ae_score: ae_scores[i],
anomaly_score: anomaly,
c2_score: c2,
packet_count: flow.packet_count() as u64,
flow_duration_us: flow.duration_us(),
});
}
results
}
/// Run AE on a padded batch, return MSE scores for the first `actual` rows.
fn run_ae_batch(
models: &MLModels,
input: &tract_ndarray::Array2<f32>,
actual: usize,
n_features: usize,
) -> TractResult<Vec<f32>> {
let result = models.deep_autoencoder.run(tvec![input.clone().into_tensor().into()])?;
let output = result[0]
.to_array_view::<f32>()?
.into_dimensionality::<tract_ndarray::Ix2>()?;
let diff = input - &output;
let sq = &diff * &diff;
let mut scores = Vec::with_capacity(actual);
let n_f = n_features as f32;
for i in 0..actual {
let mse: f32 = sq.row(i).sum() / n_f;
scores.push(mse);
}
Ok(scores)
}
/// Returns (anomaly_scores, per_class_probs, c2_scores) for the first `actual` rows.
fn run_classifier_batch(
models: &MLModels,
input: &tract_ndarray::Array2<f32>,
actual: usize,
) -> TractResult<ClassifierBatchOutput> {
let result = models.classifier.run(tvec![input.clone().into_tensor().into()])?;
// Output 0: anomaly (batch_size, 1)
let anomaly_view = result[0].to_array_view::<f32>()?;
let anomaly: Vec<f32> = (0..actual)
.map(|i| anomaly_view.as_slice().map(|s| s[i]).unwrap_or(0.0))
.collect();
// Output 1: class_probs (batch_size, n_classes)
let class_view = result[1]
.to_array_view::<f32>()?
.into_dimensionality::<tract_ndarray::Ix2>()?;
let class_probs: Vec<Vec<f32>> = (0..actual)
.map(|i| class_view.row(i).iter().copied().collect())
.collect();
// Output 2: c2_score (batch_size, 1)
let c2_view = result[2].to_array_view::<f32>()?;
let c2: Vec<f32> = (0..actual)
.map(|i| c2_view.as_slice().map(|s| s[i]).unwrap_or(0.0))
.collect();
Ok((anomaly, class_probs, c2))
}
fn build_detection_result(
/// AutoencoderOnly path. Output is reconstruction MSE; when it exceeds
/// `anomaly_threshold` the flow is marked as a generic `anomaly`. There
/// are no classifier outputs, so no per-class gating happens here.
fn infer_autoencoder_only(
&self,
flow: &FlowData,
flow_key: String,
ae_score: f32,
anomaly_score: f32,
class_probs: &[f32],
c2_score: f32,
) -> DetectionResult {
let is_attack = anomaly_score > self.config.anomaly_threshold;
model: &RunnableModel,
batch_size: usize,
n_features: usize,
flows: &[FlowData],
) -> Vec<DetectionResult> {
let n = flows.len();
let all_features: Vec<Vec<f32>> = flows.iter().map(|f| self.preprocess_ae_features(f)).collect();
let mut scores = Vec::with_capacity(n);
let (predicted_class, class_confidence) = class_probs
.iter()
.enumerate()
.max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(CmpOrdering::Equal))
.map(|(i, &p)| (i, p))
.unwrap_or((0, 0.0));
let mut attack_type = self
.config
.attack_labels
.get(&predicted_class.to_string())
.cloned()
.unwrap_or_else(|| "UNKNOWN".to_string());
let mut confidence = class_confidence;
// C2 head override
if c2_score > self.config.c2_threshold {
let c2_class_prob = self
.c2_class_idx
.and_then(|idx| class_probs.get(idx).copied())
.unwrap_or(0.0);
if c2_score > c2_class_prob {
attack_type = "C2 Communication".to_string();
confidence = c2_score;
for chunk_start in (0..n).step_by(batch_size) {
let chunk_end = (chunk_start + batch_size).min(n);
let actual = chunk_end - chunk_start;
let input = tract_ndarray::Array2::<f32>::from_shape_fn((batch_size, n_features), |(i, j)| {
if i < actual {
all_features[chunk_start + i][j]
} else {
0.0
}
});
match run_ae_batch(model, &input, actual, n_features) {
Ok(s) => scores.extend_from_slice(&s),
Err(e) => {
log!(MLLog::InferenceFailed("AutoencoderOnly".to_string(), e.to_string()));
self.record_failure();
return Vec::new();
}
}
}
// "Normal" class = benign
if Some(predicted_class) == self.normal_class_idx && c2_score <= self.config.c2_threshold {
return DetectionResult {
flow_key,
flow_key_raw: flow.flow_key.clone(),
direction: flow.direction,
is_attack: false,
attack_type: None,
confidence: class_confidence,
ae_score,
anomaly_score,
c2_score,
packet_count: flow.packet_count() as u64,
flow_duration_us: flow.duration_us(),
let thr = self.config.anomaly_threshold;
flows
.iter()
.zip(scores.iter())
.map(|(flow, &score)| {
let is_attack = score > thr;
DetectionResult {
flow_key: build_flow_key_label(flow),
flow_key_raw: flow.flow_key.clone(),
direction: flow.direction,
is_attack,
attack_type: if is_attack { Some("anomaly".to_string()) } else { None },
// AE-only has no separate classifier confidence; reuse the score.
confidence: score,
ae_score: score,
anomaly_score: score,
c2_score: 0.0,
packet_count: flow.packet_count() as u64,
flow_duration_us: flow.duration_us(),
}
})
.collect()
}
/// ClassifierOnly path. Output is per-class softmax; the manifest's
/// labels drive attack_type and the flow fires when the argmax class
/// isn't Normal and confidence ≥ `class_min_confidence`.
fn infer_classifier_only(
&self,
model: &RunnableModel,
batch_size: usize,
n_features: usize,
labels: &BTreeMap<String, LabelSpec>,
normal_idx: Option<usize>,
flows: &[FlowData],
) -> Vec<DetectionResult> {
let n = flows.len();
let all_features: Vec<Vec<f32>> = flows.iter().map(|f| self.preprocess_ae_features(f)).collect();
let class_min_conf = self.config.class_min_confidence;
let mut results = Vec::with_capacity(n);
for chunk_start in (0..n).step_by(batch_size) {
let chunk_end = (chunk_start + batch_size).min(n);
let actual = chunk_end - chunk_start;
let input = tract_ndarray::Array2::<f32>::from_shape_fn((batch_size, n_features), |(i, j)| {
if i < actual {
all_features[chunk_start + i][j]
} else {
0.0
}
});
let class_probs = match run_classifier_only_batch(model, &input, actual) {
Ok(cp) => cp,
Err(e) => {
log!(MLLog::InferenceFailed("ClassifierOnly".to_string(), e.to_string()));
self.record_failure();
return Vec::new();
}
};
for (i, probs) in class_probs.into_iter().enumerate() {
let flow = &flows[chunk_start + i];
let (predicted_class, confidence) = argmax(&probs);
let is_attack = normal_idx.is_none_or(|ni| predicted_class != ni) && confidence >= class_min_conf;
let attack_type = if is_attack {
labels
.get(&predicted_class.to_string())
.map(|l| l.name.clone())
.unwrap_or_else(|| "UNKNOWN".to_string())
} else {
"Normal".to_string()
};
results.push(DetectionResult {
flow_key: build_flow_key_label(flow),
flow_key_raw: flow.flow_key.clone(),
direction: flow.direction,
is_attack,
attack_type: if is_attack { Some(attack_type) } else { None },
confidence,
ae_score: 0.0,
anomaly_score: 0.0,
c2_score: 0.0,
packet_count: flow.packet_count() as u64,
flow_duration_us: flow.duration_us(),
});
}
}
DetectionResult {
flow_key,
flow_key_raw: flow.flow_key.clone(),
direction: flow.direction,
is_attack,
attack_type: if is_attack { Some(attack_type) } else { None },
confidence,
ae_score,
anomaly_score,
c2_score,
packet_count: flow.packet_count() as u64,
flow_duration_us: flow.duration_us(),
}
results
}
fn preprocess_ae_features(&self, flow: &FlowData) -> Vec<f32> {
@ -325,6 +446,13 @@ impl Inference {
features.features.iter().map(|&x| x as f32).collect()
}
/// Publish a rolling QPS estimate visible in `current_status().info.qps_recent`.
/// Called by the engine after each inference tick completes.
pub fn record_tick_qps(&self, flows_per_second: f32) {
self.qps_recent
.store(flows_per_second.max(0.0) as u32, Ordering::Relaxed);
}
// -- Circuit breaker ---------------------------------------------------------
fn now_secs() -> u64 {
@ -341,7 +469,6 @@ impl Inference {
}
let elapsed = Self::now_secs().saturating_sub(open_since);
if elapsed >= CIRCUIT_BREAKER_COOLDOWN_SECS {
// Cooldown elapsed — reset and allow inference
self.circuit_open_since.store(0, Ordering::Relaxed);
self.failure_count.store(0, Ordering::Relaxed);
self.failure_window_start.store(0, Ordering::Relaxed);
@ -355,7 +482,6 @@ impl Inference {
let now = Self::now_secs();
let window_start = self.failure_window_start.load(Ordering::Relaxed);
// If window has expired, start a new window
if window_start == 0 || now.saturating_sub(window_start) > CIRCUIT_BREAKER_WINDOW_SECS {
self.failure_window_start.store(now, Ordering::Relaxed);
self.failure_count.store(1, Ordering::Relaxed);
@ -369,3 +495,117 @@ impl Inference {
}
}
}
// -- Free helpers (can be unit-tested without an Inference) -------------------
fn argmax(probs: &[f32]) -> (usize, f32) {
probs
.iter()
.enumerate()
.max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(CmpOrdering::Equal))
.map(|(i, &p)| (i, p))
.unwrap_or((0, 0.0))
}
fn build_flow_key_label(flow: &FlowData) -> String {
format!(
"{}:{} -> {}:{} (proto {}) [{}]",
flow.flow_key.src_ip_string(),
flow.flow_key.src_port,
flow.flow_key.dst_ip_string(),
flow.flow_key.dst_port,
flow.flow_key.protocol,
flow.direction
)
}
/// Run an autoencoder-style model: output shape == input shape, score is
/// per-row mean squared error between input and reconstruction.
fn run_ae_batch(
model: &RunnableModel,
input: &tract_ndarray::Array2<f32>,
actual: usize,
n_features: usize,
) -> TractResult<Vec<f32>> {
let result = model.run(tvec![input.clone().into_tensor().into()])?;
let output = result[0]
.to_array_view::<f32>()?
.into_dimensionality::<tract_ndarray::Ix2>()?;
let diff = input - &output;
let sq = &diff * &diff;
let mut scores = Vec::with_capacity(actual);
let n_f = n_features as f32;
for i in 0..actual {
scores.push(sq.row(i).sum() / n_f);
}
Ok(scores)
}
/// Run a 3-output multi-task classifier: (anomaly, class_probs, c2_score).
fn run_classifier_batch(
model: &RunnableModel,
input: &tract_ndarray::Array2<f32>,
actual: usize,
) -> TractResult<ClassifierBatchOutput> {
let result = model.run(tvec![input.clone().into_tensor().into()])?;
let anomaly_view = result[0].to_array_view::<f32>()?;
let anomaly: Vec<f32> = (0..actual)
.map(|i| anomaly_view.as_slice().map(|s| s[i]).unwrap_or(0.0))
.collect();
let class_view = result[1]
.to_array_view::<f32>()?
.into_dimensionality::<tract_ndarray::Ix2>()?;
let class_probs: Vec<Vec<f32>> = (0..actual)
.map(|i| class_view.row(i).iter().copied().collect())
.collect();
let c2_view = result[2].to_array_view::<f32>()?;
let c2: Vec<f32> = (0..actual)
.map(|i| c2_view.as_slice().map(|s| s[i]).unwrap_or(0.0))
.collect();
Ok((anomaly, class_probs, c2))
}
/// Run a single-output classifier (ClassifierOnly adapter).
fn run_classifier_only_batch(
model: &RunnableModel,
input: &tract_ndarray::Array2<f32>,
actual: usize,
) -> TractResult<Vec<Vec<f32>>> {
let result = model.run(tvec![input.clone().into_tensor().into()])?;
let view = result[0]
.to_array_view::<f32>()?
.into_dimensionality::<tract_ndarray::Ix2>()?;
Ok((0..actual).map(|i| view.row(i).iter().copied().collect()).collect())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn argmax_picks_highest() {
let (idx, val) = argmax(&[0.1, 0.5, 0.3, 0.4]);
assert_eq!(idx, 1);
assert!((val - 0.5).abs() < 1e-6);
}
#[test]
fn argmax_empty_returns_zero() {
assert_eq!(argmax(&[]), (0, 0.0));
}
#[test]
fn argmax_equal_picks_last() {
// Iterator::max_by returns the LAST element when comparisons are
// equal (in contrast to min_by). Ties in softmax probabilities are
// rare in practice, and "last wins" is a consistent contract across
// this codebase.
let (idx, _) = argmax(&[0.25, 0.25, 0.25, 0.25]);
assert_eq!(idx, 3);
}
}

View File

@ -39,8 +39,7 @@ pub struct ModelPaths {
pub classifier: Option<String>,
}
/// Per-label metadata. `confirmations` and `playbook` are parsed but not yet consumed
/// by the aggregator — they belong to M2's aggregator rewrite.
/// Per-label metadata driving classifier output decoding and alert gating.
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct LabelSpec {
pub name: String,
@ -60,10 +59,11 @@ pub struct Thresholds {
pub class_min_confidence: Option<f32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub ae: Option<f32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub alert_multiplier: Option<f32>,
}
/// Pointer to a preprocessing sidecar (scaler/clip arrays). For v10 this is
/// the legacy `inference_config.json`.
/// Pointer to the JSON sidecar carrying scaler / clip / weight arrays.
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Preprocessing {
pub scaler_sidecar: String,
@ -134,6 +134,43 @@ impl ModelManifest {
}
}
}
self.validate_labels(path)?;
self.validate_thresholds(path)?;
Ok(())
}
fn validate_labels(&self, path: &Path) -> Result<(), MLError> {
let mut seen: Vec<String> = Vec::with_capacity(self.labels.len());
for spec in self.labels.values() {
if let Some(n) = spec.confirmations
&& n == 0
{
return Err(MLError::ManifestInvalid(
path.to_path_buf(),
format!("label '{}' has confirmations: 0 (must be ≥ 1)", spec.name),
));
}
let lower = spec.name.to_ascii_lowercase();
if seen.iter().any(|s| s == &lower) {
return Err(MLError::ManifestInvalid(
path.to_path_buf(),
format!("duplicate label name '{}' (case-insensitive)", spec.name),
));
}
seen.push(lower);
}
Ok(())
}
fn validate_thresholds(&self, path: &Path) -> Result<(), MLError> {
if let Some(m) = self.thresholds.alert_multiplier
&& !(m.is_finite() && m > 0.0)
{
return Err(MLError::ManifestInvalid(
path.to_path_buf(),
format!("thresholds.alert_multiplier must be finite and > 0 (got {m})"),
));
}
Ok(())
}
@ -226,4 +263,59 @@ features: []
let err = parsed.validate(path).expect_err("should reject empty features");
assert!(matches!(err, MLError::ManifestInvalid { .. }), "got {err:?}");
}
#[test]
fn rejects_zero_confirmations() {
let yaml = r#"
name: bad
adapter: classifier_only
models:
model: m.onnx
features:
- flow_duration
labels:
"0": { name: Bot, confirmations: 0 }
"#;
let path = Path::new("/tmp/test-manifest.yaml");
let parsed: ModelManifest = serde_yaml_ng::from_str(yaml).unwrap();
let err = parsed.validate(path).expect_err("should reject confirmations: 0");
assert!(matches!(err, MLError::ManifestInvalid { .. }), "got {err:?}");
}
#[test]
fn rejects_duplicate_label_name_ignoring_case() {
let yaml = r#"
name: bad
adapter: classifier_only
models:
model: m.onnx
features:
- flow_duration
labels:
"0": { name: Bot }
"1": { name: BOT }
"#;
let path = Path::new("/tmp/test-manifest.yaml");
let parsed: ModelManifest = serde_yaml_ng::from_str(yaml).unwrap();
let err = parsed.validate(path).expect_err("should reject duplicate label names");
assert!(matches!(err, MLError::ManifestInvalid { .. }), "got {err:?}");
}
#[test]
fn rejects_non_positive_alert_multiplier() {
let yaml = r#"
name: bad
adapter: classifier_only
models:
model: m.onnx
features:
- flow_duration
thresholds:
alert_multiplier: 0
"#;
let path = Path::new("/tmp/test-manifest.yaml");
let parsed: ModelManifest = serde_yaml_ng::from_str(yaml).unwrap();
let err = parsed.validate(path).expect_err("should reject alert_multiplier <= 0");
assert!(matches!(err, MLError::ManifestInvalid { .. }), "got {err:?}");
}
}

View File

@ -1,3 +1,4 @@
pub mod adapter;
pub mod aggregator;
pub mod alert;
pub mod config_loader;

View File

@ -1,4 +1,13 @@
//! Factory that builds an `MLModelAdapter` from a `ModelManifest`: inspects
//! the manifest's `adapter` field, loads the named ONNX file(s) with shape
//! validation against the inference config's feature counts, and wraps the
//! `RunnableModel`s in `Arc` so hot-reload can swap without per-tick clones.
//! A shape mismatch yields `MLError::FeatureMismatch` carrying both
//! expected and observed dims for the upload UI to render.
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Instant;
use macros::log;
@ -6,105 +15,153 @@ use tract_onnx::prelude::*;
use tract_onnx::tract_hir::infer::Factoid;
use tract_onnx::tract_hir::internal::DimLike;
use crate::infrastructure::app_config::AppConfig;
use super::adapter::MLModelAdapter;
use super::manifest::{AdapterKind, LabelSpec, ModelManifest};
use crate::model::config::constants::MODELS_DIR;
use crate::model::detection::ml_detection::RunnableModel;
use crate::model::error::ml::MLError;
use crate::model::log::ml::MLLog;
use crate::model::system::config::MLInferenceConfig;
pub struct MLModels {
pub deep_autoencoder: RunnableModel,
pub classifier: RunnableModel,
pub batch_size: usize,
/// Build an `MLModelAdapter` by loading the ONNX file(s) the manifest names,
/// validating shape against the inference config's feature counts, and
/// wrapping the underlying `RunnableModel`s in `Arc` for zero-copy swap.
pub fn build_adapter(
manifest: &ModelManifest,
manifest_path: Option<&Path>,
inference_config: &MLInferenceConfig,
batch_size: usize,
) -> Result<MLModelAdapter, MLError> {
let resolve = |rel: &str| -> PathBuf {
match manifest_path {
Some(mp) => ModelManifest::resolve_relative(mp, rel),
None => PathBuf::from(MODELS_DIR).join(rel),
}
};
match manifest.adapter {
AdapterKind::AutoencoderOnly => {
let Some(model_name) = manifest.models.model.as_deref() else {
return Err(MLError::ManifestInvalid(
manifest_path.unwrap_or_else(|| Path::new("")).to_path_buf(),
"autoencoder_only adapter requires models.model".to_string(),
));
};
let path = resolve(model_name);
let n_features = inference_config.num_ae_features();
let model = Arc::new(loader(&path, model_name, n_features, batch_size)?);
Ok(MLModelAdapter::AutoencoderOnly {
model,
batch_size,
n_features,
})
}
AdapterKind::ClassifierOnly => {
let Some(model_name) = manifest.models.model.as_deref() else {
return Err(MLError::ManifestInvalid(
manifest_path.unwrap_or_else(|| Path::new("")).to_path_buf(),
"classifier_only adapter requires models.model".to_string(),
));
};
let path = resolve(model_name);
let n_features = inference_config.num_classifier_features();
let model = Arc::new(loader(&path, model_name, n_features, batch_size)?);
let labels = manifest.labels.clone();
let normal_idx = find_label_index(&labels, "Normal");
Ok(MLModelAdapter::ClassifierOnly {
model,
batch_size,
n_features,
labels,
normal_idx,
})
}
AdapterKind::MultiTask => {
let Some(ae_name) = manifest.models.autoencoder.as_deref() else {
return Err(MLError::ManifestInvalid(
manifest_path.unwrap_or_else(|| Path::new("")).to_path_buf(),
"multi_task adapter requires models.autoencoder".to_string(),
));
};
let Some(cls_name) = manifest.models.classifier.as_deref() else {
return Err(MLError::ManifestInvalid(
manifest_path.unwrap_or_else(|| Path::new("")).to_path_buf(),
"multi_task adapter requires models.classifier".to_string(),
));
};
let n_ae = inference_config.num_ae_features();
let n_cls = inference_config.num_classifier_features();
let ae = Arc::new(loader(&resolve(ae_name), ae_name, n_ae, batch_size)?);
let classifier = Arc::new(loader(&resolve(cls_name), cls_name, n_cls, batch_size)?);
let labels = manifest.labels.clone();
let normal_idx = find_label_index(&labels, "Normal");
let c2_idx = find_label_index(&labels, "C2 Communication");
Ok(MLModelAdapter::MultiTask {
ae,
classifier,
batch_size,
n_ae,
n_cls,
labels,
normal_idx,
c2_idx,
})
}
}
}
impl MLModels {
pub fn load_models(
app_config: &Arc<AppConfig>,
inference_config: &Arc<MLInferenceConfig>,
batch_size: usize,
) -> Result<Self, MLError> {
Self::load_named(
&app_config.inference.deep_autoencoder_name,
&app_config.inference.classifier_name,
inference_config,
batch_size,
)
}
/// Locate a label index by its `name` field. Used to cache hot-path indices
/// (Normal, C2 Communication) at adapter build time rather than re-scanning
/// the label map on every inference batch.
fn find_label_index(labels: &BTreeMap<String, LabelSpec>, target: &str) -> Option<usize> {
labels
.iter()
.find(|(_, v)| v.name.eq_ignore_ascii_case(target))
.and_then(|(k, _)| k.parse::<usize>().ok())
}
/// Manifest-friendly constructor: caller supplies the AE and classifier ONNX filenames
/// directly. Used when a `ModelManifest` overrides the defaults in `AppConfig`.
pub fn load_named(
autoencoder_name: &str,
classifier_name: &str,
inference_config: &Arc<MLInferenceConfig>,
batch_size: usize,
) -> Result<Self, MLError> {
Ok(Self {
deep_autoencoder: Self::loader(autoencoder_name, inference_config.num_ae_features(), batch_size)?,
classifier: Self::loader(classifier_name, inference_config.num_classifier_features(), batch_size)?,
batch_size,
})
}
fn loader(model_path: &Path, model_name: &str, features: usize, batch_size: usize) -> Result<RunnableModel, MLError> {
log!(MLLog::ModelLoading(model_name.to_string(), features, batch_size));
fn loader(model: &str, features: usize, batch_size: usize) -> Result<RunnableModel, MLError> {
let model_path = PathBuf::from("models").join(model);
let start = Instant::now();
let result = loader_inner(model_path, model_name, features, batch_size);
let elapsed_ms = start.elapsed().as_millis() as u64;
log!(MLLog::ModelLoadComplete(model_name.to_string(), elapsed_ms));
log!(MLLog::ModelLoading(model.to_string(), features, batch_size));
result
}
let start = Instant::now();
let result = Self::loader_inner(&model_path, model, features, batch_size);
let elapsed_ms = start.elapsed().as_millis() as u64;
log!(MLLog::ModelLoadComplete(model.to_string(), elapsed_ms));
fn loader_inner(
model_path: &Path,
model_name: &str,
features: usize,
batch_size: usize,
) -> Result<RunnableModel, MLError> {
let mut onnx_model = onnx()
.model_for_path(model_path)
.map_err(|e| MLError::ModelLoadFailed(model_path.to_path_buf(), format!("parse ONNX: {e}")))?;
result
}
fn loader_inner(
model_path: &Path,
model: &str,
features: usize,
batch_size: usize,
) -> Result<RunnableModel, MLError> {
let mut onnx_model = onnx()
.model_for_path(model_path)
.map_err(|e| MLError::ModelLoadFailed(model_path.to_path_buf(), format!("parse ONNX: {e}")))?;
// Introspect the ONNX input fact and cross-check its last dim against the
// expected feature count. A dynamic/symbolic dim is skipped — we only fail
// when the ONNX graph declares a concrete integer that disagrees.
if let Some(onnx_dim) = introspect_input_features(&onnx_model) {
let matched = onnx_dim == features;
log!(MLLog::OnnxShapeChecked(model.to_string(), features, onnx_dim, matched));
if !matched {
return Err(MLError::FeatureMismatch(model_path.to_path_buf(), features, onnx_dim));
}
if let Some(onnx_dim) = introspect_input_features(&onnx_model) {
let matched = onnx_dim == features;
log!(MLLog::OnnxShapeChecked(
model_name.to_string(),
features,
onnx_dim,
matched,
));
if !matched {
return Err(MLError::FeatureMismatch(model_path.to_path_buf(), features, onnx_dim));
}
onnx_model
.set_input_fact(0, f32::fact([batch_size, features]).into())
.map_err(|e| MLError::ModelLoadFailed(model_path.to_path_buf(), format!("set_input_fact: {e}")))?;
onnx_model
.into_optimized()
.and_then(|m| m.into_runnable())
.map_err(|e| MLError::ModelLoadFailed(model_path.to_path_buf(), format!("optimize/runnable: {e}")))
}
pub fn get_model_info(&self, name: &str) -> String {
let model = match name {
"deep_autoencoder" => &self.deep_autoencoder,
"classifier" => &self.classifier,
_ => return "unknown model".to_string(),
};
onnx_model
.set_input_fact(0, f32::fact([batch_size, features]).into())
.map_err(|e| MLError::ModelLoadFailed(model_path.to_path_buf(), format!("set_input_fact: {e}")))?;
let inputs = model.model().inputs.len();
let outputs = model.model().outputs.len();
format!(
"{name}: inputs: {inputs}, outputs: {outputs}, batch_size: {}",
self.batch_size
)
}
onnx_model
.into_optimized()
.and_then(|m| m.into_runnable())
.map_err(|e| MLError::ModelLoadFailed(model_path.to_path_buf(), format!("optimize/runnable: {e}")))
}
/// Read the concrete last-dim (feature count) from an ONNX model's declared input fact.
@ -123,8 +180,8 @@ fn introspect_input_features(model: &InferenceModel) -> Option<usize> {
mod tests {
use super::*;
/// Integration test: tract must introspect the shipped v10 AE ONNX's input dim as 31.
/// Silently skipped outside the repo root (no `models/` directory).
/// Repo-root integration: the shipped v10 AE ONNX's input last-dim must
/// still be 31. A change here means the shipped manifest + sidecar drifted.
#[test]
fn introspect_v10_autoencoder_input_is_31() {
let ae_path = PathBuf::from("models/deep_autoencoder.onnx");
@ -149,24 +206,23 @@ mod tests {
assert_eq!(dim, 32, "v10 classifier ONNX input dim changed unexpectedly");
}
/// End-to-end M1 smoke: loading v10 via `load_named` with the manifest-driven
/// names + sidecar-derived inference config must succeed without shape mismatch.
/// Smoke: the shipped manifest + sidecar must load and produce a
/// `MultiTask` adapter with both underlying models.
#[test]
fn v10_load_named_with_manifest_paths_succeeds() {
fn v10_build_adapter_multitask() {
let manifest_path = Path::new("models/manifest.yaml");
if !manifest_path.exists() {
eprintln!("skipping: models/manifest.yaml absent");
return;
}
let (cfg, manifest) =
MLInferenceConfig::from_manifest_with_sidecar(manifest_path).expect("manifest + sidecar load");
let cfg = Arc::new(cfg);
let ae = manifest.models.autoencoder.as_deref().expect("multi_task has AE");
let cls = manifest
.models
.classifier
.as_deref()
.expect("multi_task has classifier");
let _models = MLModels::load_named(ae, cls, &cfg, 8).expect("v10 load via manifest names");
let (cfg, manifest) = MLInferenceConfig::from_manifest_with_sidecar(manifest_path).expect("config load");
let adapter = build_adapter(&manifest, Some(manifest_path), &cfg, 8).expect("build adapter");
match adapter {
MLModelAdapter::MultiTask { n_ae, n_cls, .. } => {
assert_eq!(n_ae, 31);
assert_eq!(n_cls, 32);
}
_ => panic!("v10 manifest should build MultiTask adapter"),
}
}
}

View File

@ -1,43 +1,41 @@
use std::path::PathBuf;
//! Filesystem watcher over `models/`. Reloads the ML source whenever the
//! manifest or ONNX files change, re-reading both the manifest and the
//! scaler sidecar so feature-changing uploads land without a restart.
//! Events inside the staging subdirectory are filtered so partial uploads
//! can't flicker the UI through transient `Error` states.
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;
use std::time::{Duration, SystemTime};
use macros::log;
use notify::{Event, EventKind, RecommendedWatcher, RecursiveMode, Watcher};
use tokio::sync::mpsc;
use tokio::time::sleep;
use super::adapter::ModelSourceState;
use super::inference::Inference;
use super::model_loader::MLModels;
use super::model_loader::build_adapter;
use crate::infrastructure::app_config::AppConfig;
use crate::model::config::constants::{MANIFEST_FILENAME, MODELS_DIR, STAGING_SUBDIR};
use crate::model::detection::model_source::ModelInfo;
use crate::model::error::ml::MLError;
use crate::model::log::ml::MLLog;
use crate::model::system::config::MLInferenceConfig;
/// Debounce window: wait for both AE + classifier to land before reloading.
/// Debounce window: wait for both manifest and ONNX to land before reloading.
const DEBOUNCE_SECS: u64 = 5;
/// Watch the models/ directory for .onnx file changes and hot-reload into the inference pipeline.
pub struct ModelWatcher {
inference: Arc<Inference>,
app_config: Arc<AppConfig>,
inference_config: Arc<MLInferenceConfig>,
}
impl ModelWatcher {
pub fn new(
inference: Arc<Inference>,
app_config: Arc<AppConfig>,
inference_config: Arc<MLInferenceConfig>,
) -> Self {
Self {
inference,
app_config,
inference_config,
}
pub fn new(inference: Arc<Inference>, app_config: Arc<AppConfig>) -> Self {
Self { inference, app_config }
}
/// Start watching in a background task.
pub fn start(self) {
tokio::spawn(async move {
if let Err(e) = self.run().await {
@ -50,7 +48,7 @@ impl ModelWatcher {
}
async fn run(self) -> Result<(), MLError> {
let models_dir = PathBuf::from("models");
let models_dir = PathBuf::from(MODELS_DIR);
if !models_dir.exists() {
log!(MLLog::InferenceFailed(
"ModelWatcher".to_string(),
@ -60,23 +58,16 @@ impl ModelWatcher {
}
let (tx, mut rx) = mpsc::channel::<()>(16);
// notify watcher runs on a blocking thread — forward events via channel
let _watcher = Self::spawn_watcher(models_dir, tx)?;
log!(MLLog::ModelWatcherStarted);
loop {
// Wait for first event
if rx.recv().await.is_none() {
break;
}
// Debounce: drain any additional events within the window
sleep(Duration::from_secs(DEBOUNCE_SECS)).await;
while rx.try_recv().is_ok() {}
// Attempt reload
self.try_reload();
}
@ -86,37 +77,135 @@ impl ModelWatcher {
fn spawn_watcher(models_dir: PathBuf, tx: mpsc::Sender<()>) -> Result<RecommendedWatcher, MLError> {
let mut watcher = notify::recommended_watcher(move |res: Result<Event, notify::Error>| {
if let Ok(event) = res {
let dominated = matches!(event.kind, EventKind::Create(_) | EventKind::Modify(_));
let has_onnx = event
.paths
.iter()
.any(|p| p.extension().is_some_and(|ext| ext == "onnx"));
if dominated && has_onnx {
let _ = tx.blocking_send(());
if !is_relevant_event(&event) {
return;
}
let _ = tx.blocking_send(());
}
})
.map_err(MLError::ModelWatcherFailed)?;
// Recursive watch so the staging filter gets exercised — otherwise a
// drop-in to `.staging/` wouldn't trigger notify at all on some FSes.
watcher
.watch(&models_dir, RecursiveMode::NonRecursive)
.watch(&models_dir, RecursiveMode::Recursive)
.map_err(MLError::ModelWatcherFailed)?;
Ok(watcher)
}
/// Full reload: manifest presence check → config re-parse → adapter build →
/// atomic state swap. Any failure lands the pipeline in `Error` rather
/// than crashing.
fn try_reload(&self) {
let batch_size = self.app_config.inference.inference_batch_size;
let manifest_path = PathBuf::from(MODELS_DIR).join(MANIFEST_FILENAME);
// Path 1 — manifest disappeared: transition to Dormant.
if !manifest_path.exists() {
log!(MLLog::ModelReloadStarting);
self.inference.swap_state(ModelSourceState::Dormant);
log!(MLLog::ModelReloadSuccess);
return;
}
// Path 2 — manifest present: re-read + rebuild adapter.
log!(MLLog::ModelReloadStarting);
match MLModels::load_models(&self.app_config, &self.inference_config, batch_size) {
Ok(new_models) => {
self.inference.swap_models(Arc::new(new_models));
log!(MLLog::ModelReloadSuccess);
}
let batch_size = self.app_config.inference.inference_batch_size;
match MLInferenceConfig::from_manifest_with_sidecar(&manifest_path) {
Ok((config, manifest)) => match build_adapter(&manifest, Some(&manifest_path), &config, batch_size) {
Ok(adapter) => {
let info = ModelInfo::new(
manifest.name.clone(),
manifest.adapter.as_str().to_string(),
manifest.features.len(),
);
self.inference.swap_state(ModelSourceState::Active { adapter, info });
log!(MLLog::ModelReloadSuccess);
}
Err(e) => {
self.record_error(e.to_string(), Some(manifest_path.clone()));
}
},
Err(e) => {
log!(MLLog::ModelReloadFailed(e.to_string()));
self.record_error(e.to_string(), Some(manifest_path.clone()));
}
}
}
fn record_error(&self, msg: String, last_attempted_path: Option<PathBuf>) {
log!(MLLog::ModelReloadFailed(msg.clone()));
self.inference.swap_state(ModelSourceState::Error {
msg,
since: SystemTime::now(),
last_attempted_path,
});
}
}
/// Inbound event filter. Ignore `.staging/` paths entirely; pass through
/// `.onnx` / `.yaml` / `.yml` / `.json` changes in `models/`.
fn is_relevant_event(event: &Event) -> bool {
// Only Create / Modify events trigger a reload; renames / removes would
// also surface but debouncing handles both equally well.
if !matches!(event.kind, EventKind::Create(_) | EventKind::Modify(_)) {
return false;
}
event.paths.iter().any(|p| {
if path_is_inside_staging(p) {
return false;
}
matches!(
p.extension().and_then(|e| e.to_str()),
Some("onnx") | Some("yaml") | Some("yml") | Some("json")
)
})
}
fn path_is_inside_staging(path: &Path) -> bool {
path.components().any(|c| c.as_os_str() == STAGING_SUBDIR)
}
#[cfg(test)]
mod tests {
use super::*;
use notify::event::{CreateKind, ModifyKind};
#[test]
fn staging_paths_are_filtered() {
assert!(path_is_inside_staging(&PathBuf::from("models/.staging/bad.onnx")));
assert!(path_is_inside_staging(&PathBuf::from(
"/tmp/models/.staging/sub/x.yaml"
)));
assert!(!path_is_inside_staging(&PathBuf::from("models/good.onnx")));
}
#[test]
fn non_relevant_extensions_rejected() {
let event = Event {
kind: EventKind::Create(CreateKind::Any),
paths: vec![PathBuf::from("models/readme.md")],
attrs: Default::default(),
};
assert!(!is_relevant_event(&event));
}
#[test]
fn onnx_outside_staging_accepted() {
let event = Event {
kind: EventKind::Create(CreateKind::Any),
paths: vec![PathBuf::from("models/foo.onnx")],
attrs: Default::default(),
};
assert!(is_relevant_event(&event));
}
#[test]
fn staging_paths_always_rejected() {
let event = Event {
kind: EventKind::Modify(ModifyKind::Any),
paths: vec![PathBuf::from("models/.staging/partial.onnx")],
attrs: Default::default(),
};
assert!(!is_relevant_event(&event));
}
}

View File

@ -20,6 +20,7 @@ use crate::interface::port::notification::AlertNotifier;
use crate::interface::port::rate_limit_api::RateLimitPort;
use crate::interface::port::secret_store::SecretStorePort;
use crate::model::config::constants::MAX_PENDING_UNBLOCK_RETRIES;
use crate::model::detection::attack_type::canonical_from_str;
use crate::model::error::Error;
use crate::model::error::soar::SoarError;
use crate::model::event::ThreatDetectedEvent;
@ -59,7 +60,9 @@ pub struct SoarEngine {
pub(super) geoip: Option<Arc<GeoIpService>>,
/// Optional rate limit config for adjust_rate_limit action.
pub(super) rate_limit: Option<Arc<dyn RateLimitPort>>,
/// Lock to serialize rate limit read-save-write sequences (Item 6: atomicity).
/// Serializes rate-limit read-save-write sequences so concurrent SOAR
/// actions can't race the `soar_rate_limit_*` DB settings into an
/// inconsistent pair.
pub(super) rate_limit_lock: TokioMutex<()>,
/// Cached enforce level: Monitor=0, MlOnly=1, Enforce=2.
pub(super) enforce_level_cache: Arc<AtomicU8>,
@ -120,7 +123,7 @@ impl SoarEngine {
// Check if this row belongs to the same playbook as the last one
let needs_new = playbooks.last().is_none_or(|last| last.id != pb_id);
if needs_new {
let _ = threshold; // stored in DB but no longer used at runtime
let _ = threshold; // persisted for schema stability; runtime gating comes from the Condition rows
playbooks.push(Playbook {
id: pb_id,
name,
@ -165,6 +168,9 @@ impl SoarEngine {
}
ConditionType::RepeatOffender => operator == "==",
ConditionType::Frequency => operator == ">=",
ConditionType::MultiSourceMin => operator == ">=",
ConditionType::SingleSourceHigh => operator == ">=",
ConditionType::FusedConfidenceAbove => matches!(operator.as_str(), ">=" | "<="),
};
if !valid {
log!(SoarLog::InvalidConditionOperator(
@ -183,6 +189,18 @@ impl SoarEngine {
}
}
// Warn on playbooks whose trigger_event isn't in the canonical
// dictionary — those won't ever match a fused event and usually
// signal a typo or stale pre-canonicalization playbook.
for pb in &playbooks {
if canonical_from_str(&pb.trigger_event).is_none() {
log!(SoarLog::NonCanonicalTriggerEvent(
pb.name.clone(),
pb.trigger_event.clone(),
));
}
}
*self.playbooks.write() = playbooks;
// Load admin whitelist
@ -334,7 +352,6 @@ impl SoarEngine {
/// Restore original rate limits if the TTL has expired.
/// Called by TTL scheduler on each sweep.
pub async fn check_rate_limit_restoration(&self) -> Result<(), Error> {
// Acquire lock to serialize rate limit read-save-write (Item 6: atomicity)
let _guard = self.rate_limit_lock.lock().await;
let expires_str = match self
@ -484,6 +501,8 @@ mod tests {
geoip_country: None,
is_repeat_offender: false,
sources: vec![DetectionSource::ML],
active_source_count: 1,
fused_confidence: 0.95,
ae_score: 0.0,
anomaly_score: 0.0,
c2_score: 0.0,
@ -519,6 +538,8 @@ mod tests {
geoip_country: None,
is_repeat_offender: false,
sources: vec![DetectionSource::ML],
active_source_count: 1,
fused_confidence: 0.95,
ae_score: 0.0,
anomaly_score: 0.0,
c2_score: 0.0,
@ -601,6 +622,8 @@ mod tests {
geoip_country: None,
is_repeat_offender: false,
sources: vec![DetectionSource::ML],
active_source_count: 1,
fused_confidence: 0.95,
ae_score: 0.0,
anomaly_score: 0.0,
c2_score: 0.0,
@ -636,6 +659,8 @@ mod tests {
geoip_country: None,
is_repeat_offender: false,
sources: vec![DetectionSource::ML],
active_source_count: 1,
fused_confidence: 0.95,
ae_score: 0.0,
anomaly_score: 0.0,
c2_score: 0.0,
@ -681,6 +706,8 @@ mod tests {
geoip_country: None,
is_repeat_offender: false,
sources: vec![DetectionSource::ML],
active_source_count: 1,
fused_confidence: 0.95,
ae_score: 0.0,
anomaly_score: 0.0,
c2_score: 0.0,
@ -757,6 +784,8 @@ mod tests {
geoip_country: country.map(|s| s.to_string()),
is_repeat_offender: repeat,
sources: vec![DetectionSource::ML],
active_source_count: 1,
fused_confidence: 0.95,
ae_score: 0.0,
anomaly_score: 0.0,
c2_score: 0.0,

View File

@ -6,17 +6,29 @@
//! but none of them touch anything outside those fields.
use std::net::IpAddr;
use std::str::FromStr;
use std::sync::atomic::Ordering;
use std::time::{Duration, Instant};
use macros::log;
use crate::core::soar::engine::SoarEngine;
use crate::model::event::ThreatDetectedEvent;
use crate::model::event::{DetectionSource, ThreatDetectedEvent};
use crate::model::log::soar::SoarLog;
use crate::model::soar::condition::{ConditionType, PlaybookCondition};
use crate::model::soar::playbook::Playbook;
/// Default frequency-condition window when the playbook omits `value2`.
const DEFAULT_FREQUENCY_WINDOW_SECS: u64 = 60;
/// Default minimum confidence for `SingleSourceHigh` when the playbook
/// omits `value2`. Conservative enough that ad-hoc solo playbooks don't
/// auto-block noisy single-source hits.
const DEFAULT_SINGLE_SOURCE_HIGH_MIN_CONFIDENCE: f32 = 0.95;
/// Default expiry for cooldown cleanup when no playbook has a cooldown set.
const DEFAULT_COOLDOWN_EXPIRY_SECS: u64 = 3600;
impl SoarEngine {
/// Find playbooks matching the event via trigger_event + multi-condition AND logic.
pub(super) fn find_matching_playbooks(&self, event: &ThreatDetectedEvent) -> Vec<Playbook> {
@ -142,7 +154,11 @@ impl SoarEngine {
Ok(v) => v,
Err(_) => return false,
};
let window_secs = cond.value2.as_ref().and_then(|s| s.parse::<u64>().ok()).unwrap_or(60);
let window_secs = cond
.value2
.as_ref()
.and_then(|s| s.parse::<u64>().ok())
.unwrap_or(DEFAULT_FREQUENCY_WINDOW_SECS);
let count = self
.frequency_tracker
.record_and_count(pb.id, &event.source_ip, window_secs);
@ -152,6 +168,70 @@ impl SoarEngine {
}
met
}
ConditionType::MultiSourceMin => {
let required = match cond.value.parse::<usize>() {
Ok(v) => v,
Err(_) => return false,
};
let met = event.active_source_count >= required;
if !met {
log!(SoarLog::ConditionNotMet(
"multi_source_min".to_string(),
pb.name.clone(),
format!("{} / {}", event.active_source_count, required),
));
}
met
}
ConditionType::SingleSourceHigh => {
// `value` names the target source (canonical Display form
// or a common alias); `value2` is the minimum confidence.
let target_source = match DetectionSource::from_str(&cond.value) {
Ok(s) => s,
Err(_) => return false,
};
let min_conf = cond
.value2
.as_ref()
.and_then(|s| s.parse::<f32>().ok())
.unwrap_or(DEFAULT_SINGLE_SOURCE_HIGH_MIN_CONFIDENCE);
// Solo = exactly one contributing source AND it matches the
// target source AND confidence clears the escape-hatch bar.
let solo_match =
event.active_source_count == 1 && event.sources.len() == 1 && event.sources[0] == target_source;
let conf_met = event.confidence >= min_conf;
let met = solo_match && conf_met;
if !met {
log!(SoarLog::ConditionNotMet(
"single_source_high".to_string(),
pb.name.clone(),
format!(
"sources={:?} count={} conf={:.3} target={} need_conf>={:.3}",
event.sources, event.active_source_count, event.confidence, cond.value, min_conf
),
));
}
met
}
ConditionType::FusedConfidenceAbove => {
let threshold = match cond.value.parse::<f32>() {
Ok(v) => v,
Err(_) => return false,
};
let met = if cond.operator == "<=" {
event.fused_confidence <= threshold
} else {
event.fused_confidence >= threshold
};
if !met {
log!(SoarLog::ConditionNotMet(
"fused_confidence_above".to_string(),
pb.name.clone(),
format!("{:.3} vs {:.3}", event.fused_confidence, threshold),
));
}
met
}
}
}
@ -178,9 +258,13 @@ impl SoarEngine {
pub fn cleanup_expired_cooldowns(&self) {
let max_cooldown_secs = {
let playbooks = self.playbooks.read();
playbooks.iter().map(|p| p.cooldown_secs as u64).max().unwrap_or(3600)
playbooks
.iter()
.map(|p| p.cooldown_secs as u64)
.max()
.unwrap_or(DEFAULT_COOLDOWN_EXPIRY_SECS)
};
let expiry = Duration::from_secs(max_cooldown_secs.saturating_mul(2).max(3600));
let expiry = Duration::from_secs(max_cooldown_secs.saturating_mul(2).max(DEFAULT_COOLDOWN_EXPIRY_SECS));
let before = self.cooldowns.len();
self.cooldowns.retain(|_, instant| instant.elapsed() < expiry);
let removed = before.saturating_sub(self.cooldowns.len());

View File

@ -1,24 +1,30 @@
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use std::time::{Duration, SystemTime};
use crossbeam::queue::SegQueue;
use macros::log;
use tokio::sync::oneshot;
use crate::core::ml::adapter::ModelSourceState;
use crate::core::ml::alert::MLAlert;
use crate::core::ml::drift_detector::DriftDetector;
use crate::core::ml::engine::Engine;
use crate::core::ml::manifest::{AdapterKind, ModelManifest};
use crate::core::ml::model_loader::MLModels;
use crate::core::ml::inference::Inference;
use crate::core::ml::manifest::ModelManifest;
use crate::core::ml::model_loader::build_adapter;
use crate::core::ml::traffic_logger::TrafficLogger;
use crate::infrastructure::app_config::AppConfig;
use crate::infrastructure::health::SystemHealth;
use crate::infrastructure::statistics::FlowStatistics;
use crate::model::config::constants::{MANIFEST_FILENAME, MODELS_DIR};
use crate::model::detection::flow_features::FlowFeatures;
use crate::model::detection::ml_detection::EngineConfig;
use crate::model::detection::model_source::ModelInfo;
use crate::model::error::Error;
use crate::model::error::misc::MiscError;
use crate::model::error::system::SystemError;
use crate::model::log::ml::MLLog;
use crate::model::log::system::SystemLog;
use crate::model::system::config::MLInferenceConfig;
use crate::model::system::health::EbpfHealth;
@ -29,7 +35,7 @@ use crate::model::system::health::EbpfHealth;
pub struct AppServices {
pub health: Arc<SystemHealth>,
pub ml_alert: Arc<MLAlert>,
pub ml_models: Arc<MLModels>,
pub ml_inference: Arc<Inference>,
pub ml_engine: Arc<Engine>,
pub flow_statistics: Arc<FlowStatistics>,
shutdowns: SegQueue<oneshot::Sender<()>>,
@ -46,25 +52,45 @@ impl AppServices {
let health = SystemHealth::new(app_config.clone(), ebpf_health)?;
let batch_size = app_config.inference.inference_batch_size;
// When a `multi_task` manifest is present, its model filenames are authoritative.
// Otherwise (no manifest, or single-model adapter in v2) fall through to the
// legacy app_config-driven path. M1 only needs `multi_task` wired end-to-end.
let ml_models = Arc::new(match ml_manifest.as_ref() {
Some(m) if m.adapter == AdapterKind::MultiTask => {
let ae = m
.models
.autoencoder
.as_deref()
.unwrap_or(&app_config.inference.deep_autoencoder_name);
let cls = m
.models
.classifier
.as_deref()
.unwrap_or(&app_config.inference.classifier_name);
MLModels::load_named(ae, cls, &inference_config, batch_size)?
let initial_state = match ml_manifest.as_ref() {
Some(manifest) => {
let manifest_path = PathBuf::from(MODELS_DIR).join(MANIFEST_FILENAME);
match build_adapter(manifest, Some(&manifest_path), &inference_config, batch_size) {
Ok(adapter) => {
let info = ModelInfo::new(
manifest.name.clone(),
manifest.adapter.as_str().to_string(),
manifest.features.len(),
);
log!(MLLog::ModelsLoaded(format!(
"{} ({}) — {} features, {} labels",
manifest.name,
manifest.adapter.as_str(),
manifest.features.len(),
manifest.labels.len()
)));
ModelSourceState::Active { adapter, info }
}
Err(e) => {
log!(MLLog::ModelReloadFailed(e.to_string()));
ModelSourceState::Error {
msg: e.to_string(),
since: SystemTime::now(),
last_attempted_path: Some(manifest_path),
}
}
}
}
_ => MLModels::load_models(&app_config, &inference_config, batch_size)?,
});
None => {
log!(MLLog::ModelsLoaded(
"no manifest present — ML source dormant".to_string()
));
ModelSourceState::Dormant
}
};
let ml_inference = Arc::new(Inference::new(initial_state, inference_config.clone()));
let ml_alert = Arc::new(MLAlert::new());
let traffic_logger = if app_config.inference.traffic_logging_mode {
@ -88,8 +114,7 @@ impl AppServices {
};
let ml_engine = Arc::new(Engine::new(
ml_models.clone(),
inference_config.clone(),
ml_inference.clone(),
ml_alert.clone(),
drift_detector,
engine_config,
@ -102,7 +127,7 @@ impl AppServices {
Ok(Self {
health: Arc::new(health),
ml_alert,
ml_models,
ml_inference,
ml_engine,
flow_statistics,
shutdowns: SegQueue::new(),

View File

@ -156,12 +156,15 @@ impl System {
pub async fn run(&mut self) -> Result<ShutdownMode, Error> {
log!(SystemLog::Initializing);
log!(MLLog::ModelsLoaded(
self.app_services.ml_models.get_model_info("deep_autoencoder")
));
log!(MLLog::ModelsLoaded(
self.app_services.ml_models.get_model_info("classifier")
));
// ML source state snapshot. Day 1 with no manifest renders as Dormant
// — the rest of the stack still runs (3-source fusion).
{
let status = self.app_services.ml_inference.current_status();
match serde_json::to_string(&status) {
Ok(s) => log!(MLLog::ModelsLoaded(s)),
Err(e) => log!(MLLog::ModelsLoaded(format!("<unserializable status: {e}>"))),
}
}
log!(MLLog::ConfigLoaded(
self.inference_config.num_ae_features(),
self.inference_config.num_attack_types(),
@ -267,11 +270,7 @@ impl System {
});
// Start model hot-reload watcher (monitors models/ for .onnx changes)
let model_watcher = ModelWatcher::new(
self.app_services.ml_engine.inference_pipeline().clone(),
self.app_config.clone(),
self.inference_config.clone(),
);
let model_watcher = ModelWatcher::new(self.app_services.ml_inference.clone(), self.app_config.clone());
model_watcher.start();
// Initialize force_https flag from DB setting

View File

@ -14,6 +14,11 @@ pub const FLOW_BULK_MIN_BYTES: u64 = 1000;
pub const FLOW_IDLE_TIMEOUT_US: u64 = 120_000_000;
pub const FLOW_TERMINATED_TIMEOUT_US: u64 = 5_000_000;
// ── ML Model Directory ─────────────────────────────────────────────
pub const MODELS_DIR: &str = "models";
pub const MANIFEST_FILENAME: &str = "manifest.yaml";
pub const STAGING_SUBDIR: &str = ".staging";
// ── Notification ───────────────────────────────────────────────────
pub const TELEGRAM_MAX_RETRIES: u32 = 2;

View File

@ -0,0 +1,327 @@
//! Canonical attack-type dictionary and cross-source translator.
//!
//! Fusion v1 requires that Suricata / ML / CV (Beaconing) / Graph (Correlation)
//! use a shared vocabulary — otherwise the orchestrator's dedup key
//! `(source_ip, attack_type)` never collides across sources, and the
//! cross-source "sources agreed" fusion signal is impossible.
//!
//! The 13 canonical types below are the v1 seed set. Each source ships a
//! translation table from its own raw labels (Suricata classtype, ML class
//! name, Beaconing tag, Correlation sub-type) into the canonical vocabulary.
//! Unknown labels land in `Unknown` — a valid dedup bucket that still
//! participates in fusion.
use std::fmt;
use serde::{Deserialize, Serialize};
use crate::model::event::DetectionSource;
/// The v1 seed dictionary — 13 canonical attack types every detection source
/// maps into. New types may be added without breaking change; renaming or
/// removing one IS a breaking change (dedup keys drift).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CanonicalAttackType {
BruteForce,
PortScan,
C2Beacon,
DnsTunnel,
SqlInjection,
Xss,
Exploit,
LateralMovement,
Reconnaissance,
Cryptomining,
DosDdos,
BotActivity,
/// Fallback bucket. Still a valid dedup key — two sources firing
/// `Unknown` on the same src_ip within the fusion window DO fuse.
Unknown,
}
impl CanonicalAttackType {
/// Stable wire-format string. Must match `DedupKey.attack_type` exactly
/// across releases — renaming breaks dedup on in-flight alerts.
pub const fn as_str(self) -> &'static str {
match self {
Self::BruteForce => "brute_force",
Self::PortScan => "port_scan",
Self::C2Beacon => "c2_beacon",
Self::DnsTunnel => "dns_tunnel",
Self::SqlInjection => "sql_injection",
Self::Xss => "xss",
Self::Exploit => "exploit",
Self::LateralMovement => "lateral_movement",
Self::Reconnaissance => "reconnaissance",
Self::Cryptomining => "cryptomining",
Self::DosDdos => "dos_ddos",
Self::BotActivity => "bot_activity",
Self::Unknown => "unknown",
}
}
/// All 13 canonical values, in declaration order. Used by BYO-contract
/// docs + CI consistency check.
pub const ALL: &'static [Self] = &[
Self::BruteForce,
Self::PortScan,
Self::C2Beacon,
Self::DnsTunnel,
Self::SqlInjection,
Self::Xss,
Self::Exploit,
Self::LateralMovement,
Self::Reconnaissance,
Self::Cryptomining,
Self::DosDdos,
Self::BotActivity,
Self::Unknown,
];
}
impl fmt::Display for CanonicalAttackType {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(self.as_str())
}
}
/// Reverse lookup: canonical wire string → enum. Returns `None` when the
/// input isn't one of the 13 seeds (caller typically treats this as an
/// upstream bug, not an `Unknown` bucket).
pub fn canonical_from_str(s: &str) -> Option<CanonicalAttackType> {
CanonicalAttackType::ALL.iter().copied().find(|c| c.as_str() == s)
}
/// Translate a source-specific raw label into the canonical dictionary.
///
/// Never fails — unknown labels map to `CanonicalAttackType::Unknown` so
/// the event still lands in a valid dedup bucket. The `raw_label`
/// comparison is case-insensitive and trims whitespace; callers that pass
/// user-visible labels verbatim don't need to pre-normalize.
pub fn translate(source: DetectionSource, raw_label: &str) -> CanonicalAttackType {
let normalized = raw_label.trim().to_ascii_lowercase();
match source {
DetectionSource::ML => translate_ml(&normalized),
DetectionSource::Beaconing => translate_beaconing(&normalized),
DetectionSource::Correlation => translate_correlation(&normalized),
DetectionSource::Suricata => translate_suricata(&normalized),
}
}
/// Map ML classifier class names into the canonical dictionary. Case-insensitive;
/// covers the baseline class vocabulary.
fn translate_ml(label: &str) -> CanonicalAttackType {
match label {
"brute force" | "brute_force" | "bruteforce" => CanonicalAttackType::BruteForce,
"reconnaissance" | "recon" | "port_scan" | "portscan" => CanonicalAttackType::Reconnaissance,
"c2 communication" | "c2" | "c2_beacon" | "command_and_control" => CanonicalAttackType::C2Beacon,
"dns tunneling" | "dns_tunnel" | "dns tunnel" => CanonicalAttackType::DnsTunnel,
"sql injection" | "sql_injection" | "sqli" => CanonicalAttackType::SqlInjection,
"xss" | "cross_site_scripting" => CanonicalAttackType::Xss,
"web attack" | "web_attack" => CanonicalAttackType::Exploit,
"exploitation" | "exploit" => CanonicalAttackType::Exploit,
"dos/ddos" | "dos_ddos" | "ddos" | "dos" => CanonicalAttackType::DosDdos,
"cryptomining" | "cryptocurrency_mining" | "mining" => CanonicalAttackType::Cryptomining,
"bot" | "bot_activity" | "botnet" | "malware" => CanonicalAttackType::BotActivity,
"lateral movement" | "lateral_movement" => CanonicalAttackType::LateralMovement,
// "Normal" is not a threat — translators should not see it, but if they do,
// fall through to Unknown rather than panicking.
_ => CanonicalAttackType::Unknown,
}
}
/// Beaconing (Layer 2 CV) only produces C2-style temporal beacons in v1.
/// Sub-tags (e.g. "c2_beacon", "heartbeat") all collapse here.
fn translate_beaconing(_label: &str) -> CanonicalAttackType {
CanonicalAttackType::C2Beacon
}
/// Correlation (Layer 3 graph) splits across scan / lateral / botnet.
fn translate_correlation(label: &str) -> CanonicalAttackType {
match label {
"scan" | "port_scan" | "reconnaissance" => CanonicalAttackType::PortScan,
"lateral" | "lateral_movement" => CanonicalAttackType::LateralMovement,
"botnet" | "bot" | "bot_activity" => CanonicalAttackType::BotActivity,
_ => CanonicalAttackType::Unknown,
}
}
/// Map Suricata classtypes (`eve.json.alert.category`, not sid) into the
/// canonical dictionary. We match on classtype because sid numbering isn't
/// stable across rule packs; classtype is part of the rule DSL and stable
/// across ET Open / Talos releases.
fn translate_suricata(label: &str) -> CanonicalAttackType {
// classtype strings come lowercase+trimmed from `translate`
match label {
// Scan / reconnaissance
"attempted-recon" | "network-scan" | "misc-activity" => CanonicalAttackType::Reconnaissance,
// Exploits / admin compromise
"attempted-admin" | "successful-admin" | "attempted-user" | "successful-user" | "shellcode-detect"
| "attempted-exploit" => CanonicalAttackType::Exploit,
// Web-application attacks
"web-application-attack" => CanonicalAttackType::Exploit,
"web-application-activity" => CanonicalAttackType::Exploit,
// SQL injection is usually emitted as web-application-attack, but some
// rule packs use "sql-injection" directly.
"sql-injection" => CanonicalAttackType::SqlInjection,
// XSS — same note as SQLi
"xss" | "cross-site-scripting" => CanonicalAttackType::Xss,
// DoS / DDoS
"attempted-dos" | "successful-dos" | "denial-of-service" => CanonicalAttackType::DosDdos,
// Trojan / malware / C2
"trojan-activity" | "malware-cnc" | "command-and-control" => CanonicalAttackType::C2Beacon,
// Credential attacks
"suspicious-login" | "unsuccessful-user" | "brute-force" => CanonicalAttackType::BruteForce,
// Policy / Crypto miner
"coin-mining" | "policy-violation" => CanonicalAttackType::Cryptomining,
// DNS tunneling detections emitted by some ET Open rules
"dns-tunnel" | "protocol-command-decode" => CanonicalAttackType::DnsTunnel,
_ => CanonicalAttackType::Unknown,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn wire_format_roundtrip() {
for canonical in CanonicalAttackType::ALL {
let wire = canonical.as_str();
let round = canonical_from_str(wire).expect("wire format must round-trip");
assert_eq!(*canonical, round, "roundtrip mismatch for {wire}");
}
}
#[test]
fn unknown_wire_string_returns_none() {
assert!(canonical_from_str("this_attack_does_not_exist").is_none());
assert!(canonical_from_str("").is_none());
}
#[test]
fn ml_v10_class_names_translate() {
// v10 ships 10 classes — every non-"Normal" class must hit a canonical entry.
// "Normal" is a legitimate benign class and IS expected to return Unknown
// because translators should not be called on benign flows in the first place.
let cases = [
("Brute Force", CanonicalAttackType::BruteForce),
("Reconnaissance", CanonicalAttackType::Reconnaissance),
("C2 Communication", CanonicalAttackType::C2Beacon),
("DoS/DDoS", CanonicalAttackType::DosDdos),
("Exploitation", CanonicalAttackType::Exploit),
("Web Attack", CanonicalAttackType::Exploit),
("Bot", CanonicalAttackType::BotActivity),
("Malware", CanonicalAttackType::BotActivity),
("Cryptomining", CanonicalAttackType::Cryptomining),
];
for (raw, expected) in cases {
assert_eq!(
translate(DetectionSource::ML, raw),
expected,
"ML class {raw:?} must translate to {expected:?}",
);
}
}
#[test]
fn beaconing_always_c2() {
assert_eq!(
translate(DetectionSource::Beaconing, "anything"),
CanonicalAttackType::C2Beacon
);
assert_eq!(
translate(DetectionSource::Beaconing, "c2_beacon"),
CanonicalAttackType::C2Beacon
);
}
#[test]
fn correlation_subtypes_split() {
assert_eq!(
translate(DetectionSource::Correlation, "scan"),
CanonicalAttackType::PortScan
);
assert_eq!(
translate(DetectionSource::Correlation, "lateral"),
CanonicalAttackType::LateralMovement
);
assert_eq!(
translate(DetectionSource::Correlation, "botnet"),
CanonicalAttackType::BotActivity
);
}
#[test]
fn suricata_classtype_mapping() {
let cases = [
("attempted-admin", CanonicalAttackType::Exploit),
("web-application-attack", CanonicalAttackType::Exploit),
("trojan-activity", CanonicalAttackType::C2Beacon),
("attempted-recon", CanonicalAttackType::Reconnaissance),
("attempted-dos", CanonicalAttackType::DosDdos),
("coin-mining", CanonicalAttackType::Cryptomining),
("brute-force", CanonicalAttackType::BruteForce),
("sql-injection", CanonicalAttackType::SqlInjection),
];
for (raw, expected) in cases {
assert_eq!(
translate(DetectionSource::Suricata, raw),
expected,
"Suricata classtype {raw:?} must translate to {expected:?}",
);
}
}
#[test]
fn unknown_label_lands_in_unknown_bucket() {
// Unknown is the fallback — MUST NOT panic, MUST be dedup-safe.
for source in [
DetectionSource::ML,
DetectionSource::Correlation,
DetectionSource::Suricata,
] {
assert_eq!(
translate(source, "this_is_not_a_real_label"),
CanonicalAttackType::Unknown,
);
}
}
#[test]
fn case_insensitive_and_whitespace_tolerant() {
assert_eq!(
translate(DetectionSource::ML, " BRUTE FORCE "),
CanonicalAttackType::BruteForce,
);
assert_eq!(
translate(DetectionSource::Suricata, "Attempted-Admin"),
CanonicalAttackType::Exploit,
);
}
#[test]
fn dedup_key_non_collision_across_sources() {
// The central invariant: two sources hitting the SAME attack on the
// SAME src_ip produce the SAME canonical wire string, so dedup collides.
let ml = translate(DetectionSource::ML, "Brute Force");
let suricata = translate(DetectionSource::Suricata, "brute-force");
assert_eq!(
ml.as_str(),
suricata.as_str(),
"cross-source dedup key MUST match for the same canonical attack",
);
}
#[test]
fn all_canonical_have_unique_wire_strings() {
use std::collections::HashSet;
let strings: HashSet<&str> = CanonicalAttackType::ALL.iter().map(|c| c.as_str()).collect();
assert_eq!(
strings.len(),
CanonicalAttackType::ALL.len(),
"duplicate wire string in CanonicalAttackType::ALL",
);
}
}

View File

@ -1,3 +1,5 @@
pub mod attack_type;
pub mod drift;
pub mod flow_features;
pub mod ml_detection;
pub mod model_source;

View File

@ -0,0 +1,113 @@
//! Wire-level types for broadcasting ML source state to the frontend.
//!
//! Two layers:
//! - `ModelInfo` — stable facts about a loaded model (name, adapter kind, feature count).
//! - `ModelSourceStatus` — current state of the ML source: Dormant / Active / Error.
//!
//! The internal `core::ml::state::ModelSourceState` holds the actual model
//! adapter plus this metadata; it converts into `ModelSourceStatus` for
//! WebSocket / HTTP responses via `From`.
use std::time::SystemTime;
use serde::{Deserialize, Serialize};
/// Public facts about a currently-loaded model. Drives the UI's ML Status panel.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelInfo {
/// Human-readable manifest name (e.g., "netguardia-v10").
pub name: String,
/// One of "autoencoder_only" | "classifier_only" | "multi_task".
pub adapter_kind: String,
/// When this model became Active (unix epoch seconds).
pub loaded_at_secs: u64,
/// Number of features the manifest declares. Useful for UI "31 features" display.
pub features_count: usize,
/// Recent inference QPS (rolling window). Zero before first tick.
pub qps_recent: f32,
}
impl ModelInfo {
pub fn new(name: String, adapter_kind: String, features_count: usize) -> Self {
let loaded_at_secs = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0);
Self {
name,
adapter_kind,
loaded_at_secs,
features_count,
qps_recent: 0.0,
}
}
}
/// Wire-format ML source status. Broadcast to frontend; returned by
/// `GET /api/models/current`. Serde-tagged so the frontend can discriminate
/// on the `state` field.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "state", rename_all = "snake_case")]
pub enum ModelSourceStatus {
/// No model loaded. Day 1 default; fusion runs with 3 sources.
Dormant,
/// Model loaded and serving inference. `info` populates the UI card.
Active { info: ModelInfo },
/// Last load attempt failed. UI shows the reason in red.
/// `since_secs` is unix epoch seconds; `last_attempted_path` is the
/// file that failed (if any).
Error {
msg: String,
since_secs: u64,
#[serde(default, skip_serializing_if = "Option::is_none")]
last_attempted_path: Option<String>,
},
}
impl ModelSourceStatus {
pub fn is_active(&self) -> bool {
matches!(self, Self::Active { .. })
}
pub fn is_dormant(&self) -> bool {
matches!(self, Self::Dormant)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn serde_roundtrip_dormant() {
let json = serde_json::to_string(&ModelSourceStatus::Dormant).unwrap();
assert!(json.contains("\"state\":\"dormant\""));
let back: ModelSourceStatus = serde_json::from_str(&json).unwrap();
assert!(back.is_dormant());
}
#[test]
fn serde_roundtrip_active() {
let status = ModelSourceStatus::Active {
info: ModelInfo::new("netguardia-v10".into(), "multi_task".into(), 31),
};
let json = serde_json::to_string(&status).unwrap();
assert!(json.contains("\"state\":\"active\""));
assert!(json.contains("\"adapter_kind\":\"multi_task\""));
let back: ModelSourceStatus = serde_json::from_str(&json).unwrap();
assert!(back.is_active());
}
#[test]
fn serde_roundtrip_error() {
let status = ModelSourceStatus::Error {
msg: "feature mismatch".into(),
since_secs: 1_700_000_000,
last_attempted_path: Some("models/.staging/bad.onnx".into()),
};
let json = serde_json::to_string(&status).unwrap();
assert!(json.contains("\"state\":\"error\""));
let back: ModelSourceStatus = serde_json::from_str(&json).unwrap();
assert!(matches!(back, ModelSourceStatus::Error { .. }));
}
}

View File

@ -1,12 +1,13 @@
use std::fmt;
use std::str::FromStr;
use crate::interface::communication::event::Event;
// -- Detection Source ---------------------------------------------------------
/// Identifies which detection subsystem produced a detection.
/// Used for attribution tracking and future cross-source deduplication.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
/// Used for attribution tracking and cross-source deduplication.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum DetectionSource {
ML,
Correlation,
@ -25,6 +26,22 @@ impl fmt::Display for DetectionSource {
}
}
impl FromStr for DetectionSource {
type Err = ();
/// Accepts the canonical `Display` form plus common aliases so playbook
/// authors can write `"CV"` for Beaconing or `"Graph"` for Correlation.
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s {
"ML" | "ml" => Ok(Self::ML),
"Suricata" | "suricata" => Ok(Self::Suricata),
"Beaconing" | "beaconing" | "CV" | "cv" => Ok(Self::Beaconing),
"Correlation" | "correlation" | "Graph" | "graph" => Ok(Self::Correlation),
_ => Err(()),
}
}
}
// -- Detection Event (internal pipeline) --------------------------------------
/// Raw detection from any source. Sent via mpsc channel to DetectionOrchestrator.
@ -68,9 +85,19 @@ pub struct ThreatDetectedEvent {
pub geoip_country: Option<String>,
/// Whether this src_ip had a block action in the past 24h
pub is_repeat_offender: bool,
/// Which detection sources contributed to this threat (for attribution)
/// Which detection sources contributed to this threat (for attribution).
/// Single-source events have length 1; fused events have 2..=4.
pub sources: Vec<DetectionSource>,
/// Per-model scores for debugging false positives
/// Number of distinct sources that contributed to this event.
/// Used by SOAR `MultiSourceMin` / `SingleSourceHigh` conditions —
/// counts unique sources, not per-source fires within the window.
pub active_source_count: usize,
/// Cross-source fused confidence (1 ∏(1 c_i)). Equal to
/// `confidence` after the fusion engine runs; retained as a separate
/// field so SOAR policies can discriminate "single source" from
/// "fused multi-source" numerically identical confidences.
pub fused_confidence: f32,
/// Per-model scores for debugging false positives.
pub ae_score: f32,
pub anomaly_score: f32,
pub c2_score: f32,

View File

@ -44,5 +44,17 @@ loggable! {
#[error("Correlation cleanup: removed {removed} expired entries")]
CorrelationCleanup { removed: usize } => tracing::Level::DEBUG,
#[error("Fusion: emitted {source_ip} {attack_type} fused={fused:.3} sources={count}")]
FusionEmitted { source_ip: String, attack_type: String, fused: f32, count: usize } => tracing::Level::DEBUG,
#[error("Fusion: window evicted under LRU pressure ({key_src} {key_type})")]
FusionWindowEvicted { key_src: String, key_type: String } => tracing::Level::WARN,
#[error("Fusion: failed to publish ThreatDetectedEvent: {err}")]
FusionPublishFailed { err: String } => tracing::Level::ERROR,
#[error("Fusion: failed to publish AuditEvent: {err}")]
FusionAuditPublishFailed { err: String } => tracing::Level::ERROR,
}
}

View File

@ -92,5 +92,8 @@ loggable! {
#[error("Invalid operator '{operator}' for condition type '{condition_type}' on playbook '{name}', condition skipped")]
InvalidConditionOperator { name: String, condition_type: String, operator: String } => tracing::Level::WARN,
#[error("Playbook '{name}' uses non-canonical trigger_event '{trigger_event}' — cross-source fusion dedup may silently miss this rule")]
NonCanonicalTriggerEvent { name: String, trigger_event: String } => tracing::Level::WARN,
}
}

View File

@ -16,6 +16,17 @@ pub enum ConditionType {
RepeatOffender,
/// Frequency: N events from same source_ip within window_secs
Frequency,
/// Multi-source agreement: `event.active_source_count >= value`.
MultiSourceMin,
/// Solo high-confidence escape hatch: exactly one contributing source,
/// and that source matches a specific name with confidence ≥ `value2`.
/// Lets a single high-confidence signature-class detection block
/// without waiting on peer agreement.
SingleSourceHigh,
/// Fused confidence above threshold — reads `event.fused_confidence`
/// rather than the per-event `confidence`, so a single high-confidence
/// event doesn't pass a threshold intended for multi-source agreement.
FusedConfidenceAbove,
}
impl fmt::Display for ConditionType {
@ -26,6 +37,9 @@ impl fmt::Display for ConditionType {
Self::IpPattern => write!(f, "ip_pattern"),
Self::RepeatOffender => write!(f, "repeat_offender"),
Self::Frequency => write!(f, "frequency"),
Self::MultiSourceMin => write!(f, "multi_source_min"),
Self::SingleSourceHigh => write!(f, "single_source_high"),
Self::FusedConfidenceAbove => write!(f, "fused_confidence_above"),
}
}
}
@ -40,6 +54,9 @@ impl FromStr for ConditionType {
"ip_pattern" => Ok(Self::IpPattern),
"repeat_offender" => Ok(Self::RepeatOffender),
"frequency" => Ok(Self::Frequency),
"multi_source_min" => Ok(Self::MultiSourceMin),
"single_source_high" => Ok(Self::SingleSourceHigh),
"fused_confidence_above" => Ok(Self::FusedConfidenceAbove),
other => Err(SoarError::UnknownConditionType(other)),
}
}

View File

@ -100,6 +100,11 @@ pub struct MLInferenceConfig {
pub c2_threshold: f32,
#[serde(default = "default_class_min_confidence")]
pub class_min_confidence: f32,
/// Multiplier applied to the confidence threshold before the aggregator
/// fires an alert. The manifest can override this via
/// `thresholds.alert_multiplier`.
#[serde(default = "default_alert_threshold_multiplier")]
pub alert_threshold_multiplier: f32,
pub model_type: String,
pub output_names: Vec<String>,
pub ae_feature_weights: HashMap<String, f64>,
@ -109,6 +114,10 @@ fn default_class_min_confidence() -> f32 {
0.4
}
fn default_alert_threshold_multiplier() -> f32 {
1.2
}
impl MLInferenceConfig {
pub fn num_ae_features(&self) -> usize {
self.ae_feature_names.len()