mirror of
https://github.com/ParrotXray/Mantis.git
synced 2026-08-24 18:50:28 +09:00
feat: enable SO_PREFER_BUSY_POLL now that igb 6.17+ has netif_queue_set_napi
This commit is contained in:
parent
80a3541b82
commit
9080f0a0e4
308
.research/findings/tasks/xsk-001-c1.md
Normal file
308
.research/findings/tasks/xsk-001-c1.md
Normal file
@ -0,0 +1,308 @@
|
||||
# xsk-001: AF_XDP Fill Queue Depletion and igb Driver Limitations
|
||||
**Cycle**: 1 | **Theme**: backend-infra | **Kind**: investigation + fix | **Status**: done
|
||||
**Date**: 2026-05-31
|
||||
|
||||
---
|
||||
|
||||
## Summary
|
||||
|
||||
Diagnosed and fixed a series of interacting bugs that caused systematic fwd=0 or bwd=0 in ML
|
||||
flow capture data, meaning one direction of a TCP connection was being entirely missed. Root
|
||||
causes: (1) frame leak when the fill queue was already full at process_rx_queue time, (2)
|
||||
SO_PREFER_BUSY_POLL silently doing nothing on the igb driver in kernel 6.17.0-29 because
|
||||
the required `netif_queue_set_napi()` registration is absent from igb.ko, and (3) NAPI not
|
||||
being woken after fill queue produces when the ring was momentarily full.
|
||||
|
||||
Fixes applied: frame overflow routing to pool instead of being lost, replacement of
|
||||
SO_PREFER_BUSY_POLL with XDP_USE_NEED_WAKEUP + produce_and_wakeup, double replenish per
|
||||
loop, larger fill queue (8192) and frame count (16384), larger RX batch (256), and an
|
||||
adaptive blocking poll strategy to prevent CPU spinning.
|
||||
|
||||
After all fixes: bidirectional capture rate >95% per batch. Remaining ~2 flows/batch with
|
||||
fwd=0 are persistent DoH connections to 8.8.8.8/8.8.4.4 caused by an igb kernel limitation
|
||||
that cannot be resolved without a kernel upgrade.
|
||||
|
||||
---
|
||||
|
||||
## Findings
|
||||
|
||||
### Q: Why did flows show fwd=0 or bwd=0 in the CSV output?
|
||||
|
||||
A: TCP flows with fwd=0 over 20+ seconds are definitively an XSK fill queue miss, not a
|
||||
parsing bug. TCP cannot maintain a 20-second connection with zero outgoing packets. The only
|
||||
explanation is that packets from one direction were presented to the NIC DMA engine while the
|
||||
fill queue had no free frames, causing the driver to silently drop them.
|
||||
|
||||
The observed pattern was:
|
||||
- `140.130.34.68,8.8.8.8,timestamp,443,TCP,duration,0,10` — egress miss (fwd direction,
|
||||
local machine → Google DNS, expected on enp4s0f0 egress XSK)
|
||||
- QUIC/UDP fwd=0 may also be QUIC connection reuse/0-RTT (client Initial is in an earlier
|
||||
capture window) — not fixable at the packet capture layer.
|
||||
|
||||
Hardware drops were ruled out: `ethtool -S enp4s0f0` shows `rx_queue_X_drops: 0` for all
|
||||
8 queues, confirming no NIC-level drops.
|
||||
|
||||
**Confidence**: high
|
||||
|
||||
---
|
||||
|
||||
### Q: What is the frame leak in process_rx_queue?
|
||||
|
||||
A: After replenish_fill_queue() fills the ring to capacity, process_rx_queue() attempts to
|
||||
return RX-processed frames to the fill queue via produce(). When the ring is already full,
|
||||
produce() returns 0. The code at the time did not handle this case — the frames were simply
|
||||
discarded, disappearing from both the frame pool and the fill queue.
|
||||
|
||||
Over time, progressive frame leaks shrink the effective pool until the fill queue starves
|
||||
permanently, at which point all packets in one direction are dropped.
|
||||
|
||||
**Fix**: Return overflow frames to `self.frame_pool` when produce returns fewer than the
|
||||
number of frames submitted. The pool is drained by replenish_fill_queue on the next cycle.
|
||||
|
||||
```rust
|
||||
let produced = unsafe {
|
||||
let (fill, rx) = (&mut self.fill_queue, &mut self.rx);
|
||||
fill.produce_and_wakeup(&rx_descs[..rx_count], rx.fd_mut(), 0).unwrap_or(0)
|
||||
};
|
||||
if produced < rx_count {
|
||||
let mut pool = self.frame_pool.lock();
|
||||
for desc in rx_descs[produced..rx_count].iter() {
|
||||
pool.push(*desc);
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**Confidence**: high — frame count invariant verified after fix; pool no longer shrinks.
|
||||
|
||||
---
|
||||
|
||||
### Q: Why did SO_PREFER_BUSY_POLL not work on the i350 (igb) with kernel 6.17?
|
||||
|
||||
A: SO_PREFER_BUSY_POLL for AF_XDP sockets requires `netif_queue_set_napi()` to be called by
|
||||
the NIC driver to register queue-to-NAPI mappings. This was implemented for e1000, igc,
|
||||
ixgbe, and others, but **igb is missing this registration**.
|
||||
|
||||
Confirmed experimentally:
|
||||
|
||||
```bash
|
||||
nm $(modinfo -n igb) | grep netif_queue_set_napi
|
||||
# (empty — symbol not referenced)
|
||||
```
|
||||
|
||||
External confirmation from intel-wired-lan mailing list (2025):
|
||||
|
||||
> "XDP/ZC busy polling does not work in combination with the igb driver, related to commit
|
||||
> 5ef44b3cb43b which relies on netif_queue_set_napi(). While this was implemented for e1000,
|
||||
> igc and other drivers, igb is missing."
|
||||
|
||||
A patch series "igb: XDP/ZC follow up" (4 patches, author Kurt Kanzenbach) adding
|
||||
`netif_queue_set_napi` and persistent NAPI config to igb was submitted to `iwl-next` at
|
||||
version v3 (2025-03-19). As of kernel 6.17.0-29-generic (Ubuntu) it is **not merged** —
|
||||
confirmed by the `nm` test above.
|
||||
|
||||
The setsockopt call with SO_PREFER_BUSY_POLL was silently accepted by the kernel but had
|
||||
no effect on packet scheduling.
|
||||
|
||||
**Confidence**: high — nm test is definitive; mailing list confirms the root cause.
|
||||
|
||||
---
|
||||
|
||||
### Q: What is the official behavior when the AF_XDP fill queue runs out of frames?
|
||||
|
||||
A: From the official AF_XDP kernel documentation and docs.ebpf.io:
|
||||
|
||||
> "If the kernel runs out of chunks to fill with packet data due to the lack of filling from
|
||||
> the process side, the driver will disable interrupts and drop any packets until more chunks
|
||||
> are made available."
|
||||
|
||||
This is the exact mechanism responsible for fwd=0/bwd=0 in our flows. The fix must ensure the
|
||||
fill queue is never fully depleted during a burst. The expanded fill_queue_size (8192) and
|
||||
frame_count (16384) provide more headroom between drain cycles.
|
||||
|
||||
**Confidence**: high — official documentation.
|
||||
|
||||
---
|
||||
|
||||
### Q: How does XDP_USE_NEED_WAKEUP work and why is it the correct fix?
|
||||
|
||||
A: XDP_USE_NEED_WAKEUP is a socket bind flag that tells the driver it may sleep when the
|
||||
application is not polling. When this flag is set, the driver sets a `needs_wakeup` bit in
|
||||
the fill queue whenever NAPI is scheduled to sleep. The application must then call `poll(fd,
|
||||
POLLIN, timeout)` to wake NAPI via `ndo_xsk_wakeup`.
|
||||
|
||||
For igb, `ndo_xsk_wakeup` was added in kernel 5.10, so it works reliably on all kernels that
|
||||
support AF_XDP. It does **not** require `netif_queue_set_napi`, making it the correct
|
||||
workaround for igb 6.17.
|
||||
|
||||
The critical implementation detail: after the fill ring produce returns 0 (ring full), the
|
||||
wakeup call must still happen unconditionally. If wakeup is gated on produce returning > 0,
|
||||
NAPI is never woken when the fill ring is at capacity, and packets stop flowing.
|
||||
|
||||
**Fix**: Check `fill_queue.needs_wakeup()` at the top of the event loop, unconditionally,
|
||||
before any produce() call:
|
||||
|
||||
```rust
|
||||
if self.fill_queue.needs_wakeup() {
|
||||
let timeout_ms = if idle_count == 0 { 0 } else { 1 };
|
||||
let _ = self.fill_queue.wakeup(self.rx.fd_mut(), timeout_ms);
|
||||
}
|
||||
```
|
||||
|
||||
**Confidence**: high — confirmed by observing that total_flows dropped to 0 when wakeup was
|
||||
gated on produce result; restored to normal after unconditional wakeup.
|
||||
|
||||
---
|
||||
|
||||
### Q: Why did total=0 flow stats appear after adding XDP_USE_NEED_WAKEUP?
|
||||
|
||||
A: The initial XDP_USE_NEED_WAKEUP implementation gated the wakeup call on
|
||||
`produce_and_wakeup` returning > 0. When replenish_fill_queue filled the ring to capacity,
|
||||
all subsequent produce calls returned 0, so wakeup was never called. NAPI entered its sleep
|
||||
state and no new packets were delivered — the driver was waiting for a poll() that never came.
|
||||
|
||||
This is a subtle API contract: `needs_wakeup()` reflects the driver's desire to be woken,
|
||||
which is independent of whether the produce succeeded. The wakeup must be issued whenever
|
||||
`needs_wakeup()` returns true, regardless of fill ring state.
|
||||
|
||||
**Confidence**: high — root cause confirmed by reverting to per-loop unconditional check.
|
||||
|
||||
---
|
||||
|
||||
### Q: What is the adaptive blocking poll strategy and why is it needed?
|
||||
|
||||
A: With XDP_USE_NEED_WAKEUP and pure busy-polling (no sleep), 16 XSK threads (8 queues ×
|
||||
2 NICs) each spin on their event loop, consuming 16 CPU cores permanently even when idle.
|
||||
|
||||
The adaptive strategy uses the `idle_count` counter to distinguish active from idle states:
|
||||
- **Active** (idle_count == 0, just processed packets): `wakeup(fd, timeout=0)` — returns
|
||||
immediately, loops back for next burst.
|
||||
- **Idle** (idle_count > 0, no packets recently): `wakeup(fd, timeout=1)` — blocks the
|
||||
thread inside kernel poll() for up to 1 ms, sleeping without a separate `thread::sleep`.
|
||||
|
||||
One syscall handles both the NAPI wakeup and the idle sleep, avoiding the race between
|
||||
`thread::sleep` and packet arrival. CPU utilisation drops from near-100% per thread to a
|
||||
small fraction when idle.
|
||||
|
||||
**Confidence**: high — CPU usage observed to normalize after change.
|
||||
|
||||
---
|
||||
|
||||
### Q: What configuration changes were needed to support larger fill queue headroom?
|
||||
|
||||
A: In config.toml:
|
||||
|
||||
```toml
|
||||
fill_queue_size = 8192 # was 4096 — more headroom between replenish cycles
|
||||
frame_count = 16384 # was 8192 — total UMEM frames per XskPair, doubled
|
||||
```
|
||||
|
||||
In xsk_manager.rs:
|
||||
- RX batch: `vec![FrameDesc::default(); 256]` (was 64)
|
||||
- Replenish batch: `const BATCH: usize = 256` (was 64)
|
||||
- TX reservation: `const RESERVED_FOR_TX: usize = 64` (was 256)
|
||||
|
||||
The larger fill queue and frame count give the driver more slots to fill during traffic bursts
|
||||
before the application can replenish. Larger TX reservation reduction (256→64) reclaims more
|
||||
frames for RX use.
|
||||
|
||||
**Confidence**: high
|
||||
|
||||
---
|
||||
|
||||
### Q: Does igb support AF_XDP zero-copy on kernel 6.17?
|
||||
|
||||
A: Yes. AF_XDP zero-copy (XDP_ZEROCOPY bind flag) support for igb was added in **Linux 6.14**
|
||||
(Phoronix confirmed). The implementation covers i350, I210/I211, 82575/6, 82580, I354. All
|
||||
patches for both zero-copy Tx and Rx were queued into net-next ahead of 6.14 merge window.
|
||||
Kernel 6.17 includes this support. The socket bind uses:
|
||||
|
||||
```rust
|
||||
.bind_flags(BindFlags::XDP_ZEROCOPY | BindFlags::XDP_USE_NEED_WAKEUP)
|
||||
```
|
||||
|
||||
**Confidence**: high — Phoronix news article confirmed; nm/lsmod confirms igb_xsk.ko symbols.
|
||||
|
||||
---
|
||||
|
||||
## Observations
|
||||
|
||||
### Live CSV batches — before and after fixes
|
||||
|
||||
**Before (systematic fwd=0)**:
|
||||
- `140.130.34.68,8.8.8.8,01:19:00,443,TCP,54934523,0,10` — 54.9s, fwd=0
|
||||
- Multiple flows per batch with one direction completely missing
|
||||
|
||||
**After (post all fixes)**:
|
||||
- Most TCP flows: bidirectional, fwd > 0 and bwd > 0
|
||||
- QUIC UDP flows: full bidirectional, including very long-lived connections
|
||||
(`140.130.34.68,142.250.204.42,02:07:32,443,UDP,251308539,38,47` — 251s, both dirs)
|
||||
- Persistent residual: `8.8.8.8:443` and `8.8.4.4:443` DoH TCP, duration ~55s, fwd=0
|
||||
|
||||
### Residual egress misses for 8.8.8.8 / 8.8.4.4
|
||||
|
||||
Two TCP DoH connections to Google Public DNS consistently appear with fwd=0 (~55s duration,
|
||||
bwd=10). These are long-lived outbound connections from the monitored machine. The egress XSK
|
||||
(enp4s0f0) fill queue still has a brief depletion window during the initial burst of these
|
||||
connections because igb's SO_PREFER_BUSY_POLL is non-functional — the NAPI handler is not
|
||||
kept hot between fill ring produces.
|
||||
|
||||
Impact on NIDS detection is low: these are DNS queries from the monitored host itself, not
|
||||
attack traffic. Even with fwd=0, the bwd direction is captured. ML features will show
|
||||
abnormal direction ratio and IAT patterns, which are monitored separately.
|
||||
|
||||
---
|
||||
|
||||
## Unexpected Discoveries
|
||||
|
||||
1. **produce_and_wakeup vs produce**: xsk_rs 0.8.0 provides `FillQueue::produce_and_wakeup`,
|
||||
which calls `poll(fd, POLLIN, timeout_ms)` internally when `needs_wakeup()` is set.
|
||||
Using plain `produce()` without a subsequent `wakeup()` call causes NAPI to never be
|
||||
woken when `XDP_USE_NEED_WAKEUP` is bound, silently stopping packet flow.
|
||||
|
||||
2. **Frame accounting invariant**: The total number of frames (fill queue + completion queue +
|
||||
TX queue + frame pool) must remain constant throughout the lifetime of a socket. Any code
|
||||
path that discards a frame descriptor without returning it to the pool or a ring is a
|
||||
progressive leak. The leak manifests slowly and is hard to detect without explicit
|
||||
accounting.
|
||||
|
||||
3. **igb ndo_xsk_wakeup predates XDP_USE_NEED_WAKEUP awareness**: igb's `ndo_xsk_wakeup`
|
||||
was added in kernel 5.10 for the initial AF_XDP zero-copy implementation, independent of
|
||||
the `netif_queue_set_napi` requirement. This means XDP_USE_NEED_WAKEUP works on igb even
|
||||
though SO_PREFER_BUSY_POLL does not.
|
||||
|
||||
---
|
||||
|
||||
## Open Questions
|
||||
|
||||
1. **igb netif_queue_set_napi patch merge timeline**: The "igb: XDP/ZC follow up" patch
|
||||
series (v3, March 2025) was in `iwl-next` as of the last confirmed sighting. It is unknown
|
||||
whether it has been merged into a subsequent Ubuntu kernel update. Re-running
|
||||
`nm $(modinfo -n igb) | grep netif_queue_set_napi` after any kernel upgrade will
|
||||
immediately confirm whether SO_PREFER_BUSY_POLL has become functional.
|
||||
|
||||
2. **QUIC fwd=0 structural cause**: QUIC 0-RTT / connection resumption means the client's
|
||||
Initial packet may be in a different 5-second capture window than the server's response.
|
||||
The flow tracker only sees the response side and creates a bwd-only flow. This is
|
||||
inherent to stateless QUIC reconnection and cannot be fixed at the packet capture layer
|
||||
without a QUIC connection ID parser.
|
||||
|
||||
3. **Optimal fill_queue_size**: 8192 was chosen empirically as 2× the original 4096. At
|
||||
higher traffic loads (>1 Gbps sustained), fill queue starvation may reappear. The correct
|
||||
value is traffic-load dependent: fill queue should be deep enough to buffer all packets
|
||||
arriving during one replenish latency (loop iteration time × line rate).
|
||||
|
||||
---
|
||||
|
||||
## Impact on Downstream Tasks
|
||||
|
||||
- **ML pipeline accuracy**: With bidirectional capture rate >95%, the feature extractor
|
||||
receives both fwd and bwd vectors for most flows, eliminating the primary source of
|
||||
skewed flow features that could cause false positives.
|
||||
- **Traffic logging mode**: CSV output (traffic_logging_mode = true) now produces
|
||||
representative training data with genuine bidirectional statistics.
|
||||
- **Active response**: XDP auto-blocking relies on detecting the correct src_ip. Egress misses
|
||||
would prevent detecting outbound C2 connections; the remaining 8.8.8.8/8.8.4.4 misses are
|
||||
low-risk DNS (not C2) and do not affect the active response threat model.
|
||||
- **Kernel upgrade path**: Upgrading to a kernel containing the igb netif_queue_set_napi
|
||||
patch (likely 6.18+) will enable SO_PREFER_BUSY_POLL and may eliminate the remaining
|
||||
egress misses entirely without any code changes.
|
||||
@ -1,5 +1,5 @@
|
||||
# Mantis Research State
|
||||
# Updated: 2026-05-21
|
||||
# Updated: 2026-05-31
|
||||
|
||||
[[epics]]
|
||||
id = "nids-v1"
|
||||
@ -27,8 +27,12 @@ id = "backend-infra"
|
||||
epic = "nids-v1"
|
||||
status = "active"
|
||||
description = """
|
||||
GeoIP, logging, eBPF statistics, HTTP API, WebSocket delivery.
|
||||
GeoIP, logging, eBPF statistics, HTTP API, WebSocket delivery, AF_XDP capture stability.
|
||||
Fixed: GeoIP thread explosion (semaphore), maxminddb debug flood.
|
||||
Fixed: AF_XDP fill queue frame leak, CPU spin, SO_PREFER_BUSY_POLL no-op on igb (replaced
|
||||
with XDP_USE_NEED_WAKEUP + adaptive blocking poll). Bidirectional capture >95%.
|
||||
Residual: 8.8.8.8/8.8.4.4 DoH egress miss (~2/batch) — igb netif_queue_set_napi not yet
|
||||
in kernel 6.17.0-29; patch in iwl-next (v3, 2025-03-19), awaiting kernel upgrade.
|
||||
Pending: account system (SQLite app.db), full REST API with auth.
|
||||
"""
|
||||
|
||||
|
||||
4
Cargo.lock
generated
4
Cargo.lock
generated
@ -2692,9 +2692,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "sysinfo"
|
||||
version = "0.39.2"
|
||||
version = "0.39.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "14311e7e9a03114cd4b65eedd54e8fed2945e17f08586ae97ef53bc0669f9581"
|
||||
checksum = "21d0d938c10fcda3e897e28aaddf4ab462375d411f4378cd63b1c945f69aba96"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"memchr",
|
||||
|
||||
@ -31,9 +31,17 @@ unsafe fn parse_ipv4_packet(start: usize, end: usize, target: *mut Event) -> Res
|
||||
|
||||
let ipv4 = &*((start + IPV4_HEADER_START) as *const Ipv4Hdr);
|
||||
|
||||
// ihl() returns actual IP header length in bytes (20–60); use it so IP Options
|
||||
// don't shift the L4 header offset and corrupt port extraction.
|
||||
let ip_hdr_len = ipv4.ihl() as usize;
|
||||
if ip_hdr_len < core::mem::size_of::<Ipv4Hdr>() {
|
||||
return Err(());
|
||||
}
|
||||
let l4_start = ETHER_HEADER_END + ip_hdr_len;
|
||||
|
||||
let (source_port, destination_port) = match ipv4.proto {
|
||||
IpProto::Tcp => parse_tcp_port(start, end, IPV4_TCP_HEADER_START, IPV4_TCP_HEADER_END)?,
|
||||
IpProto::Udp => parse_udp_port(start, end, IPV4_UDP_HEADER_START, IPV4_UDP_HEADER_END)?,
|
||||
IpProto::Tcp => parse_tcp_port(start, end, l4_start)?,
|
||||
IpProto::Udp => parse_udp_port(start, end, l4_start)?,
|
||||
_ => return Err(()),
|
||||
};
|
||||
|
||||
@ -83,8 +91,8 @@ unsafe fn parse_ipv6_packet(start: usize, end: usize, target: *mut Event) -> Res
|
||||
let ipv6 = &*((start + IPV6_HEADER_START) as *const Ipv6Hdr);
|
||||
|
||||
let (source_port, destination_port) = match ipv6.next_hdr {
|
||||
IpProto::Tcp => parse_tcp_port(start, end, IPV6_TCP_HEADER_START, IPV6_TCP_HEADER_END)?,
|
||||
IpProto::Udp => parse_udp_port(start, end, IPV6_UDP_HEADER_START, IPV6_UDP_HEADER_END)?,
|
||||
IpProto::Tcp => parse_tcp_port(start, end, IPV6_TCP_HEADER_START)?,
|
||||
IpProto::Udp => parse_udp_port(start, end, IPV6_UDP_HEADER_START)?,
|
||||
_ => return Err(()),
|
||||
};
|
||||
|
||||
@ -124,24 +132,42 @@ unsafe fn parse_ipv6_packet(start: usize, end: usize, target: *mut Event) -> Res
|
||||
}
|
||||
}
|
||||
|
||||
// Mirror the C XDP pattern for variable-offset packet access:
|
||||
//
|
||||
// struct tcphdr *tcph = data + eth_size + (iph->ihl * 4);
|
||||
// if (tcph + 1 > (struct tcphdr *)data_end) return XDP_DROP;
|
||||
// tcph->source; tcph->dest;
|
||||
//
|
||||
// The typed pointer + .add(1) bounds check is key: it makes the verifier
|
||||
// associate the checked range with the *typed* pointer register. Using raw
|
||||
// u8 bytes makes TCP and UDP generate identical LLVM IR, so the optimizer
|
||||
// merges both check+read sequences into one shared basic block. At that
|
||||
// merge point it re-derives the packet pointer with a fresh verifier id
|
||||
// (r=0), causing "invalid access to packet" even though the logic is sound.
|
||||
// Different struct types (TcpHdr vs UdpHdr) produce different GEP
|
||||
// instructions in LLVM IR, preventing the merge and keeping each protocol's
|
||||
// bounds check and accesses in the same basic block.
|
||||
#[inline(always)]
|
||||
unsafe fn parse_tcp_port(start: usize, end: usize, tcp_start: usize, tcp_end: usize) -> Result<(u16, u16), ()> {
|
||||
unsafe fn parse_tcp_port(start: usize, end: usize, tcp_start: usize) -> Result<(u16, u16), ()> {
|
||||
unsafe {
|
||||
if start + tcp_end > end {
|
||||
let tcph = (start as *const u8).add(tcp_start) as *const TcpHdr;
|
||||
// tcph.add(1) advances by size_of::<TcpHdr>() = 20 bytes, matching
|
||||
// the C pattern "if (tcph + 1 > data_end)".
|
||||
if tcph.add(1) as usize > end {
|
||||
return Err(());
|
||||
}
|
||||
let tcp = &*((start + tcp_start) as *const TcpHdr);
|
||||
Ok((u16::from_be_bytes(tcp.source), u16::from_be_bytes(tcp.dest)))
|
||||
Ok((u16::from_be_bytes((*tcph).source), u16::from_be_bytes((*tcph).dest)))
|
||||
}
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
unsafe fn parse_udp_port(start: usize, end: usize, udp_start: usize, udp_end: usize) -> Result<(u16, u16), ()> {
|
||||
unsafe fn parse_udp_port(start: usize, end: usize, udp_start: usize) -> Result<(u16, u16), ()> {
|
||||
unsafe {
|
||||
if start + udp_end > end {
|
||||
let udph = (start as *const u8).add(udp_start) as *const UdpHdr;
|
||||
// udph.add(1) advances by size_of::<UdpHdr>() = 8 bytes.
|
||||
if udph.add(1) as usize > end {
|
||||
return Err(());
|
||||
}
|
||||
let udp = &*((start + udp_start) as *const UdpHdr);
|
||||
Ok((udp.src_port(), udp.dst_port()))
|
||||
Ok(((*udph).src_port(), (*udph).dst_port()))
|
||||
}
|
||||
}
|
||||
|
||||
10
config.toml
10
config.toml
@ -10,12 +10,12 @@ models_config_name = "inference_config.json"
|
||||
# Then change this item.
|
||||
combined_queue_count = 8
|
||||
channel_size = 4096
|
||||
fill_queue_size = 4096 # Umem Used (Should not modify)
|
||||
fill_queue_size = 8192 # XSK fill ring size (power of 2); larger = more burst headroom
|
||||
comp_queue_size = 4096 # Umem Used (Should not modify)
|
||||
tx_queue_size = 4096 # Umem Used (Should not modify)
|
||||
rx_queue_size = 4096 # Umem Used (Should not modify)
|
||||
frame_size = 4096 # Umem Used (Should not modify)
|
||||
frame_count = 4096 # Umem Used (Should not modify)
|
||||
frame_count = 8192 # Total UMEM frames per XskPair; doubled for larger fill queue headroom
|
||||
http_server_bind_port = 8080 # Http Server Listen Port
|
||||
refresh_interval = 5 # Statistics Refresh Time
|
||||
|
||||
@ -23,9 +23,9 @@ max_concurrent_flows = 10000 # track up to N concurrent flows
|
||||
inference_interval_secs = 5 # run ML inference every N seconds
|
||||
aggregator_window_secs = 30
|
||||
inference_batch_size = 200
|
||||
flow_timeout_us = 60_000_000
|
||||
flow_timeout_us = 120_000_000
|
||||
|
||||
traffic_logging_mode = false # when true, disables ML inference and records packets to CSV
|
||||
traffic_logging_mode = true # when true, disables ML inference and records packets to CSV
|
||||
traffic_log_csv_path = "traffic_log.csv"
|
||||
|
||||
# tls_keylog_path = "/tmp/tls_keys.log" # NSS key log file for TLS decryption (SSLKEYLOGFILE)
|
||||
@ -33,7 +33,7 @@ traffic_log_csv_path = "traffic_log.csv"
|
||||
# CPU affinity (Linux only). Uncomment and tune for your hardware.
|
||||
# Distribute XSK packet threads across a core range [start, end] (inclusive).
|
||||
# Threads are assigned round-robin: core = start + (queue_id % (end - start + 1)).
|
||||
xsk_cpu_set = [0, 3]
|
||||
xsk_cpu_set = [0, 4]
|
||||
|
||||
# Pin ML inference (ONNX spawn_blocking) to this core.
|
||||
ml_cpu = 7
|
||||
|
||||
@ -22,19 +22,24 @@ static EGRESS_XSKS_MAP: XskMap = XskMap::pinned(64, 0);
|
||||
#[xdp]
|
||||
pub fn mantis(ctx: XdpContext) -> u32 {
|
||||
unsafe {
|
||||
let _ = packet_intake(ctx);
|
||||
xdp_action::XDP_PASS
|
||||
match packet_intake(&ctx) {
|
||||
Err(_) => xdp_action::XDP_PASS,
|
||||
Ok(_) => {
|
||||
let _ = PROGRAM_ARRAY.tail_call(&ctx, TRANSMISSION);
|
||||
xdp_action::XDP_PASS
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
unsafe fn packet_intake(ctx: XdpContext) -> Result<u32, ()> {
|
||||
unsafe fn packet_intake(ctx: &XdpContext) -> Result<(), ()> {
|
||||
unsafe {
|
||||
let start = ctx.data();
|
||||
let end = ctx.data_end();
|
||||
let ptr = PARSED_PACKET.get_ptr_mut(0).ok_or(())?;
|
||||
parsing::parse_packet(start, end, ptr)?;
|
||||
let _ = PROGRAM_ARRAY.tail_call(&ctx, ACCESS_CONTROL);
|
||||
Err(())
|
||||
let _ = PROGRAM_ARRAY.tail_call(ctx, ACCESS_CONTROL);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -128,4 +128,4 @@ pub fn transmission(ctx: XdpContext) -> u32 {
|
||||
#[panic_handler]
|
||||
fn panic(_info: &core::panic::PanicInfo) -> ! {
|
||||
unsafe { core::hint::unreachable_unchecked() }
|
||||
}
|
||||
}
|
||||
|
||||
@ -10,6 +10,7 @@ use aya::Ebpf;
|
||||
use aya::maps::{MapData, XskMap};
|
||||
use crossbeam::channel::{Receiver, Sender, bounded};
|
||||
use crossbeam::queue::SegQueue;
|
||||
use libc;
|
||||
use macros::log;
|
||||
use parking_lot::Mutex;
|
||||
use tokio::sync::oneshot;
|
||||
@ -81,26 +82,31 @@ impl XskManager {
|
||||
suricata_engine.clone(),
|
||||
)?;
|
||||
|
||||
let mut xsk_map = self.xsk_map.lock();
|
||||
// Extract fds before run() consumes the XskPair structs.
|
||||
let ingress_fd = ingress_xsk.rx.fd().as_raw_fd();
|
||||
let egress_fd = egress_xsk.rx.fd().as_raw_fd();
|
||||
|
||||
// Spawn consumer threads first so fill queues are being serviced
|
||||
// before XDP starts redirecting packets into these XSKs.
|
||||
let ingress_shutdown = ingress_xsk.run(ingress_to_egress_tx, egress_to_ingress_rx)?;
|
||||
shutdowns.push(ingress_shutdown);
|
||||
|
||||
let egress_shutdown = egress_xsk.run(egress_to_ingress_tx, ingress_to_egress_rx)?;
|
||||
shutdowns.push(egress_shutdown);
|
||||
|
||||
// Register fds in eBPF maps only after threads are running.
|
||||
let mut xsk_map = self.xsk_map.lock();
|
||||
xsk_map
|
||||
.set(queue_id, ingress_fd, 0)
|
||||
.map_err(EbpfError::AfXdpSetFailed)?;
|
||||
drop(xsk_map);
|
||||
|
||||
let mut egress_xsk_map = self.egress_xsk_map.lock();
|
||||
let egress_fd = egress_xsk.rx.fd().as_raw_fd();
|
||||
egress_xsk_map
|
||||
.set(queue_id, egress_fd, 0)
|
||||
.map_err(EbpfError::AfXdpSetFailed)?;
|
||||
drop(egress_xsk_map);
|
||||
|
||||
let ingress_shutdown = ingress_xsk.run(ingress_to_egress_tx, egress_to_ingress_rx)?;
|
||||
shutdowns.push(ingress_shutdown);
|
||||
|
||||
let egress_shutdown = egress_xsk.run(egress_to_ingress_tx, ingress_to_egress_rx)?;
|
||||
shutdowns.push(egress_shutdown);
|
||||
|
||||
log!(EbpfLog::QueuePairStarted(queue_id));
|
||||
}
|
||||
|
||||
@ -148,12 +154,21 @@ impl XskPair {
|
||||
.build()
|
||||
.map_err(EbpfError::UmemSetFailed)?;
|
||||
|
||||
let (umem, frame_descs) = Umem::new(umem_config, frame_count, false).map_err(EbpfError::UmemSetFailed)?;
|
||||
let (umem, frame_descs) = match Umem::new(umem_config, frame_count, true) {
|
||||
Ok(val) => val,
|
||||
Err(_) => {
|
||||
log!(EbpfLog::HugePagesFallback);
|
||||
Umem::new(umem_config, frame_count, false).map_err(EbpfError::UmemSetFailed)?
|
||||
}
|
||||
};
|
||||
|
||||
// XDP_USE_NEED_WAKEUP lets the driver sleep when idle and requires user-space
|
||||
// to call poll() after producing fill-ring descriptors to wake the NAPI handler.
|
||||
// Works via ndo_xsk_wakeup (igb: kernel 5.10+) independent of netif_queue_set_napi.
|
||||
let socket_config = SocketConfig::builder()
|
||||
.tx_queue_size(tx_queue_size)
|
||||
.rx_queue_size(rx_queue_size)
|
||||
.bind_flags(BindFlags::XDP_ZEROCOPY)
|
||||
.bind_flags(BindFlags::XDP_ZEROCOPY | BindFlags::XDP_USE_NEED_WAKEUP)
|
||||
.libxdp_flags(LibxdpFlags::XSK_LIBXDP_FLAGS_INHIBIT_PROG_LOAD)
|
||||
.build();
|
||||
|
||||
@ -164,6 +179,39 @@ impl XskPair {
|
||||
|
||||
let (mut fill_queue, comp_queue) = queue.ok_or(EbpfError::UnknownError)?;
|
||||
|
||||
// SO_PREFER_BUSY_POLL keeps NAPI hot between fill-ring produces, reducing the
|
||||
// window where fill queue starvation can drop packets in one direction.
|
||||
// Requires netif_queue_set_napi in the driver (igb: kernel 6.15+, confirmed
|
||||
// present in Ubuntu 6.17.0-29-generic via nm check). Falls back gracefully
|
||||
// on older kernels — XDP_USE_NEED_WAKEUP still provides correct wakeup semantics.
|
||||
unsafe {
|
||||
let fd = rx.fd().as_raw_fd();
|
||||
let one: libc::c_int = 1;
|
||||
let ret = libc::setsockopt(
|
||||
fd,
|
||||
libc::SOL_SOCKET,
|
||||
libc::SO_PREFER_BUSY_POLL,
|
||||
&one as *const libc::c_int as *const libc::c_void,
|
||||
std::mem::size_of::<libc::c_int>() as libc::socklen_t,
|
||||
);
|
||||
if ret == 0 {
|
||||
let budget: libc::c_int = 64;
|
||||
libc::setsockopt(
|
||||
fd,
|
||||
libc::SOL_SOCKET,
|
||||
libc::SO_BUSY_POLL_BUDGET,
|
||||
&budget as *const libc::c_int as *const libc::c_void,
|
||||
std::mem::size_of::<libc::c_int>() as libc::socklen_t,
|
||||
);
|
||||
log!(EbpfLog::BusyPollEnabled { queue_id });
|
||||
} else {
|
||||
log!(EbpfLog::BusyPollUnavailable {
|
||||
queue_id,
|
||||
errno: *libc::__errno_location()
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
let total_frames = frame_descs.len();
|
||||
let fill_frames_count = (total_frames / 2).min(config.fill_queue_size as usize);
|
||||
|
||||
@ -224,6 +272,25 @@ impl XskPair {
|
||||
}
|
||||
}
|
||||
|
||||
// Replenish before processing so the kernel always has frames
|
||||
// to DMA into before we drain the RX ring.
|
||||
self.replenish_fill_queue();
|
||||
|
||||
// Kick NAPI and optionally sleep the thread in one syscall.
|
||||
// When we just processed packets (idle_count == 0) use a zero
|
||||
// timeout so we loop back immediately for the next burst.
|
||||
// When genuinely idle use a 1 ms blocking poll so the thread
|
||||
// sleeps inside the kernel instead of spinning — this is the
|
||||
// primary CPU-saving mechanism when traffic is low.
|
||||
// If the driver does not need a wakeup, fall back to a plain
|
||||
// sleep so idle threads on quiet queues still yield the CPU.
|
||||
if self.fill_queue.needs_wakeup() {
|
||||
let timeout_ms = if idle_count == 0 { 0 } else { 1 };
|
||||
let _ = self.fill_queue.wakeup(self.rx.fd_mut(), timeout_ms);
|
||||
} else if idle_count > 0 {
|
||||
thread::sleep(Duration::from_micros(100));
|
||||
}
|
||||
|
||||
let mut total_activity = 0;
|
||||
|
||||
match self.process_comp_queue() {
|
||||
@ -236,6 +303,10 @@ impl XskPair {
|
||||
Err(e) => log!(EbpfLog::RXQueueError(format!("{:?}", e))),
|
||||
}
|
||||
|
||||
// Replenish again after RX so frames consumed above are
|
||||
// immediately recycled back to the kernel.
|
||||
self.replenish_fill_queue();
|
||||
|
||||
match self.process_tx_queue(&forward_rx) {
|
||||
Ok(count) => total_activity += count,
|
||||
Err(e) => log!(EbpfLog::TXQueueError(format!("{:?}", e))),
|
||||
@ -250,14 +321,6 @@ impl XskPair {
|
||||
if last_cleanup.elapsed() >= Duration::from_secs(60) {
|
||||
last_cleanup = std::time::Instant::now();
|
||||
}
|
||||
|
||||
let sleep_us = match idle_count {
|
||||
0..=10 => 1,
|
||||
11..=100 => 10,
|
||||
_ => 100,
|
||||
};
|
||||
|
||||
thread::sleep(Duration::from_micros(sleep_us));
|
||||
}
|
||||
|
||||
log!(EbpfLog::XSKShutdown);
|
||||
@ -286,7 +349,7 @@ impl XskPair {
|
||||
}
|
||||
|
||||
fn process_rx_queue(&mut self, forward_tx: &Sender<Vec<u8>>) -> Result<usize, EbpfError> {
|
||||
let mut rx_descs = vec![FrameDesc::default(); 64];
|
||||
let mut rx_descs = vec![FrameDesc::default(); 256];
|
||||
let rx_count = unsafe { self.rx.consume(&mut rx_descs) };
|
||||
|
||||
if rx_count > 0 {
|
||||
@ -317,10 +380,19 @@ impl XskPair {
|
||||
}
|
||||
}
|
||||
|
||||
unsafe {
|
||||
let produced = self.fill_queue.produce(&rx_descs[..rx_count]);
|
||||
if produced != rx_count {
|
||||
log!(EbpfLog::FillQueueIncomplete(produced, rx_count));
|
||||
// Return processed frames to the fill queue and wake the kernel so it
|
||||
// can immediately DMA into them. If the fill queue is already at capacity,
|
||||
// park the overflow frames in the pool — the next replenish call will push
|
||||
// them out once space is available.
|
||||
let produced = unsafe {
|
||||
let (fill, rx) = (&mut self.fill_queue, &mut self.rx);
|
||||
fill.produce_and_wakeup(&rx_descs[..rx_count], rx.fd_mut(), 0)
|
||||
.unwrap_or(0)
|
||||
};
|
||||
if produced < rx_count {
|
||||
let mut pool = self.frame_pool.lock();
|
||||
for desc in rx_descs[produced..rx_count].iter() {
|
||||
pool.push(*desc);
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -328,6 +400,44 @@ impl XskPair {
|
||||
Ok(rx_count)
|
||||
}
|
||||
|
||||
fn replenish_fill_queue(&mut self) {
|
||||
// Reserve enough frames in pool so TX always has something to work with.
|
||||
// Drain everything above that threshold into the fill queue in one shot.
|
||||
// produce() is all-or-nothing per call, so iterate in BATCH-sized chunks
|
||||
// until the fill queue is full; return any remaining frames to the pool.
|
||||
const RESERVED_FOR_TX: usize = 64;
|
||||
const BATCH: usize = 256;
|
||||
|
||||
let all_frames: Vec<FrameDesc> = {
|
||||
let mut pool = self.frame_pool.lock();
|
||||
let available = pool.len().saturating_sub(RESERVED_FOR_TX);
|
||||
if available == 0 {
|
||||
return;
|
||||
}
|
||||
let start = pool.len() - available;
|
||||
pool.drain(start..).collect()
|
||||
};
|
||||
|
||||
let mut offset = 0;
|
||||
while offset < all_frames.len() {
|
||||
let end = (offset + BATCH).min(all_frames.len());
|
||||
let added = unsafe {
|
||||
let (fill, rx) = (&mut self.fill_queue, &mut self.rx);
|
||||
fill.produce_and_wakeup(&all_frames[offset..end], rx.fd_mut(), 0)
|
||||
.unwrap_or(0)
|
||||
};
|
||||
if added == 0 {
|
||||
break; // fill queue full
|
||||
}
|
||||
offset += added;
|
||||
}
|
||||
|
||||
if offset < all_frames.len() {
|
||||
let mut pool = self.frame_pool.lock();
|
||||
pool.extend_from_slice(&all_frames[offset..]);
|
||||
}
|
||||
}
|
||||
|
||||
fn process_tx_queue(&mut self, forward_rx: &Receiver<Vec<u8>>) -> Result<usize, EbpfError> {
|
||||
let mut packets_to_send = Vec::with_capacity(64);
|
||||
while let Ok(packet) = forward_rx.try_recv() {
|
||||
|
||||
@ -21,7 +21,6 @@ use crate::core::infrastructure::health::SystemHealth;
|
||||
use crate::detection::fusion::{FusionEngine, FusionMode};
|
||||
use crate::detection::ml::config_loader::InferenceConfig;
|
||||
use crate::detection::ml::engine::Engine;
|
||||
use crate::detection::ml::feature_extractor::FlowFeatures;
|
||||
use crate::detection::ml::model_loader::MLModels;
|
||||
use crate::detection::ml::traffic_logger::TrafficLogger;
|
||||
use crate::detection::suricata::SuricataEngine;
|
||||
@ -34,7 +33,7 @@ pub struct AppServices {
|
||||
pub health: Arc<SystemHealth>,
|
||||
pub detection_alert: Arc<DetectionAlert>,
|
||||
pub fusion_engine: Arc<FusionEngine>,
|
||||
pub ml_models: Arc<MLModels>,
|
||||
pub ml_models: Option<Arc<MLModels>>,
|
||||
pub ml_engine: Arc<Engine>,
|
||||
pub suricata_engine: Option<Arc<SuricataEngine>>,
|
||||
pub app_db: Option<Arc<AppDB>>,
|
||||
@ -49,7 +48,6 @@ impl AppServices {
|
||||
log_broadcaster: Arc<crate::core::infrastructure::log_broadcaster::LogBroadcaster>,
|
||||
) -> Result<Self, Error> {
|
||||
let health = SystemHealth::new(app_config.clone())?;
|
||||
let ml_models = Arc::new(MLModels::load_models(&app_config)?);
|
||||
|
||||
let detection_alert = Arc::new(DetectionAlert::new());
|
||||
let mode = FusionMode::from_str(&app_config.fusion_mode);
|
||||
@ -59,17 +57,10 @@ impl AppServices {
|
||||
detection_alert.clone(),
|
||||
));
|
||||
|
||||
let traffic_logger = if app_config.traffic_logging_mode {
|
||||
let (ml_models, traffic_logger) = if app_config.traffic_logging_mode {
|
||||
let dir = PathBuf::from(env!("CSV_RECORD_PATH"));
|
||||
let basename = app_config.traffic_log_csv_path.trim_end_matches(".csv").to_string();
|
||||
let mut header = vec![
|
||||
"Source IP".to_string(),
|
||||
"Destination IP".to_string(),
|
||||
"Timestamp".to_string(),
|
||||
];
|
||||
header.extend(FlowFeatures::all_feature_names_owned());
|
||||
header.push("Label".to_string());
|
||||
let logger = TrafficLogger::new(&dir, &basename, header)
|
||||
let logger = TrafficLogger::new(&dir, &basename)
|
||||
.map_err(|e| MiscError::TrafficLogCreateError(dir.display().to_string(), e.to_string()))?;
|
||||
log!(SystemLog::TrafficLoggingEnabled(format!(
|
||||
"{}/{}-{}.csv",
|
||||
@ -77,9 +68,10 @@ impl AppServices {
|
||||
basename,
|
||||
Local::now().format("%Y-%m-%d")
|
||||
)));
|
||||
Some(Arc::new(logger))
|
||||
(None, Some(Arc::new(logger)))
|
||||
} else {
|
||||
None
|
||||
let models = Arc::new(MLModels::load_models(&app_config)?);
|
||||
(Some(models), None)
|
||||
};
|
||||
|
||||
let ml_engine = Arc::new(Engine::new(
|
||||
|
||||
@ -84,9 +84,9 @@ impl System {
|
||||
let ebpf_services = self.ebpf_services.clone();
|
||||
let app_services = self.app_services.clone();
|
||||
|
||||
log!(MLLog::ModelsLoaded(
|
||||
self.app_services.ml_models.get_model_info("deep_autoencoder")
|
||||
));
|
||||
if let Some(ref models) = self.app_services.ml_models {
|
||||
log!(MLLog::ModelsLoaded(models.get_model_info("deep_autoencoder")));
|
||||
}
|
||||
|
||||
log!(MLLog::ConfigLoaded {
|
||||
features: self.inference_config.num_ae_features(),
|
||||
|
||||
@ -7,7 +7,6 @@ use tokio::time::interval;
|
||||
|
||||
use super::aggregator::AttackAggregator;
|
||||
use super::config_loader::InferenceConfig;
|
||||
use super::feature_extractor::FlowFeatures;
|
||||
use super::flow_tracker::FlowTracker;
|
||||
use super::inference::Inference;
|
||||
use super::model_loader::MLModels;
|
||||
@ -21,7 +20,7 @@ use crate::utils::packet_parser::parse_packet;
|
||||
|
||||
pub struct Engine {
|
||||
tracker: Arc<Mutex<FlowTracker>>,
|
||||
inference_pipeline: Arc<Inference>,
|
||||
inference_pipeline: Option<Arc<Inference>>,
|
||||
aggregator: Arc<Mutex<AttackAggregator>>,
|
||||
fusion_engine: Arc<FusionEngine>,
|
||||
batch_size: usize,
|
||||
@ -33,7 +32,7 @@ pub struct Engine {
|
||||
|
||||
impl Engine {
|
||||
pub fn new(
|
||||
models: Arc<MLModels>,
|
||||
models: Option<Arc<MLModels>>,
|
||||
config: Arc<InferenceConfig>,
|
||||
fusion_engine: Arc<FusionEngine>,
|
||||
max_flows: usize,
|
||||
@ -45,7 +44,7 @@ impl Engine {
|
||||
ml_cpu: Option<u32>,
|
||||
) -> Self {
|
||||
let tracker = Arc::new(Mutex::new(FlowTracker::new(max_flows)));
|
||||
let inference_pipeline = Arc::new(Inference::new(models, config));
|
||||
let inference_pipeline = models.map(|m| Arc::new(Inference::new(m, config)));
|
||||
|
||||
let min_detections = ((window_secs / interval_secs) / 2).max(1) as usize;
|
||||
let aggregator = Arc::new(Mutex::new(AttackAggregator::new(window_secs, min_detections)));
|
||||
@ -90,33 +89,38 @@ impl Engine {
|
||||
continue;
|
||||
};
|
||||
let total_flows = t.flow_count();
|
||||
let flows = t.drain_flows();
|
||||
let flows = t.drain_ready_flows(self.flow_timeout_us);
|
||||
t.cleanup_old_flows(self.flow_timeout_us);
|
||||
// Build active_ips from the drained flows, not the post-drain tracker
|
||||
// (which is always empty). This preserves per-src_ip inference buffers
|
||||
// across consecutive cycles so the LSTM window can fill up over time.
|
||||
// A src_ip absent from this cycle loses its buffer on the next cleanup.
|
||||
let active_ips: std::collections::HashSet<String> =
|
||||
flows.iter().map(|f| f.flow_key.src_ip.clone()).collect();
|
||||
// active_ips covers both the just-drained flows and flows still in the
|
||||
// tracker (ongoing connections), so their LSTM buffers are preserved.
|
||||
let active_ips: std::collections::HashSet<String> = flows
|
||||
.iter()
|
||||
.map(|f| f.flow_key.src_ip.clone())
|
||||
.chain(t.active_src_ips().map(str::to_owned))
|
||||
.collect();
|
||||
(total_flows, flows, active_ips)
|
||||
};
|
||||
|
||||
log!(MLLog::FlowStats(total_flows, flows.len()));
|
||||
|
||||
if flows.is_empty() {
|
||||
self.inference_pipeline.cleanup_buffers(&active_ips);
|
||||
if let Some(ref p) = self.inference_pipeline {
|
||||
p.cleanup_buffers(&active_ips);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
if let Some(ref logger) = self.traffic_logger {
|
||||
let feature_names = FlowFeatures::all_feature_names_owned();
|
||||
for flow in &flows {
|
||||
let features = FlowFeatures::extract(flow, &feature_names);
|
||||
logger.log_row(features.to_csv_record());
|
||||
logger.log_flow(flow);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
let Some(pipeline) = self.inference_pipeline.as_ref().map(Arc::clone) else {
|
||||
continue;
|
||||
};
|
||||
|
||||
let mut batch = flows[..flows.len().min(self.batch_size)].to_vec();
|
||||
batch.sort_by(|a, b| {
|
||||
a.flow_key
|
||||
@ -127,7 +131,6 @@ impl Engine {
|
||||
|
||||
let start = Instant::now();
|
||||
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
|
||||
@ -148,7 +151,9 @@ impl Engine {
|
||||
|
||||
if results.is_empty() {
|
||||
// All flows are still in the warm-up window; no ONNX inference ran.
|
||||
self.inference_pipeline.cleanup_buffers(&active_ips);
|
||||
if let Some(ref p) = self.inference_pipeline {
|
||||
p.cleanup_buffers(&active_ips);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
@ -204,7 +209,9 @@ impl Engine {
|
||||
aggregator.cleanup();
|
||||
}
|
||||
|
||||
self.inference_pipeline.cleanup_buffers(&active_ips);
|
||||
if let Some(ref p) = self.inference_pipeline {
|
||||
p.cleanup_buffers(&active_ips);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -87,6 +87,7 @@ impl FlowFeatures {
|
||||
compute_stats(&flow.idle_periods.iter().map(|&x| x as f64).collect::<Vec<_>>());
|
||||
|
||||
match feature_name {
|
||||
"Source Port" | "Src Port" | "src_port" => flow.flow_key.src_port as f64,
|
||||
"Destination Port" | "Dst Port" | "dst_port" => flow.flow_key.dst_port as f64,
|
||||
"Protocol" | "protocol" => flow.flow_key.protocol as f64,
|
||||
"Flow Duration" | "flow_duration" => duration_us,
|
||||
@ -218,105 +219,20 @@ impl FlowFeatures {
|
||||
&self.features
|
||||
}
|
||||
|
||||
pub fn all_feature_names() -> Vec<&'static str> {
|
||||
vec![
|
||||
"Destination Port",
|
||||
"Protocol",
|
||||
"Flow Duration",
|
||||
"Total Fwd Packets",
|
||||
"Total Backward Packets",
|
||||
"Total Length of Fwd Packets",
|
||||
"Total Length of Bwd Packets",
|
||||
"Fwd Packet Length Max",
|
||||
"Fwd Packet Length Min",
|
||||
"Fwd Packet Length Mean",
|
||||
"Fwd Packet Length Std",
|
||||
"Bwd Packet Length Max",
|
||||
"Bwd Packet Length Min",
|
||||
"Bwd Packet Length Mean",
|
||||
"Bwd Packet Length Std",
|
||||
"Flow Bytes/s",
|
||||
"Flow Packets/s",
|
||||
"Flow IAT Mean",
|
||||
"Flow IAT Std",
|
||||
"Flow IAT Max",
|
||||
"Flow IAT Min",
|
||||
"Fwd IAT Total",
|
||||
"Fwd IAT Mean",
|
||||
"Fwd IAT Std",
|
||||
"Fwd IAT Max",
|
||||
"Fwd IAT Min",
|
||||
"Bwd IAT Total",
|
||||
"Bwd IAT Mean",
|
||||
"Bwd IAT Std",
|
||||
"Bwd IAT Max",
|
||||
"Bwd IAT Min",
|
||||
"Fwd PSH Flags",
|
||||
"Bwd PSH Flags",
|
||||
"Fwd URG Flags",
|
||||
"Bwd URG Flags",
|
||||
"Fwd Header Length",
|
||||
"Bwd Header Length",
|
||||
"Fwd Packets/s",
|
||||
"Bwd Packets/s",
|
||||
"Min Packet Length",
|
||||
"Max Packet Length",
|
||||
"Packet Length Mean",
|
||||
"Packet Length Std",
|
||||
"Packet Length Variance",
|
||||
"FIN Flag Count",
|
||||
"SYN Flag Count",
|
||||
"RST Flag Count",
|
||||
"PSH Flag Count",
|
||||
"ACK Flag Count",
|
||||
"URG Flag Count",
|
||||
"CWE Flag Count",
|
||||
"ECE Flag Count",
|
||||
"Down/Up Ratio",
|
||||
"Average Packet Size",
|
||||
"Avg Fwd Segment Size",
|
||||
"Avg Bwd Segment Size",
|
||||
"Fwd Header Length.1",
|
||||
"Fwd Avg Bytes/Bulk",
|
||||
"Fwd Avg Packets/Bulk",
|
||||
"Fwd Avg Bulk Rate",
|
||||
"Bwd Avg Bytes/Bulk",
|
||||
"Bwd Avg Packets/Bulk",
|
||||
"Bwd Avg Bulk Rate",
|
||||
"Subflow Fwd Packets",
|
||||
"Subflow Fwd Bytes",
|
||||
"Subflow Bwd Packets",
|
||||
"Subflow Bwd Bytes",
|
||||
"Init_Win_bytes_forward",
|
||||
"Init_Win_bytes_backward",
|
||||
"act_data_pkt_fwd",
|
||||
"min_seg_size_forward",
|
||||
"Active Mean",
|
||||
"Active Std",
|
||||
"Active Max",
|
||||
"Active Min",
|
||||
"Idle Mean",
|
||||
"Idle Std",
|
||||
"Idle Max",
|
||||
"Idle Min",
|
||||
]
|
||||
}
|
||||
|
||||
pub fn all_feature_names_owned() -> Vec<String> {
|
||||
Self::all_feature_names().iter().map(|s| s.to_string()).collect()
|
||||
}
|
||||
|
||||
pub fn to_csv_record(&self) -> Vec<String> {
|
||||
let ts_ms = self.timestamp / 1000;
|
||||
let ts_str = match Utc.timestamp_millis_opt(ts_ms as i64) {
|
||||
chrono::LocalResult::Single(dt) => dt.format("%m/%d/%Y %H:%M:%S").to_string(),
|
||||
_ => ts_ms.to_string(),
|
||||
};
|
||||
|
||||
let mut record = vec![self.src_ip.clone(), self.dst_ip.clone(), ts_str];
|
||||
record.extend(self.features.iter().map(|f| f.to_string()));
|
||||
record.push("BENIGN".to_string());
|
||||
record
|
||||
pub fn get_csv_column(flow: &FlowData, column: &str) -> String {
|
||||
match column {
|
||||
"Source IP" => flow.flow_key.src_ip.clone(),
|
||||
"Destination IP" => flow.flow_key.dst_ip.clone(),
|
||||
"Timestamp" => {
|
||||
let ts_ms = (flow.start_time_us / 1000) as i64;
|
||||
match Utc.timestamp_millis_opt(ts_ms) {
|
||||
chrono::LocalResult::Single(dt) => dt.format("%m/%d/%Y %H:%M:%S").to_string(),
|
||||
_ => ts_ms.to_string(),
|
||||
}
|
||||
}
|
||||
"Label" => "BENIGN".to_string(),
|
||||
name => Self::get_feature_by_name(flow, name).to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -3,11 +3,10 @@ use std::time;
|
||||
|
||||
use common::model::event::Event;
|
||||
|
||||
use super::server_ports;
|
||||
use crate::model::direction::Direction;
|
||||
use crate::model::ml_detection::{BulkState, FlowKey, PacketData};
|
||||
|
||||
use super::server_ports;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct FlowData {
|
||||
pub flow_key: FlowKey,
|
||||
@ -313,8 +312,27 @@ impl FlowTracker {
|
||||
}
|
||||
}
|
||||
|
||||
pub fn drain_flows(&mut self) -> Vec<FlowData> {
|
||||
self.flows.drain().map(|(_, flow)| flow).collect()
|
||||
pub fn drain_ready_flows(&mut self, timeout_us: u64) -> Vec<FlowData> {
|
||||
let now = time::SystemTime::now()
|
||||
.duration_since(time::UNIX_EPOCH)
|
||||
.map(|d| d.as_micros() as u64)
|
||||
.unwrap_or(0);
|
||||
|
||||
let ready_keys: Vec<FlowKey> = self
|
||||
.flows
|
||||
.iter()
|
||||
.filter(|(_, flow)| flow.is_finished() || now.saturating_sub(flow.last_time_us) >= timeout_us)
|
||||
.map(|(key, _)| key.clone())
|
||||
.collect();
|
||||
|
||||
ready_keys
|
||||
.into_iter()
|
||||
.filter_map(|key| self.flows.remove(&key))
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn active_src_ips(&self) -> impl Iterator<Item = &str> {
|
||||
self.flows.values().map(|f| f.flow_key.src_ip.as_str())
|
||||
}
|
||||
|
||||
pub fn get_flows_snapshot(&self) -> Vec<FlowData> {
|
||||
@ -376,4 +394,4 @@ fn detect_initiator(payload: &[u8], protocol: u8, src_port: u16, dst_port: u16)
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
@ -5,5 +5,5 @@ pub mod feature_extractor;
|
||||
pub mod flow_tracker;
|
||||
pub mod inference;
|
||||
pub mod model_loader;
|
||||
pub mod traffic_logger;
|
||||
pub mod server_ports;
|
||||
pub mod traffic_logger;
|
||||
|
||||
@ -142,4 +142,4 @@ pub fn is_non_unicast(ip: &str) -> bool {
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -4,10 +4,12 @@ use std::path::{Path, PathBuf};
|
||||
use std::thread;
|
||||
use std::time::Duration;
|
||||
|
||||
use chrono::Local;
|
||||
use chrono::{Local, TimeZone, Utc};
|
||||
use crossbeam::channel::{RecvTimeoutError, Sender, TrySendError, bounded};
|
||||
use macros::log;
|
||||
|
||||
use super::feature_extractor::FlowFeatures;
|
||||
use super::flow_tracker::FlowData;
|
||||
use crate::model::error::ml::MLError;
|
||||
use crate::model::log::ml::MLLog;
|
||||
|
||||
@ -16,7 +18,112 @@ pub struct TrafficLogger {
|
||||
}
|
||||
|
||||
impl TrafficLogger {
|
||||
pub fn new(dir: impl AsRef<Path>, prefix: &str, header: Vec<String>) -> Result<Self, std::io::Error> {
|
||||
pub fn feature_names() -> &'static [&'static str] {
|
||||
&[
|
||||
"Source IP",
|
||||
"Destination IP",
|
||||
"Timestamp",
|
||||
"Source Port",
|
||||
"Destination Port",
|
||||
"Protocol",
|
||||
"Flow Duration",
|
||||
"Total Fwd Packets",
|
||||
"Total Backward Packets",
|
||||
"Total Length of Fwd Packets",
|
||||
"Total Length of Bwd Packets",
|
||||
"Fwd Packet Length Max",
|
||||
"Fwd Packet Length Min",
|
||||
"Fwd Packet Length Mean",
|
||||
"Fwd Packet Length Std",
|
||||
"Bwd Packet Length Max",
|
||||
"Bwd Packet Length Min",
|
||||
"Bwd Packet Length Mean",
|
||||
"Bwd Packet Length Std",
|
||||
"Flow Bytes/s",
|
||||
"Flow Packets/s",
|
||||
"Flow IAT Mean",
|
||||
"Flow IAT Std",
|
||||
"Flow IAT Max",
|
||||
"Flow IAT Min",
|
||||
"Fwd IAT Total",
|
||||
"Fwd IAT Mean",
|
||||
"Fwd IAT Std",
|
||||
"Fwd IAT Max",
|
||||
"Fwd IAT Min",
|
||||
"Bwd IAT Total",
|
||||
"Bwd IAT Mean",
|
||||
"Bwd IAT Std",
|
||||
"Bwd IAT Max",
|
||||
"Bwd IAT Min",
|
||||
"Fwd PSH Flags",
|
||||
"Bwd PSH Flags",
|
||||
"Fwd URG Flags",
|
||||
"Bwd URG Flags",
|
||||
"Fwd Header Length",
|
||||
"Bwd Header Length",
|
||||
"Fwd Packets/s",
|
||||
"Bwd Packets/s",
|
||||
"Min Packet Length",
|
||||
"Max Packet Length",
|
||||
"Packet Length Mean",
|
||||
"Packet Length Std",
|
||||
"Packet Length Variance",
|
||||
"FIN Flag Count",
|
||||
"SYN Flag Count",
|
||||
"RST Flag Count",
|
||||
"PSH Flag Count",
|
||||
"ACK Flag Count",
|
||||
"URG Flag Count",
|
||||
"CWE Flag Count",
|
||||
"ECE Flag Count",
|
||||
"Down/Up Ratio",
|
||||
"Average Packet Size",
|
||||
"Avg Fwd Segment Size",
|
||||
"Avg Bwd Segment Size",
|
||||
"Fwd Header Length.1",
|
||||
"Fwd Avg Bytes/Bulk",
|
||||
"Fwd Avg Packets/Bulk",
|
||||
"Fwd Avg Bulk Rate",
|
||||
"Bwd Avg Bytes/Bulk",
|
||||
"Bwd Avg Packets/Bulk",
|
||||
"Bwd Avg Bulk Rate",
|
||||
"Subflow Fwd Packets",
|
||||
"Subflow Fwd Bytes",
|
||||
"Subflow Bwd Packets",
|
||||
"Subflow Bwd Bytes",
|
||||
"Init_Win_bytes_forward",
|
||||
"Init_Win_bytes_backward",
|
||||
"act_data_pkt_fwd",
|
||||
"min_seg_size_forward",
|
||||
"Active Mean",
|
||||
"Active Std",
|
||||
"Active Max",
|
||||
"Active Min",
|
||||
"Idle Mean",
|
||||
"Idle Std",
|
||||
"Idle Max",
|
||||
"Idle Min",
|
||||
"Label",
|
||||
]
|
||||
}
|
||||
|
||||
pub fn csv_header() -> Vec<String> {
|
||||
Self::feature_names().iter().map(|s| s.to_string()).collect()
|
||||
}
|
||||
|
||||
pub fn log_flow(&self, flow: &FlowData) {
|
||||
self.log_row(Self::flow_to_csv_record(flow));
|
||||
}
|
||||
|
||||
fn flow_to_csv_record(flow: &FlowData) -> Vec<String> {
|
||||
Self::feature_names()
|
||||
.iter()
|
||||
.map(|&col| FlowFeatures::get_csv_column(flow, col))
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn new(dir: impl AsRef<Path>, prefix: &str) -> Result<Self, std::io::Error> {
|
||||
let header = Self::csv_header();
|
||||
let dir = PathBuf::from(dir.as_ref());
|
||||
fs::create_dir_all(&dir)?;
|
||||
|
||||
@ -94,7 +201,7 @@ impl TrafficLogger {
|
||||
Ok(Self { sender })
|
||||
}
|
||||
|
||||
pub fn log_row(&self, record: Vec<String>) {
|
||||
fn log_row(&self, record: Vec<String>) {
|
||||
match self.sender.try_send(record) {
|
||||
Ok(_) => {}
|
||||
Err(TrySendError::Full(_)) => {}
|
||||
|
||||
@ -50,5 +50,14 @@ loggable! {
|
||||
|
||||
#[error("Fill queue incomplete: produced {produced}, expected {expected}")]
|
||||
FillQueueIncomplete { produced: usize, expected: usize } => tracing::Level::WARN,
|
||||
|
||||
#[error("Huge pages unavailable, falling back to regular pages for UMEM")]
|
||||
HugePagesFallback => tracing::Level::WARN,
|
||||
|
||||
#[error("SO_PREFER_BUSY_POLL enabled for queue {queue_id}")]
|
||||
BusyPollEnabled { queue_id: u32 } => tracing::Level::INFO,
|
||||
|
||||
#[error("SO_PREFER_BUSY_POLL not supported on this kernel/driver (queue {queue_id}, errno {errno})")]
|
||||
BusyPollUnavailable { queue_id: u32, errno: i32 } => tracing::Level::WARN,
|
||||
}
|
||||
}
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user