mirror of
https://github.com/ParrotXray/Mantis.git
synced 2026-08-24 19:00:27 +09:00
feat/optimized-inference-speed (#5)
* feat: complete rule scan, maybe * feat: improve inference speed and timely stop for ^C and solve cpu race
This commit is contained in:
GitHub
parent
5e42c137ae
commit
e686449a98
86
TODO
86
TODO
@ -1,12 +1,12 @@
|
||||
1. .csv 一開始文件格式為 traffic-yyyy-oo-zz.csv, 換天後直接變日期創新檔，例如今天啟動今天為 traffic-2026-05-11.csv 隔天變成 traffic-2026-05-12.csv [x]
|
||||
2. will change tract-onnx engine to ort-tract engine [x]
|
||||
3. 改善推論效能及速度，在大量資料時
|
||||
3. 改善推論效能及速度，在大量資料時 [x]
|
||||
4. 實現 ML + RULE 的共用 HashMap 實現共同決策結果
|
||||
5. 代碼最佳化，檢查是否除了 build.rs 以外有無 .unwarp() and eprintln() [x]
|
||||
6. 完成前端 detection 頁面
|
||||
7. 使用 sqllite 實現帳號系統、白黑名單永久記錄
|
||||
|
||||
[o]
|
||||
[x]
|
||||
Suricata
|
||||
提升偵測準確率：
|
||||
pcre 很多規則依賴正則，沒有會漏掉很多威脅
|
||||
@ -18,7 +18,7 @@ $HOME_NET / $EXTERNAL_NET 區分內外網，很多規則依賴這個
|
||||
flow:established 順便過濾掉不完整連線
|
||||
|
||||
============================================================
|
||||
PRIORITY BACKLOG (updated 2026-05-16)
|
||||
PRIORITY BACKLOG (updated 2026-05-18)
|
||||
============================================================
|
||||
|
||||
-- MUST DO (system is not usable without these) -----------
|
||||
@ -59,23 +59,69 @@ PRIORITY BACKLOG (updated 2026-05-16)
|
||||
|
||||
-- SHOULD DO (runtime management) ------------------------
|
||||
|
||||
[ ] HTTP API - rule / suppress / whitelist management
|
||||
Endpoints needed:
|
||||
GET /api/rules - list loaded rules (sid, msg, enabled)
|
||||
POST /api/rules/reload - trigger hot-reload (see below)
|
||||
GET /api/suppress - list suppress entries
|
||||
POST /api/suppress - add entry { sid, track, ip }
|
||||
[ ] Account system + persistent list DB (updated 2026-05-18)
|
||||
Two SQLite files, separated by concern:
|
||||
|
||||
static/db/rules.db (rule-related, already exists)
|
||||
suppress (id, sid INTEGER, track INTEGER, ip_net TEXT, comment TEXT, created_at)
|
||||
- track: 1=by_src 2=by_dst 4=by_either (Suricata values)
|
||||
- already populated by build.rs; API writes go here; RuleEngine reloads on change
|
||||
|
||||
static/db/app.db (account + network list, new)
|
||||
accounts (id, username, password_hash, role, created_at)
|
||||
- role: "admin" | "viewer"
|
||||
- password_hash: argon2 or bcrypt
|
||||
sessions (token TEXT PK, account_id, expires_at)
|
||||
- token: 32-byte random hex, expires in 24 h
|
||||
whitelist (id, ip_net TEXT, comment TEXT, created_at)
|
||||
- packets from whitelisted CIDRs skip rule + ML evaluation entirely
|
||||
blacklist (id, ip_net TEXT, comment TEXT, action TEXT, created_at)
|
||||
- action: "alert" | "drop"
|
||||
- packets from blacklisted CIDRs auto-alert without rule evaluation
|
||||
Migration: CREATE TABLE IF NOT EXISTS at startup in AppServices::init().
|
||||
File: core/infrastructure/app_db.rs (new)
|
||||
|
||||
[ ] HTTP API - accounts / suppress / whitelist / blacklist / rules
|
||||
All routes require session token in Authorization: Bearer <token> header
|
||||
except POST /api/auth/login.
|
||||
|
||||
Auth:
|
||||
POST /api/auth/login - { username, password } -> { token, expires_at }
|
||||
POST /api/auth/logout - invalidate current token
|
||||
GET /api/auth/me - current account info
|
||||
|
||||
Accounts (admin only):
|
||||
GET /api/accounts - list accounts
|
||||
POST /api/accounts - create account { username, password, role }
|
||||
PUT /api/accounts/:id/password - change password
|
||||
DELETE /api/accounts/:id - delete account
|
||||
|
||||
Suppress:
|
||||
GET /api/suppress - list all entries
|
||||
POST /api/suppress - add { sid, track, ip_net, comment }
|
||||
DELETE /api/suppress/:id - remove entry
|
||||
GET /api/whitelist - list whitelist IPs / CIDRs
|
||||
POST /api/whitelist - add entry { ip_net, comment }
|
||||
DELETE /api/whitelist/:id - remove entry
|
||||
GET /api/blacklist - list blacklist entries
|
||||
POST /api/blacklist - add entry
|
||||
DELETE /api/blacklist/:id - remove entry
|
||||
POST /api/system/restart - graceful restart (re-attach eBPF, reload all)
|
||||
Storage: rules.db (suppress / whitelist / blacklist tables already planned in item 7).
|
||||
Auth: session token via SQLite account system (item 7).
|
||||
Files: web/routes/rules.rs, web/routes/system.rs
|
||||
(runtime effect: RuleEngine reloads suppress table after write)
|
||||
|
||||
Whitelist:
|
||||
GET /api/whitelist - list entries
|
||||
POST /api/whitelist - add { ip_net, comment }
|
||||
DELETE /api/whitelist/:id - remove
|
||||
|
||||
Blacklist:
|
||||
GET /api/blacklist - list entries
|
||||
POST /api/blacklist - add { ip_net, action, comment }
|
||||
DELETE /api/blacklist/:id - remove
|
||||
|
||||
Rules:
|
||||
GET /api/rules - list loaded rules (sid, msg, enabled)
|
||||
POST /api/rules/reload - trigger hot-reload
|
||||
|
||||
System:
|
||||
POST /api/system/restart - graceful restart
|
||||
|
||||
Files: web/routes/auth.rs, web/routes/accounts.rs, web/routes/lists.rs,
|
||||
web/routes/rules.rs, web/routes/system.rs
|
||||
Middleware: web/middleware/auth.rs (token extractor + role check)
|
||||
|
||||
[ ] Hot-reload: rules + suppress without process restart
|
||||
Trigger: POST /api/rules/reload OR SIGHUP signal.
|
||||
@ -121,4 +167,4 @@ PRIORITY BACKLOG (updated 2026-05-16)
|
||||
[x] isdataat
|
||||
Completed 2026-05-15 (included with byte_test/jump/extract).
|
||||
ByteOp kind=3; supports negated form (!isdataat) and relative flag.
|
||||
build.rs: parse_isdataat(); rule_engine.rs: eval_byte_ops() case 3.
|
||||
build.rs: parse_isdataat(); rule_engine.rs: eval_byte_ops() case 3.
|
||||
19
config.toml
19
config.toml
@ -4,7 +4,11 @@ egress_ifname = "enp4s0f0" # Egress NIC Name
|
||||
geoip_db_name = "GeoLite2-City.mmdb"
|
||||
deep_autoencoder_name = "deep_autoencoder.onnx"
|
||||
models_config_name = "inference_config.json"
|
||||
combined_queue_count = 8 # NIC Combined Queue Count (ethtool -l <NIC>)
|
||||
|
||||
# NIC Combined Queue Count (ethtool -l <NIC>)
|
||||
# If you will edit this, please enter ethtool -L <NIC> combined <combined number> at System Terminal.
|
||||
# Then change this item.
|
||||
combined_queue_count = 8
|
||||
channel_size = 4096
|
||||
fill_queue_size = 4096 # Umem Used (Should not modify)
|
||||
comp_queue_size = 4096 # Umem Used (Should not modify)
|
||||
@ -23,9 +27,18 @@ aggregator_window_secs = 30
|
||||
inference_batch_size = 200
|
||||
flow_timeout_us = 60_000_000
|
||||
|
||||
traffic_logging_mode = true # When true, disables ML inference and records all ingress/egress packets to CSV
|
||||
traffic_logging_mode = false # When true, disables ML inference and records all ingress/egress packets to CSV
|
||||
traffic_log_csv_path = "traffic_log.csv" # Output CSV file path for traffic logging mode
|
||||
|
||||
home_net = ["140.130.34.0/24"]
|
||||
|
||||
# tls_keylog_path = "/tmp/tls_keys.log" # NSS key log file for TLS decryption (SSLKEYLOGFILE)
|
||||
# tls_keylog_path = "/tmp/tls_keys.log" # NSS key log file for TLS decryption (SSLKEYLOGFILE)
|
||||
|
||||
# CPU affinity (Linux only). Uncomment and tune for your hardware.
|
||||
# Pin XSK packet threads starting from this core (one core per queue pair).
|
||||
# Example: xsk_cpu_base=0 with combined_queue_count=4 uses cores 0-3 for packets.
|
||||
xsk_cpu_base = 0
|
||||
#
|
||||
# Pin ML inference (ONNX spawn_blocking) to this core.
|
||||
# Example: on an 8-core machine, reserve core 7 for inference.
|
||||
ml_cpu = 7
|
||||
@ -28,6 +28,7 @@ use crate::model::error::Error;
|
||||
use crate::model::log::ebpf::EbpfLog;
|
||||
use crate::model::log::ml::MLLog;
|
||||
use crate::model::log::rule::RuleLog;
|
||||
use crate::utils::cpu_affinity::set_cpu_affinity;
|
||||
|
||||
pub struct XskManager {
|
||||
app_config: Arc<AppConfig>,
|
||||
@ -124,6 +125,8 @@ impl XskManager {
|
||||
|
||||
pub struct XskPair {
|
||||
direction: Direction,
|
||||
queue_id: u32,
|
||||
xsk_cpu_base: Option<u32>,
|
||||
umem: Arc<Umem>,
|
||||
fill_queue: FillQueue,
|
||||
comp_queue: CompQueue,
|
||||
@ -193,6 +196,8 @@ impl XskPair {
|
||||
|
||||
let xsk_pair = Self {
|
||||
direction,
|
||||
queue_id,
|
||||
xsk_cpu_base: config.xsk_cpu_base,
|
||||
umem: Arc::new(umem),
|
||||
fill_queue,
|
||||
comp_queue,
|
||||
@ -219,6 +224,10 @@ impl XskPair {
|
||||
thread::Builder::new()
|
||||
.name(thread_name.clone())
|
||||
.spawn(move || {
|
||||
if let Some(base) = self.xsk_cpu_base {
|
||||
set_cpu_affinity((base + self.queue_id) as usize);
|
||||
}
|
||||
|
||||
// StreamReassembler is !Send (Rc inside protolens), so create it here.
|
||||
let min_sig_matches = self.min_signature_matches;
|
||||
let is_ingress = self.direction == Direction::Ingress;
|
||||
|
||||
@ -66,6 +66,7 @@ impl AppServices {
|
||||
app_config.aggregator_window_secs,
|
||||
app_config.flow_timeout_us,
|
||||
traffic_logger,
|
||||
app_config.ml_cpu,
|
||||
));
|
||||
|
||||
|
||||
|
||||
@ -17,6 +17,7 @@ use crate::core::infrastructure::ml_alert::MLAlert;
|
||||
use crate::model::error::ml::MLError;
|
||||
use crate::model::log::ml::MLLog;
|
||||
use crate::model::ml_detection::{EngineStats, InferenceStats};
|
||||
use crate::utils::cpu_affinity::{num_cpus, set_cpu_affinity};
|
||||
use crate::utils::packet_parser::parse_packet;
|
||||
|
||||
pub struct Engine {
|
||||
@ -29,6 +30,7 @@ pub struct Engine {
|
||||
inference_interval_secs: u64,
|
||||
flow_timeout_us: u64,
|
||||
traffic_logger: Option<Arc<TrafficLogger>>,
|
||||
ml_cpu: Option<u32>,
|
||||
}
|
||||
|
||||
impl Engine {
|
||||
@ -43,6 +45,7 @@ impl Engine {
|
||||
window_secs: u64,
|
||||
flow_timeout_us: u64,
|
||||
traffic_logger: Option<Arc<TrafficLogger>>,
|
||||
ml_cpu: Option<u32>,
|
||||
) -> Self {
|
||||
let tracker = Arc::new(Mutex::new(FlowTracker::new(max_flows)));
|
||||
let inference_pipeline = Arc::new(Inference::new(models, config));
|
||||
@ -62,6 +65,7 @@ impl Engine {
|
||||
inference_interval_secs: interval_secs,
|
||||
flow_timeout_us,
|
||||
traffic_logger,
|
||||
ml_cpu,
|
||||
}
|
||||
}
|
||||
|
||||
@ -86,21 +90,23 @@ impl Engine {
|
||||
_ = ticker.tick() => {}
|
||||
}
|
||||
|
||||
let Ok(mut t) = self.tracker.lock() else {
|
||||
log!(MLError::TrackerLockPoisoned);
|
||||
continue;
|
||||
let (total_flows, packet_counts, flows) = {
|
||||
let Ok(mut t) = self.tracker.lock() else {
|
||||
log!(MLError::TrackerLockPoisoned);
|
||||
continue;
|
||||
};
|
||||
let total_flows = t.flow_count();
|
||||
let packet_counts: Vec<usize> = t.get_flows_snapshot()
|
||||
.iter()
|
||||
.map(|f| f.packet_count())
|
||||
.collect();
|
||||
let flows = if self.traffic_logger.is_some() {
|
||||
t.drain_flows(self.min_packets)
|
||||
} else {
|
||||
t.get_flows_for_inference(self.min_packets)
|
||||
};
|
||||
(total_flows, packet_counts, flows)
|
||||
};
|
||||
let total_flows = t.flow_count();
|
||||
let packet_counts: Vec<usize> = t.get_flows_snapshot()
|
||||
.iter()
|
||||
.map(|f| f.packet_count())
|
||||
.collect();
|
||||
let flows = if self.traffic_logger.is_some() {
|
||||
t.drain_flows(self.min_packets)
|
||||
} else {
|
||||
t.get_flows_for_inference(self.min_packets)
|
||||
};
|
||||
drop(t);
|
||||
|
||||
log!(MLLog::FlowStats(
|
||||
total_flows,
|
||||
@ -149,12 +155,27 @@ impl Engine {
|
||||
log!(MLLog::RunningInference(batch.len()));
|
||||
|
||||
let start = Instant::now();
|
||||
let results = self.inference_pipeline.infer_batch(&batch);
|
||||
let batch_len = batch.len();
|
||||
let pipeline = Arc::clone(&self.inference_pipeline);
|
||||
let ml_cpu = self.ml_cpu;
|
||||
let mut handle = tokio::task::spawn_blocking(move || {
|
||||
let cpu = ml_cpu.map(|c| c as usize).unwrap_or_else(|| num_cpus().saturating_sub(1));
|
||||
set_cpu_affinity(cpu);
|
||||
pipeline.infer_batch(&batch)
|
||||
});
|
||||
|
||||
let results = tokio::select! {
|
||||
_ = &mut shutdown_rx => {
|
||||
handle.abort();
|
||||
break;
|
||||
}
|
||||
res = &mut handle => res.unwrap_or_default(),
|
||||
};
|
||||
let elapsed_us = start.elapsed().as_micros() as u64;
|
||||
let stats = InferenceStats::from_results(&results, elapsed_us);
|
||||
|
||||
if results.len() != batch.len() {
|
||||
log!(MLLog::InferenceResults(batch.len(), results.len()));
|
||||
if results.len() != batch_len {
|
||||
log!(MLLog::InferenceResults(batch_len, results.len()));
|
||||
}
|
||||
|
||||
log!(MLLog::InferenceCompleted(
|
||||
@ -226,4 +247,4 @@ impl Engine {
|
||||
};
|
||||
EngineStats { active_flows }
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -1,6 +1,7 @@
|
||||
use std::collections::VecDeque;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Instant;
|
||||
|
||||
use macros::log;
|
||||
use ndarray::Array3;
|
||||
@ -42,11 +43,12 @@ impl Inference {
|
||||
flows.iter().filter_map(|flow| self.infer_single(flow)).collect()
|
||||
}
|
||||
|
||||
pub fn infer_single(&self, flow: &FlowData) -> Option<DetectionResult> {
|
||||
fn infer_single(&self, flow: &FlowData) -> Option<DetectionResult> {
|
||||
let t0 = Instant::now();
|
||||
let features = self.preprocess_ae_features(flow);
|
||||
let t1 = Instant::now();
|
||||
let window_size = self.config.window_size;
|
||||
|
||||
// Update per-src_ip buffer
|
||||
let sequence = {
|
||||
let Ok(mut buffers) = self.flow_buffers.lock() else {
|
||||
log!(MLError::InferenceLockPoisoned);
|
||||
@ -56,43 +58,24 @@ impl Inference {
|
||||
.entry(flow.flow_key.src_ip.clone())
|
||||
.or_insert_with(VecDeque::new);
|
||||
|
||||
buf.push_back(features.clone());
|
||||
buf.push_back(features);
|
||||
if buf.len() > window_size {
|
||||
buf.pop_front();
|
||||
}
|
||||
|
||||
println!(
|
||||
"Buffer [{}->{}]: {}/{} | contents: {:?}",
|
||||
flow.flow_key.src_ip,
|
||||
flow.flow_key.dst_ip,
|
||||
buf.len(),
|
||||
window_size,
|
||||
buf.iter().map(|v| format!("{:.3}", v[0])).collect::<Vec<_>>()
|
||||
);
|
||||
if buf.len() < window_size {
|
||||
return None;
|
||||
}
|
||||
|
||||
// Left-pad with zeros if not enough history
|
||||
let mut seq: Vec<Vec<f32>> = Vec::with_capacity(window_size);
|
||||
let pad_len = window_size.saturating_sub(buf.len());
|
||||
let feat_len = self.config.num_ae_features();
|
||||
for _ in 0..pad_len {
|
||||
seq.push(vec![0.0f32; feat_len]);
|
||||
}
|
||||
for v in buf.iter() {
|
||||
seq.push(v.clone());
|
||||
}
|
||||
seq
|
||||
buf.iter().cloned().collect::<Vec<_>>()
|
||||
};
|
||||
|
||||
// Build 3D tensor (1, window_size, features)
|
||||
let t2 = Instant::now();
|
||||
let feat_len = self.config.num_ae_features();
|
||||
let ae_input = Array3::from_shape_fn(
|
||||
(1, window_size, feat_len),
|
||||
|(_, t, f)| sequence[t][f],
|
||||
);
|
||||
|
||||
let t3 = Instant::now();
|
||||
let ae_score = match self.run_autoencoder(&ae_input) {
|
||||
Ok(score) => score,
|
||||
Err(e) => {
|
||||
@ -100,9 +83,17 @@ impl Inference {
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let t4 = Instant::now();
|
||||
|
||||
log!(MLLog::InferenceTiming(
|
||||
flow.flow_key.src_ip.clone(),
|
||||
t1.duration_since(t0).as_millis() as u64,
|
||||
t2.duration_since(t1).as_millis() as u64,
|
||||
t3.duration_since(t2).as_millis() as u64,
|
||||
t4.duration_since(t3).as_millis() as u64,
|
||||
));
|
||||
|
||||
let is_attack = ae_score >= self.config.ae_threshold;
|
||||
|
||||
let flow_key = format!(
|
||||
"{}:{} -> {}:{} (proto {}) [{}]",
|
||||
flow.flow_key.src_ip,
|
||||
@ -145,7 +136,6 @@ impl Inference {
|
||||
|
||||
let output = outputs[0].try_extract_array::<f32>().map_err(|e| e.to_string())?;
|
||||
|
||||
// MSE between input and reconstructed output
|
||||
let input_dyn = input.view().into_dyn();
|
||||
let diff = &input_dyn - &output;
|
||||
let n = (self.config.window_size * self.config.num_ae_features()) as f32;
|
||||
@ -153,4 +143,4 @@ impl Inference {
|
||||
|
||||
Ok(mse)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -27,7 +27,7 @@ impl MLModels {
|
||||
|
||||
Session::builder()
|
||||
.map_err(|_| MLError::ModelLoadFailed { path: model_path.clone() })?
|
||||
.with_optimization_level(GraphOptimizationLevel::Disable)
|
||||
.with_optimization_level(GraphOptimizationLevel::All)
|
||||
.map_err(|_| MLError::ModelLoadFailed { path: model_path.clone() })?
|
||||
.commit_from_file(&model_path)
|
||||
.map_err(|_| MLError::ModelLoadFailed { path: model_path })
|
||||
|
||||
@ -12,5 +12,7 @@ async fn main() -> Result<(), Error> {
|
||||
let mut system = System::new().await?;
|
||||
system.run().await?;
|
||||
system.terminate().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
drop(system);
|
||||
std::process::exit(0);
|
||||
}
|
||||
@ -37,4 +37,11 @@ pub struct Config {
|
||||
/// Only useful in [external]->NetGuardia->[internal] deployments where the
|
||||
/// internal server can be configured to write TLS session keys.
|
||||
pub tls_keylog_path: Option<String>,
|
||||
/// First CPU core assigned to XSK packet threads. Each queue pair gets one core
|
||||
/// starting from this base (e.g. base=0 with 4 queues pins to cores 0-3).
|
||||
/// If absent, no affinity is set.
|
||||
pub xsk_cpu_base: Option<u32>,
|
||||
/// CPU core pinned to the ML inference spawn_blocking thread.
|
||||
/// If absent, defaults to the last available core.
|
||||
pub ml_cpu: Option<u32>,
|
||||
}
|
||||
@ -72,5 +72,8 @@ loggable! {
|
||||
#[error("CSV rotated to: {path}")]
|
||||
TrafficLogRotated { path: String } => tracing::Level::INFO,
|
||||
|
||||
#[error("Timing [{src}]: feature={feature_ms}ms buffer={buffer_ms}ms tensor={tensor_ms}ms onnx={onnx_ms}ms")]
|
||||
InferenceTiming { src: String, feature_ms: u64, buffer_ms: u64, tensor_ms: u64, onnx_ms: u64 } => tracing::Level::INFO,
|
||||
|
||||
}
|
||||
}
|
||||
@ -2,5 +2,6 @@ pub mod logging;
|
||||
pub mod static_files;
|
||||
pub mod boot_time;
|
||||
pub mod packet_parser;
|
||||
pub mod cpu_affinity;
|
||||
|
||||
pub mod ip_address;
|
||||
Loading…
x
Reference in New Issue
Block a user