refactor: Complete architecture overhaul and code review fixes (#14)

* refactor: Complete architecture overhaul and code review fixes

eBPF:
- Replace Event enum with flat ParsedPacket struct (56 bytes)
- Replace TcpFlags (8 bools) with u8 bitmap constants
- Remove all eBPF statistics (24 maps) — moved to userspace
- Replace port magic number with PortRule struct (match_all flag)
- Add rate_limit stage: packet/SYN/UDP/DNS per-IP rate limiting
- Implement dynamic pipeline via NEXT_STAGE map
- Flatten egress to single XSK redirect
- Fix TCP data offset validation and header bounds checks
- Fix HTTP protocol_filter: None→false, use pkt.tcp_flags
- Fix rate limit window off-by-one
- Fix verifier bounds check for packet access
- Add static assertions for header sizes

Userspace:
- Move ml/ → core/ml/, rename AppServices → MLService
- Rename service → protocol_filter
- Per-thread FlowTracker with parking_lot::Mutex
- Extract EngineConfig, remove PacketProcessor wrapper
- Add UserPacket, FlowStatistics, RateLimitConfig
- Split config.toml into sections
- FlowKey bytes, FlowData memory limits, EntryMap generics
- Fix MSE denominator, segment sizes, flow eviction
- Fix CORS, config validation, FD race, packet bounds
- Fix IPv6 RFC 5952, add ICMP support, ML config validation
- Add safety comments, DOS protection, configurable log level

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* refactor: Restructure REST API with /api/v1/ prefix

- Add /api/v1/ prefix for all REST endpoints
- Move WebSocket routes to /ws/ (health, alerts, flows)
- Rename: access_control → acl, service → filter, misc → system
- Restructure filter into /filter/http and /filter/ssh sub-scopes
- Add rate-limit config API (GET/PUT /api/v1/rate-limit/config)
- Wire flow_stats_ws to /ws/flows
- Fix all error responses to JSON format
- Fix double JSON serialization (.json(web::Json(x)) → .json(x))
- Remove old control/ directory

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* feat: WebSocket flow stats with subscription-based filtering

- Add direction and last_seen_us fields to FlowStatsEntry
- Add FlowSubscription type for client-side query filters
- WebSocket /ws/flows now supports subscription messages:
  {"direction": "ingress", "window_secs": 60, "top_n": 10, "interval_secs": 3}
- Client can update filter at any time by sending new subscription JSON
- Server responds immediately with filtered data on subscription change
- Default: all flows, no filter, 5 second push interval
- Remove broadcast channel from FlowStatistics (per-client filtering instead)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
DaLaw2 2026-03-19 18:00:03 +08:00 committed by GitHub
parent 26f2e1851b
commit 18a2d4a168
80 changed files with 2337 additions and 2592 deletions

View File

@ -1,4 +1,5 @@
pub mod offset;
pub mod other;
pub mod program_array;
pub mod pipeline;
pub mod setting;
pub mod tcp_flags;

View File

@ -23,3 +23,10 @@ pub const IPV4_UDP_HEADER_END: usize = IPV4_UDP_HEADER_START + size_of::<UdpHdr>
pub const IPV6_UDP_HEADER_START: usize = IPV6_HEADER_END;
pub const IPV6_UDP_HEADER_END: usize = IPV6_UDP_HEADER_START + size_of::<UdpHdr>();
#[cfg(not(feature = "user"))]
const _: () = {
assert!(core::mem::size_of::<EthHdr>() == 14);
assert!(core::mem::size_of::<Ipv4Hdr>() == 20);
assert!(core::mem::size_of::<Ipv6Hdr>() == 40);
};

View File

@ -0,0 +1,7 @@
pub const MAX_STAGES: u32 = 8;
pub const STAGE_NONE: u32 = u32::MAX;
pub const STAGE_ENTRY: u32 = 0;
pub const STAGE_ACCESS_CONTROL: u32 = 1;
pub const STAGE_RATE_LIMIT: u32 = 2;
pub const STAGE_SERVICE: u32 = 3;
pub const STAGE_TRANSMISSION: u32 = 7;

View File

@ -1,11 +0,0 @@
pub mod ingress {
pub const ACCESS_CONTROL: u32 = 0;
pub const SERVICE: u32 = 1;
pub const STATISTICS: u32 = 2;
pub const TRANSMISSION: u32 = 3;
}
pub mod egress {
pub const STATISTICS: u32 = 0;
pub const TRANSMISSION: u32 = 1;
}

View File

@ -0,0 +1,8 @@
pub const TCP_FIN: u8 = 0x01;
pub const TCP_SYN: u8 = 0x02;
pub const TCP_RST: u8 = 0x04;
pub const TCP_PSH: u8 = 0x08;
pub const TCP_ACK: u8 = 0x10;
pub const TCP_URG: u8 = 0x20;
pub const TCP_ECE: u8 = 0x40;
pub const TCP_CWR: u8 = 0x80;

View File

@ -5,9 +5,9 @@ use network_types::tcp::TcpHdr;
use network_types::udp::UdpHdr;
use crate::define::offset::*;
use crate::model::event::{Event, IPv4Event, IPv6Event};
use crate::model::parsed_packet::ParsedPacket;
pub fn parse_packet(start: usize, end: usize, target: *mut Event) -> Result<(), ()> {
pub fn parse_packet(start: usize, end: usize, target: *mut ParsedPacket) -> Result<(), ()> {
unsafe {
if start + ETHER_HEADER_END > end {
return Err(());
@ -23,125 +23,93 @@ pub fn parse_packet(start: usize, end: usize, target: *mut Event) -> Result<(),
}
#[inline(always)]
unsafe fn parse_ipv4_packet(start: usize, end: usize, target: *mut Event) -> Result<(), ()> {
unsafe {
if start + IPV4_HEADER_END > end {
return Err(());
}
let ipv4 = &*((start + IPV4_HEADER_START) as *const Ipv4Hdr);
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)?,
_ => return Err(()),
};
*(target as *mut u32) = 0;
let ipv4_data_ptr = (target as *mut u8).add(16);
core::ptr::write(ipv4_data_ptr as *mut IpProto, ipv4.proto);
core::ptr::copy_nonoverlapping(
ipv4.src_addr.as_ptr(),
ipv4_data_ptr.add(core::mem::offset_of!(IPv4Event, src_ip)),
4,
);
core::ptr::copy_nonoverlapping(
ipv4.dst_addr.as_ptr(),
ipv4_data_ptr.add(core::mem::offset_of!(IPv4Event, dst_ip)),
4,
);
core::ptr::write(
ipv4_data_ptr.add(core::mem::offset_of!(IPv4Event, src_port)) as *mut u16,
source_port,
);
core::ptr::write(
ipv4_data_ptr.add(core::mem::offset_of!(IPv4Event, dst_port)) as *mut u16,
destination_port,
);
core::ptr::write(
ipv4_data_ptr.add(core::mem::offset_of!(IPv4Event, packet_length)) as *mut u32,
(end - start) as u32,
);
core::ptr::write(
ipv4_data_ptr.add(core::mem::offset_of!(IPv4Event, timestamp_us)) as *mut u64,
bpf_ktime_get_ns(),
);
Ok(())
unsafe fn parse_ipv4_packet(start: usize, end: usize, target: *mut ParsedPacket) -> Result<(), ()> {
if start + IPV4_HEADER_END > end {
return Err(());
}
let ipv4 = &*((start + IPV4_HEADER_START) as *const Ipv4Hdr);
let packet_length = (end - start) as u32;
let (src_port, dst_port, tcp_flags, l4_header_len) = match ipv4.proto {
IpProto::Tcp => parse_tcp(start, end, IPV4_TCP_HEADER_START, IPV4_TCP_HEADER_END)?,
IpProto::Udp => parse_udp(start, end, IPV4_UDP_HEADER_START, IPV4_UDP_HEADER_END)?,
_ => return Err(()),
};
let t = &mut *target;
t.timestamp_ns = bpf_ktime_get_ns();
core::ptr::copy_nonoverlapping(ipv4.src_addr.as_ptr(), t.src_ip.as_mut_ptr(), 4);
core::ptr::copy_nonoverlapping(ipv4.dst_addr.as_ptr(), t.dst_ip.as_mut_ptr(), 4);
t.packet_length = packet_length;
t.payload_length = packet_length.saturating_sub((IPV4_HEADER_END + l4_header_len) as u32);
t.src_port = src_port;
t.dst_port = dst_port;
t.ip_version = 4;
t.protocol = ipv4.proto;
t.tcp_flags = tcp_flags;
Ok(())
}
#[inline(always)]
unsafe fn parse_ipv6_packet(start: usize, end: usize, target: *mut Event) -> Result<(), ()> {
unsafe {
if start + IPV6_HEADER_END > end {
return Err(());
}
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)?,
_ => return Err(()),
};
*(target as *mut u32) = 1;
let ipv6_data_ptr = (target as *mut u8).add(16);
core::ptr::write(ipv6_data_ptr as *mut IpProto, ipv6.next_hdr);
core::ptr::copy_nonoverlapping(
ipv6.src_addr.as_ptr(),
ipv6_data_ptr.add(core::mem::offset_of!(IPv6Event, src_ip)),
16,
);
core::ptr::copy_nonoverlapping(
ipv6.dst_addr.as_ptr(),
ipv6_data_ptr.add(core::mem::offset_of!(IPv6Event, dst_ip)),
16,
);
core::ptr::write(
ipv6_data_ptr.add(core::mem::offset_of!(IPv6Event, src_port)) as *mut u16,
source_port,
);
core::ptr::write(
ipv6_data_ptr.add(core::mem::offset_of!(IPv6Event, dst_port)) as *mut u16,
destination_port,
);
core::ptr::write(
ipv6_data_ptr.add(core::mem::offset_of!(IPv6Event, packet_length)) as *mut u32,
(end - start) as u32,
);
core::ptr::write(
ipv6_data_ptr.add(core::mem::offset_of!(IPv6Event, timestamp_us)) as *mut u64,
bpf_ktime_get_ns(),
);
Ok(())
unsafe fn parse_ipv6_packet(start: usize, end: usize, target: *mut ParsedPacket) -> Result<(), ()> {
if start + IPV6_HEADER_END > end {
return Err(());
}
let ipv6 = &*((start + IPV6_HEADER_START) as *const Ipv6Hdr);
let packet_length = (end - start) as u32;
let (src_port, dst_port, tcp_flags, l4_header_len) = match ipv6.next_hdr {
IpProto::Tcp => parse_tcp(start, end, IPV6_TCP_HEADER_START, IPV6_TCP_HEADER_END)?,
IpProto::Udp => parse_udp(start, end, IPV6_UDP_HEADER_START, IPV6_UDP_HEADER_END)?,
_ => return Err(()),
};
let t = &mut *target;
t.timestamp_ns = bpf_ktime_get_ns();
core::ptr::copy_nonoverlapping(ipv6.src_addr.as_ptr(), t.src_ip.as_mut_ptr(), 16);
core::ptr::copy_nonoverlapping(ipv6.dst_addr.as_ptr(), t.dst_ip.as_mut_ptr(), 16);
t.packet_length = packet_length;
t.payload_length = packet_length.saturating_sub((IPV6_HEADER_END + l4_header_len) as u32);
t.src_port = src_port;
t.dst_port = dst_port;
t.ip_version = 6;
t.protocol = ipv6.next_hdr;
t.tcp_flags = tcp_flags;
Ok(())
}
#[inline(always)]
unsafe fn parse_tcp_port(start: usize, end: usize, tcp_start: usize, tcp_end: usize) -> Result<(u16, u16), ()> {
unsafe {
if start + tcp_end > end {
return Err(());
}
let tcp = &*((start + tcp_start) as *const TcpHdr);
Ok((u16::from_be_bytes(tcp.source), u16::from_be_bytes(tcp.dest)))
unsafe fn parse_tcp(start: usize, end: usize, tcp_start: usize, tcp_end: usize) -> Result<(u16, u16, u8, usize), ()> {
if start + tcp_end > end {
return Err(());
}
let tcp = &*((start + tcp_start) as *const TcpHdr);
let data_offset = (*((start + tcp_start + 12) as *const u8) >> 4) as usize;
if data_offset < 5 || data_offset > 15 {
return Err(());
}
let header_len = data_offset * 4;
if start + tcp_start + header_len > end {
return Err(());
}
let flags = *((start + tcp_start + 13) as *const u8);
Ok((
u16::from_be_bytes(tcp.source),
u16::from_be_bytes(tcp.dest),
flags,
header_len,
))
}
#[inline(always)]
unsafe fn parse_udp_port(start: usize, end: usize, udp_start: usize, udp_end: usize) -> Result<(u16, u16), ()> {
unsafe {
if start + udp_end > end {
return Err(());
}
let udp = &*((start + udp_start) as *const UdpHdr);
Ok((udp.src_port(), udp.dst_port()))
unsafe fn parse_udp(start: usize, end: usize, udp_start: usize, udp_end: usize) -> Result<(u16, u16, u8, usize), ()> {
if start + udp_end > end {
return Err(());
}
let udp = &*((start + udp_start) as *const UdpHdr);
Ok((udp.src_port(), udp.dst_port(), 0u8, 8usize))
}

View File

@ -1,197 +0,0 @@
use network_types::ip::IpProto;
use crate::model::ip_address::{AddrPortV4, AddrPortV6};
#[repr(C, align(8))]
#[derive(Clone)]
pub enum Event {
IPv4(IPv4Event),
IPv6(IPv6Event),
}
#[repr(C, align(8))]
#[derive(Clone)]
pub enum RawIp {
V4(u32),
V6(u128),
}
impl Event {
pub fn timestamp_us(&self) -> u64 {
match self {
Event::IPv4(e) => e.timestamp_us,
Event::IPv6(e) => e.timestamp_us,
}
}
pub fn packet_length(&self) -> u32 {
match self {
Event::IPv4(e) => e.packet_length,
Event::IPv6(e) => e.packet_length,
}
}
pub fn header_length(&self) -> u16 {
match self {
Event::IPv4(e) => e.header_length,
Event::IPv6(e) => e.header_length,
}
}
pub fn payload_length(&self) -> u32 {
match self {
Event::IPv4(e) => e.payload_length,
Event::IPv6(e) => e.payload_length,
}
}
pub fn tcp_flags(&self) -> &TcpFlags {
match self {
Event::IPv4(e) => &e.tcp_flags,
Event::IPv6(e) => &e.tcp_flags,
}
}
pub fn tcp_window_size(&self) -> u16 {
match self {
Event::IPv4(e) => e.tcp_window_size,
Event::IPv6(e) => e.tcp_window_size,
}
}
pub fn is_forward(&self) -> bool {
match self {
Event::IPv4(e) => e.is_forward,
Event::IPv6(e) => e.is_forward,
}
}
pub fn protocol(&self) -> &IpProto {
match self {
Event::IPv4(e) => &e.protocol,
Event::IPv6(e) => &e.protocol,
}
}
pub fn src_ip(&self) -> RawIp {
match self {
Event::IPv4(e) => RawIp::V4(e.src_ip),
Event::IPv6(e) => RawIp::V6(e.src_ip),
}
}
pub fn dst_ip(&self) -> RawIp {
match self {
Event::IPv4(e) => RawIp::V4(e.dst_ip),
Event::IPv6(e) => RawIp::V6(e.dst_ip),
}
}
pub fn src_port(&self) -> u16 {
match self {
Event::IPv4(e) => e.src_port,
Event::IPv6(e) => e.src_port,
}
}
pub fn dst_port(&self) -> u16 {
match self {
Event::IPv4(e) => e.dst_port,
Event::IPv6(e) => e.dst_port,
}
}
pub fn set_is_forward(&mut self, value: bool) {
match self {
Event::IPv4(e) => e.is_forward = value,
Event::IPv6(e) => e.is_forward = value,
}
}
}
#[repr(C, align(8))]
#[derive(Debug, Clone)]
pub struct IPv4Event {
pub protocol: IpProto,
pub src_ip: u32,
pub dst_ip: u32,
pub src_port: u16,
pub dst_port: u16,
pub packet_length: u32,
pub payload_length: u32,
pub header_length: u16,
pub timestamp_us: u64,
pub tcp_flags: TcpFlags,
pub tcp_window_size: u16,
pub is_forward: bool,
}
impl IPv4Event {
#[inline(always)]
pub fn source_addr(&self) -> AddrPortV4 {
AddrPortV4::new(self.src_ip, self.src_port)
}
#[inline(always)]
pub fn destination_addr(&self) -> AddrPortV4 {
AddrPortV4::new(self.dst_ip, self.dst_port)
}
}
#[repr(C, align(8))]
#[derive(Clone)]
pub struct IPv6Event {
pub protocol: IpProto,
pub src_ip: u128,
pub dst_ip: u128,
pub src_port: u16,
pub dst_port: u16,
pub packet_length: u32,
pub payload_length: u32,
pub header_length: u16,
pub timestamp_us: u64,
pub tcp_flags: TcpFlags,
pub tcp_window_size: u16,
pub is_forward: bool,
}
impl IPv6Event {
#[inline(always)]
pub fn source_addr(&self) -> AddrPortV6 {
AddrPortV6::new(self.src_ip, self.src_port)
}
#[inline(always)]
pub fn destination_addr(&self) -> AddrPortV6 {
AddrPortV6::new(self.dst_ip, self.dst_port)
}
}
#[repr(C, align(8))]
#[derive(Debug, Clone, Default)]
pub struct TcpFlags {
pub fin: bool,
pub syn: bool,
pub rst: bool,
pub psh: bool,
pub ack: bool,
pub urg: bool,
pub ece: bool,
pub cwr: bool,
}
impl TcpFlags {
pub fn from_byte(flags: u8) -> Self {
Self {
fin: (flags & 0x01) != 0,
syn: (flags & 0x02) != 0,
rst: (flags & 0x04) != 0,
psh: (flags & 0x08) != 0,
ack: (flags & 0x10) != 0,
urg: (flags & 0x20) != 0,
ece: (flags & 0x40) != 0,
cwr: (flags & 0x80) != 0,
}
}
}

View File

@ -1,7 +1,8 @@
pub mod event;
pub mod flow_stats;
pub mod http_method;
pub mod ip_address;
pub mod packet;
pub mod parsed_packet;
pub mod placeholder;
pub mod port_rule;
pub mod pseudo_header;
pub mod rate_limit;

View File

@ -1,5 +0,0 @@
use crate::define::other::STANDARD_MTU;
use crate::model::event::Event;
#[repr(transparent)]
pub struct Packet(pub [u8; size_of::<Event>() + STANDARD_MTU]);

View File

@ -0,0 +1,61 @@
use network_types::ip::IpProto;
use crate::model::ip_address::{AddrPortV4, AddrPortV6};
#[repr(C, align(8))]
pub struct ParsedPacket {
pub timestamp_ns: u64,
pub src_ip: [u8; 16],
pub dst_ip: [u8; 16],
pub packet_length: u32,
pub payload_length: u32,
pub src_port: u16,
pub dst_port: u16,
pub ip_version: u8,
pub protocol: IpProto,
pub tcp_flags: u8,
/// Padding for 8-byte alignment (required by eBPF PerCpuArray)
pub _pad: u8,
}
impl ParsedPacket {
#[inline(always)]
pub fn src_ip_v4(&self) -> u32 {
u32::from_ne_bytes([self.src_ip[0], self.src_ip[1], self.src_ip[2], self.src_ip[3]])
}
#[inline(always)]
pub fn dst_ip_v4(&self) -> u32 {
u32::from_ne_bytes([self.dst_ip[0], self.dst_ip[1], self.dst_ip[2], self.dst_ip[3]])
}
#[inline(always)]
pub fn src_ip_v6(&self) -> u128 {
u128::from_ne_bytes(self.src_ip)
}
#[inline(always)]
pub fn dst_ip_v6(&self) -> u128 {
u128::from_ne_bytes(self.dst_ip)
}
#[inline(always)]
pub fn src_addr_v4(&self) -> AddrPortV4 {
AddrPortV4::new(self.src_ip_v4(), self.src_port)
}
#[inline(always)]
pub fn dst_addr_v4(&self) -> AddrPortV4 {
AddrPortV4::new(self.dst_ip_v4(), self.dst_port)
}
#[inline(always)]
pub fn src_addr_v6(&self) -> AddrPortV6 {
AddrPortV6::new(self.src_ip_v6(), self.src_port)
}
#[inline(always)]
pub fn dst_addr_v6(&self) -> AddrPortV6 {
AddrPortV6::new(self.dst_ip_v6(), self.dst_port)
}
}

View File

@ -0,0 +1,110 @@
#[cfg(feature = "user")]
use aya::Pod;
#[cfg(feature = "user")]
use std::vec::Vec;
use crate::define::setting::MAX_RULES_PORT;
use crate::model::ip_address::Port;
pub const PORT_RULE_MATCH_ALL: u8 = 1;
#[repr(C)]
#[derive(Clone, Copy)]
pub struct PortRule {
pub match_all: u8,
pub count: u8,
pub _pad: [u8; 2],
pub ports: [Port; MAX_RULES_PORT],
}
impl PortRule {
pub fn new_empty() -> Self {
Self {
match_all: 0,
count: 0,
_pad: [0; 2],
ports: [0; MAX_RULES_PORT],
}
}
pub fn new_match_all() -> Self {
Self {
match_all: PORT_RULE_MATCH_ALL,
count: 0,
_pad: [0; 2],
ports: [0; MAX_RULES_PORT],
}
}
pub fn is_match_all(&self) -> bool {
self.match_all == PORT_RULE_MATCH_ALL
}
pub fn contains(&self, port: Port) -> bool {
if self.is_match_all() {
return true;
}
for i in 0..(self.count as usize) {
if i >= MAX_RULES_PORT {
break;
}
if self.ports[i] == port {
return true;
}
}
false
}
#[cfg(feature = "user")]
pub fn add_port(&mut self, port: Port) -> bool {
if self.is_match_all() {
return true;
}
for i in 0..(self.count as usize) {
if i >= MAX_RULES_PORT {
return false;
}
if self.ports[i] == port {
return true; // already exists
}
}
if (self.count as usize) >= MAX_RULES_PORT {
return false; // full
}
self.ports[self.count as usize] = port;
self.count += 1;
true
}
#[cfg(feature = "user")]
pub fn remove_port(&mut self, port: Port) -> bool {
for i in 0..(self.count as usize) {
if i >= MAX_RULES_PORT {
break;
}
if self.ports[i] == port {
// shift remaining
for j in i..(self.count as usize - 1) {
self.ports[j] = self.ports[j + 1];
}
self.count -= 1;
self.ports[self.count as usize] = 0;
return true;
}
}
false
}
#[cfg(feature = "user")]
pub fn to_port_vec(&self) -> Vec<Port> {
self.ports[..self.count as usize].to_vec()
}
#[cfg(feature = "user")]
pub fn is_empty(&self) -> bool {
!self.is_match_all() && self.count == 0
}
}
#[cfg(feature = "user")]
unsafe impl Pod for PortRule {}

View File

@ -0,0 +1,19 @@
#[cfg(feature = "user")]
use aya::Pod;
pub const MAX_TRACKED_IPS: u32 = 65536;
pub const DEFAULT_WINDOW_NS: u64 = 1_000_000_000;
pub const DEFAULT_PACKET_RATE: u64 = 10000;
pub const DEFAULT_SYN_RATE: u64 = 100;
pub const DEFAULT_UDP_RATE: u64 = 5000;
pub const DEFAULT_DNS_RATE: u64 = 200;
#[repr(C)]
#[derive(Clone, Copy)]
pub struct RateState {
pub count: u64,
pub window_start: u64,
}
#[cfg(feature = "user")]
unsafe impl Pod for RateState {}

View File

@ -1,26 +1,34 @@
[Config]
ingress_ifname = "enp4s0f1" # Ingress NIC Name
egress_ifname = "enp4s0f0" # Egress NIC Name
geoip_db_name = "GeoLite2-City.mmdb"
[Http]
http_server_bind_port = 8080
[Network]
ingress_ifname = "ng-ext"
egress_ifname = "ng-int"
combined_queue_count = 16
channel_size = 4096
fill_queue_size = 4096
comp_queue_size = 4096
tx_queue_size = 4096
rx_queue_size = 4096
frame_size = 4096
frame_count = 4096
refresh_interval = 5
[Inference]
deep_autoencoder_name = "deep_autoencoder.onnx"
classifier_name = "classifier.onnx"
models_config_name = "inference_config.json"
combined_queue_count = 8 # NIC Combined Queue Count (ethtool -l <NIC>)
channel_size = 4096
fill_queue_size = 4096 # Umem Used (Should not modify)
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)
http_server_bind_port = 8080 # Http Server Listen Port
refresh_interval = 5 # Statistics Refresh Time
max_concurrent_flows = 10000 # max_flows: track up to 10000 concurrent flows
min_packets_for_inference = 5 # min_packets: minimum 10 packets per flow for inference
inference_interval_secs = 5 # interval_secs: run inference every 5 seconds
max_concurrent_flows = 10000
min_packets_for_inference = 5
inference_interval_secs = 5
aggregator_window_secs = 30
inference_batch_size = 200
traffic_logging_mode = true
traffic_log_csv_path = "traffic_log.csv"
traffic_logging_mode = true # 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
[Misc]
geoip_db_name = "GeoLite2-City.mmdb"
[Pipeline]
ingress = ["access_control", "rate_limit", "service"]
egress = []

View File

@ -1 +0,0 @@
pub mod statistics;

View File

@ -1,93 +0,0 @@
use aya_ebpf::macros::map;
use aya_ebpf::maps::LruHashMap;
use common::define::setting::MAX_STATS;
use common::model::event::{IPv4Event, IPv6Event};
use common::model::flow_stats::FlowStats;
use common::model::ip_address::{AddrPortV4, AddrPortV6};
#[map]
static IPV4_EGRESS_SRC_1MIN: LruHashMap<AddrPortV4, FlowStats> = LruHashMap::with_max_entries(MAX_STATS as u32, 0);
#[map]
static IPV4_EGRESS_SRC_10MIN: LruHashMap<AddrPortV4, FlowStats> = LruHashMap::with_max_entries(MAX_STATS as u32, 0);
#[map]
static IPV4_EGRESS_SRC_1HOUR: LruHashMap<AddrPortV4, FlowStats> = LruHashMap::with_max_entries(MAX_STATS as u32, 0);
#[map]
static IPV6_EGRESS_SRC_1MIN: LruHashMap<AddrPortV6, FlowStats> = LruHashMap::with_max_entries(MAX_STATS as u32, 0);
#[map]
static IPV6_EGRESS_SRC_10MIN: LruHashMap<AddrPortV6, FlowStats> = LruHashMap::with_max_entries(MAX_STATS as u32, 0);
#[map]
static IPV6_EGRESS_SRC_1HOUR: LruHashMap<AddrPortV6, FlowStats> = LruHashMap::with_max_entries(MAX_STATS as u32, 0);
#[map]
static IPV4_EGRESS_DST_1MIN: LruHashMap<AddrPortV4, FlowStats> = LruHashMap::with_max_entries(MAX_STATS as u32, 0);
#[map]
static IPV4_EGRESS_DST_10MIN: LruHashMap<AddrPortV4, FlowStats> = LruHashMap::with_max_entries(MAX_STATS as u32, 0);
#[map]
static IPV4_EGRESS_DST_1HOUR: LruHashMap<AddrPortV4, FlowStats> = LruHashMap::with_max_entries(MAX_STATS as u32, 0);
#[map]
static IPV6_EGRESS_DST_1MIN: LruHashMap<AddrPortV6, FlowStats> = LruHashMap::with_max_entries(MAX_STATS as u32, 0);
#[map]
static IPV6_EGRESS_DST_10MIN: LruHashMap<AddrPortV6, FlowStats> = LruHashMap::with_max_entries(MAX_STATS as u32, 0);
#[map]
static IPV6_EGRESS_DST_1HOUR: LruHashMap<AddrPortV6, FlowStats> = LruHashMap::with_max_entries(MAX_STATS as u32, 0);
pub fn ipv4_update_stats(event: &IPv4Event) {
unsafe {
let source = event.source_addr();
let destination = event.destination_addr();
ipv4_update_flow_stats(&IPV4_EGRESS_SRC_1MIN, &source, event);
ipv4_update_flow_stats(&IPV4_EGRESS_SRC_10MIN, &source, event);
ipv4_update_flow_stats(&IPV4_EGRESS_SRC_1HOUR, &source, event);
ipv4_update_flow_stats(&IPV4_EGRESS_DST_1MIN, &destination, event);
ipv4_update_flow_stats(&IPV4_EGRESS_DST_10MIN, &destination, event);
ipv4_update_flow_stats(&IPV4_EGRESS_DST_1HOUR, &destination, event);
}
}
pub fn ipv6_update_stats(event: &IPv6Event) {
unsafe {
let source = event.source_addr();
let destination = event.destination_addr();
ipv6_update_flow_status(&IPV6_EGRESS_SRC_1MIN, &source, event);
ipv6_update_flow_status(&IPV6_EGRESS_SRC_10MIN, &source, event);
ipv6_update_flow_status(&IPV6_EGRESS_SRC_1HOUR, &source, event);
ipv6_update_flow_status(&IPV6_EGRESS_DST_1MIN, &destination, event);
ipv6_update_flow_status(&IPV6_EGRESS_DST_10MIN, &destination, event);
ipv6_update_flow_status(&IPV6_EGRESS_DST_1HOUR, &destination, event);
}
}
#[inline(always)]
unsafe fn ipv4_update_flow_stats(map: &LruHashMap<AddrPortV4, FlowStats>, key: &AddrPortV4, event: &IPv4Event) {
unsafe {
if let Some(status) = map.get_ptr_mut(key) {
(*status).bytes += event.packet_length as u64;
(*status).packets += 1;
(*status).last_seen = event.timestamp_us;
} else {
let new_stats = FlowStats {
bytes: event.packet_length as u64,
packets: 1,
last_seen: event.timestamp_us,
};
let _ = map.insert(key, &new_stats, 0);
}
}
}
#[inline(always)]
unsafe fn ipv6_update_flow_status(map: &LruHashMap<AddrPortV6, FlowStats>, key: &AddrPortV6, event: &IPv6Event) {
unsafe {
if let Some(status) = map.get_ptr_mut(key) {
(*status).bytes += event.packet_length as u64;
(*status).packets += 1;
(*status).last_seen = event.timestamp_us;
} else {
let new_stats = FlowStats {
bytes: event.packet_length as u64,
packets: 1,
last_seen: event.timestamp_us,
};
let _ = map.insert(key, &new_stats, 0);
}
}
}

View File

@ -1,70 +1,18 @@
#![no_std]
#![no_main]
mod action;
use action::statistics;
use aya_ebpf::bindings::xdp_action;
use aya_ebpf::macros::{map, xdp};
use aya_ebpf::maps::{PerCpuArray, ProgramArray, XskMap};
use aya_ebpf::maps::XskMap;
use aya_ebpf::programs::XdpContext;
#[allow(unused_imports)]
use aya_log_ebpf::info;
use common::define::program_array::egress::*;
use common::{ebpf::parsing, model::event::Event};
#[map]
static PROGRAM_ARRAY: ProgramArray = ProgramArray::with_max_entries(8, 0);
#[map]
static PARSED_PACKET: PerCpuArray<Event> = PerCpuArray::with_max_entries(1, 0);
#[map]
static EGRESS_XSKS_MAP: XskMap = XskMap::pinned(64, 0);
#[xdp]
pub fn net_guardia(ctx: XdpContext) -> u32 {
unsafe {
let _ = packet_intake(ctx);
xdp_action::XDP_PASS
}
}
unsafe fn packet_intake(ctx: XdpContext) -> Result<u32, ()> {
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, STATISTICS);
Err(())
}
}
#[xdp]
pub fn statistics(ctx: XdpContext) -> u32 {
unsafe {
let _ = try_statistics(&ctx);
xdp_action::XDP_PASS
}
}
unsafe fn try_statistics(ctx: &XdpContext) -> Result<u32, ()> {
unsafe {
let ptr = PARSED_PACKET.get_ptr(0).ok_or(())?;
let parsed_packet = &*ptr;
match parsed_packet {
Event::IPv4(event) => {
statistics::ipv4_update_stats(event);
}
Event::IPv6(event) => {
statistics::ipv6_update_stats(event);
}
}
let _ = PROGRAM_ARRAY.tail_call(ctx, TRANSMISSION);
Ok(xdp_action::XDP_PASS)
}
}
#[xdp]
pub fn transmission(ctx: XdpContext) -> u32 {
let queue_id = unsafe { (*ctx.ctx).rx_queue_index };
match EGRESS_XSKS_MAP.redirect(queue_id, 0) {
Ok(action) => action,

View File

@ -1,35 +1,38 @@
use aya_ebpf::macros::map;
use aya_ebpf::maps::HashMap;
use common::define::setting::{MAX_RULES, MAX_RULES_PORT};
use common::model::event::{IPv4Event, IPv6Event};
use common::model::ip_address::{IPv4, IPv6, Port};
use common::define::setting::MAX_RULES;
use common::model::parsed_packet::ParsedPacket;
use common::model::ip_address::{IPv4, IPv6};
use common::model::port_rule::PortRule;
#[map]
static IPV4_SRC_WHITELIST: HashMap<IPv4, [Port; MAX_RULES_PORT]> = HashMap::with_max_entries(MAX_RULES as u32, 0);
static IPV4_SRC_WHITELIST: HashMap<IPv4, PortRule> = HashMap::with_max_entries(MAX_RULES as u32, 0);
#[map]
static IPV6_SRC_WHITELIST: HashMap<IPv6, [Port; MAX_RULES_PORT]> = HashMap::with_max_entries(MAX_RULES as u32, 0);
static IPV6_SRC_WHITELIST: HashMap<IPv6, PortRule> = HashMap::with_max_entries(MAX_RULES as u32, 0);
#[map]
static IPV4_DST_WHITELIST: HashMap<IPv4, [Port; MAX_RULES_PORT]> = HashMap::with_max_entries(MAX_RULES as u32, 0);
static IPV4_DST_WHITELIST: HashMap<IPv4, PortRule> = HashMap::with_max_entries(MAX_RULES as u32, 0);
#[map]
static IPV6_DST_WHITELIST: HashMap<IPv6, [Port; MAX_RULES_PORT]> = HashMap::with_max_entries(MAX_RULES as u32, 0);
static IPV6_DST_WHITELIST: HashMap<IPv6, PortRule> = HashMap::with_max_entries(MAX_RULES as u32, 0);
#[map]
static IPV4_SRC_BLACKLIST: HashMap<IPv4, [Port; MAX_RULES_PORT]> = HashMap::with_max_entries(MAX_RULES as u32, 0);
static IPV4_SRC_BLACKLIST: HashMap<IPv4, PortRule> = HashMap::with_max_entries(MAX_RULES as u32, 0);
#[map]
static IPV6_SRC_BLACKLIST: HashMap<IPv6, [Port; MAX_RULES_PORT]> = HashMap::with_max_entries(MAX_RULES as u32, 0);
static IPV6_SRC_BLACKLIST: HashMap<IPv6, PortRule> = HashMap::with_max_entries(MAX_RULES as u32, 0);
#[map]
static IPV4_DST_BLACKLIST: HashMap<IPv4, [Port; MAX_RULES_PORT]> = HashMap::with_max_entries(MAX_RULES as u32, 0);
static IPV4_DST_BLACKLIST: HashMap<IPv4, PortRule> = HashMap::with_max_entries(MAX_RULES as u32, 0);
#[map]
static IPV6_DST_BLACKLIST: HashMap<IPv6, [Port; MAX_RULES_PORT]> = HashMap::with_max_entries(MAX_RULES as u32, 0);
static IPV6_DST_BLACKLIST: HashMap<IPv6, PortRule> = HashMap::with_max_entries(MAX_RULES as u32, 0);
pub fn ipv4_is_whitelisted(event: &IPv4Event) -> bool {
pub fn ipv4_is_whitelisted(pkt: &ParsedPacket) -> bool {
let src_ip = pkt.src_ip_v4();
let dst_ip = pkt.dst_ip_v4();
unsafe {
if let Some(ports) = IPV4_SRC_WHITELIST.get(&event.src_ip) {
if is_port_exist(ports, event.src_port) {
if let Some(rule) = IPV4_SRC_WHITELIST.get(&src_ip) {
if rule.contains(pkt.src_port) {
return true;
}
}
if let Some(ports) = IPV4_DST_WHITELIST.get(&event.dst_ip) {
if is_port_exist(ports, event.dst_port) {
if let Some(rule) = IPV4_DST_WHITELIST.get(&dst_ip) {
if rule.contains(pkt.dst_port) {
return true;
}
}
@ -37,15 +40,17 @@ pub fn ipv4_is_whitelisted(event: &IPv4Event) -> bool {
false
}
pub fn ipv6_is_whitelisted(event: &IPv6Event) -> bool {
pub fn ipv6_is_whitelisted(pkt: &ParsedPacket) -> bool {
let src_ip = pkt.src_ip_v6();
let dst_ip = pkt.dst_ip_v6();
unsafe {
if let Some(ports) = IPV6_SRC_WHITELIST.get(&event.src_ip) {
if is_port_exist(ports, event.src_port) {
if let Some(rule) = IPV6_SRC_WHITELIST.get(&src_ip) {
if rule.contains(pkt.src_port) {
return true;
}
}
if let Some(ports) = IPV6_DST_WHITELIST.get(&event.dst_ip) {
if is_port_exist(ports, event.dst_port) {
if let Some(rule) = IPV6_DST_WHITELIST.get(&dst_ip) {
if rule.contains(pkt.dst_port) {
return true;
}
}
@ -53,15 +58,17 @@ pub fn ipv6_is_whitelisted(event: &IPv6Event) -> bool {
false
}
pub fn ipv4_is_blacklisted(event: &IPv4Event) -> bool {
pub fn ipv4_is_blacklisted(pkt: &ParsedPacket) -> bool {
let src_ip = pkt.src_ip_v4();
let dst_ip = pkt.dst_ip_v4();
unsafe {
if let Some(ports) = IPV4_SRC_BLACKLIST.get(&event.src_ip) {
if is_port_exist(ports, event.src_port) {
if let Some(rule) = IPV4_SRC_BLACKLIST.get(&src_ip) {
if rule.contains(pkt.src_port) {
return true;
}
}
if let Some(ports) = IPV4_DST_BLACKLIST.get(&event.dst_ip) {
if is_port_exist(ports, event.dst_port) {
if let Some(rule) = IPV4_DST_BLACKLIST.get(&dst_ip) {
if rule.contains(pkt.dst_port) {
return true;
}
}
@ -69,34 +76,20 @@ pub fn ipv4_is_blacklisted(event: &IPv4Event) -> bool {
false
}
pub fn ipv6_is_blacklisted(event: &IPv6Event) -> bool {
pub fn ipv6_is_blacklisted(pkt: &ParsedPacket) -> bool {
let src_ip = pkt.src_ip_v6();
let dst_ip = pkt.dst_ip_v6();
unsafe {
if let Some(ports) = IPV6_SRC_BLACKLIST.get(&event.src_ip) {
if is_port_exist(ports, event.src_port) {
if let Some(rule) = IPV6_SRC_BLACKLIST.get(&src_ip) {
if rule.contains(pkt.src_port) {
return true;
}
}
if let Some(ports) = IPV6_DST_BLACKLIST.get(&event.dst_ip) {
if is_port_exist(ports, event.dst_port) {
if let Some(rule) = IPV6_DST_BLACKLIST.get(&dst_ip) {
if rule.contains(pkt.dst_port) {
return true;
}
}
}
false
}
#[inline(always)]
fn is_port_exist(ports: &[Port; MAX_RULES_PORT], target_port: Port) -> bool {
if ports.get(0) == Some(&0) {
return true;
}
for &port in ports.iter() {
if port == 0 {
break;
}
if port == target_port {
return true;
}
}
false
}

View File

@ -1,3 +1,3 @@
pub mod access_control;
pub mod service;
pub mod statistics;
pub mod rate_limit;
pub mod protocol_filter;

View File

@ -0,0 +1,137 @@
use aya_ebpf::macros::map;
use aya_ebpf::maps::{Array, HashMap};
use common::define::setting::MAX_RULES;
use common::define::tcp_flags::*;
use common::model::http_method::HttpMethodBitmap;
use common::model::ip_address::*;
use common::model::parsed_packet::ParsedPacket;
use common::model::placeholder::PlaceHolder;
use network_types::ip::IpProto;
#[map]
static IPV4_HTTP_SERVICE: HashMap<AddrPortV4, HttpMethodBitmap> = HashMap::with_max_entries(MAX_RULES as u32, 0);
#[map]
static IPV6_HTTP_SERVICE: HashMap<AddrPortV6, HttpMethodBitmap> = HashMap::with_max_entries(MAX_RULES as u32, 0);
#[map]
static SSH_WHITE_LIST_ENABLE: Array<PlaceHolder> = Array::with_max_entries(1, 0);
#[map]
static IPV4_SSH_SERVICE: HashMap<AddrPortV4, PlaceHolder> = HashMap::with_max_entries(MAX_RULES as u32, 0);
#[map]
static IPV6_SSH_SERVICE: HashMap<AddrPortV6, PlaceHolder> = HashMap::with_max_entries(MAX_RULES as u32, 0);
#[map]
static IPV4_SSH_WHITE_LIST: HashMap<IPv4, PlaceHolder> = HashMap::with_max_entries(MAX_RULES as u32, 0);
#[map]
static IPV6_SSH_WHITE_LIST: HashMap<IPv6, PlaceHolder> = HashMap::with_max_entries(MAX_RULES as u32, 0);
#[map]
static IPV4_SSH_BLACK_LIST: HashMap<IPv4, PlaceHolder> = HashMap::with_max_entries(MAX_RULES as u32, 0);
#[map]
static IPV6_SSH_BLACK_LIST: HashMap<IPv6, PlaceHolder> = HashMap::with_max_entries(MAX_RULES as u32, 0);
pub fn ipv4_service_rule_violation(start: usize, end: usize, pkt: &ParsedPacket) -> bool {
let source = pkt.src_addr_v4();
let destination = pkt.dst_addr_v4();
http_service_violation(start, end, pkt, &IPV4_HTTP_SERVICE, &destination)
|| ipv4_ssh_service_violation(&source, &destination)
}
pub fn ipv6_service_rule_violation(start: usize, end: usize, pkt: &ParsedPacket) -> bool {
let source = pkt.src_addr_v6();
let destination = pkt.dst_addr_v6();
http_service_violation(start, end, pkt, &IPV6_HTTP_SERVICE, &destination)
|| ipv6_ssh_service_violation(&source, &destination)
}
#[inline(always)]
fn http_service_violation<K>(
start: usize,
end: usize,
pkt: &ParsedPacket,
map: &HashMap<K, HttpMethodBitmap>,
destination: &K,
) -> bool {
match map.get_ptr_mut(destination) {
Some(allow_method) => {
if !matches!(pkt.protocol, IpProto::Tcp) {
return false;
}
if pkt.tcp_flags & (TCP_SYN | TCP_RST | TCP_FIN) != 0 {
return false;
}
if pkt.tcp_flags & (TCP_PSH | TCP_ACK) != (TCP_PSH | TCP_ACK) {
return false;
}
// Bounds check before reading IHL — verifier needs to see this
if start + 15 > end {
return false;
}
let l4_offset = match pkt.ip_version {
4 => 14 + ((unsafe { *((start + 14) as *const u8) } & 0x0F) as usize) * 4,
6 => 14 + 40,
_ => return false,
};
if start + l4_offset + 13 > end {
return false;
}
let doff = (unsafe { *((start + l4_offset + 12) as *const u8) } >> 4) as usize;
if doff < 5 || doff > 15 {
return false;
}
let payload_offset = l4_offset + doff * 4;
match get_http_request_method(start, end, payload_offset) {
Some(http_method) => unsafe { *allow_method & http_method == 0 },
None => false,
}
}
None => false,
}
}
#[inline(always)]
fn get_http_request_method(start: usize, end: usize, offset: usize) -> Option<HttpMethodBitmap> {
if start + offset + 8 > end {
return None;
}
let data = unsafe { core::slice::from_raw_parts((start + offset) as *const u8, 8) };
match &data[..4] {
b"GET " => Some(1 << 0),
b"POST" if &data[4..5] == b" " => Some(1 << 1),
b"PUT " => Some(1 << 2),
b"DELE" if &data[4..7] == b"TE " => Some(1 << 3),
b"HEAD" if &data[4..5] == b" " => Some(1 << 4),
b"OPTI" if &data[4..8] == b"ONS " => Some(1 << 5),
b"PATC" if &data[4..6] == b"H " => Some(1 << 6),
b"TRAC" if &data[4..6] == b"E " => Some(1 << 7),
b"CONN" if &data[4..8] == b"ECT " => Some(1 << 8),
_ => None,
}
}
#[inline(always)]
fn ipv4_ssh_service_violation(source: &AddrPortV4, destination: &AddrPortV4) -> bool {
unsafe {
if IPV4_SSH_SERVICE.get(destination).is_some() {
if SSH_WHITE_LIST_ENABLE.get(0).is_some() {
IPV4_SSH_WHITE_LIST.get(&source.ip()).is_none()
} else {
IPV4_SSH_BLACK_LIST.get(&source.ip()).is_some()
}
} else {
false
}
}
}
#[inline(always)]
fn ipv6_ssh_service_violation(source: &AddrPortV6, destination: &AddrPortV6) -> bool {
unsafe {
if IPV6_SSH_SERVICE.get(destination).is_some() {
if SSH_WHITE_LIST_ENABLE.get(0).is_some() {
IPV6_SSH_WHITE_LIST.get(&source.ip()).is_none()
} else {
IPV6_SSH_BLACK_LIST.get(&source.ip()).is_some()
}
} else {
false
}
}
}

View File

@ -0,0 +1,154 @@
use aya_ebpf::helpers::bpf_ktime_get_ns;
use aya_ebpf::macros::map;
use aya_ebpf::maps::{Array, LruHashMap};
use common::define::tcp_flags::*;
use common::model::ip_address::{IPv4, IPv6};
use common::model::parsed_packet::ParsedPacket;
use common::model::rate_limit::*;
use network_types::ip::IpProto;
const CFG_PACKET_RATE: u32 = 0;
const CFG_SYN_RATE: u32 = 1;
const CFG_UDP_RATE: u32 = 2;
const CFG_DNS_RATE: u32 = 3;
const CFG_WINDOW_NS: u32 = 4;
#[map]
static RATE_LIMIT_CONFIG: Array<u64> = Array::with_max_entries(5, 0);
#[map]
static IPV4_PACKET_RATE_MAP: LruHashMap<IPv4, RateState> = LruHashMap::with_max_entries(MAX_TRACKED_IPS, 0);
#[map]
static IPV6_PACKET_RATE_MAP: LruHashMap<IPv6, RateState> = LruHashMap::with_max_entries(MAX_TRACKED_IPS, 0);
#[map]
static IPV4_SYN_RATE_MAP: LruHashMap<IPv4, RateState> = LruHashMap::with_max_entries(MAX_TRACKED_IPS, 0);
#[map]
static IPV6_SYN_RATE_MAP: LruHashMap<IPv6, RateState> = LruHashMap::with_max_entries(MAX_TRACKED_IPS, 0);
#[map]
static IPV4_UDP_RATE_MAP: LruHashMap<IPv4, RateState> = LruHashMap::with_max_entries(MAX_TRACKED_IPS, 0);
#[map]
static IPV6_UDP_RATE_MAP: LruHashMap<IPv6, RateState> = LruHashMap::with_max_entries(MAX_TRACKED_IPS, 0);
#[map]
static IPV4_DNS_RATE_MAP: LruHashMap<IPv4, RateState> = LruHashMap::with_max_entries(MAX_TRACKED_IPS, 0);
#[map]
static IPV6_DNS_RATE_MAP: LruHashMap<IPv6, RateState> = LruHashMap::with_max_entries(MAX_TRACKED_IPS, 0);
pub fn should_drop(pkt: &ParsedPacket) -> bool {
match pkt.ip_version {
4 => ipv4_should_drop(pkt),
6 => ipv6_should_drop(pkt),
_ => false,
}
}
#[inline(always)]
fn get_config(index: u32, default: u64) -> u64 {
unsafe {
RATE_LIMIT_CONFIG
.get(index)
.copied()
.filter(|&v| v > 0)
.unwrap_or(default)
}
}
/// Returns true if the packet is a TCP SYN-only (no ACK) packet.
/// For non-TCP packets (e.g. UDP), tcp_flags is 0, so this safely returns false.
#[inline(always)]
fn is_syn_only(pkt: &ParsedPacket) -> bool {
matches!(pkt.protocol, IpProto::Tcp)
&& (pkt.tcp_flags & TCP_SYN != 0)
&& (pkt.tcp_flags & TCP_ACK == 0)
}
#[inline(always)]
fn check_rate<K>(
map: &LruHashMap<K, RateState>,
key: &K,
now: u64,
window: u64,
limit: u64,
) -> bool {
unsafe {
if let Some(state) = map.get_ptr_mut(key) {
if now - (*state).window_start >= window {
(*state).count = 1;
(*state).window_start = now;
} else {
(*state).count += 1;
if (*state).count > limit {
return true;
}
}
} else {
let new_state = RateState {
count: 1,
window_start: now,
};
let _ = map.insert(key, &new_state, 0);
}
}
false
}
#[inline(always)]
fn ipv4_should_drop(pkt: &ParsedPacket) -> bool {
let now = unsafe { bpf_ktime_get_ns() };
let window = get_config(CFG_WINDOW_NS, DEFAULT_WINDOW_NS);
let src_ip = pkt.src_ip_v4();
if check_rate(&IPV4_PACKET_RATE_MAP, &src_ip, now, window, get_config(CFG_PACKET_RATE, DEFAULT_PACKET_RATE)) {
return true;
}
if is_syn_only(pkt) {
if check_rate(&IPV4_SYN_RATE_MAP, &src_ip, now, window, get_config(CFG_SYN_RATE, DEFAULT_SYN_RATE)) {
return true;
}
}
if matches!(pkt.protocol, IpProto::Udp) {
if check_rate(&IPV4_UDP_RATE_MAP, &src_ip, now, window, get_config(CFG_UDP_RATE, DEFAULT_UDP_RATE)) {
return true;
}
}
if pkt.dst_port == 53 {
if check_rate(&IPV4_DNS_RATE_MAP, &src_ip, now, window, get_config(CFG_DNS_RATE, DEFAULT_DNS_RATE)) {
return true;
}
}
false
}
#[inline(always)]
fn ipv6_should_drop(pkt: &ParsedPacket) -> bool {
let now = unsafe { bpf_ktime_get_ns() };
let window = get_config(CFG_WINDOW_NS, DEFAULT_WINDOW_NS);
let src_ip = pkt.src_ip_v6();
if check_rate(&IPV6_PACKET_RATE_MAP, &src_ip, now, window, get_config(CFG_PACKET_RATE, DEFAULT_PACKET_RATE)) {
return true;
}
if is_syn_only(pkt) {
if check_rate(&IPV6_SYN_RATE_MAP, &src_ip, now, window, get_config(CFG_SYN_RATE, DEFAULT_SYN_RATE)) {
return true;
}
}
if matches!(pkt.protocol, IpProto::Udp) {
if check_rate(&IPV6_UDP_RATE_MAP, &src_ip, now, window, get_config(CFG_UDP_RATE, DEFAULT_UDP_RATE)) {
return true;
}
}
if pkt.dst_port == 53 {
if check_rate(&IPV6_DNS_RATE_MAP, &src_ip, now, window, get_config(CFG_DNS_RATE, DEFAULT_DNS_RATE)) {
return true;
}
}
false
}

View File

@ -1,163 +0,0 @@
use aya_ebpf::macros::map;
use aya_ebpf::maps::{Array, HashMap};
use common::define::offset::*;
use common::define::setting::MAX_RULES;
use common::model::event::{IPv4Event, IPv6Event};
use common::model::http_method::HttpMethodBitmap;
use common::model::ip_address::*;
use common::model::placeholder::PlaceHolder;
use network_types::ip::IpProto;
use network_types::tcp::TcpHdr;
#[map]
static IPV4_HTTP_SERVICE: HashMap<AddrPortV4, HttpMethodBitmap> = HashMap::with_max_entries(MAX_RULES as u32, 0);
#[map]
static IPV6_HTTP_SERVICE: HashMap<AddrPortV6, HttpMethodBitmap> = HashMap::with_max_entries(MAX_RULES as u32, 0);
#[map]
static SSH_WHITE_LIST_ENABLE: Array<PlaceHolder> = Array::with_max_entries(1, 0);
#[map]
static IPV4_SSH_SERVICE: HashMap<AddrPortV4, PlaceHolder> = HashMap::with_max_entries(MAX_RULES as u32, 0);
#[map]
static IPV6_SSH_SERVICE: HashMap<AddrPortV6, PlaceHolder> = HashMap::with_max_entries(MAX_RULES as u32, 0);
#[map]
static IPV4_SSH_WHITE_LIST: HashMap<IPv4, PlaceHolder> = HashMap::with_max_entries(MAX_RULES as u32, 0);
#[map]
static IPV6_SSH_WHITE_LIST: HashMap<IPv6, PlaceHolder> = HashMap::with_max_entries(MAX_RULES as u32, 0);
#[map]
static IPV4_SSH_BLACK_LIST: HashMap<IPv4, PlaceHolder> = HashMap::with_max_entries(MAX_RULES as u32, 0);
#[map]
static IPV6_SSH_BLACK_LIST: HashMap<IPv6, PlaceHolder> = HashMap::with_max_entries(MAX_RULES as u32, 0);
pub fn ipv4_service_rule_violation(start: usize, end: usize, event: &IPv4Event) -> bool {
let protocol = event.protocol;
let source = event.source_addr();
let destination = event.destination_addr();
ipv4_http_service_violation(start, end, &protocol, &destination)
|| ipv4_ssh_service_violation(&source, &destination)
}
pub fn ipv6_service_rule_violation(start: usize, end: usize, event: &IPv6Event) -> bool {
let protocol = event.protocol;
let source = event.source_addr();
let destination = event.destination_addr();
ipv6_http_service_violation(start, end, &protocol, &destination)
|| ipv6_ssh_service_violation(&source, &destination)
}
#[inline(always)]
fn ipv4_http_service_violation(start: usize, end: usize, protocol: &IpProto, destination: &AddrPortV4) -> bool {
match IPV4_HTTP_SERVICE.get_ptr_mut(destination) {
Some(allow_method) => {
if !matches!(protocol, IpProto::Tcp) {
return false;
}
unsafe {
if start + IPV4_TCP_HEADER_END > end {
return false;
}
let tcp_header = &*((start + IPV4_TCP_HEADER_START) as *const TcpHdr);
if tcp_header.syn() != 0 || tcp_header.rst() != 0 || tcp_header.fin() != 0 {
return false;
}
if tcp_header.psh() == 0 || tcp_header.ack() == 0 {
return false;
}
let doff = tcp_header.doff();
if doff < 5 || doff > 15 {
return false;
}
let tcp_header_len = (doff * 4) as usize;
let tcp_payload_start = IPV4_TCP_HEADER_END + tcp_header_len;
match get_http_request_method(start, end, tcp_payload_start) {
Some(http_method) => *allow_method & http_method == 0,
None => true,
}
}
}
None => false,
}
}
#[inline(always)]
fn ipv6_http_service_violation(start: usize, end: usize, protocol: &IpProto, destination: &AddrPortV6) -> bool {
match IPV6_HTTP_SERVICE.get_ptr_mut(destination) {
Some(allow_method) => {
if !matches!(protocol, IpProto::Tcp) {
return false;
}
unsafe {
if start + IPV6_TCP_HEADER_END > end {
return false;
}
let tcp_header = &*((start + IPV6_TCP_HEADER_START) as *const TcpHdr);
if tcp_header.syn() != 0 || tcp_header.rst() != 0 || tcp_header.fin() != 0 {
return false;
}
if tcp_header.psh() == 0 || tcp_header.ack() == 0 {
return false;
}
let doff = tcp_header.doff();
if doff < 5 || doff > 15 {
return false;
}
let tcp_header_len = (doff * 4) as usize;
let tcp_payload_start = IPV6_TCP_HEADER_END + tcp_header_len;
match get_http_request_method(start, end, tcp_payload_start) {
Some(http_method) => *allow_method & http_method == 0,
None => true,
}
}
}
None => false,
}
}
#[inline(always)]
fn get_http_request_method(start: usize, end: usize, offset: usize) -> Option<HttpMethodBitmap> {
if start + offset + 8 > end {
return None;
}
let data = unsafe { core::slice::from_raw_parts((start + offset) as *const u8, 8) };
match &data[..4] {
b"GET " => Some(1 << 0),
b"POST" if &data[4..5] == b" " => Some(1 << 1),
b"PUT " => Some(1 << 2),
b"DELE" if &data[4..7] == b"TE " => Some(1 << 3),
b"HEAD" if &data[4..5] == b" " => Some(1 << 4),
b"OPTI" if &data[4..8] == b"ONS " => Some(1 << 5),
b"PATC" if &data[4..6] == b"H " => Some(1 << 6),
b"TRAC" if &data[4..6] == b"E " => Some(1 << 7),
b"CONN" if &data[4..8] == b"ECT " => Some(1 << 8),
_ => None,
}
}
#[inline(always)]
fn ipv4_ssh_service_violation(source: &AddrPortV4, destination: &AddrPortV4) -> bool {
unsafe {
if IPV4_SSH_SERVICE.get(destination).is_some() {
if SSH_WHITE_LIST_ENABLE.get(0).is_some() {
IPV4_SSH_WHITE_LIST.get(&source.ip()).is_none()
} else {
IPV4_SSH_BLACK_LIST.get(&source.ip()).is_some()
}
} else {
false
}
}
}
#[inline(always)]
fn ipv6_ssh_service_violation(source_ip: &AddrPortV6, destination: &AddrPortV6) -> bool {
unsafe {
if IPV6_SSH_SERVICE.get(destination).is_some() {
if SSH_WHITE_LIST_ENABLE.get(0).is_some() {
IPV6_SSH_WHITE_LIST.get(&source_ip.ip()).is_none()
} else {
IPV6_SSH_BLACK_LIST.get(&source_ip.ip()).is_some()
}
} else {
false
}
}
}

View File

@ -1,93 +0,0 @@
use aya_ebpf::macros::map;
use aya_ebpf::maps::LruHashMap;
use common::define::setting::MAX_STATS;
use common::model::event::{IPv4Event, IPv6Event};
use common::model::flow_stats::FlowStats;
use common::model::ip_address::{AddrPortV4, AddrPortV6};
#[map]
static IPV4_INGRESS_SRC_1MIN: LruHashMap<AddrPortV4, FlowStats> = LruHashMap::with_max_entries(MAX_STATS as u32, 0);
#[map]
static IPV4_INGRESS_SRC_10MIN: LruHashMap<AddrPortV4, FlowStats> = LruHashMap::with_max_entries(MAX_STATS as u32, 0);
#[map]
static IPV4_INGRESS_SRC_1HOUR: LruHashMap<AddrPortV4, FlowStats> = LruHashMap::with_max_entries(MAX_STATS as u32, 0);
#[map]
static IPV6_INGRESS_SRC_1MIN: LruHashMap<AddrPortV6, FlowStats> = LruHashMap::with_max_entries(MAX_STATS as u32, 0);
#[map]
static IPV6_INGRESS_SRC_10MIN: LruHashMap<AddrPortV6, FlowStats> = LruHashMap::with_max_entries(MAX_STATS as u32, 0);
#[map]
static IPV6_INGRESS_SRC_1HOUR: LruHashMap<AddrPortV6, FlowStats> = LruHashMap::with_max_entries(MAX_STATS as u32, 0);
#[map]
static IPV4_INGRESS_DST_1MIN: LruHashMap<AddrPortV4, FlowStats> = LruHashMap::with_max_entries(MAX_STATS as u32, 0);
#[map]
static IPV4_INGRESS_DST_10MIN: LruHashMap<AddrPortV4, FlowStats> = LruHashMap::with_max_entries(MAX_STATS as u32, 0);
#[map]
static IPV4_INGRESS_DST_1HOUR: LruHashMap<AddrPortV4, FlowStats> = LruHashMap::with_max_entries(MAX_STATS as u32, 0);
#[map]
static IPV6_INGRESS_DST_1MIN: LruHashMap<AddrPortV6, FlowStats> = LruHashMap::with_max_entries(MAX_STATS as u32, 0);
#[map]
static IPV6_INGRESS_DST_10MIN: LruHashMap<AddrPortV6, FlowStats> = LruHashMap::with_max_entries(MAX_STATS as u32, 0);
#[map]
static IPV6_INGRESS_DST_1HOUR: LruHashMap<AddrPortV6, FlowStats> = LruHashMap::with_max_entries(MAX_STATS as u32, 0);
pub fn ipv4_update_stats(event: &IPv4Event) {
unsafe {
let source = event.source_addr();
let destination = event.destination_addr();
ipv4_update_flow_stats(&IPV4_INGRESS_SRC_1MIN, &source, event);
ipv4_update_flow_stats(&IPV4_INGRESS_SRC_10MIN, &source, event);
ipv4_update_flow_stats(&IPV4_INGRESS_SRC_1HOUR, &source, event);
ipv4_update_flow_stats(&IPV4_INGRESS_DST_1MIN, &destination, event);
ipv4_update_flow_stats(&IPV4_INGRESS_DST_10MIN, &destination, event);
ipv4_update_flow_stats(&IPV4_INGRESS_DST_1HOUR, &destination, event);
}
}
pub fn ipv6_update_stats(event: &IPv6Event) {
unsafe {
let source = event.source_addr();
let destination = event.destination_addr();
ipv6_update_flow_status(&IPV6_INGRESS_SRC_1MIN, &source, event);
ipv6_update_flow_status(&IPV6_INGRESS_SRC_10MIN, &source, event);
ipv6_update_flow_status(&IPV6_INGRESS_SRC_1HOUR, &source, event);
ipv6_update_flow_status(&IPV6_INGRESS_DST_1MIN, &destination, event);
ipv6_update_flow_status(&IPV6_INGRESS_DST_10MIN, &destination, event);
ipv6_update_flow_status(&IPV6_INGRESS_DST_1HOUR, &destination, event);
}
}
#[inline(always)]
unsafe fn ipv4_update_flow_stats(map: &LruHashMap<AddrPortV4, FlowStats>, key: &AddrPortV4, event: &IPv4Event) {
unsafe {
if let Some(status) = map.get_ptr_mut(key) {
(*status).bytes += event.packet_length as u64;
(*status).packets += 1;
(*status).last_seen = event.timestamp_us;
} else {
let new_stats = FlowStats {
bytes: event.packet_length as u64,
packets: 1,
last_seen: event.timestamp_us,
};
let _ = map.insert(key, &new_stats, 0);
}
}
}
#[inline(always)]
unsafe fn ipv6_update_flow_status(map: &LruHashMap<AddrPortV6, FlowStats>, key: &AddrPortV6, event: &IPv6Event) {
unsafe {
if let Some(status) = map.get_ptr_mut(key) {
(*status).bytes += event.packet_length as u64;
(*status).packets += 1;
(*status).last_seen = event.timestamp_us;
} else {
let new_stats = FlowStats {
bytes: event.packet_length as u64,
packets: 1,
last_seen: event.timestamp_us,
};
let _ = map.insert(key, &new_stats, 0);
}
}
}

View File

@ -4,41 +4,51 @@ mod action;
use aya_ebpf::bindings::xdp_action;
use aya_ebpf::macros::{map, xdp};
use aya_ebpf::maps::{PerCpuArray, ProgramArray, XskMap};
use aya_ebpf::maps::{Array, PerCpuArray, ProgramArray, XskMap};
use aya_ebpf::programs::XdpContext;
#[allow(unused_imports)]
use aya_log_ebpf::info;
use common::define::program_array::ingress::*;
use common::ebpf::parsing;
use common::model::event::Event;
use common::define::pipeline::*;
use common::model::parsed_packet::ParsedPacket;
use crate::action::{access_control, service, statistics};
use crate::action::{access_control, rate_limit, protocol_filter};
#[map]
static PROGRAM_ARRAY: ProgramArray = ProgramArray::with_max_entries(8, 0);
static PROGRAM_ARRAY: ProgramArray = ProgramArray::with_max_entries(MAX_STAGES, 0);
#[map]
static PARSED_PACKET: PerCpuArray<Event> = PerCpuArray::with_max_entries(1, 0);
static NEXT_STAGE: Array<u32> = Array::with_max_entries(MAX_STAGES, 0);
#[map]
static PARSED_PACKET: PerCpuArray<ParsedPacket> = PerCpuArray::with_max_entries(1, 0);
#[map]
static INGRESS_XSKS_MAP: XskMap = XskMap::pinned(64, 0);
#[inline(always)]
unsafe fn chain_next(ctx: &XdpContext, current_id: u32) {
unsafe {
if let Some(&next_slot) = NEXT_STAGE.get(current_id) {
if next_slot != STAGE_NONE {
let _ = PROGRAM_ARRAY.tail_call(ctx, next_slot);
}
}
let _ = PROGRAM_ARRAY.tail_call(ctx, STAGE_TRANSMISSION);
}
}
#[xdp]
pub fn net_guardia(ctx: XdpContext) -> u32 {
unsafe {
let _ = packet_intake(&ctx);
let _ = PROGRAM_ARRAY.tail_call(&ctx, TRANSMISSION);
packet_intake(&ctx);
let _ = PROGRAM_ARRAY.tail_call(&ctx, STAGE_TRANSMISSION);
xdp_action::XDP_PASS
}
}
#[inline(always)]
unsafe fn packet_intake(ctx: &XdpContext) -> Result<u32, ()> {
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(())
unsafe fn packet_intake(ctx: &XdpContext) {
let Some(ptr) = PARSED_PACKET.get_ptr_mut(0) else { return };
if parsing::parse_packet(ctx.data(), ctx.data_end(), ptr).is_ok() {
chain_next(ctx, STAGE_ENTRY);
}
}
@ -48,7 +58,7 @@ pub fn access_control(ctx: XdpContext) -> u32 {
match try_access_control(&ctx) {
Ok(action) => action,
Err(_) => {
let _ = PROGRAM_ARRAY.tail_call(&ctx, TRANSMISSION);
chain_next(&ctx, STAGE_ACCESS_CONTROL);
xdp_action::XDP_PASS
}
}
@ -59,39 +69,66 @@ pub fn access_control(ctx: XdpContext) -> u32 {
unsafe fn try_access_control(ctx: &XdpContext) -> Result<u32, ()> {
unsafe {
let ptr = PARSED_PACKET.get_ptr(0).ok_or(())?;
let parsed_packet = &*ptr;
match parsed_packet {
Event::IPv4(event) => {
if access_control::ipv4_is_whitelisted(event) {
let _ = PROGRAM_ARRAY.tail_call(ctx, STATISTICS);
let pkt = &*ptr;
match pkt.ip_version {
4 => {
if access_control::ipv4_is_whitelisted(pkt) {
let _ = PROGRAM_ARRAY.tail_call(ctx, STAGE_TRANSMISSION);
return Err(());
}
if access_control::ipv4_is_blacklisted(event) {
if access_control::ipv4_is_blacklisted(pkt) {
return Ok(xdp_action::XDP_DROP);
}
}
Event::IPv6(event) => {
if access_control::ipv6_is_whitelisted(event) {
let _ = PROGRAM_ARRAY.tail_call(ctx, STATISTICS);
6 => {
if access_control::ipv6_is_whitelisted(pkt) {
let _ = PROGRAM_ARRAY.tail_call(ctx, STAGE_TRANSMISSION);
return Err(());
}
if access_control::ipv6_is_blacklisted(event) {
if access_control::ipv6_is_blacklisted(pkt) {
return Ok(xdp_action::XDP_DROP);
}
}
_ => {}
}
let _ = PROGRAM_ARRAY.tail_call(ctx, SERVICE);
chain_next(ctx, STAGE_ACCESS_CONTROL);
Err(())
}
}
#[xdp]
pub fn service(ctx: XdpContext) -> u32 {
pub fn rate_limit(ctx: XdpContext) -> u32 {
unsafe {
match try_rate_limit(&ctx) {
Ok(action) => action,
Err(_) => {
chain_next(&ctx, STAGE_RATE_LIMIT);
xdp_action::XDP_PASS
}
}
}
}
#[inline(always)]
unsafe fn try_rate_limit(ctx: &XdpContext) -> Result<u32, ()> {
unsafe {
let ptr = PARSED_PACKET.get_ptr(0).ok_or(())?;
let pkt = &*ptr;
if rate_limit::should_drop(pkt) {
return Ok(xdp_action::XDP_DROP);
}
chain_next(ctx, STAGE_RATE_LIMIT);
Err(())
}
}
#[xdp]
pub fn protocol_filter(ctx: XdpContext) -> u32 {
unsafe {
match try_service(&ctx) {
Ok(action) => action,
Err(_) => {
let _ = PROGRAM_ARRAY.tail_call(&ctx, TRANSMISSION);
chain_next(&ctx, STAGE_SERVICE);
xdp_action::XDP_PASS
}
}
@ -104,46 +141,21 @@ unsafe fn try_service(ctx: &XdpContext) -> Result<u32, ()> {
let start = ctx.data();
let end = ctx.data_end();
let ptr = PARSED_PACKET.get_ptr(0).ok_or(())?;
let parsed_packet = &*ptr;
match parsed_packet {
Event::IPv4(event) => {
if service::ipv4_service_rule_violation(start, end, event) {
let pkt = &*ptr;
match pkt.ip_version {
4 => {
if protocol_filter::ipv4_service_rule_violation(start, end, pkt) {
return Ok(xdp_action::XDP_DROP);
}
}
Event::IPv6(event) => {
if service::ipv6_service_rule_violation(start, end, event) {
6 => {
if protocol_filter::ipv6_service_rule_violation(start, end, pkt) {
return Ok(xdp_action::XDP_DROP);
}
}
_ => {}
}
let _ = PROGRAM_ARRAY.tail_call(ctx, STATISTICS);
Err(())
}
}
#[xdp]
pub fn statistics(ctx: XdpContext) -> u32 {
unsafe {
let _ = try_statistics(&ctx);
xdp_action::XDP_PASS
}
}
#[inline(always)]
unsafe fn try_statistics(ctx: &XdpContext) -> Result<u32, ()> {
unsafe {
let ptr = PARSED_PACKET.get_ptr(0).ok_or(())?;
let parsed_packet = &*ptr;
match parsed_packet {
Event::IPv4(event) => {
statistics::ipv4_update_stats(&event);
}
Event::IPv6(event) => {
statistics::ipv6_update_stats(&event);
}
}
let _ = PROGRAM_ARRAY.tail_call(ctx, TRANSMISSION);
chain_next(ctx, STAGE_SERVICE);
Err(())
}
}

@ -1 +1 @@
Subproject commit 55da05e710cb00bbc9bef85804667624b0ddf377
Subproject commit 4e8b39bbb93641926bba18899b6c63cee187564d

View File

@ -3,8 +3,8 @@ use std::net::{Ipv4Addr, Ipv6Addr, SocketAddrV4, SocketAddrV6};
use aya::maps::{HashMap as AyaHashMap, MapData};
use aya::{Ebpf, Pod};
use common::define::setting::MAX_RULES_PORT;
use common::model::ip_address::{IPv4, IPv6, Port};
use common::model::port_rule::PortRule;
use tokio::sync::RwLock;
use crate::model::direction::FlowDirection;
@ -12,7 +12,6 @@ use crate::model::error::ebpf::EbpfError;
use crate::model::error::Error;
use crate::model::ip_address::NativeConvert;
use crate::model::list_type::ListType;
use crate::utils::ip_address::convert_ports_to_vec;
pub struct AccessControl {
ipv4_src_whitelist: RwLock<MapWrapper<IPv4>>,
@ -130,7 +129,7 @@ impl AccessControl {
}
struct MapWrapper<T> {
map: AyaHashMap<MapData, T, [Port; MAX_RULES_PORT]>,
map: AyaHashMap<MapData, T, PortRule>,
}
impl<T: NativeConvert + Pod> MapWrapper<T> {
@ -144,63 +143,59 @@ impl<T: NativeConvert + Pod> MapWrapper<T> {
self.map
.iter()
.filter_map(Result::ok)
.map(|(key, value)| (key.into_native(), convert_ports_to_vec(value)))
.map(|(key, rule)| (key.into_native(), rule.to_port_vec()))
.collect()
}
fn add(&mut self, ip: T, port: Port) -> Result<(), Error> {
let mut new_ports = [0_u16; MAX_RULES_PORT];
if port == 0 {
new_ports[0] = 0;
} else if let Ok(ports) = self.map.get(&ip, 0) {
if ports[0] == 0 {
return Ok(());
}
let mut index = None;
for (i, &value) in ports.iter().enumerate() {
if value == port {
return Ok(());
}
if index.is_none() && value == 0 {
index = Some(i);
}
}
if index.is_none() {
Err(EbpfError::RuleReachLimit)?;
}
new_ports.copy_from_slice(&ports);
new_ports[index.unwrap()] = port;
} else {
new_ports[0] = port;
// port 0 in API = match all ports
self.map
.insert(ip, PortRule::new_match_all(), 0)
.map_err(EbpfError::MapOperationError)?;
return Ok(());
}
let mut rule = self.map.get(&ip, 0).unwrap_or_else(|_| PortRule::new_empty());
if rule.is_match_all() {
return Ok(()); // already matching all
}
if !rule.add_port(port) {
Err(EbpfError::RuleReachLimit)?;
}
self.map
.insert(ip, new_ports, 0)
.insert(ip, rule, 0)
.map_err(EbpfError::MapOperationError)?;
Ok(())
}
fn remove(&mut self, ip: T, port: Port) -> Result<(), Error> {
if let Ok(mut ports) = self.map.get(&ip, 0) {
if port == 0 {
self.map.remove(&ip).map_err(EbpfError::MapOperationError)?;
return Ok(());
}
if let Some(index) = ports.iter().position(|&x| x == port) {
for i in index..(MAX_RULES_PORT - 1) {
ports[i] = ports[i + 1];
}
ports[MAX_RULES_PORT - 1] = 0;
if ports[0] == 0 {
self.map.remove(&ip).map_err(EbpfError::MapOperationError)?;
} else {
self.map.insert(ip, ports, 0).map_err(EbpfError::MapOperationError)?;
}
}
Ok(())
} else {
Err(EbpfError::IpDoesNotExist)?
if port == 0 {
// port 0 in API = remove entire IP
self.map.remove(&ip).map_err(EbpfError::MapOperationError)?;
return Ok(());
}
let mut rule = self.map.get(&ip, 0).map_err(|_| EbpfError::IpDoesNotExist)?;
if rule.is_match_all() {
// Can't remove a single port from match_all — remove the whole IP
self.map.remove(&ip).map_err(EbpfError::MapOperationError)?;
return Ok(());
}
rule.remove_port(port);
if rule.is_empty() {
self.map.remove(&ip).map_err(EbpfError::MapOperationError)?;
} else {
self.map
.insert(ip, rule, 0)
.map_err(EbpfError::MapOperationError)?;
}
Ok(())
}
}

View File

@ -1,6 +1,6 @@
pub mod access_control;
pub mod service;
pub mod statistics;
pub mod rate_limit;
pub mod protocol_filter;
pub mod xsk_manager;
use std::sync::Arc;
@ -11,51 +11,40 @@ use macros::log;
use tokio::sync::oneshot;
use crate::core::ebpf::access_control::AccessControl;
use crate::core::ebpf::service::Service;
use crate::core::ebpf::statistics::Statistics;
use crate::core::ebpf::rate_limit::RateLimitConfig;
use crate::core::ebpf::protocol_filter::ProtocolFilter;
use crate::core::ebpf::xsk_manager::XskManager;
use crate::core::infrastructure::app_config::AppConfig;
use crate::core::ml::engine::Engine;
use crate::model::error::system::SystemError;
use crate::model::error::Error;
use crate::ml::engine::Engine;
pub struct EbpfServices {
pub xsk_manager: Arc<XskManager>,
pub access_control: Arc<AccessControl>,
pub service: Arc<Service>,
pub statistics: Arc<Statistics>,
pub protocol_filter: Arc<ProtocolFilter>,
pub rate_limit: Arc<RateLimitConfig>,
pub shutdowns: SegQueue<oneshot::Sender<()>>,
}
impl EbpfServices {
pub fn new(
app_config: Arc<AppConfig>,
ingress_ebpf: &mut Ebpf,
egress_ebpf: &mut Ebpf,
) -> Result<Self, Error> {
pub fn new(app_config: Arc<AppConfig>, ingress_ebpf: &mut Ebpf, egress_ebpf: &mut Ebpf) -> Result<Self, Error> {
let xsk_manager = XskManager::new(app_config.clone(), ingress_ebpf, egress_ebpf)?;
let access_control = AccessControl::new(ingress_ebpf)?;
let service = Service::new(ingress_ebpf)?;
let statistics = Statistics::new(app_config.clone(), ingress_ebpf, egress_ebpf)?;
let ebpf_services = Self {
let protocol_filter = ProtocolFilter::new(ingress_ebpf)?;
let rate_limit = RateLimitConfig::new(ingress_ebpf)?;
Ok(Self {
xsk_manager: Arc::new(xsk_manager),
access_control: Arc::new(access_control),
service: Arc::new(service),
statistics: Arc::new(statistics),
protocol_filter: Arc::new(protocol_filter),
rate_limit: Arc::new(rate_limit),
shutdowns: SegQueue::new(),
};
Ok(ebpf_services)
})
}
pub async fn run(self: Arc<Self>, ml_engine: Arc<Engine>) -> Result<(), Error> {
let xsk_manager = self.xsk_manager.clone();
let statistics = self.statistics.clone();
xsk_manager.run(Some(ml_engine), &self.shutdowns)?;
let statistics_shutdown = statistics.run().await;
self.shutdowns.push(statistics_shutdown);
Ok(())
}
@ -66,4 +55,4 @@ impl EbpfServices {
}
}
}
}
}

View File

@ -12,30 +12,30 @@ use crate::model::error::ebpf::EbpfError;
use crate::model::error::Error;
use crate::model::ip_address::NativeConvert;
pub struct Service {
pub struct ProtocolFilter {
ipv4_http_service: RwLock<HttpServiceWrapper<AddrPortV4>>,
ipv6_http_service: RwLock<HttpServiceWrapper<AddrPortV6>>,
ssh_white_list_enable: RwLock<WhiteListControl>,
ipv4_ssh_service: RwLock<SshServiceWrapper<AddrPortV4>>,
ipv6_ssh_service: RwLock<SshServiceWrapper<AddrPortV6>>,
ipv4_ssh_white_list: RwLock<SshListWrapper<IPv4>>,
ipv6_ssh_white_list: RwLock<SshListWrapper<IPv6>>,
ipv4_ssh_black_list: RwLock<SshListWrapper<IPv4>>,
ipv6_ssh_black_list: RwLock<SshListWrapper<IPv6>>,
ipv4_ssh_service: RwLock<EntryMap<AddrPortV4>>,
ipv6_ssh_service: RwLock<EntryMap<AddrPortV6>>,
ipv4_ssh_white_list: RwLock<EntryMap<IPv4>>,
ipv6_ssh_white_list: RwLock<EntryMap<IPv6>>,
ipv4_ssh_black_list: RwLock<EntryMap<IPv4>>,
ipv6_ssh_black_list: RwLock<EntryMap<IPv6>>,
}
impl Service {
impl ProtocolFilter {
pub fn new(ebpf: &mut Ebpf) -> Result<Self, Error> {
let service = Self {
ipv4_http_service: RwLock::new(HttpServiceWrapper::new(ebpf, "IPV4_HTTP_SERVICE")?),
ipv6_http_service: RwLock::new(HttpServiceWrapper::new(ebpf, "IPV6_HTTP_SERVICE")?),
ssh_white_list_enable: RwLock::new(WhiteListControl::new(ebpf, "SSH_WHITE_LIST_ENABLE")?),
ipv4_ssh_service: RwLock::new(SshServiceWrapper::new(ebpf, "IPV4_SSH_SERVICE")?),
ipv6_ssh_service: RwLock::new(SshServiceWrapper::new(ebpf, "IPV6_SSH_SERVICE")?),
ipv4_ssh_white_list: RwLock::new(SshListWrapper::new(ebpf, "IPV4_SSH_WHITE_LIST")?),
ipv6_ssh_white_list: RwLock::new(SshListWrapper::new(ebpf, "IPV6_SSH_WHITE_LIST")?),
ipv4_ssh_black_list: RwLock::new(SshListWrapper::new(ebpf, "IPV4_SSH_BLACK_LIST")?),
ipv6_ssh_black_list: RwLock::new(SshListWrapper::new(ebpf, "IPV6_SSH_BLACK_LIST")?),
ipv4_ssh_service: RwLock::new(EntryMap::new(ebpf, "IPV4_SSH_SERVICE")?),
ipv6_ssh_service: RwLock::new(EntryMap::new(ebpf, "IPV6_SSH_SERVICE")?),
ipv4_ssh_white_list: RwLock::new(EntryMap::new(ebpf, "IPV4_SSH_WHITE_LIST")?),
ipv6_ssh_white_list: RwLock::new(EntryMap::new(ebpf, "IPV6_SSH_WHITE_LIST")?),
ipv4_ssh_black_list: RwLock::new(EntryMap::new(ebpf, "IPV4_SSH_BLACK_LIST")?),
ipv6_ssh_black_list: RwLock::new(EntryMap::new(ebpf, "IPV6_SSH_BLACK_LIST")?),
};
Ok(service)
}
@ -105,75 +105,75 @@ impl Service {
}
pub async fn get_ipv4_ssh_service(&self) -> Vec<SocketAddrV4> {
self.ipv4_ssh_service.read().await.get_ssh_service()
self.ipv4_ssh_service.read().await.get_all()
}
pub async fn get_ipv6_ssh_service(&self) -> Vec<SocketAddrV6> {
self.ipv6_ssh_service.read().await.get_ssh_service()
self.ipv6_ssh_service.read().await.get_all()
}
pub async fn add_ipv4_ssh_service(&self, address: SocketAddrV4) -> Result<(), Error> {
self.ipv4_ssh_service.write().await.add_ssh_service(address)
self.ipv4_ssh_service.write().await.add(address)
}
pub async fn add_ipv6_ssh_service(&self, address: SocketAddrV6) -> Result<(), Error> {
self.ipv6_ssh_service.write().await.add_ssh_service(address)
self.ipv6_ssh_service.write().await.add(address)
}
pub async fn remove_ipv4_ssh_service(&self, address: SocketAddrV4) -> Result<(), Error> {
self.ipv4_ssh_service.write().await.remove_ssh_service(address)
self.ipv4_ssh_service.write().await.remove(address)
}
pub async fn remove_ipv6_ssh_service(&self, address: SocketAddrV6) -> Result<(), Error> {
self.ipv6_ssh_service.write().await.remove_ssh_service(address)
self.ipv6_ssh_service.write().await.remove(address)
}
pub async fn get_ipv4_ssh_white_list(&self) -> Vec<Ipv4Addr> {
self.ipv4_ssh_white_list.read().await.get_list()
self.ipv4_ssh_white_list.read().await.get_all()
}
pub async fn get_ipv6_ssh_white_list(&self) -> Vec<Ipv6Addr> {
self.ipv6_ssh_white_list.read().await.get_list()
self.ipv6_ssh_white_list.read().await.get_all()
}
pub async fn add_ipv4_ssh_white_list(&self, ip: Ipv4Addr) -> Result<(), Error> {
self.ipv4_ssh_white_list.write().await.add_list(ip)
self.ipv4_ssh_white_list.write().await.add(ip)
}
pub async fn add_ipv6_ssh_white_list(&self, ip: Ipv6Addr) -> Result<(), Error> {
self.ipv6_ssh_white_list.write().await.add_list(ip)
self.ipv6_ssh_white_list.write().await.add(ip)
}
pub async fn remove_ipv4_ssh_white_list(&self, ip: Ipv4Addr) -> Result<(), Error> {
self.ipv4_ssh_white_list.write().await.remove_list(ip)
self.ipv4_ssh_white_list.write().await.remove(ip)
}
pub async fn remove_ipv6_ssh_white_list(&self, ip: Ipv6Addr) -> Result<(), Error> {
self.ipv6_ssh_white_list.write().await.remove_list(ip)
self.ipv6_ssh_white_list.write().await.remove(ip)
}
pub async fn get_ipv4_ssh_black_list(&self) -> Vec<Ipv4Addr> {
self.ipv4_ssh_black_list.read().await.get_list()
self.ipv4_ssh_black_list.read().await.get_all()
}
pub async fn get_ipv6_ssh_black_list(&self) -> Vec<Ipv6Addr> {
self.ipv6_ssh_black_list.read().await.get_list()
self.ipv6_ssh_black_list.read().await.get_all()
}
pub async fn add_ipv4_ssh_black_list(&self, ip: Ipv4Addr) -> Result<(), Error> {
self.ipv4_ssh_black_list.write().await.add_list(ip)
self.ipv4_ssh_black_list.write().await.add(ip)
}
pub async fn add_ipv6_ssh_black_list(&self, ip: Ipv6Addr) -> Result<(), Error> {
self.ipv6_ssh_black_list.write().await.add_list(ip)
self.ipv6_ssh_black_list.write().await.add(ip)
}
pub async fn remove_ipv4_ssh_black_list(&self, ip: Ipv4Addr) -> Result<(), Error> {
self.ipv4_ssh_black_list.write().await.remove_list(ip)
self.ipv4_ssh_black_list.write().await.remove(ip)
}
pub async fn remove_ipv6_ssh_black_list(&self, ip: Ipv6Addr) -> Result<(), Error> {
self.ipv6_ssh_black_list.write().await.remove_list(ip)
self.ipv6_ssh_black_list.write().await.remove(ip)
}
}
@ -239,7 +239,7 @@ impl<T: NativeConvert + Pod> HttpServiceWrapper<T> {
let ebpf_method = HttpMethod::convert_to_bitmap(http_method);
self.map
.insert(address, ebpf_method, 0)
.map_err(|_| EbpfError::RuleReachLimit)?;
.map_err(EbpfError::MapOperationError)?;
Ok(())
}
@ -263,18 +263,18 @@ impl<T: NativeConvert + Pod> HttpServiceWrapper<T> {
}
}
struct SshServiceWrapper<T> {
struct EntryMap<T> {
map: AyaHashMap<MapData, T, PlaceHolder>,
}
impl<T: NativeConvert + Pod> SshServiceWrapper<T> {
impl<T: NativeConvert + Pod> EntryMap<T> {
fn new(ebpf: &mut Ebpf, map_name: &str) -> Result<Self, Error> {
let map = ebpf.take_map(map_name).ok_or(EbpfError::MapNotFound)?;
let map = AyaHashMap::try_from(map).map_err(EbpfError::MapOperationError)?;
Ok(Self { map })
}
fn get_ssh_service(&self) -> Vec<T::Native> {
fn get_all(&self) -> Vec<T::Native> {
self.map
.keys()
.filter_map(Result::ok)
@ -282,51 +282,17 @@ impl<T: NativeConvert + Pod> SshServiceWrapper<T> {
.collect()
}
fn add_ssh_service(&mut self, address: T::Native) -> Result<(), Error> {
let address = T::from_native(address);
fn add(&mut self, key: T::Native) -> Result<(), Error> {
let key = T::from_native(key);
self.map
.insert(address, 0_u8, 0)
.map_err(|_| EbpfError::RuleReachLimit)?;
.insert(key, 0_u8, 0)
.map_err(EbpfError::MapOperationError)?;
Ok(())
}
fn remove_ssh_service(&mut self, address: T::Native) -> Result<(), Error> {
let address = T::from_native(address);
self.map.remove(&address).map_err(|_| EbpfError::IpDoesNotExist)?;
Ok(())
}
}
struct SshListWrapper<T> {
map: AyaHashMap<MapData, T, PlaceHolder>,
}
impl<T: NativeConvert + Pod> SshListWrapper<T> {
fn new(ebpf: &mut Ebpf, map_name: &str) -> Result<Self, Error> {
let map = ebpf.take_map(map_name).ok_or(EbpfError::MapNotFound)?;
let map = AyaHashMap::try_from(map).map_err(EbpfError::MapOperationError)?;
Ok(Self { map })
}
fn get_list(&self) -> Vec<T::Native> {
self.map
.keys()
.filter_map(Result::ok)
.map(|key| key.into_native())
.collect()
}
fn add_list(&mut self, address: T::Native) -> Result<(), Error> {
let address = T::from_native(address);
self.map
.insert(address, 0_u8, 0)
.map_err(|_| EbpfError::RuleReachLimit)?;
Ok(())
}
fn remove_list(&mut self, address: T::Native) -> Result<(), Error> {
let address = T::from_native(address);
self.map.remove(&address).map_err(|_| EbpfError::IpDoesNotExist)?;
fn remove(&mut self, key: T::Native) -> Result<(), Error> {
let key = T::from_native(key);
self.map.remove(&key).map_err(EbpfError::MapOperationError)?;
Ok(())
}
}

View File

@ -0,0 +1,43 @@
use aya::maps::{Array, MapData};
use aya::Ebpf;
use parking_lot::Mutex;
use crate::model::error::ebpf::EbpfError;
use crate::model::error::Error;
pub struct RateLimitConfig {
config_map: Mutex<Array<MapData, u64>>,
}
impl RateLimitConfig {
pub fn new(ebpf: &mut Ebpf) -> Result<Self, Error> {
let map = ebpf.take_map("RATE_LIMIT_CONFIG").ok_or(EbpfError::MapNotFound)?;
let config_map = Array::try_from(map).map_err(EbpfError::MapOperationError)?;
Ok(Self { config_map: Mutex::new(config_map) })
}
pub fn set_packet_rate(&self, rate: u64) -> Result<(), Error> {
self.config_map.lock().set(0, rate, 0).map_err(EbpfError::MapOperationError)?;
Ok(())
}
pub fn set_syn_rate(&self, rate: u64) -> Result<(), Error> {
self.config_map.lock().set(1, rate, 0).map_err(EbpfError::MapOperationError)?;
Ok(())
}
pub fn set_udp_rate(&self, rate: u64) -> Result<(), Error> {
self.config_map.lock().set(2, rate, 0).map_err(EbpfError::MapOperationError)?;
Ok(())
}
pub fn set_dns_rate(&self, rate: u64) -> Result<(), Error> {
self.config_map.lock().set(3, rate, 0).map_err(EbpfError::MapOperationError)?;
Ok(())
}
pub fn set_window_ns(&self, ns: u64) -> Result<(), Error> {
self.config_map.lock().set(4, ns, 0).map_err(EbpfError::MapOperationError)?;
Ok(())
}
}

View File

@ -1,258 +0,0 @@
use std::collections::HashMap;
use std::net::{IpAddr, SocketAddrV4, SocketAddrV6};
use std::sync::Arc;
use aya::maps::{HashMap as AyaHashMap, MapData};
use aya::{Ebpf, Pod};
use common::model::flow_stats::FlowStats;
use common::model::ip_address::{AddrPortV4, AddrPortV6};
use futures::future::join_all;
use macros::log;
use tokio::select;
use tokio::sync::{RwLock, oneshot};
use tokio::time::{Duration, sleep};
use crate::core::infrastructure::app_config::AppConfig;
use crate::core::infrastructure::geoip::GeoIpService;
use crate::model::direction::{Direction, FlowDirection};
use crate::model::error::Error;
use crate::model::error::ebpf::EbpfError;
use crate::model::error::misc::MiscError;
use crate::model::geo_stats::FlowStatsWithGeo;
use crate::model::ip_address::NativeConvert;
use crate::model::time_type::TimeType;
use crate::utils::boot_time::boot_time;
pub struct Statistics {
app_config: Arc<AppConfig>,
boot_time: u64,
geo_ip: Option<Arc<GeoIpService>>,
ipv4_maps: HashMap<(Direction, FlowDirection, TimeType), RwLock<FlowMap<AddrPortV4>>>,
ipv6_maps: HashMap<(Direction, FlowDirection, TimeType), RwLock<FlowMap<AddrPortV6>>>,
}
impl Statistics {
const INGRESS_MAPS: [((Direction, FlowDirection, TimeType), (&'static str, &'static str)); 6] = [
(
(Direction::Ingress, FlowDirection::Source, TimeType::_1Min),
("IPV4_INGRESS_SRC_1MIN", "IPV6_INGRESS_SRC_1MIN"),
),
(
(Direction::Ingress, FlowDirection::Source, TimeType::_10Min),
("IPV4_INGRESS_SRC_10MIN", "IPV6_INGRESS_SRC_10MIN"),
),
(
(Direction::Ingress, FlowDirection::Source, TimeType::_1Hour),
("IPV4_INGRESS_SRC_1HOUR", "IPV6_INGRESS_SRC_1HOUR"),
),
(
(Direction::Ingress, FlowDirection::Destination, TimeType::_1Min),
("IPV4_INGRESS_DST_1MIN", "IPV6_INGRESS_DST_1MIN"),
),
(
(Direction::Ingress, FlowDirection::Destination, TimeType::_10Min),
("IPV4_INGRESS_DST_10MIN", "IPV6_INGRESS_DST_10MIN"),
),
(
(Direction::Ingress, FlowDirection::Destination, TimeType::_1Hour),
("IPV4_INGRESS_DST_1HOUR", "IPV6_INGRESS_DST_1HOUR"),
),
];
const EGRESS_MAPS: [((Direction, FlowDirection, TimeType), (&'static str, &'static str)); 6] = [
(
(Direction::Egress, FlowDirection::Source, TimeType::_1Min),
("IPV4_EGRESS_SRC_1MIN", "IPV6_EGRESS_SRC_1MIN"),
),
(
(Direction::Egress, FlowDirection::Source, TimeType::_10Min),
("IPV4_EGRESS_SRC_10MIN", "IPV6_EGRESS_SRC_10MIN"),
),
(
(Direction::Egress, FlowDirection::Source, TimeType::_1Hour),
("IPV4_EGRESS_SRC_1HOUR", "IPV6_EGRESS_SRC_1HOUR"),
),
(
(Direction::Egress, FlowDirection::Destination, TimeType::_1Min),
("IPV4_EGRESS_DST_1MIN", "IPV6_EGRESS_DST_1MIN"),
),
(
(Direction::Egress, FlowDirection::Destination, TimeType::_10Min),
("IPV4_EGRESS_DST_10MIN", "IPV6_EGRESS_DST_10MIN"),
),
(
(Direction::Egress, FlowDirection::Destination, TimeType::_1Hour),
("IPV4_EGRESS_DST_1HOUR", "IPV6_EGRESS_DST_1HOUR"),
),
];
pub fn new(
app_config: Arc<AppConfig>,
ingress_ebpf: &mut Ebpf,
egress_ebpf: &mut Ebpf,
) -> Result<Statistics, Error> {
let boot_time = boot_time();
let mut ipv4_maps = HashMap::new();
let mut ipv6_maps = HashMap::new();
let geo_ip = match GeoIpService::new(&app_config.geoip_db_name) {
Ok(service) => Some(Arc::new(service)),
Err(err) => {
log!(MiscError::InvalidGeoIPConfiguration(err));
None
}
};
for (key, (ipv4_name, ipv6_name)) in Self::INGRESS_MAPS {
ipv4_maps.insert(key, RwLock::new(FlowMap::new(ingress_ebpf, ipv4_name)?));
ipv6_maps.insert(key, RwLock::new(FlowMap::new(ingress_ebpf, ipv6_name)?));
}
for (key, (ipv4_name, ipv6_name)) in Self::EGRESS_MAPS {
ipv4_maps.insert(key, RwLock::new(FlowMap::new(egress_ebpf, ipv4_name)?));
ipv6_maps.insert(key, RwLock::new(FlowMap::new(egress_ebpf, ipv6_name)?));
}
let statistics = Statistics {
app_config,
boot_time,
geo_ip,
ipv4_maps,
ipv6_maps,
};
Ok(statistics)
}
pub async fn run(self: Arc<Self>) -> oneshot::Sender<()> {
let refresh_interval = self.app_config.refresh_interval;
let (sender, receiver) = oneshot::channel();
tokio::spawn(async move {
let mut receiver = receiver;
loop {
select! {
biased;
_ = &mut receiver => break,
_ = sleep(Duration::from_secs(refresh_interval)) => {
self.cleanup_expired_flows().await;
},
}
}
});
sender
}
pub async fn cleanup_expired_flows(self: &Arc<Self>) {
let boot_time = self.boot_time;
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos() as u64;
for ((_, _, time_type), map) in self.ipv4_maps.iter() {
map.write().await.cleanup(boot_time, now, time_type.duration())
}
for ((_, _, time_type), map) in self.ipv6_maps.iter() {
map.write().await.cleanup(boot_time, now, time_type.duration())
}
}
pub async fn get_ipv4_flow_data(
&self,
direction: Direction,
flow_direction: FlowDirection,
time_type: TimeType,
) -> HashMap<SocketAddrV4, FlowStatsWithGeo> {
let flow_data = self
.ipv4_maps
.get(&(direction, flow_direction, time_type))
.unwrap()
.write()
.await
.get_map();
if let Some(ref geo_ip) = self.geo_ip {
let futures: Vec<_> = flow_data
.into_iter()
.map(|(addr, stats)| {
let geo_ip = geo_ip.clone();
async move {
let ip = IpAddr::V4(*addr.ip());
let geo = geo_ip.lookup(ip).await.ok().flatten();
(addr, FlowStatsWithGeo { stats, geo })
}
})
.collect();
join_all(futures).await.into_iter().collect()
} else {
flow_data
.into_iter()
.map(|(addr, stats)| (addr, FlowStatsWithGeo { stats, geo: None }))
.collect()
}
}
pub async fn get_ipv6_flow_data(
&self,
direction: Direction,
flow_direction: FlowDirection,
time_type: TimeType,
) -> HashMap<SocketAddrV6, FlowStatsWithGeo> {
let flow_data = self
.ipv6_maps
.get(&(direction, flow_direction, time_type))
.unwrap()
.write()
.await
.get_map();
if let Some(ref geo_ip) = self.geo_ip {
let futures: Vec<_> = flow_data
.into_iter()
.map(|(addr, stats)| {
let geo_ip = geo_ip.clone();
async move {
let ip = IpAddr::V6(*addr.ip());
let geo = geo_ip.lookup(ip).await.ok().flatten();
(addr, FlowStatsWithGeo { stats, geo })
}
})
.collect();
join_all(futures).await.into_iter().collect()
} else {
flow_data
.into_iter()
.map(|(addr, stats)| (addr, FlowStatsWithGeo { stats, geo: None }))
.collect()
}
}
}
struct FlowMap<T> {
map: AyaHashMap<MapData, T, FlowStats>,
}
impl<T: NativeConvert + Pod> FlowMap<T> {
fn new(ebpf: &mut Ebpf, map_name: &str) -> Result<Self, Error> {
let map = ebpf.take_map(map_name).ok_or(EbpfError::MapNotFound)?;
let map = AyaHashMap::try_from(map).map_err(EbpfError::MapOperationError)?;
Ok(Self { map })
}
fn get_map(&self) -> HashMap<T::Native, FlowStats> {
self.map
.iter()
.filter_map(Result::ok)
.map(|(key, value)| (key.into_native(), FlowStats::from(value)))
.collect()
}
fn cleanup(&mut self, boot_time: u64, now: u64, window: u64) {
let expired_keys: Vec<T> = self
.map
.iter()
.filter_map(|result| {
result
.ok()
.and_then(|(key, stats)| (now - stats.last_seen - boot_time > window).then_some(key))
})
.collect();
expired_keys.iter().for_each(|key| {
let _ = self.map.remove(key);
});
}
}

View File

@ -17,8 +17,9 @@ use xsk_rs::config::{BindFlags, FrameSize, Interface, LibxdpFlags, QueueSize, So
use xsk_rs::{CompQueue, FillQueue, FrameDesc, RxQueue, Socket, TxQueue, Umem};
use crate::core::infrastructure::app_config::AppConfig;
use crate::ml::engine::{Engine, PacketProcessor};
use crate::model::config::Config;
use crate::core::ml::engine::Engine;
use crate::core::ml::flow_tracker::FlowTracker;
use crate::model::config::NetworkConfig;
use crate::model::direction::Direction;
use crate::model::error::ebpf::EbpfError;
use crate::model::error::system::SystemError;
@ -49,45 +50,47 @@ impl XskManager {
}
pub fn run(&self, ml_engine: Option<Arc<Engine>>, shutdowns: &SegQueue<oneshot::Sender<()>>) -> Result<(), Error> {
let config = self.app_config.config.clone();
let combined_queue_count = config.combined_queue_count;
let packet_processor = ml_engine.map(|engine| Arc::new(PacketProcessor::new(engine)));
let network = self.app_config.network.clone();
let combined_queue_count = network.combined_queue_count;
for queue_id in 0..combined_queue_count {
let (ingress_to_egress_tx, ingress_to_egress_rx) = bounded(config.channel_size);
let (egress_to_ingress_tx, egress_to_ingress_rx) = bounded(config.channel_size);
let (ingress_to_egress_tx, ingress_to_egress_rx) = bounded(network.channel_size);
let (egress_to_ingress_tx, egress_to_ingress_rx) = bounded(network.channel_size);
let tracker = ml_engine.as_ref().map(|engine| engine.tracker(queue_id).clone());
let ingress_xsk = XskPair::new(
config.clone(),
network.clone(),
queue_id,
&config.ingress_ifname,
&config.egress_ifname,
&network.ingress_ifname,
&network.egress_ifname,
Direction::Ingress,
packet_processor.clone(),
tracker.clone(),
)?;
let egress_xsk = XskPair::new(
config.clone(),
network.clone(),
queue_id,
&config.egress_ifname,
&config.ingress_ifname,
&network.egress_ifname,
&network.ingress_ifname,
Direction::Egress,
packet_processor.clone(),
tracker,
)?;
let mut xsk_map = self.xsk_map.lock();
let mut egress_xsk_map = self.egress_xsk_map.lock();
let ingress_fd = ingress_xsk.rx.fd().as_raw_fd();
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(xsk_map);
drop(egress_xsk_map);
let ingress_shutdown = ingress_xsk.run(ingress_to_egress_tx, egress_to_ingress_rx)?;
@ -110,18 +113,18 @@ pub struct XskPair {
comp_queue: CompQueue,
tx: TxQueue,
rx: RxQueue,
frame_pool: Arc<Mutex<Vec<FrameDesc>>>, // SegQueue
packet_processor: Option<Arc<PacketProcessor>>,
frame_pool: Arc<Mutex<Vec<FrameDesc>>>,
tracker: Option<Arc<Mutex<FlowTracker>>>,
}
impl XskPair {
pub fn new(
config: Config,
config: NetworkConfig,
queue_id: u32,
rx_ifname: &str,
tx_ifname: &str,
_tx_ifname: &str,
direction: Direction,
packet_processor: Option<Arc<PacketProcessor>>,
tracker: Option<Arc<Mutex<FlowTracker>>>,
) -> Result<Self, Error> {
let rx_ifname_c = CString::new(rx_ifname).map_err(|_| SystemError::UnknownError)?;
@ -145,12 +148,13 @@ impl XskPair {
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_COPY)
.libxdp_flags(LibxdpFlags::XSK_LIBXDP_FLAGS_INHIBIT_PROG_LOAD)
.build();
let interface = Interface::new(rx_ifname_c);
// SAFETY: Interface and umem are valid and outlive the socket
let (tx, rx, queue) =
unsafe { Socket::new(socket_config, &umem, &interface, queue_id).map_err(EbpfError::SocketSetFailed)? };
@ -161,9 +165,10 @@ impl XskPair {
let fill_frames: Vec<FrameDesc> = frame_descs.iter().take(fill_frames_count).copied().collect();
// SAFETY: Frame descriptors are valid and owned by this UMEM
let submitted = unsafe { fill_queue.produce(&fill_frames) };
if submitted != fill_frames.len() {
log!(EbpfLog::QueueInitIncomplete);
return Err(EbpfError::FillQueueInitFailed.into());
}
let pool_frames: Vec<FrameDesc> = frame_descs.iter().skip(fill_frames_count).copied().collect();
@ -176,7 +181,7 @@ impl XskPair {
tx,
rx,
frame_pool: Arc::new(Mutex::new(pool_frames)),
packet_processor,
tracker,
};
Ok(xsk_pair)
@ -198,6 +203,9 @@ impl XskPair {
let mut idle_count: u32 = 0;
loop {
// Non-blocking shutdown check: try_recv avoids blocking the hot loop.
// The idle backoff below (sleep_us) ensures we don't busy-spin when idle,
// which also bounds how quickly we detect shutdown to at most 100us.
if let Some(ref mut rx) = shutdown_rx {
match rx.try_recv() {
Ok(_) | Err(oneshot::error::TryRecvError::Closed) => {
@ -266,6 +274,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];
// SAFETY: rx_descs buffer is large enough for consume
let rx_count = unsafe { self.rx.consume(&mut rx_descs) };
if rx_count > 0 {
@ -273,11 +282,17 @@ impl XskPair {
let lengths = rx_desc.lengths();
let packet_len = lengths.data() as usize;
// SAFETY: rx_desc is valid and belongs to this UMEM
let data = unsafe { self.umem.data(rx_desc) };
let packet_data = data.contents()[..packet_len].to_vec();
let contents = data.contents();
if packet_len > contents.len() {
log!(EbpfLog::InvalidPacketLength);
continue;
}
let packet_data = contents[..packet_len].to_vec();
if let Some(ref processor) = self.packet_processor {
processor.process(&packet_data, self.direction == Direction::Ingress);
if let Some(ref tracker) = self.tracker {
Engine::process_packet(tracker, &packet_data, self.direction == Direction::Ingress);
}
if let Err(e) = forward_tx.try_send(packet_data) {
@ -316,7 +331,9 @@ impl XskPair {
return Ok(0);
}
let _ = self.process_comp_queue();
if let Err(e) = self.process_comp_queue() {
log!(EbpfLog::CompQueueError(format!("{:?}", e)));
}
let pool_size = {
let pool = self.frame_pool.lock();
@ -355,6 +372,7 @@ impl XskPair {
}
}
// SAFETY: Frames contain valid packet data written above
let nb_submitted = unsafe { self.tx.produce(&frames) };
if let Err(e) = self.tx.wakeup() {

View File

@ -1,39 +1,51 @@
use std::fs;
use std::ops::Deref;
use crate::model::config::{Config, ConfigTable};
use crate::model::config::{
AppConfigTable, HttpConfig, InferenceConfig as InfConfig, MiscConfig, NetworkConfig, PipelineConfig,
};
use crate::model::error::system::SystemError;
use crate::model::error::Error;
pub struct AppConfig {
pub config: Config,
pub http: HttpConfig,
pub network: NetworkConfig,
pub inference: InfConfig,
#[allow(dead_code)]
pub misc: MiscConfig,
pub pipeline: PipelineConfig,
}
impl AppConfig {
pub fn new() -> Result<Self, Error> {
let toml_string = fs::read_to_string("./config.toml").map_err(SystemError::ConfigNotFound)?;
let config_table = toml::from_str::<ConfigTable>(&toml_string).map_err(|_| SystemError::InvalidConfig)?;
let config = config_table.config;
if !Self::validate(&config) {
let table = toml::from_str::<AppConfigTable>(&toml_string).map_err(|_| SystemError::InvalidConfig)?;
if !Self::validate(&table) {
Err(SystemError::InvalidConfig)?
} else {
Ok(Self { config })
}
Ok(Self {
http: table.http,
network: table.network,
inference: table.inference,
misc: table.misc,
pipeline: table.pipeline,
})
}
fn validate(config: &Config) -> bool {
Self::validate_second(config.refresh_interval)
}
fn validate_second(second: u64) -> bool {
second <= 3600
}
}
impl Deref for AppConfig {
type Target = Config;
fn deref(&self) -> &Self::Target {
&self.config
fn validate(table: &AppConfigTable) -> bool {
let net = &table.network;
let inf = &table.inference;
net.refresh_interval <= 3600
&& net.combined_queue_count > 0
&& net.fill_queue_size > 0
&& net.comp_queue_size > 0
&& net.tx_queue_size > 0
&& net.rx_queue_size > 0
&& net.frame_size > 0
&& net.frame_count > 0
&& table.http.http_server_bind_port > 0
&& inf.max_concurrent_flows > 0
&& inf.min_packets_for_inference > 0
&& inf.inference_interval_secs > 0
&& inf.inference_batch_size > 0
}
}

View File

@ -8,14 +8,27 @@ use lru::LruCache;
use std::num::NonZeroUsize;
use tokio::task;
use crate::model::geo_stats::GeoLocation;
use crate::utils::ip_address;
// TODO: Wire into statistics endpoint when GeoIP enrichment is enabled
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
#[allow(dead_code)]
pub struct GeoLocation {
pub country: Option<String>,
pub country_code: Option<String>,
pub city: Option<String>,
pub latitude: Option<f64>,
pub longitude: Option<f64>,
pub timezone: Option<String>,
}
#[allow(dead_code)]
pub struct GeoIpService {
reader: Arc<Reader<Vec<u8>>>,
cache: Arc<RwLock<LruCache<IpAddr, Option<GeoLocation>>>>,
}
#[allow(dead_code)]
impl GeoIpService {
pub fn new(db_name: &str) -> Result<Self, MaxMindDbError> {
let db_path = PathBuf::from("net-guardia/static/geo").join(db_name);
@ -90,8 +103,8 @@ impl GeoIpService {
let city_name = city.city.names.english
.map(|s| s.to_string());
let latitude = city.location.latitude.or(Some(0.0));
let longitude = city.location.longitude.or(Some(0.0));
let latitude = city.location.latitude;
let longitude = city.location.longitude;
let timezone = city.location.time_zone.map(|s| s.to_string());
GeoLocation {

View File

@ -42,8 +42,8 @@ impl SystemHealth {
networks: RwLock::new(Networks::new_with_refreshed_list()),
components: RwLock::new(Components::new_with_refreshed_list()),
broadcast_tx,
ingress_interface: config.ingress_ifname.clone(),
egress_interface: config.egress_ifname.clone(),
ingress_interface: config.network.ingress_ifname.clone(),
egress_interface: config.network.egress_ifname.clone(),
// management_interface: config.management_ifindex.clone(),
};
@ -112,8 +112,8 @@ impl SystemHealth {
) -> SystemHealthMetrics {
let timestamp = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_secs();
.map(|d| d.as_secs())
.unwrap_or(0);
let boot_time = System::boot_time();
let uptime_seconds = timestamp - boot_time;
@ -254,10 +254,8 @@ impl SystemHealth {
}
pub async fn get_current_metrics(&self) -> SystemHealthMetrics {
self.system.write().await.refresh_all();
self.networks.write().await.refresh(true);
self.components.write().await.refresh(true);
// Read last cached metrics from background task, don't refresh here
// to avoid racing with the background refresh_and_broadcast task
let system = self.system.read().await;
let networks = self.networks.read().await;
let components = self.components.read().await;

View File

@ -1,8 +1,7 @@
pub mod app_config;
pub mod health;
pub mod geoip;
pub mod ml_alert;
pub mod health;
pub mod statistics;
use std::sync::Arc;
use std::time::Duration;
@ -13,34 +12,37 @@ use tokio::sync::oneshot;
use crate::core::infrastructure::app_config::AppConfig;
use crate::core::infrastructure::health::SystemHealth;
use crate::core::infrastructure::ml_alert::MLAlert;
use crate::ml::config_loader::InferenceConfig;
use crate::ml::engine::Engine;
use crate::ml::feature_extractor::FlowFeatures;
use crate::ml::model_loader::MLModels;
use crate::core::ml::alert::MLAlert;
use crate::core::infrastructure::statistics::FlowStatistics;
use crate::core::ml::config_loader::InferenceConfig;
use crate::core::ml::engine::Engine;
use crate::model::ml_detection::EngineConfig;
use crate::core::ml::feature_extractor::FlowFeatures;
use crate::core::ml::model_loader::MLModels;
use crate::model::error::misc::MiscError;
use crate::model::error::system::SystemError;
use crate::model::error::Error;
use crate::model::log::system::SystemLog;
use crate::ml::traffic_logger::TrafficLogger;
use crate::core::ml::traffic_logger::TrafficLogger;
pub struct AppServices {
pub struct MLService {
pub health: Arc<SystemHealth>,
pub ml_alert: Arc<MLAlert>,
pub ml_models: Arc<MLModels>,
pub ml_engine: Arc<Engine>,
pub flow_statistics: Arc<FlowStatistics>,
shutdowns: SegQueue<oneshot::Sender<()>>,
}
impl AppServices {
impl MLService {
pub fn new(app_config: Arc<AppConfig>, inference_config: Arc<InferenceConfig>) -> Result<Self, Error> {
let health = SystemHealth::new(app_config.clone())?;
let ml_models = Arc::new(MLModels::load_models(&app_config, &inference_config)?);
let ml_alert = Arc::new(MLAlert::new());
let traffic_logger = if app_config.traffic_logging_mode {
let csv_path = app_config.traffic_log_csv_path.clone();
let traffic_logger = if app_config.inference.traffic_logging_mode {
let csv_path = app_config.inference.traffic_log_csv_path.clone();
let mut header = FlowFeatures::all_feature_names_owned();
header.push("Label".to_string());
let logger = TrafficLogger::new(&csv_path, header)
@ -51,23 +53,32 @@ impl AppServices {
None
};
let engine_config = EngineConfig {
max_flows: app_config.inference.max_concurrent_flows,
min_packets: app_config.inference.min_packets_for_inference,
batch_size: app_config.inference.inference_batch_size,
inference_interval_secs: app_config.inference.inference_interval_secs,
aggregator_window_secs: app_config.inference.aggregator_window_secs,
flow_timeout_us: 60_000_000,
};
let ml_engine = Arc::new(Engine::new(
ml_models.clone(),
inference_config.clone(),
ml_alert.clone(),
app_config.max_concurrent_flows,
app_config.min_packets_for_inference,
app_config.inference_batch_size,
app_config.inference_interval_secs,
app_config.aggregator_window_secs,
engine_config,
traffic_logger,
app_config.network.combined_queue_count,
));
let flow_statistics = Arc::new(FlowStatistics::new(ml_engine.clone()));
Ok(Self {
health: Arc::new(health),
ml_alert,
ml_models,
ml_engine,
flow_statistics,
shutdowns: SegQueue::new(),
})
}
@ -92,4 +103,4 @@ impl AppServices {
}
}
}
}
}

View File

@ -0,0 +1,78 @@
use std::sync::Arc;
use std::time;
use crate::core::ml::engine::Engine;
use crate::model::direction::Direction;
use crate::model::flow_stats::{FlowStatsEntry, FlowSubscription, StatsSummary};
pub struct FlowStatistics {
engine: Arc<Engine>,
}
impl FlowStatistics {
pub fn new(engine: Arc<Engine>) -> Self {
Self { engine }
}
pub fn get_all_flows(&self) -> Vec<FlowStatsEntry> {
let mut entries = Vec::new();
for tracker in self.engine.trackers() {
let t = tracker.lock();
entries.extend(t.get_flows().iter().map(FlowStatsEntry::from));
}
entries
}
pub fn get_filtered_flows(&self, sub: &FlowSubscription) -> Vec<FlowStatsEntry> {
let now_us = time::SystemTime::now()
.duration_since(time::UNIX_EPOCH)
.map(|d| d.as_micros() as u64)
.unwrap_or(0);
let mut flows = self.get_all_flows();
// Filter by direction
if let Some(dir) = &sub.direction {
flows.retain(|f| &f.direction == dir);
}
// Filter by time window
if let Some(window_secs) = sub.window_secs {
let cutoff = now_us.saturating_sub(window_secs * 1_000_000);
flows.retain(|f| f.last_seen_us >= cutoff);
}
// Sort by total bytes descending
flows.sort_by(|a, b| {
(b.fwd_bytes + b.bwd_bytes).cmp(&(a.fwd_bytes + a.bwd_bytes))
});
// Limit
if let Some(n) = sub.top_n {
flows.truncate(n.min(10000));
}
flows
}
pub fn get_top_flows(&self, n: usize) -> Vec<FlowStatsEntry> {
self.get_filtered_flows(&FlowSubscription {
direction: None,
window_secs: None,
top_n: Some(n),
interval_secs: None,
})
}
pub fn get_summary(&self) -> StatsSummary {
let flows = self.get_all_flows();
let total_flows = flows.len();
let total_bytes: u64 = flows.iter().map(|f| f.fwd_bytes + f.bwd_bytes).sum();
let total_packets: usize = flows.iter().map(|f| f.fwd_packets + f.bwd_packets).sum();
StatsSummary {
total_flows,
total_bytes,
total_packets,
}
}
}

View File

@ -45,7 +45,4 @@ impl AttackAggregator {
});
}
pub fn tracked_flows(&self) -> usize {
self.detections.len()
}
}

View File

@ -4,6 +4,8 @@ use tracing::error;
use crate::model::ml_detection::DetectionResult;
const ALERT_CHANNEL_CAPACITY: usize = 100;
#[derive(Debug, Clone, Serialize)]
pub struct AlertMessage {
pub timestamp: u64,
@ -23,14 +25,14 @@ impl AlertMessage {
pub fn from_detection_result(result: &DetectionResult) -> Self {
let timestamp = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_secs();
.map(|d| d.as_secs())
.unwrap_or(0);
Self {
timestamp,
flow_key: result.flow_key.clone(),
src_ip: result.flow_key_raw.src_ip.clone(),
dst_ip: result.flow_key_raw.dst_ip.clone(),
src_ip: result.flow_key_raw.src_ip_string(),
dst_ip: result.flow_key_raw.dst_ip_string(),
src_port: result.flow_key_raw.src_port,
dst_port: result.flow_key_raw.dst_port,
protocol: result.flow_key_raw.protocol,
@ -48,7 +50,7 @@ pub struct MLAlert {
impl MLAlert {
pub fn new() -> Self {
let (broadcast_tx, _) = broadcast::channel(100);
let (broadcast_tx, _) = broadcast::channel(ALERT_CHANNEL_CAPACITY);
MLAlert {
broadcast_tx,
@ -68,9 +70,6 @@ impl MLAlert {
}
}
pub fn has_subscribers(&self) -> bool {
self.broadcast_tx.receiver_count() > 0
}
}
impl Default for MLAlert {

View File

@ -4,7 +4,7 @@ use std::path::PathBuf;
use serde::{Deserialize, Serialize};
use crate::model::error::ml::MLError;
use crate::model::ml_detection::{AENormalization, ClipParams, PrecisionLevels};
use crate::model::ml_detection::ClipParams;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InferenceConfig {
@ -26,6 +26,15 @@ impl InferenceConfig {
.map_err(|_| MLError::ConfigLoadFailed { path: path.to_path_buf() })?;
let config: InferenceConfig = serde_json::from_str(&content)
.map_err(|e| MLError::ConfigParseFailed { reason: e.to_string() })?;
if config.ae_feature_names.is_empty() {
return Err(MLError::ConfigParseFailed { reason: "ae_feature_names is empty".into() });
}
if config.ae_scaler_mean.len() != config.ae_feature_names.len() {
return Err(MLError::ConfigParseFailed { reason: "scaler mean length mismatch".into() });
}
if config.ae_scaler_std.len() != config.ae_feature_names.len() {
return Err(MLError::ConfigParseFailed { reason: "scaler std length mismatch".into() });
}
Ok(config)
}
@ -41,6 +50,7 @@ impl InferenceConfig {
self.attack_labels.len()
}
#[allow(dead_code)]
pub fn get_attack_label(&self, id: usize) -> Option<&String> {
self.attack_labels.get(&id.to_string())
}

View File

@ -0,0 +1,203 @@
use std::sync::Arc;
use std::time::{Duration, Instant};
use parking_lot::Mutex;
use macros::log;
use tokio::sync::oneshot;
use tokio::time::interval;
use super::aggregator::AttackAggregator;
use super::config_loader::InferenceConfig;
use super::feature_extractor::FlowFeatures;
use super::flow_tracker::{FlowData, FlowTracker};
use super::inference::Inference;
use super::model_loader::MLModels;
use super::traffic_logger::TrafficLogger;
use super::alert::MLAlert;
use crate::model::log::ml::MLLog;
use crate::model::ml_detection::{EngineConfig, InferenceStats};
use crate::utils::packet_parser::parse_packet;
/// Per-thread tracker. Mutex is only contested during inference tick (every N seconds).
/// Hot path (process_packet): lock is uncontended → ~15ns.
pub type ThreadTracker = Arc<Mutex<FlowTracker>>;
pub struct Engine {
trackers: Vec<ThreadTracker>,
inference_pipeline: Arc<Inference>,
aggregator: Mutex<AttackAggregator>,
ml_alert: Arc<MLAlert>,
max_flows_per_thread: usize,
min_packets: usize,
batch_size: usize,
inference_interval_secs: u64,
flow_timeout_us: u64,
traffic_logger: Option<Arc<TrafficLogger>>,
}
impl Engine {
pub fn new(
models: Arc<MLModels>,
config: Arc<InferenceConfig>,
ml_alert: Arc<MLAlert>,
engine_config: EngineConfig,
traffic_logger: Option<Arc<TrafficLogger>>,
num_threads: u32,
) -> Self {
let inference_pipeline = Arc::new(Inference::new(models, config));
let min_detections = ((engine_config.aggregator_window_secs / engine_config.inference_interval_secs) / 2).max(1) as usize;
let aggregator = Mutex::new(AttackAggregator::new(engine_config.aggregator_window_secs, min_detections));
let max_flows_per_thread = engine_config.max_flows / num_threads.max(1) as usize;
let trackers: Vec<ThreadTracker> = (0..num_threads)
.map(|_| Arc::new(Mutex::new(FlowTracker::new(max_flows_per_thread))))
.collect();
Self {
trackers,
inference_pipeline,
aggregator,
ml_alert,
max_flows_per_thread,
min_packets: engine_config.min_packets,
batch_size: engine_config.batch_size,
inference_interval_secs: engine_config.inference_interval_secs,
flow_timeout_us: engine_config.flow_timeout_us,
traffic_logger,
}
}
pub fn tracker(&self, queue_id: u32) -> &ThreadTracker {
&self.trackers[queue_id as usize]
}
pub fn trackers(&self) -> &[ThreadTracker] {
&self.trackers
}
pub async fn run(self: Arc<Self>) -> oneshot::Sender<()> {
let (shutdown_tx, shutdown_rx) = oneshot::channel();
tokio::spawn(async move {
self.run_inference_loop(shutdown_rx).await;
});
shutdown_tx
}
async fn run_inference_loop(&self, mut shutdown_rx: oneshot::Receiver<()>) {
let mut ticker = interval(Duration::from_secs(self.inference_interval_secs));
loop {
tokio::select! {
_ = &mut shutdown_rx => break,
_ = ticker.tick() => {}
}
self.run_inference_tick();
}
}
fn run_inference_tick(&self) {
// Collect flows from all per-thread trackers.
// Each lock is held only for the duration of get_flows_for_inference (~microseconds).
// XSK threads are barely impacted since they process on different trackers.
let mut all_flows = Vec::new();
let mut total_count = 0;
for tracker in &self.trackers {
let t = tracker.lock();
total_count += t.flow_count();
all_flows.extend(t.get_flows_for_inference(self.min_packets));
}
log!(MLLog::FlowStats(
total_count,
all_flows.len(),
self.min_packets,
format!("{} trackers", self.trackers.len())
));
if all_flows.is_empty() {
return;
}
if let Some(ref logger) = self.traffic_logger {
self.log_traffic(&all_flows, logger);
} else {
self.run_inference(&all_flows);
}
for tracker in &self.trackers {
tracker.lock().cleanup_old_flows(self.flow_timeout_us);
}
}
fn log_traffic(&self, flows: &[FlowData], logger: &TrafficLogger) {
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());
}
}
fn run_inference(&self, flows: &[FlowData]) {
let batch = &flows[..flows.len().min(self.batch_size)];
log!(MLLog::RunningInference(batch.len()));
let start = Instant::now();
let results = self.inference_pipeline.infer_batch(batch);
let elapsed_us = start.elapsed().as_micros() as u64;
let stats = InferenceStats::from_results(&results, elapsed_us);
if results.len() != batch.len() {
log!(MLLog::InferenceResults(batch.len(), results.len()));
}
log!(MLLog::InferenceCompleted(
stats.total_flows,
stats.malicious_flows,
stats.benign_flows,
(elapsed_us as f64 / 1000.0) as u32,
stats.flows_per_second
));
{
let mut aggregator = self.aggregator.lock();
for result in &results {
if result.is_attack {
let should_alert =
aggregator.should_alert(&result.flow_key_raw, result.ae_score, result.threshold);
if should_alert {
log!(MLLog::ThreatDetected(
format!("{:?}", result.direction),
result.flow_key.clone(),
result.attack_type.clone().unwrap_or_else(|| "UNKNOWN".to_string()),
result.confidence,
result.ae_score,
));
self.ml_alert.broadcast_alert(result);
}
}
}
aggregator.cleanup();
}
}
/// Called by XSK threads. Lock is per-thread, uncontended on hot path.
pub fn process_packet(tracker: &Mutex<FlowTracker>, packet_data: &[u8], is_ingress: bool) {
match parse_packet(packet_data) {
Some((packet_info, payload_start)) => {
let payload = packet_data.get(payload_start..).unwrap_or(&[]);
tracker.lock().process_packet(packet_info, is_ingress, payload);
}
None => log!(MLLog::ParsePacketFailed(packet_data.len())),
}
}
}

View File

@ -1,5 +1,7 @@
use std::collections::HashMap;
use common::define::tcp_flags::*;
use super::flow_tracker::FlowData;
use crate::model::ml_detection::{ClipParams, PacketData};
@ -60,10 +62,10 @@ impl FlowFeatures {
let (bwd_iat_max, bwd_iat_min, bwd_iat_mean, bwd_iat_std) = compute_stats(&bwd_iats);
// 30-37
let fwd_psh = flow.fwd_packets.iter().filter(|p| p.flags.psh).count() as f64;
let bwd_psh = flow.bwd_packets.iter().filter(|p| p.flags.psh).count() as f64;
let fwd_urg = flow.fwd_packets.iter().filter(|p| p.flags.urg).count() as f64;
let bwd_urg = flow.bwd_packets.iter().filter(|p| p.flags.urg).count() as f64;
let fwd_psh = flow.fwd_packets.iter().filter(|p| p.flags & TCP_PSH != 0).count() as f64;
let bwd_psh = flow.bwd_packets.iter().filter(|p| p.flags & TCP_PSH != 0).count() as f64;
let fwd_urg = flow.fwd_packets.iter().filter(|p| p.flags & TCP_URG != 0).count() as f64;
let bwd_urg = flow.bwd_packets.iter().filter(|p| p.flags & TCP_URG != 0).count() as f64;
// 38-55
let all_lengths: Vec<f64> = flow
@ -84,7 +86,7 @@ impl FlowFeatures {
.fwd_packets
.iter()
.filter(|p| p.payload_length > 0)
.map(|p| p.header_length as f64)
.map(|p| p.payload_length as f64)
.collect();
// 70-73
@ -216,10 +218,6 @@ impl FlowFeatures {
}
}
pub fn get_features_content(&self) -> &Vec<f64> {
&self.features
}
pub fn all_feature_names() -> Vec<&'static str> {
vec![
"Destination Port",

View File

@ -1,11 +1,11 @@
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time;
use common::model::event::Event;
use common::define::tcp_flags::*;
use crate::model::direction::Direction;
use crate::model::ml_detection::{BulkState, FlowKey, PacketData};
use crate::model::user_packet::UserPacket;
#[derive(Debug, Clone)]
pub struct FlowData {
@ -37,12 +37,12 @@ pub struct FlowData {
}
impl FlowData {
pub fn new(flow_key: FlowKey, first_packet: &Event, direction: Direction) -> Self {
pub fn new(flow_key: FlowKey, first_packet: &UserPacket, direction: Direction) -> Self {
Self {
flow_key,
direction,
start_time_us: first_packet.timestamp_us(),
last_time_us: first_packet.timestamp_us(),
start_time_us: first_packet.timestamp_us,
last_time_us: first_packet.timestamp_us,
fwd_packets: Vec::new(),
fwd_total_bytes: 0,
fwd_header_bytes: 0,
@ -57,89 +57,64 @@ impl FlowData {
urg_count: 0,
cwe_count: 0,
ece_count: 0,
init_win_bytes_fwd: if first_packet.is_forward() {
first_packet.tcp_window_size()
} else {
0
},
init_win_bytes_bwd: if !first_packet.is_forward() {
first_packet.tcp_window_size()
} else {
0
},
init_win_bytes_fwd: if first_packet.is_forward { first_packet.tcp_window_size } else { 0 },
init_win_bytes_bwd: if !first_packet.is_forward { first_packet.tcp_window_size } else { 0 },
active_periods: Vec::new(),
idle_periods: Vec::new(),
last_packet_time: first_packet.timestamp_us(),
last_packet_time: first_packet.timestamp_us,
fwd_bulk_state: BulkState::default(),
bwd_bulk_state: BulkState::default(),
}
}
pub fn add_packet(&mut self, packet: &Event) {
pub fn add_packet(&mut self, packet: &UserPacket) {
const MAX_PACKETS_PER_DIRECTION: usize = 1000;
const MAX_PERIODS: usize = 10000;
let packet_data = PacketData {
timestamp_us: packet.timestamp_us(),
length: packet.packet_length(),
header_length: packet.header_length(),
payload_length: packet.payload_length(),
flags: packet.tcp_flags().clone(),
timestamp_us: packet.timestamp_us,
length: packet.packet_length,
header_length: packet.header_length,
payload_length: packet.payload_length,
flags: packet.tcp_flags,
};
if packet.tcp_flags().fin {
self.fin_count += 1;
}
if packet.tcp_flags().syn {
self.syn_count += 1;
}
if packet.tcp_flags().rst {
self.rst_count += 1;
}
if packet.tcp_flags().psh {
self.psh_count += 1;
}
if packet.tcp_flags().ack {
self.ack_count += 1;
}
if packet.tcp_flags().urg {
self.urg_count += 1;
}
if packet.tcp_flags().cwr {
self.cwe_count += 1;
}
if packet.tcp_flags().ece {
self.ece_count += 1;
}
if packet.tcp_flags & TCP_FIN != 0 { self.fin_count += 1; }
if packet.tcp_flags & TCP_SYN != 0 { self.syn_count += 1; }
if packet.tcp_flags & TCP_RST != 0 { self.rst_count += 1; }
if packet.tcp_flags & TCP_PSH != 0 { self.psh_count += 1; }
if packet.tcp_flags & TCP_ACK != 0 { self.ack_count += 1; }
if packet.tcp_flags & TCP_URG != 0 { self.urg_count += 1; }
if packet.tcp_flags & TCP_CWR != 0 { self.cwe_count += 1; }
if packet.tcp_flags & TCP_ECE != 0 { self.ece_count += 1; }
let iat = packet.timestamp_us().saturating_sub(self.last_packet_time);
let iat = packet.timestamp_us.saturating_sub(self.last_packet_time);
const IDLE_THRESHOLD_US: u64 = 1_000_000;
if iat > IDLE_THRESHOLD_US {
self.idle_periods.push(iat);
if self.idle_periods.len() < MAX_PERIODS { self.idle_periods.push(iat); }
} else if iat > 0 {
self.active_periods.push(iat);
if self.active_periods.len() < MAX_PERIODS { self.active_periods.push(iat); }
}
self.last_packet_time = packet.timestamp_us();
self.last_time_us = packet.timestamp_us();
self.last_packet_time = packet.timestamp_us;
self.last_time_us = packet.timestamp_us;
if packet.is_forward() {
self.fwd_packets.push(packet_data.clone());
self.fwd_total_bytes += packet.packet_length() as u64;
self.fwd_header_bytes += packet.header_length() as u64;
if self.init_win_bytes_fwd == 0 {
self.init_win_bytes_fwd = packet.tcp_window_size();
if packet.is_forward {
if self.fwd_packets.len() < MAX_PACKETS_PER_DIRECTION {
self.fwd_packets.push(packet_data.clone());
}
self.fwd_total_bytes += packet.packet_length as u64;
self.fwd_header_bytes += packet.header_length as u64;
if self.init_win_bytes_fwd == 0 { self.init_win_bytes_fwd = packet.tcp_window_size; }
Self::update_bulk_state(&mut self.fwd_bulk_state, &packet_data);
} else {
self.bwd_packets.push(packet_data.clone());
self.bwd_total_bytes += packet.packet_length() as u64;
self.bwd_header_bytes += packet.header_length() as u64;
if self.init_win_bytes_bwd == 0 {
self.init_win_bytes_bwd = packet.tcp_window_size();
if self.bwd_packets.len() < MAX_PACKETS_PER_DIRECTION {
self.bwd_packets.push(packet_data.clone());
}
self.bwd_total_bytes += packet.packet_length as u64;
self.bwd_header_bytes += packet.header_length as u64;
if self.init_win_bytes_bwd == 0 { self.init_win_bytes_bwd = packet.tcp_window_size; }
Self::update_bulk_state(&mut self.bwd_bulk_state, &packet_data);
}
}
@ -189,55 +164,39 @@ impl FlowData {
}
}
/// Per-thread flow tracker. No locks — each XSK thread owns one.
/// RSS guarantees the same flow always goes to the same thread.
pub struct FlowTracker {
flows: Arc<Mutex<HashMap<FlowKey, FlowData>>>,
flows: HashMap<FlowKey, FlowData>,
max_flows: usize,
}
impl FlowTracker {
pub fn new(max_flows: usize) -> Self {
Self {
flows: Arc::new(Mutex::new(HashMap::new())),
flows: HashMap::new(),
max_flows,
}
}
pub fn process_packet(&self, mut packet: Event, is_ingress: bool, payload: &[u8]) {
let direction = if is_ingress {
Direction::Ingress
} else {
Direction::Egress
};
pub fn process_packet(&mut self, mut packet: UserPacket, is_ingress: bool, payload: &[u8]) {
let direction = if is_ingress { Direction::Ingress } else { Direction::Egress };
let packet_key = FlowKey::from_packet(&packet);
let proto = packet_key.protocol;
let src_port = packet_key.src_port;
let dst_port = packet_key.dst_port;
let reversed_key = packet_key.clone().reverse();
let Ok(mut flows) = self.flows.lock() else {
return;
};
// Try-both: canonical key is whichever orientation already exists in the flow table.
// For new flows, identify the initiator using (in priority order):
// 1. TCP SYN / SYN+ACK flags
// 2. DPI: TLS ClientHello/ServerHello, HTTP request/response, DNS QR bit
// 3. Best effort: use packet as-is
let (actual_key, is_forward) = if flows.contains_key(&packet_key) {
let (actual_key, is_forward) = if self.flows.contains_key(&packet_key) {
(packet_key, true)
} else if flows.contains_key(&reversed_key) {
} else if self.flows.contains_key(&reversed_key) {
(reversed_key, false)
} else {
let flags = packet.tcp_flags();
if flags.syn && flags.ack {
// Normal: Server (egress side) sends SYN+ACK, packet arrives on ingress → reverse
// Bot attack: Client (egress side) sends SYN+ACK, packet arrives on egress → keep as-is
if is_ingress {
(reversed_key, false)
} else {
(packet_key, true)
}
} else if flags.syn {
let has_syn = packet.tcp_flags & TCP_SYN != 0;
let has_ack = packet.tcp_flags & TCP_ACK != 0;
if has_syn && has_ack {
if is_ingress { (reversed_key, false) } else { (packet_key, true) }
} else if has_syn {
(packet_key, true)
} else {
match detect_initiator(payload, proto, src_port, dst_port) {
@ -248,74 +207,60 @@ impl FlowTracker {
}
};
packet.set_is_forward(is_forward);
// `direction` should reflect the initiator's interface.
// If this packet is backward (is_forward = false), the initiator is on the opposite side.
packet.is_forward = is_forward;
let initiator_direction = if is_forward { direction } else { direction.flip() };
let flow = flows
let flow = self.flows
.entry(actual_key.clone())
.or_insert_with(|| FlowData::new(actual_key, &packet, initiator_direction));
flow.add_packet(&packet);
if flows.len() > self.max_flows {
if let Some(key) = flows.keys().next().cloned() {
flows.remove(&key);
if self.flows.len() > self.max_flows {
if let Some(oldest_key) = self.flows.iter()
.min_by_key(|(_, flow)| flow.last_time_us)
.map(|(k, _)| k.clone())
{
self.flows.remove(&oldest_key);
}
}
}
pub fn get_flows_snapshot(&self) -> Vec<FlowData> {
let Ok(flows) = self.flows.lock() else {
return Vec::new();
};
flows.values().cloned().collect()
/// Take all flows out, leaving this tracker empty. Lock-free.
pub fn drain_flows(&mut self) -> Vec<FlowData> {
self.flows.drain().map(|(_, v)| v).collect()
}
/// Get a snapshot without draining.
pub fn get_flows(&self) -> Vec<FlowData> {
self.flows.values().cloned().collect()
}
pub fn get_flows_for_inference(&self, min_packets: usize) -> Vec<FlowData> {
let Ok(flows) = self.flows.lock() else {
return Vec::new();
};
flows
self.flows
.values()
.filter(|flow| flow.packet_count() >= min_packets)
.cloned()
.collect()
}
pub fn cleanup_old_flows(&self, max_age_us: u64) {
pub fn cleanup_old_flows(&mut self, max_age_us: u64) {
let now = time::SystemTime::now()
.duration_since(time::UNIX_EPOCH)
.map(|d| d.as_micros() as u64)
.unwrap_or(0);
let Ok(mut flows) = self.flows.lock() else {
return;
};
flows.retain(|_, flow| now.saturating_sub(flow.last_time_us) < max_age_us);
self.flows.retain(|_, flow| now.saturating_sub(flow.last_time_us) < max_age_us);
}
pub fn flow_count(&self) -> usize {
let Ok(flows) = self.flows.lock() else {
return 0;
};
flows.len()
self.flows.len()
}
}
/// Inspect payload bytes to determine which side is the flow initiator.
/// Returns Some(true) if this packet is from the initiator, Some(false) if from the responder,
/// or None if the payload gives no useful signal.
fn detect_initiator(payload: &[u8], protocol: u8, src_port: u16, dst_port: u16) -> Option<bool> {
if payload.is_empty() {
return None;
}
if payload.is_empty() { return None; }
// TLS: record type 0x16 (Handshake), byte 5 = handshake type
// 0x01 = ClientHello → this side is the initiator
// 0x02 = ServerHello → this side is the responder
if payload.len() >= 6 && payload[0] == 0x16 {
return match payload[5] {
0x01 => Some(true),
@ -324,8 +269,6 @@ fn detect_initiator(payload: &[u8], protocol: u8, src_port: u16, dst_port: u16)
};
}
// HTTP: request line starts with a method verb (initiator),
// response starts with "HTTP/" (responder)
if payload.len() >= 5 {
if payload.starts_with(b"GET ")
|| payload.starts_with(b"POST ")
@ -342,8 +285,6 @@ fn detect_initiator(payload: &[u8], protocol: u8, src_port: u16, dst_port: u16)
}
}
// DNS over UDP (port 53): flags byte 2, MSB = QR bit
// 0 = query (initiator), 1 = response (responder)
if protocol == 17 && (src_port == 53 || dst_port == 53) && payload.len() >= 3 {
return Some((payload[2] >> 7) == 0);
}

View File

@ -1,4 +1,4 @@
use std::sync::{Arc, Mutex};
use std::sync::Arc;
use macros::log;
use tract_onnx::prelude::*;
@ -40,7 +40,7 @@ impl Inference {
let cls_input = self.build_classifier_input(&ae_features, ae_score);
let (attack_type, confidence) = match self.run_classifier(&cls_input) {
let (attack_type, confidence) = match self.run_classifier(cls_input) {
Ok(result) => result,
Err(e) => {
log!(MLLog::InferenceFailed("LightGBM".to_string(), e.to_string()));
@ -52,9 +52,9 @@ impl Inference {
let flow_key = format!(
"{}:{} -> {}:{} (proto {}) [{}]",
flow.flow_key.src_ip,
flow.flow_key.src_ip_string(),
flow.flow_key.src_port,
flow.flow_key.dst_ip,
flow.flow_key.dst_ip_string(),
flow.flow_key.dst_port,
flow.flow_key.protocol,
flow.direction
@ -77,6 +77,7 @@ impl Inference {
features.winsorize(&self.config.ae_clip_params, &self.config.ae_feature_names);
features.normalize(&self.config.ae_scaler_mean, &self.config.ae_scaler_std);
features.clip(self.config.ae_post_clip_min, self.config.ae_post_clip_max);
// Note: f64->f32 precision loss is acceptable for ML inference
features.features.iter().map(|&x| x as f32).collect()
}
@ -97,24 +98,24 @@ impl Inference {
}
fn run_autoencoder(&self, input: &tract_ndarray::Array2<f32>) -> TractResult<f32> {
let input_tensor = input.clone().into_tensor();
let result = self
.models
.deep_autoencoder
.run(tvec![input.clone().into_tensor().into()])?;
.run(tvec![input_tensor.into()])?;
let output = result[0]
.to_array_view::<f32>()?
.into_dimensionality::<tract_ndarray::Ix2>()?;
let diff = input - &output;
let mse = (&diff * &diff).sum() / self.config.ae_feature_names.len() as f32;
let mse = (&diff * &diff).sum() / output.len() as f32;
Ok(mse)
}
fn run_classifier(&self, input: &tract_ndarray::Array2<f32>) -> TractResult<(String, f32)> {
let input_tensor = input.clone().into_tensor();
let result = self.models.classifier.run(tvec![input_tensor.into()])?;
fn run_classifier(&self, input: tract_ndarray::Array2<f32>) -> TractResult<(String, f32)> {
let result = self.models.classifier.run(tvec![input.into_tensor().into()])?;
let output = result[0].to_array_view::<f32>()?;

View File

@ -1,3 +1,4 @@
pub mod alert;
pub mod model_loader;
pub mod config_loader;
pub mod flow_tracker;

View File

@ -14,36 +14,21 @@ pub struct MLModels {
impl MLModels {
pub fn load_models(app_config: &Arc<AppConfig>, inference_config: &Arc<InferenceConfig>) -> Result<Self, MLError> {
Ok(Self {
deep_autoencoder: Self::loader(&app_config.deep_autoencoder_name, inference_config.num_ae_features())?,
classifier: Self::loader(&app_config.classifier_name, inference_config.num_classifier_features())?
deep_autoencoder: Self::loader(&app_config.inference.deep_autoencoder_name, inference_config.num_ae_features())?,
classifier: Self::loader(&app_config.inference.classifier_name, inference_config.num_classifier_features())?
})
}
pub fn loader(model: &str, features: usize) -> Result<RunnableModel, MLError> {
let model_path = PathBuf::from("models").join(model);
let mut model = onnx()
.model_for_path(&model_path)
.map_err(|_| {
MLError::ModelLoadFailed { path: model_path.clone() }
})?;
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(&[1, features]).into())?;
Ok(model.into_optimized()?.into_runnable()?)
};
model.set_input_fact(0, f32::fact(&[1, features]).into())
.map_err(|_| {
MLError::ModelLoadFailed { path: model_path.clone() }
})?;
let runnable_model = model
.into_optimized()
.map_err(|_| {
MLError::ModelLoadFailed { path: model_path.clone() }
})?
.into_runnable()
.map_err(|_| {
MLError::ModelLoadFailed { path: model_path }
})?;
Ok(runnable_model)
load().map_err(|_| MLError::ModelLoadFailed { path: model_path })
}
pub fn get_model_info(&self, name: &str) -> String {

View File

@ -37,10 +37,9 @@ impl TrafficLogger {
}
pub fn log_row(&self, record: Vec<String>) {
match self.sender.try_send(record) {
Ok(_) => {}
Err(TrySendError::Full(_)) => {}
Err(TrySendError::Disconnected(_)) => {}
if let Err(TrySendError::Disconnected(_)) = self.sender.try_send(record) {
eprintln!("[traffic-logger] channel disconnected");
}
// Full is ok - just drop the record
}
}

View File

@ -1,3 +1,4 @@
pub mod ebpf;
pub mod infrastructure;
pub mod ml;
pub mod system;

View File

@ -1,18 +1,19 @@
use std::collections::HashMap;
use std::sync::Arc;
use actix_web::web::route;
use actix_web::{web, App, HttpServer};
use aya::maps::{MapData, ProgramArray};
use aya::maps::{Array, MapData, ProgramArray};
use aya::programs::{Xdp, XdpFlags};
use aya::Ebpf;
use aya_log::EbpfLogger;
use common::define::program_array::*;
use common::define::pipeline::*;
use macros::log;
use crate::core::ebpf::EbpfServices;
use crate::core::infrastructure::app_config::AppConfig;
use crate::core::infrastructure::AppServices;
use crate::ml::config_loader::InferenceConfig;
use crate::core::infrastructure::MLService;
use crate::core::ml::config_loader::InferenceConfig;
use crate::model::error::ebpf::EbpfError;
use crate::model::error::http::HttpError;
use crate::model::error::misc::MiscError;
@ -20,28 +21,40 @@ use crate::model::error::Error;
use crate::model::log::ml::MLLog;
use crate::model::log::system::SystemLog;
use crate::utils::logging::Logging;
use crate::web::api::{control, default, health, misc, ml_alert};
use crate::web::api::{acl, filter, rate_limit as rate_limit_api, stats, health as health_api, ml, system as system_api, default, ws};
/// Maps stage name (from config.toml) to (function_name, stage_id)
fn stage_registry() -> HashMap<&'static str, (&'static str, u32)> {
HashMap::from([
("access_control", ("access_control", STAGE_ACCESS_CONTROL)),
("rate_limit", ("rate_limit", STAGE_RATE_LIMIT)),
("service", ("protocol_filter", STAGE_SERVICE)),
])
}
pub struct System {
pub app_config: Arc<AppConfig>,
pub inference_config: Arc<InferenceConfig>,
pub ebpf_services: Arc<EbpfServices>,
pub app_services: Arc<AppServices>,
pub app_services: Arc<MLService>,
pub ingress_ebpf: Ebpf,
pub egress_ebpf: Ebpf,
#[allow(dead_code)]
ingress_program_array: ProgramArray<MapData>,
#[allow(dead_code)]
egress_program_array: ProgramArray<MapData>,
}
impl System {
pub async fn new() -> Result<Self, Error> {
let (mut ingress_ebpf, ingress_program_array) = System::get_ingress_ebpf()?;
let (mut egress_ebpf, egress_program_array) = System::get_egress_ebpf()?;
let mut ingress_ebpf = Self::load_ebpf("ingress")?;
let mut egress_ebpf = Self::load_ebpf("egress")?;
let app_config = Arc::new(AppConfig::new()?);
let inference_config = Arc::new(InferenceConfig::load_file(&app_config.models_config_name)?);
let ingress_program_array = Self::configure_ingress_pipeline(
&mut ingress_ebpf,
&app_config.pipeline.ingress,
)?;
let inference_config = Arc::new(InferenceConfig::load_file(&app_config.inference.models_config_name)?);
let ebpf_services = Arc::new(EbpfServices::new(
app_config.clone(),
@ -49,9 +62,9 @@ impl System {
&mut egress_ebpf,
)?);
let app_services = Arc::new(AppServices::new(app_config.clone(), inference_config.clone())?);
let app_services = Arc::new(MLService::new(app_config.clone(), inference_config.clone())?);
let system = System {
Ok(System {
app_config,
inference_config,
ebpf_services,
@ -59,9 +72,7 @@ impl System {
ingress_ebpf,
egress_ebpf,
ingress_program_array,
egress_program_array,
};
Ok(system)
})
}
pub async fn run(&mut self) -> Result<(), Error> {
@ -110,29 +121,25 @@ impl System {
}
fn attach_ebpf(&mut self) -> Result<(), Error> {
let config = self.app_config.config.clone();
let ingress_ifname = config.ingress_ifname;
let egress_ifname = config.egress_ifname;
let ingress_ifname = self.app_config.network.ingress_ifname.clone();
let egress_ifname = self.app_config.network.egress_ifname.clone();
Self::set_memory_limit()?;
let ingress_xdp: &mut Xdp = self
.ingress_ebpf
Self::attach_xdp(&mut self.ingress_ebpf, &ingress_ifname, true)?; // already loaded in configure_ingress_pipeline
Self::attach_xdp(&mut self.egress_ebpf, &egress_ifname, false)?; // load now
Ok(())
}
fn attach_xdp(ebpf: &mut Ebpf, ifname: &str, already_loaded: bool) -> Result<(), Error> {
let xdp: &mut Xdp = ebpf
.program_mut("net_guardia")
.ok_or(EbpfError::ProgramNotFound)?
.try_into()
.map_err(EbpfError::GetProgramFailed)?;
let egress_xdp: &mut Xdp = self
.egress_ebpf
.program_mut("net_guardia")
.ok_or(EbpfError::ProgramNotFound)?
.try_into()
.map_err(EbpfError::GetProgramFailed)?;
ingress_xdp.load().map_err(EbpfError::LoadProgramFailed)?;
ingress_xdp
.attach(&ingress_ifname, XdpFlags::DRV_MODE)
.map_err(EbpfError::AttachProgramFailed)?;
egress_xdp.load().map_err(EbpfError::LoadProgramFailed)?;
egress_xdp
.attach(&egress_ifname, XdpFlags::DRV_MODE)
if !already_loaded {
xdp.load().map_err(EbpfError::LoadProgramFailed)?;
}
xdp.attach(ifname, XdpFlags::DRV_MODE)
.map_err(EbpfError::AttachProgramFailed)?;
Ok(())
}
@ -141,14 +148,16 @@ impl System {
let app_config = self.app_config.clone();
let inference_config = self.inference_config.clone();
let access_control = self.ebpf_services.access_control.clone();
let service = self.ebpf_services.service.clone();
let statistics = self.ebpf_services.statistics.clone();
let protocol_filter = self.ebpf_services.protocol_filter.clone();
let rate_limit = self.ebpf_services.rate_limit.clone();
let health = self.app_services.health.clone();
let ml_alert = self.app_services.ml_alert.clone();
let port = self.app_config.http_server_bind_port;
let flow_statistics = self.app_services.flow_statistics.clone();
let port = self.app_config.http.http_server_bind_port;
HttpServer::new(move || {
let cors = actix_cors::Cors::default()
.allow_any_origin()
.allowed_origin("http://localhost:8080")
.allowed_origin("http://127.0.0.1:8080")
.allow_any_method()
.allow_any_header()
.max_age(3600);
@ -157,14 +166,22 @@ impl System {
.app_data(web::Data::from(app_config.clone()))
.app_data(web::Data::from(inference_config.clone()))
.app_data(web::Data::from(access_control.clone()))
.app_data(web::Data::from(service.clone()))
.app_data(web::Data::from(statistics.clone()))
.app_data(web::Data::from(protocol_filter.clone()))
.app_data(web::Data::from(rate_limit.clone()))
.app_data(web::Data::from(health.clone()))
.app_data(web::Data::from(ml_alert.clone()))
.service(control::initialize())
.service(ml_alert::initialize())
.service(health::initialize())
.service(misc::initialize())
.app_data(web::Data::from(flow_statistics.clone()))
.service(
web::scope("/api")
.service(acl::initialize())
.service(filter::initialize())
.service(rate_limit_api::initialize())
.service(stats::initialize())
.service(health_api::initialize())
.service(ml::initialize())
.service(system_api::initialize())
)
.service(ws::initialize())
.default_service(route().to(default::default_route))
})
.bind(format!("0.0.0.0:{}", port))
@ -175,54 +192,86 @@ impl System {
Ok(())
}
fn get_ingress_ebpf() -> Result<(Ebpf, ProgramArray<MapData>), Error> {
let mut ingress_ebpf = Ebpf::load(aya::include_bytes_aligned!(concat!(
env!("OUT_DIR"),
"/net-guardia-ingress"
)))
.map_err(EbpfError::EbpfNotFound)?;
let program_array = ingress_ebpf.take_map("PROGRAM_ARRAY").ok_or(EbpfError::MapNotFound)?;
let mut program_array = ProgramArray::try_from(program_array).map_err(EbpfError::MapOperationError)?;
Self::load_program(
&mut ingress_ebpf,
&mut program_array,
"access_control",
ingress::ACCESS_CONTROL,
)?;
Self::load_program(&mut ingress_ebpf, &mut program_array, "service", ingress::SERVICE)?;
Self::load_program(&mut ingress_ebpf, &mut program_array, "statistics", ingress::STATISTICS)?;
Self::load_program(
&mut ingress_ebpf,
&mut program_array,
"transmission",
ingress::TRANSMISSION,
)?;
Ok((ingress_ebpf, program_array))
fn load_ebpf(name: &str) -> Result<Ebpf, Error> {
let bytes = match name {
"ingress" => aya::include_bytes_aligned!(concat!(env!("OUT_DIR"), "/net-guardia-ingress")),
"egress" => aya::include_bytes_aligned!(concat!(env!("OUT_DIR"), "/net-guardia-egress")),
_ => return Err(EbpfError::ProgramNotFound.into()),
};
Ok(Ebpf::load(bytes).map_err(EbpfError::EbpfNotFound)?)
}
fn get_egress_ebpf() -> Result<(Ebpf, ProgramArray<MapData>), Error> {
let mut egress_ebpf = Ebpf::load(aya::include_bytes_aligned!(concat!(
env!("OUT_DIR"),
"/net-guardia-egress"
)))
.map_err(EbpfError::EbpfNotFound)?;
let program_array = egress_ebpf.take_map("PROGRAM_ARRAY").ok_or(EbpfError::MapNotFound)?;
let mut program_array = ProgramArray::try_from(program_array).map_err(EbpfError::MapOperationError)?;
Self::load_program(&mut egress_ebpf, &mut program_array, "statistics", egress::STATISTICS)?;
Self::load_program(
&mut egress_ebpf,
&mut program_array,
"transmission",
egress::TRANSMISSION,
)?;
Ok((egress_ebpf, program_array))
/// Configure the ingress pipeline based on config.toml [Pipeline] section.
/// Loads each stage program into ProgramArray and wires NEXT_STAGE map.
fn configure_ingress_pipeline(
ebpf: &mut Ebpf,
stages: &[String],
) -> Result<ProgramArray<MapData>, Error> {
let registry = stage_registry();
// Load entry point program BEFORE taking maps — verifier needs map fds at load time
let entry: &mut Xdp = ebpf
.program_mut("net_guardia")
.ok_or(EbpfError::ProgramNotFound)?
.try_into()
.map_err(EbpfError::GetProgramFailed)?;
entry.load().map_err(EbpfError::LoadProgramFailed)?;
// Take maps
let pa_map = ebpf.take_map("PROGRAM_ARRAY").ok_or(EbpfError::MapNotFound)?;
let mut program_array = ProgramArray::try_from(pa_map).map_err(EbpfError::MapOperationError)?;
let ns_map = ebpf.take_map("NEXT_STAGE").ok_or(EbpfError::MapNotFound)?;
let mut next_stage = Array::<MapData, u32>::try_from(ns_map).map_err(EbpfError::MapOperationError)?;
// Load transmission (always present at STAGE_TRANSMISSION)
Self::load_program(ebpf, &mut program_array, "transmission", STAGE_TRANSMISSION)?;
if stages.is_empty() {
// Empty pipeline: entry → transmission
next_stage
.set(STAGE_ENTRY as u32, STAGE_TRANSMISSION, 0)
.map_err(EbpfError::MapOperationError)?;
return Ok(program_array);
}
// Load each stage and assign a slot (starting from slot 1)
let mut slots: Vec<(u32, u32)> = Vec::new(); // (stage_id, slot_index)
for (i, stage_name) in stages.iter().enumerate() {
let (func_name, stage_id) = registry
.get(stage_name.as_str())
.ok_or(EbpfError::ProgramNotFound)?;
let slot = (i + 1) as u32; // slots 1, 2, 3, ...
Self::load_program(ebpf, &mut program_array, func_name, slot)?;
slots.push((*stage_id, slot));
}
// Wire NEXT_STAGE: entry → first slot
next_stage
.set(STAGE_ENTRY as u32, slots[0].1, 0)
.map_err(EbpfError::MapOperationError)?;
// Wire each stage to the next
for i in 0..slots.len() {
let (stage_id, _) = slots[i];
let next_slot = if i + 1 < slots.len() {
slots[i + 1].1
} else {
STAGE_TRANSMISSION
};
next_stage
.set(stage_id as u32, next_slot, 0)
.map_err(EbpfError::MapOperationError)?;
}
Ok(program_array)
}
fn load_program(
ebpf: &mut Ebpf,
program_array: &mut ProgramArray<MapData>,
function_name: &str,
index: u32,
slot: u32,
) -> Result<(), Error> {
let program: &mut Xdp = ebpf
.program_mut(function_name)
@ -231,7 +280,9 @@ impl System {
.map_err(EbpfError::MapOperationError)?;
program.load().map_err(EbpfError::AttachProgramFailed)?;
let fd = program.fd().map_err(|_| EbpfError::UnknownError)?;
program_array.set(index, fd, 0).map_err(EbpfError::MapOperationError)?;
program_array
.set(slot, fd, 0)
.map_err(EbpfError::MapOperationError)?;
Ok(())
}

View File

@ -2,7 +2,6 @@ mod core;
mod model;
mod utils;
mod web;
mod ml;
use crate::core::system::System;
use crate::model::error::Error;

View File

@ -1,198 +0,0 @@
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use macros::log;
use tokio::sync::oneshot;
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;
use super::traffic_logger::TrafficLogger;
use crate::core::infrastructure::ml_alert::MLAlert;
use crate::model::log::ml::MLLog;
use crate::model::ml_detection::{EngineStats, InferenceStats};
use crate::utils::packet_parser::parse_packet;
pub struct Engine {
flow_tracker: Arc<FlowTracker>,
inference_pipeline: Arc<Inference>,
aggregator: Arc<Mutex<AttackAggregator>>,
ml_alert: Arc<MLAlert>,
min_packets: usize,
batch_size: usize,
inference_interval_secs: u64,
traffic_logger: Option<Arc<TrafficLogger>>,
}
impl Engine {
pub fn new(
models: Arc<MLModels>,
config: Arc<InferenceConfig>,
ml_alert: Arc<MLAlert>,
max_flows: usize,
min_packets: usize,
batch_size: usize,
interval_secs: u64,
window_secs: u64,
traffic_logger: Option<Arc<TrafficLogger>>,
) -> Self {
let flow_tracker = Arc::new(FlowTracker::new(max_flows));
let inference_pipeline = Arc::new(Inference::new(models, 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)));
Self {
flow_tracker,
inference_pipeline,
aggregator,
ml_alert,
min_packets,
batch_size,
inference_interval_secs: interval_secs,
traffic_logger,
}
}
pub async fn run(self: Arc<Self>) -> oneshot::Sender<()> {
let (shutdown_tx, shutdown_rx) = oneshot::channel();
tokio::spawn(async move {
self.run_inference_loop(shutdown_rx).await;
});
shutdown_tx
}
pub fn get_flow_tracker(&self) -> Arc<FlowTracker> {
self.flow_tracker.clone()
}
async fn run_inference_loop(&self, mut shutdown_rx: oneshot::Receiver<()>) {
let mut ticker = interval(Duration::from_secs(self.inference_interval_secs));
loop {
tokio::select! {
_ = &mut shutdown_rx => break,
_ = ticker.tick() => {}
}
let total_flows = self.flow_tracker.flow_count();
let all_flows = self.flow_tracker.get_flows_snapshot();
let packet_counts: Vec<usize> = all_flows.iter().map(|f| f.packet_count()).collect();
let flows = self.flow_tracker.get_flows_for_inference(self.min_packets);
log!(MLLog::FlowStats(
total_flows,
flows.len(),
self.min_packets,
format!("{:?}", packet_counts)
));
if flows.is_empty() {
log!(MLLog::InferenceSkipped(format!(
"No flows with sufficient packets (total flows: {}, min packets: {})",
total_flows, self.min_packets
)));
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());
}
self.flow_tracker.cleanup_old_flows(60_000_000);
continue;
}
let batch = &flows[..flows.len().min(self.batch_size)];
log!(MLLog::RunningInference(batch.len()));
let start = Instant::now();
let results = self.inference_pipeline.infer_batch(batch);
let elapsed_us = start.elapsed().as_micros() as u64;
let stats = InferenceStats::from_results(&results, elapsed_us);
if results.len() != batch.len() {
log!(MLLog::InferenceResults(batch.len(), results.len()));
}
log!(MLLog::InferenceCompleted(
stats.total_flows,
stats.malicious_flows,
stats.benign_flows,
(elapsed_us as f64 / 1000.0) as u32,
stats.flows_per_second
));
if let Ok(mut aggregator) = self.aggregator.lock() {
for result in &results {
if result.is_attack {
let should_alert =
aggregator.should_alert(&result.flow_key_raw, result.ae_score, result.threshold);
if should_alert {
log!(MLLog::ThreatDetected(
format!("{:?}", result.direction),
result.flow_key.clone(),
result.attack_type.clone().unwrap_or_else(|| "UNKNOWN".to_string()),
result.confidence,
result.ae_score,
));
self.ml_alert.broadcast_alert(result);
}
}
}
aggregator.cleanup();
}
self.flow_tracker.cleanup_old_flows(60_000_000);
}
}
pub fn process_packet(&self, packet_data: &[u8], is_ingress: bool) {
match parse_packet(packet_data) {
Some((packet_info, payload_start)) => {
let payload = packet_data.get(payload_start..).unwrap_or(&[]);
self.flow_tracker.process_packet(packet_info, is_ingress, payload);
}
None => log!(MLLog::ParsePacketFailed(packet_data.len())),
}
}
pub fn get_stats(&self) -> EngineStats {
EngineStats {
active_flows: self.flow_tracker.flow_count(),
}
}
}
pub struct PacketProcessor {
ml_engine: Arc<Engine>,
}
impl PacketProcessor {
pub fn new(ml_engine: Arc<Engine>) -> Self {
Self { ml_engine }
}
pub fn process(&self, packet_data: &[u8], is_ingress: bool) {
self.ml_engine.process_packet(packet_data, is_ingress);
}
pub fn process_batch(&self, packets: &[Vec<u8>], is_ingress: bool) {
for packet in packets {
self.process(packet, is_ingress);
}
}
}

View File

@ -1,19 +1,28 @@
use serde::{Deserialize, Serialize};
#[derive(Debug, Deserialize)]
pub struct ConfigTable {
#[serde(rename = "Config")]
pub config: Config,
pub struct AppConfigTable {
#[serde(rename = "Http")]
pub http: HttpConfig,
#[serde(rename = "Network")]
pub network: NetworkConfig,
#[serde(rename = "Inference")]
pub inference: InferenceConfig,
#[serde(rename = "Misc")]
pub misc: MiscConfig,
#[serde(rename = "Pipeline")]
pub pipeline: PipelineConfig,
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Config {
pub struct HttpConfig {
pub http_server_bind_port: u16,
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct NetworkConfig {
pub ingress_ifname: String,
pub egress_ifname: String,
pub geoip_db_name: String,
pub deep_autoencoder_name: String,
pub classifier_name: String,
pub models_config_name: String,
pub combined_queue_count: u32,
pub channel_size: usize,
pub fill_queue_size: u32,
@ -23,7 +32,13 @@ pub struct Config {
pub frame_size: u32,
pub frame_count: u32,
pub refresh_interval: u64,
pub http_server_bind_port: u16,
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct InferenceConfig {
pub deep_autoencoder_name: String,
pub classifier_name: String,
pub models_config_name: String,
pub max_concurrent_flows: usize,
pub min_packets_for_inference: usize,
pub inference_interval_secs: u64,
@ -31,4 +46,15 @@ pub struct Config {
pub inference_batch_size: usize,
pub traffic_logging_mode: bool,
pub traffic_log_csv_path: String,
}
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct MiscConfig {
pub geoip_db_name: String,
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct PipelineConfig {
pub ingress: Vec<String>,
pub egress: Vec<String>,
}

View File

@ -49,6 +49,10 @@ traceable! {
#[error("Amount of rules has reached the upper limit")]
RuleReachLimit => tracing::Level::ERROR,
#[no_source]
#[error("Fill queue initialization failed: submitted fewer frames than expected")]
FillQueueInitFailed => tracing::Level::ERROR,
#[no_source]
#[error("Unknown error")]
UnknownError => tracing::Level::ERROR,

View File

@ -0,0 +1,59 @@
use serde::{Deserialize, Serialize};
use crate::core::ml::flow_tracker::FlowData;
use crate::model::direction::Direction;
#[derive(Debug, Clone, Serialize)]
pub struct FlowStatsEntry {
pub direction: Direction,
pub src_ip: String,
pub dst_ip: String,
pub src_port: u16,
pub dst_port: u16,
pub protocol: u8,
pub fwd_packets: usize,
pub bwd_packets: usize,
pub fwd_bytes: u64,
pub bwd_bytes: u64,
pub duration_us: u64,
pub last_seen_us: u64,
}
impl From<&FlowData> for FlowStatsEntry {
fn from(flow: &FlowData) -> Self {
Self {
direction: flow.direction,
src_ip: flow.flow_key.src_ip_string(),
dst_ip: flow.flow_key.dst_ip_string(),
src_port: flow.flow_key.src_port,
dst_port: flow.flow_key.dst_port,
protocol: flow.flow_key.protocol,
fwd_packets: flow.fwd_packets.len(),
bwd_packets: flow.bwd_packets.len(),
fwd_bytes: flow.fwd_total_bytes,
bwd_bytes: flow.bwd_total_bytes,
duration_us: flow.duration_us(),
last_seen_us: flow.last_time_us,
}
}
}
#[derive(Debug, Clone, Serialize)]
pub struct StatsSummary {
pub total_flows: usize,
pub total_bytes: u64,
pub total_packets: usize,
}
/// Client subscription filter for WebSocket flow stats
#[derive(Debug, Clone, Deserialize)]
pub struct FlowSubscription {
/// Filter by direction: "ingress", "egress", or null for both
pub direction: Option<Direction>,
/// Time window in seconds: only flows with last_seen within this window
pub window_secs: Option<u64>,
/// Max number of flows to return (sorted by bytes desc)
pub top_n: Option<usize>,
/// Push interval in seconds (default 5)
pub interval_secs: Option<u64>,
}

View File

@ -1,19 +0,0 @@
use serde::{Deserialize, Serialize};
use common::model::flow_stats::FlowStats;
#[derive(Debug, Clone, Serialize)]
pub struct FlowStatsWithGeo {
#[serde(flatten)]
pub stats: FlowStats,
pub geo: Option<GeoLocation>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GeoLocation {
pub country: Option<String>,
pub country_code: Option<String>,
pub city: Option<String>,
pub latitude: Option<f64>,
pub longitude: Option<f64>,
pub timezone: Option<String>,
}

View File

@ -3,8 +3,8 @@ use serde::{Deserialize, Serialize};
#[derive(Serialize, Deserialize, Copy, Clone, Eq, PartialEq, Hash)]
#[serde(rename_all = "lowercase")]
pub enum ListType {
#[serde(rename = "white_list")]
#[serde(rename = "whitelist")]
White,
#[serde(rename = "black_list")]
#[serde(rename = "blacklist")]
Black,
}

View File

@ -50,5 +50,8 @@ loggable! {
#[error("Fill queue incomplete: produced {produced}, expected {expected}")]
FillQueueIncomplete { produced: usize, expected: usize } => tracing::Level::WARN,
#[error("Invalid packet length exceeds buffer")]
InvalidPacketLength => tracing::Level::WARN,
}
}

View File

@ -1,73 +1,92 @@
use common::model::event::{Event, TcpFlags};
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
use serde::{Deserialize, Serialize};
use tract_onnx::prelude::{Graph, SimplePlan, TypedFact, TypedOp};
use crate::model::direction::Direction;
use crate::utils::packet_parser::{format_ipv4, format_ipv6};
use crate::model::user_packet::UserPacket;
pub type RunnableModel = SimplePlan<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>;
pub struct EngineConfig {
pub max_flows: usize,
pub min_packets: usize,
pub batch_size: usize,
pub inference_interval_secs: u64,
pub aggregator_window_secs: u64,
pub flow_timeout_us: u64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClipParams {
pub lower: f64,
pub upper: f64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AENormalization {
pub min: f64,
pub max: f64,
pub norm_max: f64,
pub mean: f64,
pub std: f64,
pub median: f64,
pub p90: f64,
pub p95: f64,
pub p99: f64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PrecisionLevels {
pub threshold: f64,
pub precision: f64,
}
#[derive(Debug, Clone, Hash, Eq, PartialEq, Serialize, Deserialize)]
pub struct FlowKey {
pub src_ip: String,
pub dst_ip: String,
pub src_ip: [u8; 16],
pub dst_ip: [u8; 16],
pub src_port: u16,
pub dst_port: u16,
pub protocol: u8,
pub ip_version: u8,
}
impl FlowKey {
pub fn from_packet(packet: &Event) -> Self {
match packet {
Event::IPv4(ipv4) => Self {
src_ip: format_ipv4(ipv4.src_ip),
dst_ip: format_ipv4(ipv4.dst_ip),
src_port: ipv4.src_port,
dst_port: ipv4.dst_port,
protocol: ipv4.protocol as u8,
},
Event::IPv6(ipv6) => Self {
src_ip: format_ipv6(ipv6.src_ip),
dst_ip: format_ipv6(ipv6.dst_ip),
src_port: ipv6.src_port,
dst_port: ipv6.dst_port,
protocol: ipv6.protocol as u8,
},
pub fn from_packet(packet: &UserPacket) -> Self {
let (src_ip, ip_version) = Self::parse_ip_to_bytes(&packet.src_ip);
let (dst_ip, _) = Self::parse_ip_to_bytes(&packet.dst_ip);
Self {
src_ip,
dst_ip,
src_port: packet.src_port,
dst_port: packet.dst_port,
protocol: packet.protocol,
ip_version,
}
}
pub fn reverse(&self) -> Self {
Self {
src_ip: self.dst_ip.clone(),
dst_ip: self.src_ip.clone(),
src_ip: self.dst_ip,
dst_ip: self.src_ip,
src_port: self.dst_port,
dst_port: self.src_port,
protocol: self.protocol,
ip_version: self.ip_version,
}
}
pub fn src_ip_string(&self) -> String {
self.ip_bytes_to_string(&self.src_ip)
}
pub fn dst_ip_string(&self) -> String {
self.ip_bytes_to_string(&self.dst_ip)
}
fn parse_ip_to_bytes(ip_str: &str) -> ([u8; 16], u8) {
if let Ok(addr) = ip_str.parse::<IpAddr>() {
match addr {
IpAddr::V4(v4) => {
let mut buf = [0u8; 16];
buf[..4].copy_from_slice(&v4.octets());
(buf, 4)
}
IpAddr::V6(v6) => (v6.octets(), 6),
}
} else {
([0u8; 16], 4)
}
}
fn ip_bytes_to_string(&self, bytes: &[u8; 16]) -> String {
if self.ip_version == 6 {
Ipv6Addr::from(*bytes).to_string()
} else {
let octets: [u8; 4] = [bytes[0], bytes[1], bytes[2], bytes[3]];
Ipv4Addr::from(octets).to_string()
}
}
}
@ -78,7 +97,7 @@ pub struct PacketData {
pub length: u32,
pub header_length: u16,
pub payload_length: u32,
pub flags: TcpFlags,
pub flags: u8,
}
#[derive(Debug, Clone, Default)]
@ -111,6 +130,7 @@ pub struct InferenceStats {
pub total_flows: usize,
pub malicious_flows: usize,
pub benign_flows: usize,
#[allow(dead_code)]
pub inference_time_us: u64,
pub flows_per_second: f32,
}
@ -136,8 +156,3 @@ impl InferenceStats {
}
}
}
#[derive(Debug, Clone)]
pub struct EngineStats {
pub active_flows: usize,
}

View File

@ -1,10 +1,10 @@
pub mod config;
pub mod direction;
pub mod error;
pub mod geo_stats;
pub mod flow_stats;
pub mod health;
pub mod ip_address;
pub mod list_type;
pub mod log;
pub mod time_type;
pub mod ml_detection;
pub mod health;
pub mod user_packet;

View File

@ -1,19 +0,0 @@
use serde::{Deserialize, Serialize};
#[derive(Serialize, Deserialize, Copy, Clone, Eq, PartialEq, Hash)]
#[serde(rename_all = "lowercase")]
pub enum TimeType {
#[serde(rename = "1min")]
_1Min = 60 * 1_000_000_000,
#[serde(rename = "10min")]
_10Min = 600 * 1_000_000_000,
#[serde(rename = "1hour")]
_1Hour = 3600 * 1_000_000_000,
}
impl TimeType {
#[inline]
pub fn duration(&self) -> u64 {
*self as u64
}
}

View File

@ -0,0 +1,16 @@
pub struct UserPacket {
#[allow(dead_code)]
pub ip_version: u8,
pub protocol: u8,
pub tcp_flags: u8,
pub src_ip: String,
pub dst_ip: String,
pub src_port: u16,
pub dst_port: u16,
pub packet_length: u32,
pub payload_length: u32,
pub header_length: u16,
pub tcp_window_size: u16,
pub timestamp_us: u64,
pub is_forward: bool,
}

View File

@ -1,16 +1,6 @@
use std::net::IpAddr;
use common::define::setting::MAX_RULES_PORT;
use common::model::ip_address::Port;
pub fn convert_ports_to_vec(ports: [u16; MAX_RULES_PORT]) -> Vec<Port> {
let mut filtered_ports: Vec<Port> = ports.into_iter().filter(|&port| port != 0).collect();
if filtered_ports.is_empty() {
filtered_ports.push(0);
}
filtered_ports
}
#[allow(dead_code)]
pub fn is_private_ip(ip: &IpAddr) -> bool {
match ip {
IpAddr::V4(v4) => {
@ -26,4 +16,4 @@ pub fn is_private_ip(ip: &IpAddr) -> bool {
|| v6.is_multicast()
}
}
}
}

View File

@ -33,11 +33,10 @@ impl Logging {
.with_ansi(false)
.with_writer(file_appender);
let level = if cfg!(debug_assertions) {
Level::DEBUG
} else {
Level::INFO
};
let level = std::env::var("RUST_LOG")
.ok()
.and_then(|s| s.parse::<Level>().ok())
.unwrap_or(if cfg!(debug_assertions) { Level::DEBUG } else { Level::INFO });
tracing_subscriber::registry()
.with(stdout_layer)

View File

@ -1,10 +1,8 @@
use std::mem;
use std::time;
use common::model::event::{Event, IPv4Event, IPv6Event, TcpFlags};
use network_types::ip::IpProto;
use crate::model::user_packet::UserPacket;
pub fn parse_packet(packet_data: &[u8]) -> Option<(Event, usize)> {
pub fn parse_packet(packet_data: &[u8]) -> Option<(UserPacket, usize)> {
if packet_data.len() < 14 {
return None;
}
@ -23,7 +21,7 @@ pub fn parse_packet(packet_data: &[u8]) -> Option<(Event, usize)> {
}
}
fn parse_ipv4(packet_data: &[u8], timestamp_us: u64) -> Option<(Event, usize)> {
fn parse_ipv4(packet_data: &[u8], timestamp_us: u64) -> Option<(UserPacket, usize)> {
if packet_data.len() < 34 {
return None;
}
@ -31,18 +29,33 @@ fn parse_ipv4(packet_data: &[u8], timestamp_us: u64) -> Option<(Event, usize)> {
let ip_header = &packet_data[14..];
let protocol_byte = ip_header[9];
let protocol = unsafe { mem::transmute::<u8, IpProto>(protocol_byte) };
let src_ip = u32::from_be_bytes([ip_header[12], ip_header[13], ip_header[14], ip_header[15]]);
let dst_ip = u32::from_be_bytes([ip_header[16], ip_header[17], ip_header[18], ip_header[19]]);
if protocol_byte != 6 && protocol_byte != 17 {
return None;
}
let src_ip_raw = u32::from_be_bytes([ip_header[12], ip_header[13], ip_header[14], ip_header[15]]);
let dst_ip_raw = u32::from_be_bytes([ip_header[16], ip_header[17], ip_header[18], ip_header[19]]);
let ihl = (ip_header[0] & 0x0F) as usize * 4;
let total_len = u16::from_be_bytes([ip_header[2], ip_header[3]]) as u32;
if protocol_byte != 6 && protocol_byte != 17 {
// Still track non-TCP/UDP packets (e.g. ICMP) for flow statistics
let packet = UserPacket {
ip_version: 4,
protocol: protocol_byte,
tcp_flags: 0,
src_ip: format_ipv4(src_ip_raw),
dst_ip: format_ipv4(dst_ip_raw),
src_port: 0,
dst_port: 0,
packet_length: total_len,
payload_length: 0,
header_length: 0,
tcp_window_size: 0,
timestamp_us,
is_forward: false,
};
return Some((packet, 14 + ihl));
}
if packet_data.len() < 14 + ihl + 4 {
return None;
}
@ -57,38 +70,39 @@ fn parse_ipv4(packet_data: &[u8], timestamp_us: u64) -> Option<(Event, usize)> {
}
let data_offset = (transport_header[12] >> 4) as u16 * 4;
let flags = TcpFlags::from_byte(transport_header[13]);
let flags = transport_header[13];
let window = u16::from_be_bytes([transport_header[14], transport_header[15]]);
(flags, window, data_offset)
} else if protocol_byte == 17 {
(TcpFlags::default(), 0, 8)
(0u8, 0, 8)
} else {
(TcpFlags::default(), 0, 0)
(0u8, 0, 0)
};
let payload_length = total_len.saturating_sub(ihl as u32 + header_length as u32);
let payload_start = 14 + ihl + header_length as usize;
let event = IPv4Event {
protocol,
src_ip,
dst_ip,
let packet = UserPacket {
ip_version: 4,
protocol: protocol_byte,
tcp_flags,
src_ip: format_ipv4(src_ip_raw),
dst_ip: format_ipv4(dst_ip_raw),
src_port,
dst_port,
packet_length: total_len,
payload_length,
header_length,
timestamp_us,
tcp_flags,
tcp_window_size,
timestamp_us,
is_forward: false,
};
Some((Event::IPv4(event), payload_start))
Some((packet, payload_start))
}
fn parse_ipv6(packet_data: &[u8], timestamp_us: u64) -> Option<(Event, usize)> {
fn parse_ipv6(packet_data: &[u8], timestamp_us: u64) -> Option<(UserPacket, usize)> {
if packet_data.len() < 54 {
return None;
}
@ -96,23 +110,38 @@ fn parse_ipv6(packet_data: &[u8], timestamp_us: u64) -> Option<(Event, usize)> {
let ip_header = &packet_data[14..];
let protocol_byte = ip_header[6];
let protocol = unsafe { mem::transmute::<u8, IpProto>(protocol_byte) };
let mut source_ip_bytes = [0u8; 16];
source_ip_bytes.copy_from_slice(&ip_header[8..24]);
let src_ip = u128::from_be_bytes(source_ip_bytes);
let src_ip_raw = u128::from_be_bytes(source_ip_bytes);
let mut dest_ip_bytes = [0u8; 16];
dest_ip_bytes.copy_from_slice(&ip_header[24..40]);
let dst_ip = u128::from_be_bytes(dest_ip_bytes);
if protocol_byte != 6 && protocol_byte != 17 {
return None;
}
let dst_ip_raw = u128::from_be_bytes(dest_ip_bytes);
let payload_len = u16::from_be_bytes([ip_header[4], ip_header[5]]) as u32;
let total_len = payload_len + 40;
if protocol_byte != 6 && protocol_byte != 17 {
// Still track non-TCP/UDP packets (e.g. ICMPv6) for flow statistics
let packet = UserPacket {
ip_version: 6,
protocol: protocol_byte,
tcp_flags: 0,
src_ip: format_ipv6(src_ip_raw),
dst_ip: format_ipv6(dst_ip_raw),
src_port: 0,
dst_port: 0,
packet_length: total_len,
payload_length: 0,
header_length: 0,
tcp_window_size: 0,
timestamp_us,
is_forward: false,
};
return Some((packet, 14 + 40));
}
if packet_data.len() < 54 + 4 {
return None;
}
@ -127,35 +156,36 @@ fn parse_ipv6(packet_data: &[u8], timestamp_us: u64) -> Option<(Event, usize)> {
}
let data_offset = (transport_header[12] >> 4) as u16 * 4;
let flags = TcpFlags::from_byte(transport_header[13]);
let flags = transport_header[13];
let window = u16::from_be_bytes([transport_header[14], transport_header[15]]);
(flags, window, data_offset)
} else if protocol_byte == 17 {
(TcpFlags::default(), 0, 8)
(0u8, 0, 8)
} else {
(TcpFlags::default(), 0, 0)
(0u8, 0, 0)
};
let payload_length = total_len.saturating_sub(40 + header_length as u32);
let payload_start = 14 + 40 + header_length as usize;
let event = IPv6Event {
protocol,
src_ip,
dst_ip,
let packet = UserPacket {
ip_version: 6,
protocol: protocol_byte,
tcp_flags,
src_ip: format_ipv6(src_ip_raw),
dst_ip: format_ipv6(dst_ip_raw),
src_port,
dst_port,
packet_length: total_len,
payload_length,
header_length,
timestamp_us,
tcp_flags,
tcp_window_size,
timestamp_us,
is_forward: false,
};
Some((Event::IPv6(event), payload_start))
Some((packet, payload_start))
}
pub fn format_ipv4(addr: u32) -> String {
@ -165,23 +195,5 @@ pub fn format_ipv4(addr: u32) -> String {
pub fn format_ipv6(addr: u128) -> String {
let bytes = addr.to_be_bytes();
format!(
"{:02x}{:02x}:{:02x}{:02x}:{:02x}{:02x}:{:02x}{:02x}:{:02x}{:02x}:{:02x}{:02x}:{:02x}{:02x}:{:02x}{:02x}",
bytes[0],
bytes[1],
bytes[2],
bytes[3],
bytes[4],
bytes[5],
bytes[6],
bytes[7],
bytes[8],
bytes[9],
bytes[10],
bytes[11],
bytes[12],
bytes[13],
bytes[14],
bytes[15]
)
std::net::Ipv6Addr::from(bytes).to_string()
}

View File

@ -1,22 +1,21 @@
use std::net::{SocketAddrV4, SocketAddrV6};
use actix_web::{delete, get, put, web, HttpResponse, Responder, Scope};
use actix_web::{web, HttpResponse, Responder, Scope};
use crate::core::ebpf::access_control::AccessControl;
use crate::model::direction::FlowDirection;
use crate::model::list_type::ListType;
pub fn initialize() -> Scope {
web::scope("/access_control")
.service(get_ipv4_list)
.service(get_ipv6_list)
.service(add_ipv4_list)
.service(add_ipv6_list)
.service(remove_ipv4_list)
.service(remove_ipv6_list)
web::scope("/acl")
.route("/ipv4/{direction}/{list_type}", web::get().to(get_ipv4_list))
.route("/ipv6/{direction}/{list_type}", web::get().to(get_ipv6_list))
.route("/ipv4/{direction}/{list_type}", web::put().to(add_ipv4_list))
.route("/ipv6/{direction}/{list_type}", web::put().to(add_ipv6_list))
.route("/ipv4/{direction}/{list_type}", web::delete().to(remove_ipv4_list))
.route("/ipv6/{direction}/{list_type}", web::delete().to(remove_ipv6_list))
}
#[get("/ipv4/{direction}/{list_type}")]
async fn get_ipv4_list(
path: web::Path<(FlowDirection, ListType)>,
access_control: web::Data<AccessControl>,
@ -26,7 +25,6 @@ async fn get_ipv4_list(
HttpResponse::Ok().json(list)
}
#[get("/ipv6/{direction}/{list_type}")]
async fn get_ipv6_list(
path: web::Path<(FlowDirection, ListType)>,
access_control: web::Data<AccessControl>,
@ -36,7 +34,6 @@ async fn get_ipv6_list(
HttpResponse::Ok().json(list)
}
#[put("/ipv4/{direction}/{list_type}")]
async fn add_ipv4_list(
address: web::Json<SocketAddrV4>,
path: web::Path<(FlowDirection, ListType)>,
@ -46,11 +43,10 @@ async fn add_ipv4_list(
let (direction, list_type) = path.into_inner();
match access_control.add_ipv4_list(direction, list_type, address).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().body(e.to_string()),
Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})),
}
}
#[put("/ipv6/{direction}/{list_type}")]
async fn add_ipv6_list(
address: web::Json<SocketAddrV6>,
path: web::Path<(FlowDirection, ListType)>,
@ -60,11 +56,10 @@ async fn add_ipv6_list(
let (direction, list_type) = path.into_inner();
match access_control.add_ipv6_list(direction, list_type, address).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().body(e.to_string()),
Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})),
}
}
#[delete("/ipv4/{direction}/{list_type}")]
async fn remove_ipv4_list(
address: web::Json<SocketAddrV4>,
path: web::Path<(FlowDirection, ListType)>,
@ -74,11 +69,10 @@ async fn remove_ipv4_list(
let (direction, list_type) = path.into_inner();
match access_control.remove_ipv4_list(direction, list_type, address).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().body(e.to_string()),
Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})),
}
}
#[delete("/ipv6/{direction}/{list_type}")]
async fn remove_ipv6_list(
address: web::Json<SocketAddrV6>,
path: web::Path<(FlowDirection, ListType)>,
@ -88,6 +82,6 @@ async fn remove_ipv6_list(
let (direction, list_type) = path.into_inner();
match access_control.remove_ipv6_list(direction, list_type, address).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().body(e.to_string()),
Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})),
}
}

View File

@ -1,12 +0,0 @@
pub mod access_control;
pub mod service;
pub mod statistics;
use actix_web::{web, Scope};
pub fn initialize() -> Scope {
web::scope("/ebpf")
.service(access_control::initialize())
.service(service::initialize())
.service(statistics::initialize())
}

View File

@ -1,251 +0,0 @@
use std::net::{Ipv4Addr, Ipv6Addr, SocketAddrV4, SocketAddrV6};
use actix_web::{delete, get, post, put, web, HttpResponse, Responder, Scope};
use common::model::http_method::HttpMethod;
use crate::core::ebpf::service::Service;
pub fn initialize() -> Scope {
web::scope("/service")
.service(get_ipv4_http_service)
.service(get_ipv6_http_service)
.service(add_ipv4_http_service)
.service(add_ipv6_http_service)
.service(remove_ipv4_http_service)
.service(remove_ipv6_http_service)
.service(is_ssh_white_list_enable)
.service(enable_ssh_white_list)
.service(disable_ssh_white_list)
.service(get_ipv4_ssh_service)
.service(get_ipv6_ssh_service)
.service(add_ipv4_ssh_service)
.service(add_ipv6_ssh_service)
.service(remove_ipv4_ssh_service)
.service(remove_ipv6_ssh_service)
.service(get_ipv4_ssh_white_list)
.service(get_ipv6_ssh_white_list)
.service(add_ipv4_ssh_white_list)
.service(add_ipv6_ssh_white_list)
.service(remove_ipv4_ssh_white_list)
.service(remove_ipv6_ssh_white_list)
.service(get_ipv4_ssh_black_list)
.service(get_ipv6_ssh_black_list)
.service(add_ipv4_ssh_black_list)
.service(add_ipv6_ssh_black_list)
.service(remove_ipv4_ssh_black_list)
.service(remove_ipv6_ssh_black_list)
}
#[get("/ipv4/http_service")]
async fn get_ipv4_http_service(service: web::Data<Service>) -> impl Responder {
let list = service.get_ipv4_http_service().await;
HttpResponse::Ok().json(web::Json(list))
}
#[get("/ipv6/http_service")]
async fn get_ipv6_http_service(service: web::Data<Service>) -> impl Responder {
let list = service.get_ipv6_http_service().await;
HttpResponse::Ok().json(web::Json(list))
}
#[put("/ipv4/http_service")]
async fn add_ipv4_http_service(
payload: web::Json<(SocketAddrV4, Vec<HttpMethod>)>,
service: web::Data<Service>,
) -> impl Responder {
let (addr, methods) = payload.into_inner();
match service.add_ipv4_http_service(addr, methods).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().body(e.to_string()),
}
}
#[put("/ipv6/http_service")]
async fn add_ipv6_http_service(
payload: web::Json<(SocketAddrV6, Vec<HttpMethod>)>,
service: web::Data<Service>,
) -> impl Responder {
let (addr, methods) = payload.into_inner();
match service.add_ipv6_http_service(addr, methods).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().body(e.to_string()),
}
}
#[delete("/ipv4/http_service")]
async fn remove_ipv4_http_service(
payload: web::Json<(SocketAddrV4, Vec<HttpMethod>)>,
service: web::Data<Service>,
) -> impl Responder {
let (addr, methods) = payload.into_inner();
match service.remove_ipv4_http_service(addr, methods).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().body(e.to_string()),
}
}
#[delete("/ipv6/http_service")]
async fn remove_ipv6_http_service(
payload: web::Json<(SocketAddrV6, Vec<HttpMethod>)>,
service: web::Data<Service>,
) -> impl Responder {
let (addr, methods) = payload.into_inner();
match service.remove_ipv6_http_service(addr, methods).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().body(e.to_string()),
}
}
#[get("/ssh_white_list")]
async fn is_ssh_white_list_enable(service: web::Data<Service>) -> impl Responder {
let enabled = service.is_ssh_white_list_enable().await;
HttpResponse::Ok().json(enabled)
}
#[post("/ssh_white_list/enable")]
async fn enable_ssh_white_list(service: web::Data<Service>) -> impl Responder {
match service.enable_ssh_white_list().await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().body(e.to_string()),
}
}
#[post("/ssh_white_list/disable")]
async fn disable_ssh_white_list(service: web::Data<Service>) -> impl Responder {
match service.disable_ssh_white_list().await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().body(e.to_string()),
}
}
#[get("/ipv4/ssh_service")]
async fn get_ipv4_ssh_service(service: web::Data<Service>) -> impl Responder {
let list = service.get_ipv4_ssh_service().await;
HttpResponse::Ok().json(web::Json(list))
}
#[get("/ipv6/ssh_service")]
async fn get_ipv6_ssh_service(service: web::Data<Service>) -> impl Responder {
let list = service.get_ipv6_ssh_service().await;
HttpResponse::Ok().json(web::Json(list))
}
#[put("/ipv4/ssh_service")]
async fn add_ipv4_ssh_service(ip_addr: web::Json<SocketAddrV4>, service: web::Data<Service>) -> impl Responder {
match service.add_ipv4_ssh_service(ip_addr.into_inner()).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().body(e.to_string()),
}
}
#[put("/ipv6/ssh_service")]
async fn add_ipv6_ssh_service(ip_addr: web::Json<SocketAddrV6>, service: web::Data<Service>) -> impl Responder {
match service.add_ipv6_ssh_service(ip_addr.into_inner()).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().body(e.to_string()),
}
}
#[delete("/ipv4/ssh_service")]
async fn remove_ipv4_ssh_service(ip_addr: web::Json<SocketAddrV4>, service: web::Data<Service>) -> impl Responder {
match service.remove_ipv4_ssh_service(ip_addr.into_inner()).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().body(e.to_string()),
}
}
#[delete("/ipv6/ssh_service")]
async fn remove_ipv6_ssh_service(ip_addr: web::Json<SocketAddrV6>, service: web::Data<Service>) -> impl Responder {
match service.remove_ipv6_ssh_service(ip_addr.into_inner()).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().body(e.to_string()),
}
}
#[get("/ipv4/ssh_white_list")]
async fn get_ipv4_ssh_white_list(service: web::Data<Service>) -> impl Responder {
let list = service.get_ipv4_ssh_white_list().await;
HttpResponse::Ok().json(web::Json(list))
}
#[get("/ipv6/ssh_white_list")]
async fn get_ipv6_ssh_white_list(service: web::Data<Service>) -> impl Responder {
let list = service.get_ipv6_ssh_white_list().await;
HttpResponse::Ok().json(web::Json(list))
}
#[put("/ipv4/ssh_white_list")]
async fn add_ipv4_ssh_white_list(ip_addr: web::Json<Ipv4Addr>, service: web::Data<Service>) -> impl Responder {
match service.add_ipv4_ssh_white_list(ip_addr.into_inner()).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().body(e.to_string()),
}
}
#[put("/ipv6/ssh_white_list")]
async fn add_ipv6_ssh_white_list(ip_addr: web::Json<Ipv6Addr>, service: web::Data<Service>) -> impl Responder {
match service.add_ipv6_ssh_white_list(ip_addr.into_inner()).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().body(e.to_string()),
}
}
#[delete("/ipv4/ssh_white_list")]
async fn remove_ipv4_ssh_white_list(ip_addr: web::Json<Ipv4Addr>, service: web::Data<Service>) -> impl Responder {
match service.remove_ipv4_ssh_white_list(ip_addr.into_inner()).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().body(e.to_string()),
}
}
#[delete("/ipv6/ssh_white_list")]
async fn remove_ipv6_ssh_white_list(ip_addr: web::Json<Ipv6Addr>, service: web::Data<Service>) -> impl Responder {
match service.remove_ipv6_ssh_white_list(ip_addr.into_inner()).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().body(e.to_string()),
}
}
#[get("/ipv4/ssh_black_list")]
async fn get_ipv4_ssh_black_list(service: web::Data<Service>) -> impl Responder {
let list = service.get_ipv4_ssh_black_list().await;
HttpResponse::Ok().json(web::Json(list))
}
#[get("/ipv6/ssh_black_list")]
async fn get_ipv6_ssh_black_list(service: web::Data<Service>) -> impl Responder {
let list = service.get_ipv6_ssh_black_list().await;
HttpResponse::Ok().json(web::Json(list))
}
#[put("/ipv4/ssh_black_list")]
async fn add_ipv4_ssh_black_list(ip_addr: web::Json<Ipv4Addr>, service: web::Data<Service>) -> impl Responder {
match service.add_ipv4_ssh_black_list(ip_addr.into_inner()).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().body(e.to_string()),
}
}
#[put("/ipv6/ssh_black_list")]
async fn add_ipv6_ssh_black_list(ip_addr: web::Json<Ipv6Addr>, service: web::Data<Service>) -> impl Responder {
match service.add_ipv6_ssh_black_list(ip_addr.into_inner()).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().body(e.to_string()),
}
}
#[delete("/ipv4/ssh_black_list")]
async fn remove_ipv4_ssh_black_list(ip_addr: web::Json<Ipv4Addr>, service: web::Data<Service>) -> impl Responder {
match service.remove_ipv4_ssh_black_list(ip_addr.into_inner()).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().body(e.to_string()),
}
}
#[delete("/ipv6/ssh_black_list")]
async fn remove_ipv6_ssh_black_list(ip_addr: web::Json<Ipv6Addr>, service: web::Data<Service>) -> impl Responder {
match service.remove_ipv6_ssh_black_list(ip_addr.into_inner()).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().body(e.to_string()),
}
}

View File

@ -1,69 +0,0 @@
use std::sync::Arc;
use actix_web::{get, web, HttpRequest, HttpResponse, Responder, Scope};
use crate::core::ebpf::statistics::Statistics;
use crate::core::infrastructure::app_config::AppConfig;
use crate::model::direction::{Direction, FlowDirection};
use crate::model::time_type::TimeType;
use crate::web::websocket::flow_websocket;
pub fn initialize() -> Scope {
web::scope("/statistics")
.service(get_ipv4_flow)
.service(get_ipv6_flow)
.service(websocket_ipv4)
.service(websocket_ipv6)
}
#[get("/get/ipv4/{direction}/{flow_direction}/{time_type}")]
async fn get_ipv4_flow(
path: web::Path<(Direction, FlowDirection, TimeType)>,
statistics: web::Data<Arc<Statistics>>,
) -> impl Responder {
let (direction, flow_direction, time_type) = path.into_inner();
let flow_data = statistics
.get_ipv4_flow_data(direction, flow_direction, time_type)
.await;
HttpResponse::Ok().json(web::Json(flow_data))
}
#[get("/get/ipv6/{direction}/{flow_direction}/{time_type}")]
async fn get_ipv6_flow(
path: web::Path<(Direction, FlowDirection, TimeType)>,
statistics: web::Data<Arc<Statistics>>,
) -> impl Responder {
let (direction, flow_direction, time_type) = path.into_inner();
let flow_data = statistics
.get_ipv6_flow_data(direction, flow_direction, time_type)
.await;
HttpResponse::Ok().json(web::Json(flow_data))
}
#[get("/websocket/ipv4/{direction}/{flow_direction}/{time_type}")]
async fn websocket_ipv4(
req: HttpRequest,
stream: web::Payload,
path: web::Path<(Direction, FlowDirection, TimeType)>,
app_config: web::Data<AppConfig>,
statistics: web::Data<Statistics>,
) -> impl Responder {
match flow_websocket::websocket_ipv4_flow(req, stream, path, app_config, statistics).await {
Ok(response) => response,
Err(err) => HttpResponse::InternalServerError().body(format!("WebSocket error: {}", err)),
}
}
#[get("/websocket/ipv6/{direction}/{flow_direction}/{time_type}")]
async fn websocket_ipv6(
req: HttpRequest,
stream: web::Payload,
path: web::Path<(Direction, FlowDirection, TimeType)>,
app_config: web::Data<AppConfig>,
statistics: web::Data<Statistics>,
) -> impl Responder {
match flow_websocket::websocket_ipv6_flow(req, stream, path, app_config, statistics).await {
Ok(response) => response,
Err(err) => HttpResponse::InternalServerError().body(format!("WebSocket error: {}", err)),
}
}

View File

@ -0,0 +1,288 @@
use std::net::{Ipv4Addr, Ipv6Addr, SocketAddrV4, SocketAddrV6};
use actix_web::{web, HttpResponse, Responder, Scope};
use common::model::http_method::HttpMethod;
use crate::core::ebpf::protocol_filter::ProtocolFilter;
pub fn initialize() -> Scope {
web::scope("/filter")
.service(http_scope())
.service(ssh_scope())
}
fn http_scope() -> Scope {
web::scope("/http")
.route("/ipv4", web::get().to(get_ipv4_http_service))
.route("/ipv6", web::get().to(get_ipv6_http_service))
.route("/ipv4", web::put().to(add_ipv4_http_service))
.route("/ipv6", web::put().to(add_ipv6_http_service))
.route("/ipv4", web::delete().to(remove_ipv4_http_service))
.route("/ipv6", web::delete().to(remove_ipv6_http_service))
}
fn ssh_scope() -> Scope {
web::scope("/ssh")
.route("/ipv4", web::get().to(get_ipv4_ssh_service))
.route("/ipv6", web::get().to(get_ipv6_ssh_service))
.route("/ipv4", web::put().to(add_ipv4_ssh_service))
.route("/ipv6", web::put().to(add_ipv6_ssh_service))
.route("/ipv4", web::delete().to(remove_ipv4_ssh_service))
.route("/ipv6", web::delete().to(remove_ipv6_ssh_service))
.service(ssh_whitelist_scope())
.service(ssh_blacklist_scope())
}
fn ssh_whitelist_scope() -> Scope {
web::scope("/whitelist")
.route("/status", web::get().to(is_ssh_white_list_enable))
.route("/enable", web::post().to(enable_ssh_white_list))
.route("/disable", web::post().to(disable_ssh_white_list))
.route("/ipv4", web::get().to(get_ipv4_ssh_white_list))
.route("/ipv6", web::get().to(get_ipv6_ssh_white_list))
.route("/ipv4", web::put().to(add_ipv4_ssh_white_list))
.route("/ipv6", web::put().to(add_ipv6_ssh_white_list))
.route("/ipv4", web::delete().to(remove_ipv4_ssh_white_list))
.route("/ipv6", web::delete().to(remove_ipv6_ssh_white_list))
}
fn ssh_blacklist_scope() -> Scope {
web::scope("/blacklist")
.route("/ipv4", web::get().to(get_ipv4_ssh_black_list))
.route("/ipv6", web::get().to(get_ipv6_ssh_black_list))
.route("/ipv4", web::put().to(add_ipv4_ssh_black_list))
.route("/ipv6", web::put().to(add_ipv6_ssh_black_list))
.route("/ipv4", web::delete().to(remove_ipv4_ssh_black_list))
.route("/ipv6", web::delete().to(remove_ipv6_ssh_black_list))
}
// --- HTTP service handlers ---
async fn get_ipv4_http_service(service: web::Data<ProtocolFilter>) -> impl Responder {
let list = service.get_ipv4_http_service().await;
HttpResponse::Ok().json(list)
}
async fn get_ipv6_http_service(service: web::Data<ProtocolFilter>) -> impl Responder {
let list = service.get_ipv6_http_service().await;
HttpResponse::Ok().json(list)
}
async fn add_ipv4_http_service(
payload: web::Json<(SocketAddrV4, Vec<HttpMethod>)>,
service: web::Data<ProtocolFilter>,
) -> impl Responder {
let (addr, methods) = payload.into_inner();
match service.add_ipv4_http_service(addr, methods).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})),
}
}
async fn add_ipv6_http_service(
payload: web::Json<(SocketAddrV6, Vec<HttpMethod>)>,
service: web::Data<ProtocolFilter>,
) -> impl Responder {
let (addr, methods) = payload.into_inner();
match service.add_ipv6_http_service(addr, methods).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})),
}
}
async fn remove_ipv4_http_service(
payload: web::Json<(SocketAddrV4, Vec<HttpMethod>)>,
service: web::Data<ProtocolFilter>,
) -> impl Responder {
let (addr, methods) = payload.into_inner();
match service.remove_ipv4_http_service(addr, methods).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})),
}
}
async fn remove_ipv6_http_service(
payload: web::Json<(SocketAddrV6, Vec<HttpMethod>)>,
service: web::Data<ProtocolFilter>,
) -> impl Responder {
let (addr, methods) = payload.into_inner();
match service.remove_ipv6_http_service(addr, methods).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})),
}
}
// --- SSH service handlers ---
async fn get_ipv4_ssh_service(service: web::Data<ProtocolFilter>) -> impl Responder {
let list = service.get_ipv4_ssh_service().await;
HttpResponse::Ok().json(list)
}
async fn get_ipv6_ssh_service(service: web::Data<ProtocolFilter>) -> impl Responder {
let list = service.get_ipv6_ssh_service().await;
HttpResponse::Ok().json(list)
}
async fn add_ipv4_ssh_service(
ip_addr: web::Json<SocketAddrV4>,
service: web::Data<ProtocolFilter>,
) -> impl Responder {
match service.add_ipv4_ssh_service(ip_addr.into_inner()).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})),
}
}
async fn add_ipv6_ssh_service(
ip_addr: web::Json<SocketAddrV6>,
service: web::Data<ProtocolFilter>,
) -> impl Responder {
match service.add_ipv6_ssh_service(ip_addr.into_inner()).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})),
}
}
async fn remove_ipv4_ssh_service(
ip_addr: web::Json<SocketAddrV4>,
service: web::Data<ProtocolFilter>,
) -> impl Responder {
match service.remove_ipv4_ssh_service(ip_addr.into_inner()).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})),
}
}
async fn remove_ipv6_ssh_service(
ip_addr: web::Json<SocketAddrV6>,
service: web::Data<ProtocolFilter>,
) -> impl Responder {
match service.remove_ipv6_ssh_service(ip_addr.into_inner()).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})),
}
}
// --- SSH whitelist handlers ---
async fn is_ssh_white_list_enable(service: web::Data<ProtocolFilter>) -> impl Responder {
let enabled = service.is_ssh_white_list_enable().await;
HttpResponse::Ok().json(enabled)
}
async fn enable_ssh_white_list(service: web::Data<ProtocolFilter>) -> impl Responder {
match service.enable_ssh_white_list().await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})),
}
}
async fn disable_ssh_white_list(service: web::Data<ProtocolFilter>) -> impl Responder {
match service.disable_ssh_white_list().await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})),
}
}
async fn get_ipv4_ssh_white_list(service: web::Data<ProtocolFilter>) -> impl Responder {
let list = service.get_ipv4_ssh_white_list().await;
HttpResponse::Ok().json(list)
}
async fn get_ipv6_ssh_white_list(service: web::Data<ProtocolFilter>) -> impl Responder {
let list = service.get_ipv6_ssh_white_list().await;
HttpResponse::Ok().json(list)
}
async fn add_ipv4_ssh_white_list(
ip_addr: web::Json<Ipv4Addr>,
service: web::Data<ProtocolFilter>,
) -> impl Responder {
match service.add_ipv4_ssh_white_list(ip_addr.into_inner()).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})),
}
}
async fn add_ipv6_ssh_white_list(
ip_addr: web::Json<Ipv6Addr>,
service: web::Data<ProtocolFilter>,
) -> impl Responder {
match service.add_ipv6_ssh_white_list(ip_addr.into_inner()).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})),
}
}
async fn remove_ipv4_ssh_white_list(
ip_addr: web::Json<Ipv4Addr>,
service: web::Data<ProtocolFilter>,
) -> impl Responder {
match service.remove_ipv4_ssh_white_list(ip_addr.into_inner()).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})),
}
}
async fn remove_ipv6_ssh_white_list(
ip_addr: web::Json<Ipv6Addr>,
service: web::Data<ProtocolFilter>,
) -> impl Responder {
match service.remove_ipv6_ssh_white_list(ip_addr.into_inner()).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})),
}
}
// --- SSH blacklist handlers ---
async fn get_ipv4_ssh_black_list(service: web::Data<ProtocolFilter>) -> impl Responder {
let list = service.get_ipv4_ssh_black_list().await;
HttpResponse::Ok().json(list)
}
async fn get_ipv6_ssh_black_list(service: web::Data<ProtocolFilter>) -> impl Responder {
let list = service.get_ipv6_ssh_black_list().await;
HttpResponse::Ok().json(list)
}
async fn add_ipv4_ssh_black_list(
ip_addr: web::Json<Ipv4Addr>,
service: web::Data<ProtocolFilter>,
) -> impl Responder {
match service.add_ipv4_ssh_black_list(ip_addr.into_inner()).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})),
}
}
async fn add_ipv6_ssh_black_list(
ip_addr: web::Json<Ipv6Addr>,
service: web::Data<ProtocolFilter>,
) -> impl Responder {
match service.add_ipv6_ssh_black_list(ip_addr.into_inner()).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})),
}
}
async fn remove_ipv4_ssh_black_list(
ip_addr: web::Json<Ipv4Addr>,
service: web::Data<ProtocolFilter>,
) -> impl Responder {
match service.remove_ipv4_ssh_black_list(ip_addr.into_inner()).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})),
}
}
async fn remove_ipv6_ssh_black_list(
ip_addr: web::Json<Ipv6Addr>,
service: web::Data<ProtocolFilter>,
) -> impl Responder {
match service.remove_ipv6_ssh_black_list(ip_addr.into_inner()).await {
Ok(_) => HttpResponse::Ok().finish(),
Err(e) => HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})),
}
}

View File

@ -1,35 +1,19 @@
use actix_web::{get, web, HttpRequest, HttpResponse, Responder, Scope};
use actix_web::{web, HttpResponse, Responder, Scope};
use crate::core::infrastructure::health::SystemHealth;
use crate::web::websocket::health_websocket;
pub fn initialize() -> Scope {
web::scope("/health")
.service(get_current_metrics)
.service(get_health_status)
.service(websocket_metrics)
.route("/metrics", web::get().to(get_current_metrics))
.route("/status", web::get().to(get_health_status))
}
#[get("/metrics")]
async fn get_current_metrics(health: web::Data<SystemHealth>) -> impl Responder {
let metrics = health.get_current_metrics().await;
HttpResponse::Ok().json(metrics)
}
#[get("/status")]
async fn get_health_status(health: web::Data<SystemHealth>) -> impl Responder {
let status = health.is_system_healthy().await;
HttpResponse::Ok().json(status)
}
#[get("/websocket/metrics")]
async fn websocket_metrics(
req: HttpRequest,
stream: web::Payload,
health: web::Data<SystemHealth>,
) -> impl Responder {
match health_websocket::websocket_system_health(req, stream, health).await {
Ok(response) => response,
Err(err) => HttpResponse::InternalServerError().body(format!("WebSocket error: {}", err)),
}
}

View File

@ -1,14 +0,0 @@
use actix_web::{get, web, HttpResponse, Responder, Scope};
use crate::utils::boot_time::boot_time;
pub fn initialize() -> Scope {
web::scope("/misc")
.service(get_boot_time)
}
#[get("/boot_time")]
async fn get_boot_time() -> impl Responder {
let boot_time = boot_time();
HttpResponse::Ok().json(boot_time)
}

View File

@ -0,0 +1,12 @@
use actix_web::{web, HttpResponse, Responder, Scope};
pub fn initialize() -> Scope {
web::scope("/ml")
.route("/status", web::get().to(get_status))
}
async fn get_status() -> impl Responder {
HttpResponse::Ok().json(serde_json::json!({
"active": true
}))
}

View File

@ -1,22 +0,0 @@
use crate::core::infrastructure::ml_alert::MLAlert;
use crate::web::websocket::alert_websocket;
use actix_web::{get, web, HttpRequest, HttpResponse, Responder, Scope};
pub fn initialize() -> Scope {
web::scope("/ml")
.service(websocket_alert)
}
#[get("/websocket/alert")]
async fn websocket_alert(
req: HttpRequest,
stream: web::Payload,
ai: web::Data<MLAlert>,
) -> impl Responder {
match alert_websocket::websocket_alert(req, stream, ai).await {
Ok(response) => response,
Err(err) => {
HttpResponse::InternalServerError().body(format!("WebSocket error: {}", err))
}
}
}

View File

@ -1,5 +1,9 @@
pub mod control;
pub mod default;
pub mod misc;
pub mod ml_alert;
pub mod acl;
pub mod filter;
pub mod rate_limit;
pub mod stats;
pub mod health;
pub mod ml;
pub mod system;
pub mod default;
pub mod ws;

View File

@ -0,0 +1,42 @@
use actix_web::{web, HttpResponse, Responder, Scope};
use serde::{Deserialize, Serialize};
use crate::core::ebpf::rate_limit::RateLimitConfig;
#[derive(Serialize, Deserialize)]
pub struct RateLimitSettings {
pub packet_rate: Option<u64>,
pub syn_rate: Option<u64>,
pub udp_rate: Option<u64>,
pub dns_rate: Option<u64>,
pub window_ns: Option<u64>,
}
pub fn initialize() -> Scope {
web::scope("/rate-limit")
.route("/config", web::get().to(get_config))
.route("/config", web::put().to(set_config))
}
async fn get_config() -> impl Responder {
HttpResponse::Ok().json(RateLimitSettings {
packet_rate: Some(common::model::rate_limit::DEFAULT_PACKET_RATE),
syn_rate: Some(common::model::rate_limit::DEFAULT_SYN_RATE),
udp_rate: Some(common::model::rate_limit::DEFAULT_UDP_RATE),
dns_rate: Some(common::model::rate_limit::DEFAULT_DNS_RATE),
window_ns: Some(common::model::rate_limit::DEFAULT_WINDOW_NS),
})
}
async fn set_config(
settings: web::Json<RateLimitSettings>,
config: web::Data<RateLimitConfig>,
) -> impl Responder {
let s = settings.into_inner();
if let Some(v) = s.packet_rate { let _ = config.set_packet_rate(v); }
if let Some(v) = s.syn_rate { let _ = config.set_syn_rate(v); }
if let Some(v) = s.udp_rate { let _ = config.set_udp_rate(v); }
if let Some(v) = s.dns_rate { let _ = config.set_dns_rate(v); }
if let Some(v) = s.window_ns { let _ = config.set_window_ns(v); }
HttpResponse::Ok().json(serde_json::json!({"status": "ok"}))
}

View File

@ -0,0 +1,26 @@
use actix_web::{web, HttpResponse, Responder, Scope};
use crate::core::infrastructure::statistics::FlowStatistics;
pub fn initialize() -> Scope {
web::scope("/stats")
.route("/flows", web::get().to(get_all_flows))
.route("/flows/top/{n}", web::get().to(get_top_flows))
.route("/summary", web::get().to(get_summary))
}
async fn get_all_flows(stats: web::Data<FlowStatistics>) -> impl Responder {
HttpResponse::Ok().json(stats.get_all_flows())
}
async fn get_top_flows(
stats: web::Data<FlowStatistics>,
path: web::Path<usize>,
) -> impl Responder {
let n = path.into_inner();
HttpResponse::Ok().json(stats.get_top_flows(n))
}
async fn get_summary(stats: web::Data<FlowStatistics>) -> impl Responder {
HttpResponse::Ok().json(stats.get_summary())
}

View File

@ -0,0 +1,10 @@
use actix_web::{web, HttpResponse, Responder, Scope};
pub fn initialize() -> Scope {
web::scope("/system")
.route("/boot-time", web::get().to(get_boot_time))
}
async fn get_boot_time() -> impl Responder {
HttpResponse::Ok().json(crate::utils::boot_time::boot_time())
}

View File

@ -0,0 +1,46 @@
use actix_web::{web, HttpRequest, HttpResponse, Responder, Scope};
use crate::core::infrastructure::health::SystemHealth;
use crate::core::infrastructure::statistics::FlowStatistics;
use crate::core::ml::alert::MLAlert;
use crate::web::websocket::{alert_websocket, flow_websocket, health_websocket};
pub fn initialize() -> Scope {
web::scope("/ws")
.route("/health", web::get().to(health_ws))
.route("/alerts", web::get().to(alerts_ws))
.route("/flows", web::get().to(flows_ws))
}
async fn health_ws(
req: HttpRequest,
stream: web::Payload,
health: web::Data<SystemHealth>,
) -> impl Responder {
match health_websocket::websocket_system_health(req, stream, health).await {
Ok(response) => response,
Err(err) => HttpResponse::InternalServerError().json(serde_json::json!({"error": format!("WebSocket error: {}", err)})),
}
}
async fn alerts_ws(
req: HttpRequest,
stream: web::Payload,
ai: web::Data<MLAlert>,
) -> impl Responder {
match alert_websocket::websocket_alert(req, stream, ai).await {
Ok(response) => response,
Err(err) => HttpResponse::InternalServerError().json(serde_json::json!({"error": format!("WebSocket error: {}", err)})),
}
}
async fn flows_ws(
req: HttpRequest,
stream: web::Payload,
stats: web::Data<FlowStatistics>,
) -> impl Responder {
match flow_websocket::flow_stats_ws(req, stream, stats).await {
Ok(response) => response,
Err(err) => HttpResponse::InternalServerError().json(serde_json::json!({"error": format!("WebSocket error: {}", err)})),
}
}

View File

@ -4,7 +4,7 @@ use futures_util::StreamExt;
use macros::log;
use tokio::sync::broadcast;
use crate::core::infrastructure::ml_alert::{MLAlert, AlertMessage};
use crate::core::ml::alert::{MLAlert, AlertMessage};
use crate::model::error::http::HttpError;
use crate::model::error::misc::MiscError;
use crate::model::log::http::HttpLog;

View File

@ -1,190 +1,87 @@
use std::sync::Arc;
use std::time::Duration;
use actix_web::{web, HttpRequest, HttpResponse, Result};
use actix_ws::{handle, Message, MessageStream, Session};
use actix_web::{web, HttpRequest, HttpResponse};
use actix_ws::Message;
use futures_util::StreamExt;
use macros::log;
use tokio::time::{interval, Duration};
use tokio::time::interval;
use crate::core::ebpf::statistics::Statistics;
use crate::core::infrastructure::app_config::AppConfig;
use crate::model::direction::{Direction, FlowDirection};
use crate::model::error::http::HttpError;
use crate::model::error::misc::MiscError;
use crate::model::time_type::TimeType;
use crate::core::infrastructure::statistics::FlowStatistics;
use crate::model::flow_stats::FlowSubscription;
use crate::model::log::http::HttpLog;
pub async fn websocket_ipv4_flow(
/// Default subscription: all flows, no filter, 5 second interval
fn default_subscription() -> FlowSubscription {
FlowSubscription {
direction: None,
window_secs: None,
top_n: None,
interval_secs: Some(5),
}
}
pub async fn flow_stats_ws(
req: HttpRequest,
body: web::Payload,
path: web::Path<(Direction, FlowDirection, TimeType)>,
app_config: web::Data<AppConfig>,
statistics: web::Data<Statistics>,
) -> Result<HttpResponse> {
let app_config = app_config.into_inner();
let statistics = statistics.into_inner();
let (direction, flow_direction, time_type) = path.into_inner();
let (response, session, msg_stream) = handle(&req, body)?;
stats: web::Data<FlowStatistics>,
) -> Result<HttpResponse, actix_web::Error> {
let (response, mut session, mut msg_stream) = actix_ws::handle(&req, body)?;
actix_web::rt::spawn(async move {
handle_ipv4_flow_connection(
app_config,
statistics,
session,
msg_stream,
direction,
flow_direction,
time_type,
)
.await;
let mut subscription = default_subscription();
let mut ticker = interval(Duration::from_secs(
subscription.interval_secs.unwrap_or(5),
));
loop {
tokio::select! {
_ = ticker.tick() => {
let flows = stats.get_filtered_flows(&subscription);
if let Ok(json) = serde_json::to_string(&flows) {
if session.text(json).await.is_err() {
break;
}
}
}
msg = msg_stream.next() => {
match msg {
Some(Ok(Message::Text(text))) => {
// Client sends subscription query as JSON
match serde_json::from_str::<FlowSubscription>(&text) {
Ok(new_sub) => {
let new_interval = new_sub.interval_secs.unwrap_or(5).max(1);
subscription = new_sub;
subscription.interval_secs = Some(new_interval);
ticker = interval(Duration::from_secs(new_interval));
// Send immediate response with new filter
let flows = stats.get_filtered_flows(&subscription);
if let Ok(json) = serde_json::to_string(&flows) {
if session.text(json).await.is_err() {
break;
}
}
}
Err(e) => {
let err_msg = serde_json::json!({"error": format!("Invalid subscription: {}", e)});
if session.text(err_msg.to_string()).await.is_err() {
break;
}
}
}
}
Some(Ok(Message::Ping(bytes))) => {
if session.pong(&bytes).await.is_err() {
break;
}
}
Some(Ok(Message::Close(_))) | None => break,
_ => {}
}
}
}
}
});
Ok(response)
}
pub async fn websocket_ipv6_flow(
req: HttpRequest,
body: web::Payload,
path: web::Path<(Direction, FlowDirection, TimeType)>,
app_config: web::Data<AppConfig>,
statistics: web::Data<Statistics>,
) -> Result<HttpResponse> {
let app_config = app_config.into_inner();
let statistics = statistics.into_inner();
let (direction, flow_direction, time_type) = path.into_inner();
let (response, session, msg_stream) = handle(&req, body)?;
actix_web::rt::spawn(async move {
handle_ipv6_flow_connection(
app_config,
statistics,
session,
msg_stream,
direction,
flow_direction,
time_type,
)
.await;
});
Ok(response)
}
async fn handle_ipv4_flow_connection(
app_config: Arc<AppConfig>,
statistics: Arc<Statistics>,
mut session: Session,
mut msg_stream: MessageStream,
direction: Direction,
flow_direction: FlowDirection,
time_type: TimeType,
) {
let config = app_config.config.clone();
let refresh_interval = Duration::from_secs(config.refresh_interval);
let mut data_interval = interval(refresh_interval);
loop {
tokio::select! {
msg_result = msg_stream.next() => {
if !handle_client_message(&mut session, msg_result).await {
break;
}
},
_ = data_interval.tick() => {
if !send_ipv4_flow_data(&statistics, &mut session, direction, flow_direction, time_type).await {
break;
}
},
}
}
let _ = session.close(None).await;
}
async fn handle_ipv6_flow_connection(
app_config: Arc<AppConfig>,
statistics: Arc<Statistics>,
mut session: Session,
mut msg_stream: MessageStream,
direction: Direction,
flow_direction: FlowDirection,
time_type: TimeType,
) {
let config = app_config.config.clone();
let refresh_interval = Duration::from_secs(config.refresh_interval);
let mut data_interval = interval(refresh_interval);
loop {
tokio::select! {
msg_result = msg_stream.next() => {
if !handle_client_message(&mut session, msg_result).await {
break;
}
},
_ = data_interval.tick() => {
if !send_ipv6_flow_data(&statistics, &mut session, direction, flow_direction, time_type).await {
break;
}
},
}
}
let _ = session.close(None).await;
}
async fn handle_client_message(
session: &mut Session,
msg_result: Option<Result<Message, actix_ws::ProtocolError>>,
) -> bool {
match msg_result {
Some(Ok(Message::Text(_))) => true,
Some(Ok(Message::Ping(bytes))) => session.pong(&bytes).await.is_ok(),
Some(Ok(Message::Close(reason))) => {
let _ = (session.clone()).close(reason).await;
false
}
Some(Err(err)) => {
log!(HttpError::WebSocketError(err));
false
}
None => false,
_ => true,
}
}
async fn send_ipv4_flow_data(
statistics: &Arc<Statistics>,
session: &mut Session,
direction: Direction,
flow_direction: FlowDirection,
time_type: TimeType,
) -> bool {
let flow_data = statistics
.get_ipv4_flow_data(direction, flow_direction, time_type)
.await;
match serde_json::to_string(&flow_data) {
Ok(json) => session.text(json).await.is_ok(),
Err(err) => {
log!(MiscError::SerializeError(err));
true
}
}
}
async fn send_ipv6_flow_data(
statistics: &Arc<Statistics>,
session: &mut Session,
direction: Direction,
flow_direction: FlowDirection,
time_type: TimeType,
) -> bool {
let flow_data = statistics
.get_ipv6_flow_data(direction, flow_direction, time_type)
.await;
match serde_json::to_string(&flow_data) {
Ok(json) => session.text(json).await.is_ok(),
Err(err) => {
log!(MiscError::SerializeError(err));
true
}
}
}

View File

@ -65,6 +65,7 @@ async fn handle_client_message(
msg_result: Option<Result<Message, actix_ws::ProtocolError>>,
) -> bool {
match msg_result {
// Text messages are intentionally ignored; no client commands are supported
Some(Ok(Message::Text(_))) => true,
Some(Ok(Message::Ping(bytes))) => session.pong(&bytes).await.is_ok(),
Some(Ok(Message::Close(reason))) => {

View File

@ -1,3 +1,3 @@
pub mod alert_websocket;
pub mod flow_websocket;
pub mod health_websocket;
pub mod alert_websocket;