feat(v12-m1): BYO-model foundation — manifest loader + FEATURE_REGISTRY

Land v12 M1. Promote the 88-arm feature match in PrecomputedStats::get()
to a FEATURE_REGISTRY (name → fn pointer) keyed by every legacy alias;
manifest validation consults it via feature_is_known(). Add a YAML
ModelManifest parser with AdapterKind { classifier_only, autoencoder_only,
multi_task }, per-label metadata (confirmations/playbook — parsed, wired
in M2), thresholds, and a scaler sidecar pointer. Introspect ONNX input
shape via tract's InferenceFact and cross-check against the manifest-
declared feature count; mismatches surface as MLError::FeatureMismatch
with both declared and onnx dims. Compose runtime MLInferenceConfig from
(manifest + scaler sidecar JSON) with order-sensitive feature reconcili-
ation — any drift between the two is rejected, not silently overridden.

service_factory prefers models/manifest.yaml when present; the legacy
JSON-only path remains as fallback. For multi_task manifests, model
filenames come from the manifest, not AppConfig. Ship models/manifest.yaml
describing the current v10 multi_task model (31 AE features, 10 labels).

Incidental clippy fixes bundled: reorder mod tests / impl blocks in
adapter/http/setup.rs and infrastructure/secret_store.rs to silence
items_after_test_module under --tests; rustfmt polish on unrelated drift.

New error variants: ManifestInvalid, FeatureMismatch, UnknownFeature.
New log variants: ManifestLoaded, OnnxShapeChecked.

Tests: +13 ML-layer tests — registry coverage of v10 features, alias
consistency, safe_div zero-denominator, manifest parse/validate, sidecar
feature-count rejection, real v10 ONNX introspection (AE=31,
classifier=32), end-to-end load_named via manifest paths. Suite 157/157;
cargo clippy --package net-guardia --tests -- -D warnings clean.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
DaLaw2 2026-04-16 21:50:30 +08:00
parent c6a77c9611
commit 86358074b9
21 changed files with 1169 additions and 226 deletions

20
Cargo.lock generated
View File

@ -2381,6 +2381,7 @@ dependencies = [
"sd-notify",
"serde",
"serde_json",
"serde_yaml_ng",
"sha2",
"sysinfo",
"thiserror 2.0.18",
@ -3437,6 +3438,19 @@ dependencies = [
"serde",
]
[[package]]
name = "serde_yaml_ng"
version = "0.10.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7b4db627b98b36d4203a7b458cf3573730f2bb591b28871d916dfa9efabfd41f"
dependencies = [
"indexmap",
"itoa",
"ryu",
"serde",
"unsafe-libyaml",
]
[[package]]
name = "sha1"
version = "0.10.6"
@ -4213,6 +4227,12 @@ dependencies = [
"subtle",
]
[[package]]
name = "unsafe-libyaml"
version = "0.2.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "673aac59facbab8a9007c7f6108d11f63b603f7cabff99fabf650fea5c32b861"
[[package]]
name = "untrusted"
version = "0.9.0"

View File

@ -20,6 +20,7 @@ libxdp-sys = { version = "0.2.4", features = ["use_cc_build", "use_precompiled_b
# Serialization
serde = { version = "1.0.228", features = ["derive"] }
serde_json = "1.0.149"
serde_yaml_ng = "0.10.0"
# Async runtime
tokio = { version = "1.50.0", features = ["rt-multi-thread", "macros", "sync", "time", "process", "io-util", "fs", "signal"] }

69
models/manifest.yaml Normal file
View File

@ -0,0 +1,69 @@
# 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.
name: netguardia-v10
adapter: multi_task
models:
autoencoder: deep_autoencoder.onnx
classifier: classifier.onnx
# 31 AE-input features. Order matters — must match ONNX input column order
# and inference_config.json `ae_feature_names`. The classifier takes these
# plus `ae_anomaly_score` appended as the 32nd input (handled in code).
features:
- flow_duration
- fwd_packets
- bwd_packets
- fwd_bytes
- bwd_bytes
- flow_bytes_per_sec
- flow_pkts_per_sec
- fwd_win_bytes
- bwd_win_bytes
- fwd_pkt_len_mean
- bwd_pkt_len_mean
- fwd_iat_mean
- bwd_iat_mean
- flow_iat_mean
- pkt_len_mean
- dst_port
- protocol
- psh_flag_cnt
- ack_flag_cnt
- syn_flag_cnt
- fin_flag_cnt
- rst_flag_cnt
- pkt_len_std
- fwd_pkt_len_std
- bwd_pkt_len_std
- fwd_seg_size_min
- fwd_act_data_pkts
- fwd_iat_std
- bwd_iat_std
- fwd_bwd_bytes_ratio
- iat_cv
# `confirmations` / `playbook` are parsed today; M2 wires them into the aggregator.
labels:
"0": { name: Bot }
"1": { name: Brute Force }
"2": { name: C2 Communication, confirmations: 1 }
"3": { name: DNS Tunneling, confirmations: 1 }
"4": { name: DoS/DDoS, confirmations: 2 }
"5": { name: Exploitation, confirmations: 1 }
"6": { name: Malware }
"7": { name: Normal }
"8": { name: Reconnaissance }
"9": { name: Web Attack }
thresholds:
anomaly: 0.9179317355155945
c2: 0.9085615873336792
class_min_confidence: 0.4
ae: 0.23011694848537445
preprocessing:
scaler_sidecar: inference_config.json

View File

@ -28,6 +28,7 @@ tokio-tungstenite = "0.28.0"
# Serialization
serde = { workspace = true }
serde_json = { workspace = true }
serde_yaml_ng = { workspace = true }
toml = "1.0.7"
# Async

View File

@ -172,47 +172,6 @@ async fn complete_setup(
}))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_valid_interface_names() {
assert!(is_valid_interface_name("eth0"));
assert!(is_valid_interface_name("ens33"));
assert!(is_valid_interface_name("br-lan"));
assert!(is_valid_interface_name("wlan0.1"));
}
#[test]
fn test_invalid_interface_empty() {
assert!(!is_valid_interface_name(""));
}
#[test]
fn test_invalid_interface_too_long() {
let long = "a".repeat(17);
assert!(!is_valid_interface_name(&long));
// Exactly 16 should be valid
let exact = "a".repeat(16);
assert!(is_valid_interface_name(&exact));
}
#[test]
fn test_invalid_interface_path_traversal() {
assert!(!is_valid_interface_name("../etc"));
assert!(!is_valid_interface_name("../../shadow"));
assert!(!is_valid_interface_name("/sys/class"));
}
#[test]
fn test_invalid_interface_special_chars() {
assert!(!is_valid_interface_name("eth0;rm"));
assert!(!is_valid_interface_name("lo&&cat"));
assert!(!is_valid_interface_name("eth0 space"));
}
}
fn save_config(
db: &Database,
secrets: &dyn SecretStorePort,
@ -258,3 +217,44 @@ fn save_config(
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_valid_interface_names() {
assert!(is_valid_interface_name("eth0"));
assert!(is_valid_interface_name("ens33"));
assert!(is_valid_interface_name("br-lan"));
assert!(is_valid_interface_name("wlan0.1"));
}
#[test]
fn test_invalid_interface_empty() {
assert!(!is_valid_interface_name(""));
}
#[test]
fn test_invalid_interface_too_long() {
let long = "a".repeat(17);
assert!(!is_valid_interface_name(&long));
// Exactly 16 should be valid
let exact = "a".repeat(16);
assert!(is_valid_interface_name(&exact));
}
#[test]
fn test_invalid_interface_path_traversal() {
assert!(!is_valid_interface_name("../etc"));
assert!(!is_valid_interface_name("../../shadow"));
assert!(!is_valid_interface_name("/sys/class"));
}
#[test]
fn test_invalid_interface_special_chars() {
assert!(!is_valid_interface_name("eth0;rm"));
assert!(!is_valid_interface_name("lo&&cat"));
assert!(!is_valid_interface_name("eth0 space"));
}
}

View File

@ -1615,11 +1615,9 @@ impl Database {
let tx = conn.transaction()?;
let prev_hash: String = tx
.query_row(
"SELECT row_hash FROM audit_log ORDER BY id DESC LIMIT 1",
[],
|row| row.get(0),
)
.query_row("SELECT row_hash FROM audit_log ORDER BY id DESC LIMIT 1", [], |row| {
row.get(0)
})
.unwrap_or_default();
let row_hash = Self::audit_row_hash(&ts, actor, action, detail, &prev_hash);

View File

@ -306,10 +306,7 @@ impl<T: NativeConvert + Pod> EntryMap<T> {
let Some(map) = self.map.as_ref() else {
return Vec::new();
};
map.keys()
.filter_map(Result::ok)
.map(|key| key.into_native())
.collect()
map.keys().filter_map(Result::ok).map(|key| key.into_native()).collect()
}
fn add(&mut self, key: T::Native) -> Result<(), Error> {

View File

@ -1,29 +1,199 @@
use std::collections::HashMap;
use std::fs;
use std::path::PathBuf;
use std::path::{Path, PathBuf};
use super::manifest::{LabelSpec, ModelManifest};
use crate::model::error::ml::MLError;
use crate::model::system::config::MLInferenceConfig;
impl MLInferenceConfig {
pub fn load_file(file: &str) -> Result<Self, MLError> {
let path = PathBuf::from("models").join(file);
let content = fs::read_to_string(&path).map_err(|_| MLError::ConfigLoadFailed(path.to_path_buf()))?;
Self::load_file_at(&path)
}
fn load_file_at(path: &Path) -> Result<Self, MLError> {
let content = fs::read_to_string(path).map_err(|_| MLError::ConfigLoadFailed(path.to_path_buf()))?;
let config: MLInferenceConfig =
serde_json::from_str(&content).map_err(|e| MLError::ConfigParseFailed(e.to_string()))?;
if config.ae_feature_names.is_empty() {
return Err(MLError::ConfigParseFailed("ae_feature_names is empty"));
}
if config.ae_scaler_mean.len() != config.ae_feature_names.len() {
return Err(MLError::ConfigParseFailed("scaler mean length mismatch"));
}
if config.ae_scaler_std.len() != config.ae_feature_names.len() {
return Err(MLError::ConfigParseFailed("scaler std length mismatch"));
}
if config.output_names.len() != 3 {
return Err(MLError::ConfigParseFailed(
"MultiTaskModel requires exactly 3 output_names (anomaly, class_probs, c2_score)",
));
}
validate(&config)?;
Ok(config)
}
/// Load a `MLInferenceConfig` by combining a manifest with its scaler sidecar.
/// The manifest supplies the authoritative feature list and per-label metadata;
/// the sidecar JSON supplies the numeric preprocessing arrays (scaler / clip / weights).
///
/// Consistency is enforced: manifest `features` must match the sidecar's
/// `ae_feature_names` exactly (order-sensitive). Any drift between the two is a
/// deployment bug, not a silent override.
pub fn from_manifest_with_sidecar(manifest_path: &Path) -> Result<(Self, ModelManifest), MLError> {
let manifest = ModelManifest::load(manifest_path)?;
let sidecar_rel = manifest
.preprocessing
.as_ref()
.map(|p| p.scaler_sidecar.as_str())
.ok_or_else(|| {
MLError::ManifestInvalid(
manifest_path.to_path_buf(),
"preprocessing.scaler_sidecar is required for v1 (scaler arrays live there)".to_string(),
)
})?;
let sidecar_path = ModelManifest::resolve_relative(manifest_path, sidecar_rel);
let mut config = Self::load_file_at(&sidecar_path)?;
reconcile_features(&manifest, &config, manifest_path)?;
apply_manifest_overrides(&manifest, &mut config);
Ok((config, manifest))
}
}
fn validate(config: &MLInferenceConfig) -> Result<(), MLError> {
if config.ae_feature_names.is_empty() {
return Err(MLError::ConfigParseFailed("ae_feature_names is empty"));
}
if config.ae_scaler_mean.len() != config.ae_feature_names.len() {
return Err(MLError::ConfigParseFailed("scaler mean length mismatch"));
}
if config.ae_scaler_std.len() != config.ae_feature_names.len() {
return Err(MLError::ConfigParseFailed("scaler std length mismatch"));
}
if config.output_names.len() != 3 {
return Err(MLError::ConfigParseFailed(
"MultiTaskModel requires exactly 3 output_names (anomaly, class_probs, c2_score)",
));
}
Ok(())
}
fn reconcile_features(
manifest: &ModelManifest,
config: &MLInferenceConfig,
manifest_path: &Path,
) -> Result<(), MLError> {
if manifest.features.len() != config.ae_feature_names.len() {
return Err(MLError::ManifestInvalid(
manifest_path.to_path_buf(),
format!(
"feature count mismatch with sidecar: manifest declares {}, sidecar lists {}",
manifest.features.len(),
config.ae_feature_names.len()
),
));
}
for (i, (mf, sf)) in manifest.features.iter().zip(config.ae_feature_names.iter()).enumerate() {
if mf != sf {
return Err(MLError::ManifestInvalid(
manifest_path.to_path_buf(),
format!("feature[{i}] mismatch: manifest='{mf}' vs sidecar='{sf}'"),
));
}
}
Ok(())
}
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;
}
if let Some(v) = manifest.thresholds.c2 {
config.c2_threshold = v;
}
if let Some(v) = manifest.thresholds.class_min_confidence {
config.class_min_confidence = v;
}
if let Some(v) = manifest.thresholds.ae {
config.ae_threshold = v;
}
}
fn manifest_labels_to_map(labels: &std::collections::BTreeMap<String, LabelSpec>) -> HashMap<String, String> {
labels.iter().map(|(k, v)| (k.clone(), v.name.clone())).collect()
}
#[cfg(test)]
mod tests {
use super::*;
/// Integration test: the shipped `models/manifest.yaml` must successfully pair
/// with its scaler sidecar to yield a valid `MLInferenceConfig`. Skipped silently
/// when run outside the repo root (no `models/` directory).
#[test]
fn v10_manifest_and_sidecar_load_successfully() {
let manifest_path = Path::new("models/manifest.yaml");
if !manifest_path.exists() {
eprintln!("skipping: models/manifest.yaml absent (not in repo root?)");
return;
}
let (cfg, manifest) = MLInferenceConfig::from_manifest_with_sidecar(manifest_path)
.expect("v10 manifest + sidecar should load cleanly");
assert_eq!(manifest.name, "netguardia-v10");
assert_eq!(cfg.ae_feature_names.len(), 31);
assert_eq!(cfg.classifier_feature_names.len(), 32);
// Label map came from manifest, not sidecar.
assert_eq!(cfg.attack_labels.get("0").map(String::as_str), Some("Bot"));
assert_eq!(cfg.attack_labels.get("7").map(String::as_str), Some("Normal"));
}
#[test]
fn feature_mismatch_between_manifest_and_sidecar_is_rejected() {
use crate::model::detection::ml_detection::ClipParams;
use std::io::Write;
// Build a minimal sidecar JSON with 2 features.
let sidecar = MLInferenceConfig {
ae_feature_names: vec!["flow_duration".into(), "fwd_packets".into()],
ae_clip_params: HashMap::from([
("flow_duration".into(), ClipParams { lower: 0.0, upper: 1.0 }),
("fwd_packets".into(), ClipParams { lower: 0.0, upper: 1.0 }),
]),
ae_scaler_mean: vec![0.0, 0.0],
ae_scaler_std: vec![1.0, 1.0],
ae_post_clip_min: -5.0,
ae_post_clip_max: 5.0,
ae_threshold: 0.5,
classifier_feature_names: vec!["flow_duration".into(), "fwd_packets".into(), "ae_anomaly_score".into()],
attack_labels: HashMap::new(),
anomaly_threshold: 0.5,
c2_threshold: 0.5,
class_min_confidence: 0.4,
model_type: "MultiTaskModel".into(),
output_names: vec!["anomaly".into(), "class_probs".into(), "c2_score".into()],
ae_feature_weights: HashMap::new(),
};
let tmp = std::env::temp_dir().join("netguardia-m1-mismatch-test");
std::fs::create_dir_all(&tmp).unwrap();
let sidecar_path = tmp.join("sidecar.json");
let manifest_path = tmp.join("manifest.yaml");
let mut f = std::fs::File::create(&sidecar_path).unwrap();
f.write_all(serde_json::to_string(&sidecar).unwrap().as_bytes())
.unwrap();
// Manifest lists 3 features, sidecar has 2 — must fail.
let manifest_yaml = r#"
name: test
adapter: multi_task
models:
autoencoder: ae.onnx
classifier: c.onnx
features:
- flow_duration
- fwd_packets
- dst_port
preprocessing:
scaler_sidecar: sidecar.json
"#;
std::fs::write(&manifest_path, manifest_yaml).unwrap();
let err =
MLInferenceConfig::from_manifest_with_sidecar(&manifest_path).expect_err("should reject count mismatch");
assert!(matches!(err, MLError::ManifestInvalid { .. }), "got {err:?}");
}
}

View File

@ -1,9 +1,22 @@
use std::collections::HashMap;
use std::sync::LazyLock;
use common::define::tcp_flags::*;
use super::flow_tracker::FlowData;
use crate::model::detection::flow_features::FlowFeatures;
use crate::model::detection::ml_detection::PacketData;
use crate::model::detection::flow_features::FlowFeatures;
/// Signature of a feature getter — takes precomputed flow statistics and returns
/// a single f64 feature value. Must be pure (no I/O, no allocation).
/// Kept module-private because `PrecomputedStats` is an implementation detail.
type FeatureGetter = fn(&PrecomputedStats) -> f64;
/// Returns true if `name` (canonical or alias) is present in FEATURE_REGISTRY.
/// Used by ModelManifest validation at load time.
pub fn feature_is_known(name: &str) -> bool {
FEATURE_REGISTRY.contains_key(name)
}
impl FlowFeatures {
pub fn extract(flow: &FlowData, feature_names: &[String]) -> Self {
@ -286,107 +299,206 @@ impl PrecomputedStats {
}
fn get(&self, feature_name: &str) -> f64 {
let safe_div = |a: f64, b: f64| if b > 0.0 { a / b } else { 0.0 };
match feature_name {
"Destination Port" | "Dst Port" | "dst_port" => self.dst_port,
"Protocol" | "protocol" => self.protocol,
"Flow Duration" | "flow_duration" => self.duration_us,
"Total Fwd Packets" | "Tot Fwd Pkts" | "fwd_packets" => self.fwd_count,
"Total Backward Packets" | "Tot Bwd Pkts" | "bwd_packets" => self.bwd_count,
"Total Length of Fwd Packets" | "TotLen Fwd Pkts" | "fwd_bytes" => self.fwd_total_bytes,
"Total Length of Bwd Packets" | "TotLen Bwd Pkts" | "bwd_bytes" => self.bwd_total_bytes,
"Fwd Packet Length Max" => self.fwd_len_max,
"Fwd Packet Length Min" => self.fwd_len_min,
"Fwd Packet Length Mean" | "Fwd Pkt Len Mean" | "fwd_pkt_len_mean" => self.fwd_len_mean,
"Fwd Packet Length Std" | "Fwd Pkt Len Std" | "fwd_pkt_len_std" => self.fwd_len_std,
"Bwd Packet Length Max" => self.bwd_len_max,
"Bwd Packet Length Min" => self.bwd_len_min,
"Bwd Packet Length Mean" | "Bwd Pkt Len Mean" | "bwd_pkt_len_mean" => self.bwd_len_mean,
"Bwd Packet Length Std" | "Bwd Pkt Len Std" | "bwd_pkt_len_std" => self.bwd_len_std,
"Flow Bytes/s" | "Flow Byts/s" | "flow_bytes_per_sec" => safe_div(self.total_bytes, self.duration_s),
"Flow Packets/s" | "Flow Pkts/s" | "flow_pkts_per_sec" => safe_div(self.total_count, self.duration_s),
"Flow IAT Mean" | "flow_iat_mean" => self.flow_iat_mean,
"Flow IAT Std" => self.flow_iat_std,
"Flow IAT Max" => self.flow_iat_max,
"Flow IAT Min" => self.flow_iat_min,
"Fwd IAT Total" => self.fwd_iat_total,
"Fwd IAT Mean" | "fwd_iat_mean" => self.fwd_iat_mean,
"Fwd IAT Std" => self.fwd_iat_std,
"Fwd IAT Max" => self.fwd_iat_max,
"Fwd IAT Min" => self.fwd_iat_min,
"Bwd IAT Total" => self.bwd_iat_total,
"Bwd IAT Mean" | "bwd_iat_mean" => self.bwd_iat_mean,
"Bwd IAT Std" => self.bwd_iat_std,
"Bwd IAT Max" => self.bwd_iat_max,
"Bwd IAT Min" => self.bwd_iat_min,
"Fwd PSH Flags" => self.fwd_psh,
"Bwd PSH Flags" => self.bwd_psh,
"Fwd URG Flags" => self.fwd_urg,
"Bwd URG Flags" => self.bwd_urg,
"Fwd Header Length" => self.fwd_header_bytes,
"Bwd Header Length" => self.bwd_header_bytes,
"Fwd Packets/s" => safe_div(self.fwd_count, self.duration_s),
"Bwd Packets/s" => safe_div(self.bwd_count, self.duration_s),
"Min Packet Length" => self.all_len_min,
"Max Packet Length" => self.all_len_max,
"Packet Length Mean" | "Pkt Len Mean" | "pkt_len_mean" => self.all_len_mean,
"Packet Length Std" | "Pkt Len Std" | "pkt_len_std" => self.all_len_std,
"Packet Length Variance" => self.all_len_std * self.all_len_std,
"FIN Flag Count" | "FIN Flag Cnt" | "fin_flag_cnt" => self.fin_count,
"SYN Flag Count" | "SYN Flag Cnt" | "syn_flag_cnt" => self.syn_count,
"RST Flag Count" | "RST Flag Cnt" | "rst_flag_cnt" => self.rst_count,
"PSH Flag Count" | "PSH Flag Cnt" | "psh_flag_cnt" => self.psh_count,
"ACK Flag Count" | "ACK Flag Cnt" | "ack_flag_cnt" => self.ack_count,
"URG Flag Count" => self.urg_count,
"CWE Flag Count" => self.cwe_count,
"ECE Flag Count" => self.ece_count,
"Down/Up Ratio" => safe_div(self.bwd_count, self.fwd_count),
"Average Packet Size" => safe_div(self.total_bytes, self.total_count),
"Avg Fwd Segment Size" => safe_div(self.fwd_total_bytes, self.fwd_count),
"Avg Bwd Segment Size" => safe_div(self.bwd_total_bytes, self.bwd_count),
"Fwd Header Length.1" => self.fwd_header_bytes,
"Fwd Avg Bytes/Bulk" => self.fwd_avg_bytes_bulk,
"Fwd Avg Packets/Bulk" => self.fwd_avg_packets_bulk,
"Fwd Avg Bulk Rate" => self.fwd_avg_bulk_rate,
"Bwd Avg Bytes/Bulk" => self.bwd_avg_bytes_bulk,
"Bwd Avg Packets/Bulk" => self.bwd_avg_packets_bulk,
"Bwd Avg Bulk Rate" => self.bwd_avg_bulk_rate,
"Subflow Fwd Packets" => self.fwd_count,
"Subflow Fwd Bytes" => self.fwd_total_bytes,
"Subflow Bwd Packets" => self.bwd_count,
"Subflow Bwd Bytes" => self.bwd_total_bytes,
"fwd_win_bytes" => self.init_win_bytes_fwd,
"bwd_win_bytes" => self.init_win_bytes_bwd,
"fwd_act_data_pkts" => self.act_data_pkt_fwd,
"fwd_seg_size_min" => self.min_seg_size_forward,
"Active Mean" => self.active_mean,
"Active Std" => self.active_std,
"Active Max" => self.active_max,
"Active Min" => self.active_min,
"Idle Mean" => self.idle_mean,
"Idle Std" => self.idle_std,
"Idle Max" => self.idle_max,
"Idle Min" => self.idle_min,
// Phase 2: unified names for IAT std (already computed, add aliases)
"fwd_iat_std" => self.fwd_iat_std,
"bwd_iat_std" => self.bwd_iat_std,
"flow_iat_std" => self.flow_iat_std,
// Phase 2: new features for C2/Bot detection
"fwd_bwd_bytes_ratio" => self.fwd_bwd_bytes_ratio,
"pkt_len_variance" => self.all_len_std * self.all_len_std,
"fwd_iat_skewness" => self.fwd_iat_skewness,
// Model lists this in ae_feature_names but had no arm — was receiving post-scaler (0-μ)/σ.
"iat_cv" => safe_div(self.flow_iat_std, self.flow_iat_mean),
_ => 0.0,
}
FEATURE_REGISTRY.get(feature_name).map(|g| g(self)).unwrap_or(0.0)
}
}
fn reg_safe_div(a: f64, b: f64) -> f64 {
if b > 0.0 { a / b } else { 0.0 }
}
fn reg_insert(m: &mut HashMap<&'static str, FeatureGetter>, names: &[&'static str], g: FeatureGetter) {
for n in names {
m.insert(*n, g);
}
}
/// Central name → getter table. Every name the system recognizes for a feature
/// lives here. Manifest validation refuses any name not present in this map.
/// Aliases (long-form CICFlowMeter names, short snake_case) map to the same getter.
static FEATURE_REGISTRY: LazyLock<HashMap<&'static str, FeatureGetter>> = LazyLock::new(|| {
let mut m: HashMap<&'static str, FeatureGetter> = HashMap::new();
reg_insert(&mut m, &["Destination Port", "Dst Port", "dst_port"], |s| s.dst_port);
reg_insert(&mut m, &["Protocol", "protocol"], |s| s.protocol);
reg_insert(&mut m, &["Flow Duration", "flow_duration"], |s| s.duration_us);
reg_insert(
&mut m,
&[
"Total Fwd Packets",
"Tot Fwd Pkts",
"fwd_packets",
"Subflow Fwd Packets",
],
|s| s.fwd_count,
);
reg_insert(
&mut m,
&[
"Total Backward Packets",
"Tot Bwd Pkts",
"bwd_packets",
"Subflow Bwd Packets",
],
|s| s.bwd_count,
);
reg_insert(
&mut m,
&[
"Total Length of Fwd Packets",
"TotLen Fwd Pkts",
"fwd_bytes",
"Subflow Fwd Bytes",
],
|s| s.fwd_total_bytes,
);
reg_insert(
&mut m,
&[
"Total Length of Bwd Packets",
"TotLen Bwd Pkts",
"bwd_bytes",
"Subflow Bwd Bytes",
],
|s| s.bwd_total_bytes,
);
reg_insert(&mut m, &["Fwd Packet Length Max"], |s| s.fwd_len_max);
reg_insert(&mut m, &["Fwd Packet Length Min"], |s| s.fwd_len_min);
reg_insert(
&mut m,
&["Fwd Packet Length Mean", "Fwd Pkt Len Mean", "fwd_pkt_len_mean"],
|s| s.fwd_len_mean,
);
reg_insert(
&mut m,
&["Fwd Packet Length Std", "Fwd Pkt Len Std", "fwd_pkt_len_std"],
|s| s.fwd_len_std,
);
reg_insert(&mut m, &["Bwd Packet Length Max"], |s| s.bwd_len_max);
reg_insert(&mut m, &["Bwd Packet Length Min"], |s| s.bwd_len_min);
reg_insert(
&mut m,
&["Bwd Packet Length Mean", "Bwd Pkt Len Mean", "bwd_pkt_len_mean"],
|s| s.bwd_len_mean,
);
reg_insert(
&mut m,
&["Bwd Packet Length Std", "Bwd Pkt Len Std", "bwd_pkt_len_std"],
|s| s.bwd_len_std,
);
reg_insert(&mut m, &["Flow Bytes/s", "Flow Byts/s", "flow_bytes_per_sec"], |s| {
reg_safe_div(s.total_bytes, s.duration_s)
});
reg_insert(&mut m, &["Flow Packets/s", "Flow Pkts/s", "flow_pkts_per_sec"], |s| {
reg_safe_div(s.total_count, s.duration_s)
});
reg_insert(&mut m, &["Flow IAT Mean", "flow_iat_mean"], |s| s.flow_iat_mean);
reg_insert(&mut m, &["Flow IAT Std", "flow_iat_std"], |s| s.flow_iat_std);
reg_insert(&mut m, &["Flow IAT Max"], |s| s.flow_iat_max);
reg_insert(&mut m, &["Flow IAT Min"], |s| s.flow_iat_min);
reg_insert(&mut m, &["Fwd IAT Total"], |s| s.fwd_iat_total);
reg_insert(&mut m, &["Fwd IAT Mean", "fwd_iat_mean"], |s| s.fwd_iat_mean);
reg_insert(&mut m, &["Fwd IAT Std", "fwd_iat_std"], |s| s.fwd_iat_std);
reg_insert(&mut m, &["Fwd IAT Max"], |s| s.fwd_iat_max);
reg_insert(&mut m, &["Fwd IAT Min"], |s| s.fwd_iat_min);
reg_insert(&mut m, &["Bwd IAT Total"], |s| s.bwd_iat_total);
reg_insert(&mut m, &["Bwd IAT Mean", "bwd_iat_mean"], |s| s.bwd_iat_mean);
reg_insert(&mut m, &["Bwd IAT Std", "bwd_iat_std"], |s| s.bwd_iat_std);
reg_insert(&mut m, &["Bwd IAT Max"], |s| s.bwd_iat_max);
reg_insert(&mut m, &["Bwd IAT Min"], |s| s.bwd_iat_min);
reg_insert(&mut m, &["Fwd PSH Flags"], |s| s.fwd_psh);
reg_insert(&mut m, &["Bwd PSH Flags"], |s| s.bwd_psh);
reg_insert(&mut m, &["Fwd URG Flags"], |s| s.fwd_urg);
reg_insert(&mut m, &["Bwd URG Flags"], |s| s.bwd_urg);
// "Fwd Header Length" and "Fwd Header Length.1" are legacy CICFlowMeter aliases.
reg_insert(&mut m, &["Fwd Header Length", "Fwd Header Length.1"], |s| {
s.fwd_header_bytes
});
reg_insert(&mut m, &["Bwd Header Length"], |s| s.bwd_header_bytes);
reg_insert(&mut m, &["Fwd Packets/s"], |s| reg_safe_div(s.fwd_count, s.duration_s));
reg_insert(&mut m, &["Bwd Packets/s"], |s| reg_safe_div(s.bwd_count, s.duration_s));
reg_insert(&mut m, &["Min Packet Length"], |s| s.all_len_min);
reg_insert(&mut m, &["Max Packet Length"], |s| s.all_len_max);
reg_insert(&mut m, &["Packet Length Mean", "Pkt Len Mean", "pkt_len_mean"], |s| {
s.all_len_mean
});
reg_insert(&mut m, &["Packet Length Std", "Pkt Len Std", "pkt_len_std"], |s| {
s.all_len_std
});
reg_insert(&mut m, &["Packet Length Variance", "pkt_len_variance"], |s| {
s.all_len_std * s.all_len_std
});
reg_insert(&mut m, &["FIN Flag Count", "FIN Flag Cnt", "fin_flag_cnt"], |s| {
s.fin_count
});
reg_insert(&mut m, &["SYN Flag Count", "SYN Flag Cnt", "syn_flag_cnt"], |s| {
s.syn_count
});
reg_insert(&mut m, &["RST Flag Count", "RST Flag Cnt", "rst_flag_cnt"], |s| {
s.rst_count
});
reg_insert(&mut m, &["PSH Flag Count", "PSH Flag Cnt", "psh_flag_cnt"], |s| {
s.psh_count
});
reg_insert(&mut m, &["ACK Flag Count", "ACK Flag Cnt", "ack_flag_cnt"], |s| {
s.ack_count
});
reg_insert(&mut m, &["URG Flag Count"], |s| s.urg_count);
reg_insert(&mut m, &["CWE Flag Count"], |s| s.cwe_count);
reg_insert(&mut m, &["ECE Flag Count"], |s| s.ece_count);
reg_insert(&mut m, &["Down/Up Ratio"], |s| reg_safe_div(s.bwd_count, s.fwd_count));
reg_insert(&mut m, &["Average Packet Size"], |s| {
reg_safe_div(s.total_bytes, s.total_count)
});
reg_insert(&mut m, &["Avg Fwd Segment Size"], |s| {
reg_safe_div(s.fwd_total_bytes, s.fwd_count)
});
reg_insert(&mut m, &["Avg Bwd Segment Size"], |s| {
reg_safe_div(s.bwd_total_bytes, s.bwd_count)
});
reg_insert(&mut m, &["Fwd Avg Bytes/Bulk"], |s| s.fwd_avg_bytes_bulk);
reg_insert(&mut m, &["Fwd Avg Packets/Bulk"], |s| s.fwd_avg_packets_bulk);
reg_insert(&mut m, &["Fwd Avg Bulk Rate"], |s| s.fwd_avg_bulk_rate);
reg_insert(&mut m, &["Bwd Avg Bytes/Bulk"], |s| s.bwd_avg_bytes_bulk);
reg_insert(&mut m, &["Bwd Avg Packets/Bulk"], |s| s.bwd_avg_packets_bulk);
reg_insert(&mut m, &["Bwd Avg Bulk Rate"], |s| s.bwd_avg_bulk_rate);
reg_insert(&mut m, &["fwd_win_bytes"], |s| s.init_win_bytes_fwd);
reg_insert(&mut m, &["bwd_win_bytes"], |s| s.init_win_bytes_bwd);
reg_insert(&mut m, &["fwd_act_data_pkts"], |s| s.act_data_pkt_fwd);
reg_insert(&mut m, &["fwd_seg_size_min"], |s| s.min_seg_size_forward);
reg_insert(&mut m, &["Active Mean"], |s| s.active_mean);
reg_insert(&mut m, &["Active Std"], |s| s.active_std);
reg_insert(&mut m, &["Active Max"], |s| s.active_max);
reg_insert(&mut m, &["Active Min"], |s| s.active_min);
reg_insert(&mut m, &["Idle Mean"], |s| s.idle_mean);
reg_insert(&mut m, &["Idle Std"], |s| s.idle_std);
reg_insert(&mut m, &["Idle Max"], |s| s.idle_max);
reg_insert(&mut m, &["Idle Min"], |s| s.idle_min);
// Phase 2: C2/Bot-oriented features
reg_insert(&mut m, &["fwd_bwd_bytes_ratio"], |s| s.fwd_bwd_bytes_ratio);
reg_insert(&mut m, &["fwd_iat_skewness"], |s| s.fwd_iat_skewness);
reg_insert(&mut m, &["iat_cv"], |s| reg_safe_div(s.flow_iat_std, s.flow_iat_mean));
m
});
fn compute_stats(values: &[f64]) -> (f64, f64, f64, f64) {
if values.is_empty() {
return (0.0, 0.0, 0.0, 0.0);
@ -469,6 +581,188 @@ mod tests {
assert_eq!(compute_bowley_skewness(&[1.0, 1.0, 1.0, 1.0, 1.0, 1.0]), 0.0);
}
/// Hand-crafted PrecomputedStats with distinctive sentinel values per field.
/// Lets us verify registry getter dispatch without constructing a real FlowData.
fn sample_stats() -> PrecomputedStats {
PrecomputedStats {
dst_port: 443.0,
protocol: 6.0,
duration_us: 1_000_000.0,
fwd_count: 10.0,
bwd_count: 4.0,
total_count: 14.0,
fwd_total_bytes: 2000.0,
bwd_total_bytes: 800.0,
total_bytes: 2800.0,
duration_s: 1.0,
fwd_len_max: 1500.0,
fwd_len_min: 40.0,
fwd_len_mean: 200.0,
fwd_len_std: 300.0,
bwd_len_max: 1200.0,
bwd_len_min: 60.0,
bwd_len_mean: 200.0,
bwd_len_std: 250.0,
all_len_max: 1500.0,
all_len_min: 40.0,
all_len_mean: 200.0,
all_len_std: 280.0,
flow_iat_max: 50_000.0,
flow_iat_min: 100.0,
flow_iat_mean: 10_000.0,
flow_iat_std: 5_000.0,
fwd_iat_total: 90_000.0,
fwd_iat_max: 40_000.0,
fwd_iat_min: 200.0,
fwd_iat_mean: 10_000.0,
fwd_iat_std: 6_000.0,
bwd_iat_total: 30_000.0,
bwd_iat_max: 15_000.0,
bwd_iat_min: 300.0,
bwd_iat_mean: 7_500.0,
bwd_iat_std: 4_000.0,
fwd_psh: 2.0,
bwd_psh: 1.0,
fwd_urg: 0.0,
bwd_urg: 0.0,
fwd_header_bytes: 200.0,
bwd_header_bytes: 80.0,
fin_count: 1.0,
syn_count: 1.0,
rst_count: 0.0,
psh_count: 3.0,
ack_count: 10.0,
urg_count: 0.0,
cwe_count: 0.0,
ece_count: 0.0,
fwd_avg_bytes_bulk: 500.0,
fwd_avg_packets_bulk: 5.0,
fwd_avg_bulk_rate: 5000.0,
bwd_avg_bytes_bulk: 400.0,
bwd_avg_packets_bulk: 4.0,
bwd_avg_bulk_rate: 4000.0,
init_win_bytes_fwd: 65535.0,
init_win_bytes_bwd: 65000.0,
act_data_pkt_fwd: 8.0,
min_seg_size_forward: 40.0,
active_max: 1000.0,
active_min: 50.0,
active_mean: 300.0,
active_std: 200.0,
idle_max: 500.0,
idle_min: 10.0,
idle_mean: 100.0,
idle_std: 80.0,
fwd_bwd_bytes_ratio: 0.71,
fwd_iat_skewness: 0.15,
}
}
#[test]
fn registry_unknown_name_returns_zero() {
let s = sample_stats();
assert_eq!(s.get("not_a_feature"), 0.0);
}
#[test]
fn registry_aliases_resolve_identically() {
// Long-form, short-form, and "Subflow" aliases must all map to the same getter.
let s = sample_stats();
for group in [
[
"Total Fwd Packets",
"Tot Fwd Pkts",
"fwd_packets",
"Subflow Fwd Packets",
],
[
"Total Length of Fwd Packets",
"TotLen Fwd Pkts",
"fwd_bytes",
"Subflow Fwd Bytes",
],
["Flow IAT Std", "flow_iat_std", "Flow IAT Std", "Flow IAT Std"], // pad to 4
["Fwd IAT Std", "fwd_iat_std", "Fwd IAT Std", "Fwd IAT Std"],
["Bwd IAT Std", "bwd_iat_std", "Bwd IAT Std", "Bwd IAT Std"],
[
"Packet Length Variance",
"pkt_len_variance",
"Packet Length Variance",
"Packet Length Variance",
],
[
"Fwd Header Length",
"Fwd Header Length.1",
"Fwd Header Length",
"Fwd Header Length",
],
] {
let expected = s.get(group[0]);
for name in &group[1..] {
assert_eq!(
s.get(name),
expected,
"alias '{name}' should resolve to same value as '{}'",
group[0]
);
}
}
}
#[test]
fn registry_safe_div_returns_zero_on_zero_denominator() {
let mut s = sample_stats();
s.flow_iat_mean = 0.0;
s.flow_iat_std = 500.0;
// iat_cv = std / mean, but mean=0 → safe_div → 0.0
assert_eq!(s.get("iat_cv"), 0.0);
s.duration_s = 0.0;
assert_eq!(s.get("flow_bytes_per_sec"), 0.0);
assert_eq!(s.get("flow_pkts_per_sec"), 0.0);
}
#[test]
fn registry_covers_v10_manifest_features() {
// Every feature the shipped v10 manifest references must be registered.
// A missing name here means the match → registry refactor dropped a binding.
const V10_FEATURES: &[&str] = &[
"flow_duration",
"fwd_packets",
"bwd_packets",
"fwd_bytes",
"bwd_bytes",
"flow_bytes_per_sec",
"flow_pkts_per_sec",
"fwd_win_bytes",
"bwd_win_bytes",
"fwd_pkt_len_mean",
"bwd_pkt_len_mean",
"fwd_iat_mean",
"bwd_iat_mean",
"flow_iat_mean",
"pkt_len_mean",
"dst_port",
"protocol",
"psh_flag_cnt",
"ack_flag_cnt",
"syn_flag_cnt",
"fin_flag_cnt",
"rst_flag_cnt",
"pkt_len_std",
"fwd_pkt_len_std",
"bwd_pkt_len_std",
"fwd_seg_size_min",
"fwd_act_data_pkts",
"fwd_iat_std",
"bwd_iat_std",
"fwd_bwd_bytes_ratio",
"iat_cv",
];
for f in V10_FEATURES {
assert!(feature_is_known(f), "v10 feature '{f}' missing from FEATURE_REGISTRY");
}
}
#[test]
fn bowley_skewness_known_output() {
// Symmetric distribution: [1, 2, 3, 4, 5, 6, 7, 8] (n=8)

View File

@ -0,0 +1,229 @@
use std::collections::BTreeMap;
use std::fs;
use std::path::{Path, PathBuf};
use serde::{Deserialize, Serialize};
use super::feature_extractor::feature_is_known;
use crate::model::error::ml::MLError;
/// The three fixed adapter shapes a v1 manifest may declare.
/// The user picks one string; never sees the internal `ModelAdapter` trait.
#[derive(Debug, Clone, Copy, Deserialize, Serialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum AdapterKind {
ClassifierOnly,
AutoencoderOnly,
MultiTask,
}
impl AdapterKind {
pub fn as_str(self) -> &'static str {
match self {
AdapterKind::ClassifierOnly => "classifier_only",
AdapterKind::AutoencoderOnly => "autoencoder_only",
AdapterKind::MultiTask => "multi_task",
}
}
}
/// ONNX file references within a manifest. `classifier_only` / `autoencoder_only`
/// adapters use `model`; `multi_task` uses `autoencoder` + `classifier`.
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct ModelPaths {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub model: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub autoencoder: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
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.
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct LabelSpec {
pub name: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub confirmations: Option<usize>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub playbook: Option<String>,
}
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct Thresholds {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub anomaly: Option<f32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub c2: Option<f32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub class_min_confidence: Option<f32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub ae: Option<f32>,
}
/// Pointer to a preprocessing sidecar (scaler/clip arrays). For v10 this is
/// the legacy `inference_config.json`.
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Preprocessing {
pub scaler_sidecar: String,
}
/// The user-authored YAML description of a model. Validated at load — unknown
/// feature names are rejected against `FEATURE_REGISTRY`.
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ModelManifest {
pub name: String,
pub adapter: AdapterKind,
#[serde(default)]
pub models: ModelPaths,
pub features: Vec<String>,
#[serde(default)]
pub labels: BTreeMap<String, LabelSpec>,
#[serde(default)]
pub thresholds: Thresholds,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub preprocessing: Option<Preprocessing>,
}
impl ModelManifest {
/// Parse and validate a manifest YAML file from disk.
pub fn load(path: impl AsRef<Path>) -> Result<Self, MLError> {
let path = path.as_ref();
let content = fs::read_to_string(path)
.map_err(|e| MLError::ManifestInvalid(path.to_path_buf(), format!("read failed: {e}")))?;
let manifest: ModelManifest = serde_yaml_ng::from_str(&content)
.map_err(|e| MLError::ManifestInvalid(path.to_path_buf(), format!("YAML parse: {e}")))?;
manifest.validate(path)?;
Ok(manifest)
}
fn validate(&self, path: &Path) -> Result<(), MLError> {
if self.name.trim().is_empty() {
return Err(MLError::ManifestInvalid(
path.to_path_buf(),
"name is empty".to_string(),
));
}
if self.features.is_empty() {
return Err(MLError::ManifestInvalid(
path.to_path_buf(),
"features is empty".to_string(),
));
}
for f in &self.features {
if !feature_is_known(f) {
return Err(MLError::UnknownFeature(f.clone()));
}
}
match self.adapter {
AdapterKind::MultiTask => {
if self.models.autoencoder.is_none() || self.models.classifier.is_none() {
return Err(MLError::ManifestInvalid(
path.to_path_buf(),
"multi_task adapter requires both models.autoencoder and models.classifier".to_string(),
));
}
}
AdapterKind::ClassifierOnly | AdapterKind::AutoencoderOnly => {
if self.models.model.is_none() {
return Err(MLError::ManifestInvalid(
path.to_path_buf(),
format!("{} adapter requires models.model", self.adapter.as_str()),
));
}
}
}
Ok(())
}
/// Resolve a manifest-relative filename (e.g. "deep_autoencoder.onnx") to an
/// absolute-ish path rooted at the manifest's parent directory.
pub fn resolve_relative(manifest_path: &Path, relative: &str) -> PathBuf {
manifest_path.parent().unwrap_or_else(|| Path::new(".")).join(relative)
}
}
#[cfg(test)]
mod tests {
use super::*;
const V10_MANIFEST: &str = r#"
name: netguardia-v10
adapter: multi_task
models:
autoencoder: deep_autoencoder.onnx
classifier: classifier.onnx
features:
- flow_duration
- fwd_packets
- bwd_packets
labels:
"0": { name: Bot }
"7": { name: Normal }
thresholds:
anomaly: 0.9
c2: 0.85
preprocessing:
scaler_sidecar: inference_config.json
"#;
#[test]
fn parses_minimal_multitask() {
let m: ModelManifest = serde_yaml_ng::from_str(V10_MANIFEST).expect("parse");
assert_eq!(m.name, "netguardia-v10");
assert_eq!(m.adapter, AdapterKind::MultiTask);
assert_eq!(m.features.len(), 3);
assert_eq!(m.models.autoencoder.as_deref(), Some("deep_autoencoder.onnx"));
assert_eq!(m.models.classifier.as_deref(), Some("classifier.onnx"));
assert_eq!(m.labels.len(), 2);
assert_eq!(m.labels.get("0").map(|l| l.name.as_str()), Some("Bot"));
}
#[test]
fn rejects_unknown_feature() {
let yaml = r#"
name: bad
adapter: classifier_only
models:
model: m.onnx
features:
- this_feature_does_not_exist
"#;
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 unknown feature");
assert!(matches!(err, MLError::UnknownFeature { .. }), "got {err:?}");
}
#[test]
fn rejects_multitask_missing_ae() {
let yaml = r#"
name: bad
adapter: multi_task
models:
classifier: c.onnx
features:
- flow_duration
"#;
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 require autoencoder");
assert!(matches!(err, MLError::ManifestInvalid { .. }), "got {err:?}");
}
#[test]
fn rejects_empty_features() {
let yaml = r#"
name: bad
adapter: classifier_only
models:
model: m.onnx
features: []
"#;
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 empty features");
assert!(matches!(err, MLError::ManifestInvalid { .. }), "got {err:?}");
}
}

View File

@ -6,6 +6,7 @@ pub mod engine;
pub mod feature_extractor;
pub mod flow_tracker;
pub mod inference;
pub mod manifest;
pub mod model_loader;
pub mod model_watcher;
pub mod traffic_logger;

View File

@ -3,6 +3,8 @@ use std::time::Instant;
use macros::log;
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 crate::model::detection::ml_detection::RunnableModel;
@ -21,18 +23,26 @@ impl MLModels {
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,
)
}
/// 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(
&app_config.inference.deep_autoencoder_name,
inference_config.num_ae_features(),
batch_size,
)?,
classifier: Self::loader(
&app_config.inference.classifier_name,
inference_config.num_classifier_features(),
batch_size,
)?,
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,
})
}
@ -43,21 +53,44 @@ impl MLModels {
log!(MLLog::ModelLoading(model.to_string(), features, batch_size));
let start = Instant::now();
let load = || -> Result<RunnableModel, Box<dyn std::error::Error>> {
let mut model = onnx().model_for_path(&model_path)?;
model.set_input_fact(0, f32::fact([batch_size, features]).into())?;
Ok(model.into_optimized()?.into_runnable()?)
};
let result = load().map_err(|_| MLError::ModelLoadFailed(model_path));
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));
result
}
fn loader_inner(
model_path: &PathBuf,
model: &str,
features: usize,
batch_size: usize,
) -> Result<RunnableModel, MLError> {
let mut onnx_model = onnx()
.model_for_path(model_path)
.map_err(|_| MLError::ModelLoadFailed(model_path.clone()))?;
// 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.clone(), features, onnx_dim));
}
}
onnx_model
.set_input_fact(0, f32::fact([batch_size, features]).into())
.map_err(|_| MLError::ModelLoadFailed(model_path.clone()))?;
onnx_model
.into_optimized()
.and_then(|m| m.into_runnable())
.map_err(|_| MLError::ModelLoadFailed(model_path.clone()))
}
pub fn get_model_info(&self, name: &str) -> String {
let model = match name {
"deep_autoencoder" => &self.deep_autoencoder,
@ -73,3 +106,67 @@ impl MLModels {
)
}
}
/// Read the concrete last-dim (feature count) from an ONNX model's declared input fact.
/// Returns None when the dim is dynamic/symbolic or when the model has no input 0.
fn introspect_input_features(model: &InferenceModel) -> Option<usize> {
let fact = model.input_fact(0).ok()?;
let rank = fact.shape.rank().concretize()? as usize;
if rank == 0 {
return None;
}
let last = fact.shape.dim(rank - 1)?;
last.concretize()?.to_usize().ok()
}
#[cfg(test)]
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).
#[test]
fn introspect_v10_autoencoder_input_is_31() {
let ae_path = PathBuf::from("models/deep_autoencoder.onnx");
if !ae_path.exists() {
eprintln!("skipping: models/deep_autoencoder.onnx absent");
return;
}
let model = onnx().model_for_path(&ae_path).expect("load AE onnx");
let dim = introspect_input_features(&model).expect("v10 AE should expose a concrete final-dim");
assert_eq!(dim, 31, "v10 AE ONNX input dim changed unexpectedly");
}
#[test]
fn introspect_v10_classifier_input_is_32() {
let cls_path = PathBuf::from("models/classifier.onnx");
if !cls_path.exists() {
eprintln!("skipping: models/classifier.onnx absent");
return;
}
let model = onnx().model_for_path(&cls_path).expect("load classifier onnx");
let dim = introspect_input_features(&model).expect("v10 classifier should expose a concrete final-dim");
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.
#[test]
fn v10_load_named_with_manifest_paths_succeeds() {
let manifest_path = std::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");
}
}

View File

@ -359,11 +359,8 @@ impl System {
// Start Suricata eve.json monitor — tails the log file, translates
// alert events into DetectionEvent on the shared mpsc. No-op if the
// bridge is disabled in config.
crate::infrastructure::suricata_monitor::SuricataMonitor::new(
self.app_config.clone(),
suricata_detection_tx,
)
.start();
crate::infrastructure::suricata_monitor::SuricataMonitor::new(self.app_config.clone(), suricata_detection_tx)
.start();
// Wait for shutdown signal (ctrl-c OR API-triggered)
tokio::select! {

View File

@ -8,6 +8,7 @@ use tokio::sync::oneshot;
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::traffic_logger::TrafficLogger;
use crate::infrastructure::app_config::AppConfig;
@ -38,13 +39,32 @@ impl AppServices {
pub fn new(
app_config: Arc<AppConfig>,
inference_config: Arc<MLInferenceConfig>,
ml_manifest: Option<ModelManifest>,
drift_detector: Arc<parking_lot::Mutex<DriftDetector>>,
ebpf_health: Arc<parking_lot::RwLock<EbpfHealth>>,
) -> Result<Self, Error> {
let health = SystemHealth::new(app_config.clone(), ebpf_health)?;
let batch_size = app_config.inference.inference_batch_size;
let ml_models = Arc::new(MLModels::load_models(&app_config, &inference_config, 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)?
}
_ => MLModels::load_models(&app_config, &inference_config, batch_size)?,
});
let ml_alert = Arc::new(MLAlert::new());
let traffic_logger = if app_config.inference.traffic_logging_mode {

View File

@ -155,8 +155,14 @@ mod tests {
#[test]
fn categorizes_permission_errors() {
assert!(matches!(categorize("Permission denied (os error 13)"), EbpfFailCategory::Permission));
assert!(matches!(categorize("Operation not permitted"), EbpfFailCategory::Permission));
assert!(matches!(
categorize("Permission denied (os error 13)"),
EbpfFailCategory::Permission
));
assert!(matches!(
categorize("Operation not permitted"),
EbpfFailCategory::Permission
));
}
#[test]

View File

@ -147,7 +147,20 @@ impl SecretStore {
.into()),
}
}
}
impl SecretStorePort for SecretStore {
fn get_secret(&self, key: &str) -> Result<Option<String>, Error> {
match self.db.get_app_secret(key)? {
Some(envelope_json) => Ok(Some(self.decrypt(&envelope_json)?)),
None => Ok(None),
}
}
fn set_secret(&self, key: &str, plaintext: &str) -> Result<(), Error> {
let envelope = self.encrypt(plaintext)?;
self.db.set_app_secret(key, &envelope)
}
}
#[cfg(test)]
@ -255,17 +268,3 @@ mod tests {
assert!(store.decrypt(&envelope).is_err());
}
}
impl SecretStorePort for SecretStore {
fn get_secret(&self, key: &str) -> Result<Option<String>, Error> {
match self.db.get_app_secret(key)? {
Some(envelope_json) => Ok(Some(self.decrypt(&envelope_json)?)),
None => Ok(None),
}
}
fn set_secret(&self, key: &str, plaintext: &str) -> Result<(), Error> {
let envelope = self.encrypt(plaintext)?;
self.db.set_app_secret(key, &envelope)
}
}

View File

@ -20,6 +20,7 @@ use crate::core::dns_filter_service::DnsFilterService;
use crate::core::ebpf::EbpfServices;
use crate::core::email::scheduler::ReportScheduler;
use crate::core::ml::drift_detector::DriftDetector;
use crate::core::ml::manifest::ModelManifest;
use crate::core::notification_service::NotificationService;
use crate::core::playbook_service::PlaybookService;
use crate::core::rate_limit_service::RateLimitService;
@ -44,6 +45,7 @@ use crate::model::error::Error;
use crate::model::error::ebpf::EbpfError;
use crate::model::error::misc::MiscError;
use crate::model::log::ebpf::EbpfLog;
use crate::model::log::ml::MLLog;
use crate::model::log::system::SystemLog;
use crate::model::monitoring::direction::FlowDirection;
use crate::model::system::config::MLInferenceConfig;
@ -104,7 +106,26 @@ impl ServiceFactory {
AppConfig::seed_defaults(&db)?;
let app_config = Arc::new(AppConfig::new(&db)?);
let inference_config = Arc::new(MLInferenceConfig::load_file(&app_config.inference.models_config_name)?);
// Prefer `models/manifest.yaml` when present (v12 BYO-model path). The manifest
// is the user-authored source of truth for features, labels, thresholds, and
// model filenames; the legacy JSON-only path is the fallback.
let manifest_path = std::path::PathBuf::from("models/manifest.yaml");
let (inference_config, ml_manifest): (Arc<MLInferenceConfig>, Option<ModelManifest>) = if manifest_path.exists()
{
let (cfg, manifest) = MLInferenceConfig::from_manifest_with_sidecar(&manifest_path)?;
log!(MLLog::ManifestLoaded(
manifest.name.clone(),
manifest.adapter.as_str().to_string(),
manifest.features.len(),
manifest.labels.len(),
));
(Arc::new(cfg), Some(manifest))
} else {
(
Arc::new(MLInferenceConfig::load_file(&app_config.inference.models_config_name)?),
None,
)
};
// Shared eBPF health handle. Initialized Healthy; downgraded to
// Unavailable with a classified reason if any stage below fails.
@ -116,16 +137,21 @@ impl ServiceFactory {
// rest of the system (HTTP API, SOAR, ML engine, auth) is built
// regardless so the operator can still reach the frontend and see
// the reason.
let (ingress_ebpf, egress_ebpf, ingress_program_array, ebpf_services) =
match Self::try_build_ebpf(&app_config) {
Ok((ingress, egress, pa, services)) => (Some(ingress), Some(egress), Some(pa), Arc::new(services)),
Err((stage, err)) => {
let health = crate::infrastructure::ebpf_preflight::classify(stage, &err, None);
log!(SystemLog::EbpfBringupFailed(format!("{:?}", health)));
*ebpf_health.write() = health;
(None, None, None, Arc::new(EbpfServices::unavailable(app_config.clone())))
}
};
let (ingress_ebpf, egress_ebpf, ingress_program_array, ebpf_services) = match Self::try_build_ebpf(&app_config)
{
Ok((ingress, egress, pa, services)) => (Some(ingress), Some(egress), Some(pa), Arc::new(services)),
Err((stage, err)) => {
let health = crate::infrastructure::ebpf_preflight::classify(stage, &err, None);
log!(SystemLog::EbpfBringupFailed(format!("{:?}", health)));
*ebpf_health.write() = health;
(
None,
None,
None,
Arc::new(EbpfServices::unavailable(app_config.clone())),
)
}
};
// Ensure enforce_mode setting exists (default: monitor)
if db.get_setting("enforce_mode")?.is_none() {
@ -153,6 +179,7 @@ impl ServiceFactory {
let app_services = Arc::new(AppServices::new(
app_config.clone(),
inference_config.clone(),
ml_manifest.clone(),
drift_detector.clone(),
ebpf_health.clone(),
)?);
@ -267,8 +294,7 @@ impl ServiceFactory {
secret_store_port,
));
let suricata_manager =
crate::infrastructure::suricata_manager::SuricataManager::new(app_config.clone());
let suricata_manager = crate::infrastructure::suricata_manager::SuricataManager::new(app_config.clone());
Ok(AppState {
app_config,

View File

@ -71,18 +71,14 @@ impl SuricataManager {
loop {
// Pre-flight: validate binary + config exist before spawning.
if let Err(e) = Self::preflight(&self.config) {
*self.health.write() = SuricataHealth::Stopped {
reason: e.to_string(),
};
*self.health.write() = SuricataHealth::Stopped { reason: e.to_string() };
return;
}
let mut child = match self.spawn_child() {
Ok(c) => c,
Err(e) => {
*self.health.write() = SuricataHealth::Stopped {
reason: e.to_string(),
};
*self.health.write() = SuricataHealth::Stopped { reason: e.to_string() };
return;
}
};

View File

@ -144,7 +144,11 @@ impl SuricataMonitor {
let alert = v.get("alert")?;
let sid = alert.get("signature_id").and_then(|x| x.as_u64()).unwrap_or(0) as u32;
let signature = alert.get("signature").and_then(|x| x.as_str()).unwrap_or("").to_string();
let signature = alert
.get("signature")
.and_then(|x| x.as_str())
.unwrap_or("")
.to_string();
// Suricata severity: 1=high, 2=medium, 3=low, 4=informational.
// Map to confidence in [0.5, 1.0] so high-severity alerts tend to trip
// SOAR condition thresholds.

View File

@ -22,5 +22,17 @@ traceable! {
#[error("Failed to flush traffic log: {err}")]
TrafficLogFlushFailed => tracing::Level::ERROR,
#[no_source]
#[error("Model manifest at {path:?} is invalid: {reason}")]
ManifestInvalid { path: PathBuf, reason: String } => tracing::Level::ERROR,
#[no_source]
#[error("Feature count mismatch for {model:?}: manifest declares {declared}, ONNX input expects {onnx_dim}")]
FeatureMismatch { model: PathBuf, declared: usize, onnx_dim: usize } => tracing::Level::ERROR,
#[no_source]
#[error("Unknown feature '{name}' — not registered in FEATURE_REGISTRY")]
UnknownFeature { name: String } => tracing::Level::ERROR,
}
}

View File

@ -62,5 +62,11 @@ loggable! {
#[error("Model reload failed, keeping current models: {error}")]
ModelReloadFailed { error: String } => tracing::Level::ERROR,
#[error("Model manifest loaded: name='{name}', adapter={adapter}, features={features}, labels={labels}")]
ManifestLoaded { name: String, adapter: String, features: usize, labels: usize } => tracing::Level::INFO,
#[error("ONNX input shape introspected for {model}: declared={declared}, onnx_dim={onnx_dim}, matched={matched}")]
OnnxShapeChecked { model: String, declared: usize, onnx_dim: usize, matched: bool } => tracing::Level::DEBUG,
}
}