Optimized statistics structure and web APIs

This commit is contained in:
DaLaw2 2024-12-13 14:30:11 +08:00
parent 8491c6b95e
commit 4b1912be7b
8 changed files with 198 additions and 213 deletions

View File

@ -26,6 +26,11 @@ panic = "abort"
[profile.release]
panic = "abort"
opt-level = 3
lto = true
strip = true
debug = false
overflow-checks = false
[profile.release.package.net-guardia-ebpf]
debug = 2

View File

@ -1,6 +1,8 @@
use crate::core::system::System;
use crate::model::direction::Direction;
use crate::model::flow_stats::FlowStats;
use crate::model::flow_type::FlowType;
use crate::model::ip_address::SocketAddressType;
use crate::model::time_type::TimeType;
use crate::utils::log_entry::system::SystemEntry;
use aya::maps::{HashMap as AyaHashMap, MapData};
use aya::Pod;
@ -16,39 +18,62 @@ static STATISTICS: OnceLock<RwLock<Statistics>> = OnceLock::new();
pub struct Statistics {
terminate: bool,
ipv4_src_1min: AyaHashMap<MapData, EbpfAddrPortV4, EbpfFlowStats>,
ipv4_src_10min: AyaHashMap<MapData, EbpfAddrPortV4, EbpfFlowStats>,
ipv4_src_1hour: AyaHashMap<MapData, EbpfAddrPortV4, EbpfFlowStats>,
ipv6_src_1min: AyaHashMap<MapData, EbpfAddrPortV6, EbpfFlowStats>,
ipv6_src_10min: AyaHashMap<MapData, EbpfAddrPortV6, EbpfFlowStats>,
ipv6_src_1hour: AyaHashMap<MapData, EbpfAddrPortV6, EbpfFlowStats>,
ipv4_dst_1min: AyaHashMap<MapData, EbpfAddrPortV4, EbpfFlowStats>,
ipv4_dst_10min: AyaHashMap<MapData, EbpfAddrPortV4, EbpfFlowStats>,
ipv4_dst_1hour: AyaHashMap<MapData, EbpfAddrPortV4, EbpfFlowStats>,
ipv6_dst_1min: AyaHashMap<MapData, EbpfAddrPortV6, EbpfFlowStats>,
ipv6_dst_10min: AyaHashMap<MapData, EbpfAddrPortV6, EbpfFlowStats>,
ipv6_dst_1hour: AyaHashMap<MapData, EbpfAddrPortV6, EbpfFlowStats>,
ipv4_maps: StdHashMap<(Direction, TimeType), FlowMap<EbpfAddrPortV4>>,
ipv6_maps: StdHashMap<(Direction, TimeType), FlowMap<EbpfAddrPortV6>>,
}
impl Statistics {
const MAP_CONFIGS: [((Direction, TimeType), (&'static str, &'static str)); 6] = [
(
(Direction::Source, TimeType::_1Min),
("IPV4_SRC_1MIN", "IPV6_SRC_1MIN"),
),
(
(Direction::Source, TimeType::_10Min),
("IPV4_SRC_10MIN", "IPV6_SRC_10MIN"),
),
(
(Direction::Source, TimeType::_1Hour),
("IPV4_SRC_1HOUR", "IPV6_SRC_1HOUR"),
),
(
(Direction::Destination, TimeType::_1Min),
("IPV4_DST_1MIN", "IPV6_DST_1MIN"),
),
(
(Direction::Destination, TimeType::_10Min),
("IPV4_DST_10MIN", "IPV6_DST_10MIN"),
),
(
(Direction::Destination, TimeType::_1Hour),
("IPV4_DST_1HOUR", "IPV6_DST_1HOUR"),
),
];
pub async fn initialize() -> anyhow::Result<()> {
info!("{}", SystemEntry::Initializing);
let mut system = System::instance_mut().await;
let ebpf = &mut system.ebpf;
let mut ipv4_maps = StdHashMap::new();
let mut ipv6_maps = StdHashMap::new();
for (key, (ipv4_name, ipv6_name)) in Self::MAP_CONFIGS {
ipv4_maps.insert(
key,
FlowMap {
map: AyaHashMap::try_from(ebpf.take_map(ipv4_name).unwrap())?,
},
);
ipv6_maps.insert(
key,
FlowMap {
map: AyaHashMap::try_from(ebpf.take_map(ipv6_name).unwrap())?,
},
);
}
let monitor = Statistics {
terminate: false,
ipv4_src_1min: AyaHashMap::try_from(ebpf.take_map("IPV4_SRC_1MIN").unwrap())?,
ipv4_src_10min: AyaHashMap::try_from(ebpf.take_map("IPV4_SRC_10MIN").unwrap())?,
ipv4_src_1hour: AyaHashMap::try_from(ebpf.take_map("IPV4_SRC_1HOUR").unwrap())?,
ipv6_src_1min: AyaHashMap::try_from(ebpf.take_map("IPV6_SRC_1MIN").unwrap())?,
ipv6_src_10min: AyaHashMap::try_from(ebpf.take_map("IPV6_SRC_10MIN").unwrap())?,
ipv6_src_1hour: AyaHashMap::try_from(ebpf.take_map("IPV6_SRC_1HOUR").unwrap())?,
ipv4_dst_1min: AyaHashMap::try_from(ebpf.take_map("IPV4_DST_1MIN").unwrap())?,
ipv4_dst_10min: AyaHashMap::try_from(ebpf.take_map("IPV4_DST_10MIN").unwrap())?,
ipv4_dst_1hour: AyaHashMap::try_from(ebpf.take_map("IPV4_DST_1HOUR").unwrap())?,
ipv6_dst_1min: AyaHashMap::try_from(ebpf.take_map("IPV6_DST_1MIN").unwrap())?,
ipv6_dst_10min: AyaHashMap::try_from(ebpf.take_map("IPV6_DST_10MIN").unwrap())?,
ipv6_dst_1hour: AyaHashMap::try_from(ebpf.take_map("IPV6_DST_1HOUR").unwrap())?,
ipv4_maps,
ipv6_maps,
};
STATISTICS.get_or_init(|| RwLock::new(monitor));
info!("{}", SystemEntry::InitializeComplete);
@ -86,89 +111,72 @@ impl Statistics {
}
pub async fn cleanup_expired_flows() {
const ONE_MIN: u64 = 60 * 1_000_000_000;
const TEN_MIN: u64 = 10 * ONE_MIN;
const ONE_HOUR: u64 = 60 * ONE_MIN;
let mut monitor = Statistics::instance_mut().await;
let boot_time = System::boot_time().await;
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos() as u64;
Statistics::cleanup_map(&mut monitor.ipv4_src_1min, now, ONE_MIN).await;
Statistics::cleanup_map(&mut monitor.ipv4_dst_1min, now, ONE_MIN).await;
Statistics::cleanup_map(&mut monitor.ipv4_src_10min, now, TEN_MIN).await;
Statistics::cleanup_map(&mut monitor.ipv4_dst_10min, now, TEN_MIN).await;
Statistics::cleanup_map(&mut monitor.ipv4_src_1hour, now, ONE_HOUR).await;
Statistics::cleanup_map(&mut monitor.ipv4_dst_1hour, now, ONE_HOUR).await;
Statistics::cleanup_map(&mut monitor.ipv6_src_1min, now, ONE_MIN).await;
Statistics::cleanup_map(&mut monitor.ipv6_dst_1min, now, ONE_MIN).await;
Statistics::cleanup_map(&mut monitor.ipv6_src_10min, now, TEN_MIN).await;
Statistics::cleanup_map(&mut monitor.ipv6_dst_10min, now, TEN_MIN).await;
Statistics::cleanup_map(&mut monitor.ipv6_src_1hour, now, ONE_HOUR).await;
Statistics::cleanup_map(&mut monitor.ipv6_dst_1hour, now, ONE_HOUR).await;
monitor
.ipv4_maps
.iter_mut()
.for_each(|((_, time_type), map)| map.cleanup(boot_time, now, time_type.duration()));
monitor
.ipv6_maps
.iter_mut()
.for_each(|((_, time_type), map)| map.cleanup(boot_time, now, time_type.duration()));
}
async fn cleanup_map<K>(map: &mut AyaHashMap<MapData, K, EbpfFlowStats>, now: u64, window: u64)
where
K: Pod,
{
let boot_time = System::boot_time().await;
let expired_keys: Vec<K> = map
.iter()
.filter_map(|result| {
if let Ok((key, stats)) = result {
if now - stats[2] - boot_time > window {
Some(key)
} else {
None
}
} else {
None
}
})
.collect();
for key in expired_keys {
let _ = map.remove(&key);
}
}
pub async fn get_ipv4_flow_data(flow_type: FlowType) -> StdHashMap<SocketAddrV4, FlowStats> {
pub async fn get_ipv4_flow_data(
direction: Direction,
time_type: TimeType,
) -> StdHashMap<SocketAddrV4, FlowStats> {
let monitor = Statistics::instance().await;
let iter = match flow_type {
FlowType::Src1Min => monitor.ipv4_src_1min.iter(),
FlowType::Src10Min => monitor.ipv4_src_10min.iter(),
FlowType::Src1Hour => monitor.ipv4_src_1hour.iter(),
FlowType::Dst1Min => monitor.ipv4_dst_1min.iter(),
FlowType::Dst10Min => monitor.ipv4_dst_10min.iter(),
FlowType::Dst1Hour => monitor.ipv4_dst_1hour.iter(),
};
iter.filter_map(Result::ok)
.map(|(key, value)| {
let ip = Ipv4Addr::from(key[0]);
let port = key[1] as u16;
(SocketAddrV4::new(ip, port), FlowStats::from(value))
})
.collect()
monitor
.ipv4_maps
.get(&(direction, time_type))
.map(|map| map.get_map())
.unwrap()
}
pub async fn get_ipv6_flow_data(flow_type: FlowType) -> StdHashMap<SocketAddrV6, FlowStats> {
pub async fn get_ipv6_flow_data(
direction: Direction,
time_type: TimeType,
) -> StdHashMap<SocketAddrV6, FlowStats> {
let monitor = Statistics::instance().await;
let iter = match flow_type {
FlowType::Src1Min => monitor.ipv6_src_1min.iter(),
FlowType::Src10Min => monitor.ipv6_src_10min.iter(),
FlowType::Src1Hour => monitor.ipv6_src_1hour.iter(),
FlowType::Dst1Min => monitor.ipv6_dst_1min.iter(),
FlowType::Dst10Min => monitor.ipv6_dst_10min.iter(),
FlowType::Dst1Hour => monitor.ipv6_dst_1hour.iter(),
};
iter.filter_map(Result::ok)
.map(|(key, value)| {
let ip = Ipv6Addr::from(key[0]);
let port = key[1] as u16;
(SocketAddrV6::new(ip, port, 0, 0), FlowStats::from(value))
})
.collect()
monitor
.ipv6_maps
.get(&(direction, time_type))
.map(|map| map.get_map())
.unwrap()
}
}
struct FlowMap<T> {
map: AyaHashMap<MapData, T, EbpfFlowStats>,
}
impl<T: SocketAddressType + Pod> FlowMap<T> {
fn get_map(&self) -> StdHashMap<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[2] - boot_time > window).then_some(key))
})
.collect();
expired_keys.iter().for_each(|key| {
let _ = self.map.remove(key);
});
}
}

View File

@ -1,18 +0,0 @@
use serde::{Deserialize, Serialize};
#[derive(Serialize, Deserialize, Debug, Copy, Clone)]
#[serde(rename_all = "lowercase")]
pub enum FlowType {
#[serde(rename = "src_1min")]
Src1Min,
#[serde(rename = "src_10min")]
Src10Min,
#[serde(rename = "src_1hour")]
Src1Hour,
#[serde(rename = "dst_1min")]
Dst1Min,
#[serde(rename = "dst_10min")]
Dst10Min,
#[serde(rename = "dst_1hour")]
Dst1Hour,
}

View File

@ -1,6 +1,6 @@
use net_guardia_common::model::ip_address::{IPv4, IPv6};
use net_guardia_common::model::ip_address::{IPv4, IPv6, EbpfAddrPortV4, EbpfAddrPortV6};
use std::hash::Hash;
use std::net::{Ipv4Addr, Ipv6Addr};
use std::net::{Ipv4Addr, Ipv6Addr, SocketAddrV4, SocketAddrV6};
pub trait IpAddressType: Copy {
type Native: Eq + PartialEq + Hash;
@ -31,3 +31,33 @@ impl IpAddressType for IPv6 {
native.into()
}
}
pub trait SocketAddressType: Copy {
type Native: Eq + PartialEq + Hash;
fn into_native(self) -> Self::Native;
fn from_native(native: Self::Native) -> Self;
}
impl SocketAddressType for EbpfAddrPortV4 {
type Native = SocketAddrV4;
fn into_native(self) -> Self::Native {
SocketAddrV4::new(Ipv4Addr::from(self[0]), self[1] as u16)
}
fn from_native(native: Self::Native) -> Self {
[(*native.ip()).into(), native.port() as u32]
}
}
impl SocketAddressType for EbpfAddrPortV6 {
type Native = SocketAddrV6;
fn into_native(self) -> Self::Native {
SocketAddrV6::new(Ipv6Addr::from(self[0]), self[1] as u16, 0, 0)
}
fn from_native(native: Self::Native) -> Self {
[(*native.ip()).into(), native.port() as u128]
}
}

View File

@ -1,7 +1,7 @@
pub mod config;
pub mod direction;
pub mod flow_stats;
pub mod flow_type;
pub mod http_method;
pub mod ip_address;
pub mod list_type;
pub mod time_type;

View File

@ -0,0 +1,19 @@
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

@ -1,122 +1,58 @@
use crate::core::statistics::Statistics;
use crate::model::flow_type::FlowType;
use crate::web::utils::flow_websocket::{IPv4FlowWebSocket, IPv6FlowWebSocket};
use actix_web::{get, web, Error, HttpRequest, HttpResponse, Responder, Scope};
use actix_web_actors::ws::start;
use crate::model::direction::Direction;
use crate::model::time_type::TimeType;
pub fn initialize() -> Scope {
web::scope("/statistics")
.service(get_ipv4_src_1min)
.service(get_ipv4_src_10min)
.service(get_ipv4_src_1hour)
.service(get_ipv6_src_1min)
.service(get_ipv6_src_10min)
.service(get_ipv6_src_1hour)
.service(get_ipv4_dst_1min)
.service(get_ipv4_dst_10min)
.service(get_ipv4_dst_1hour)
.service(get_ipv6_dst_1min)
.service(get_ipv6_dst_10min)
.service(get_ipv6_dst_1hour)
.service(get_ipv4_flow)
.service(get_ipv6_flow)
.service(websocket_ipv4)
.service(websocket_ipv6)
}
#[get("/get/ipv4/src/1min")]
async fn get_ipv4_src_1min() -> impl Responder {
let flow_data = Statistics::get_ipv4_flow_data(FlowType::Src1Min).await;
#[get("/get/ipv4/{direction}/{time_type}")]
async fn get_ipv4_flow(path: web::Path<(Direction, TimeType)>) -> impl Responder {
let (direction, time_type) = path.into_inner();
let flow_data = Statistics::get_ipv4_flow_data(direction, time_type).await;
HttpResponse::Ok().json(web::Json(flow_data))
}
#[get("/get/ipv4/src/10min")]
async fn get_ipv4_src_10min() -> impl Responder {
let flow_data = Statistics::get_ipv4_flow_data(FlowType::Src10Min).await;
#[get("/get/ipv6/{direction}/{time_type}")]
async fn get_ipv6_flow(path: web::Path<(Direction, TimeType)>) -> impl Responder {
let (direction, time_type) = path.into_inner();
let flow_data = Statistics::get_ipv6_flow_data(direction, time_type).await;
HttpResponse::Ok().json(web::Json(flow_data))
}
#[get("/get/ipv4/src/1hour")]
async fn get_ipv4_src_1hour() -> impl Responder {
let flow_data = Statistics::get_ipv4_flow_data(FlowType::Src1Hour).await;
HttpResponse::Ok().json(web::Json(flow_data))
}
#[get("/get/ipv6/src/1min")]
async fn get_ipv6_src_1min() -> impl Responder {
let flow_data = Statistics::get_ipv6_flow_data(FlowType::Src1Min).await;
HttpResponse::Ok().json(web::Json(flow_data))
}
#[get("/get/ipv6/src/10min")]
async fn get_ipv6_src_10min() -> impl Responder {
let flow_data = Statistics::get_ipv6_flow_data(FlowType::Src10Min).await;
HttpResponse::Ok().json(web::Json(flow_data))
}
#[get("/get/ipv6/src/1hour")]
async fn get_ipv6_src_1hour() -> impl Responder {
let flow_data = Statistics::get_ipv6_flow_data(FlowType::Src1Hour).await;
HttpResponse::Ok().json(web::Json(flow_data))
}
#[get("/get/ipv4/dst/1min")]
async fn get_ipv4_dst_1min() -> impl Responder {
let flow_data = Statistics::get_ipv4_flow_data(FlowType::Dst1Min).await;
HttpResponse::Ok().json(web::Json(flow_data))
}
#[get("/get/ipv4/dst/10min")]
async fn get_ipv4_dst_10min() -> impl Responder {
let flow_data = Statistics::get_ipv4_flow_data(FlowType::Dst10Min).await;
HttpResponse::Ok().json(web::Json(flow_data))
}
#[get("/get/ipv4/dst/1hour")]
async fn get_ipv4_dst_1hour() -> impl Responder {
let flow_data = Statistics::get_ipv4_flow_data(FlowType::Dst1Hour).await;
HttpResponse::Ok().json(web::Json(flow_data))
}
#[get("/get/ipv6/dst/1min")]
async fn get_ipv6_dst_1min() -> impl Responder {
let flow_data = Statistics::get_ipv6_flow_data(FlowType::Src1Min).await;
HttpResponse::Ok().json(web::Json(flow_data))
}
#[get("/get/ipv6/dst/10min")]
async fn get_ipv6_dst_10min() -> impl Responder {
let flow_data = Statistics::get_ipv6_flow_data(FlowType::Src10Min).await;
HttpResponse::Ok().json(web::Json(flow_data))
}
#[get("/get/ipv6/dst/1hour")]
async fn get_ipv6_dst_1hour() -> impl Responder {
let flow_data = Statistics::get_ipv6_flow_data(FlowType::Src1Hour).await;
HttpResponse::Ok().json(web::Json(flow_data))
}
#[get("/websocket/ipv4/{flow_type}")]
#[get("/websocket/ipv4/{direction}/{time_type}")]
async fn websocket_ipv4(
req: HttpRequest,
stream: web::Payload,
path: web::Path<FlowType>,
path: web::Path<(Direction, TimeType)>,
) -> Result<HttpResponse, Error> {
let flow_type = path.into_inner();
let (direction, time_type) = path.into_inner();
let websocket = IPv4FlowWebSocket {
flow_type,
direction,
time_type,
interval: None,
};
start(websocket, &req, stream)
}
#[get("/websocket/ipv6/{flow_type}")]
#[get("/websocket/ipv6/{direction}/{time_type}")]
async fn websocket_ipv6(
req: HttpRequest,
stream: web::Payload,
path: web::Path<FlowType>,
path: web::Path<(Direction, TimeType)>,
) -> Result<HttpResponse, Error> {
let flow_type = path.into_inner();
let (direction, time_type) = path.into_inner();
let websocket = IPv6FlowWebSocket {
flow_type,
direction,
time_type,
interval: None,
};
start(websocket, &req, stream)

View File

@ -1,12 +1,14 @@
use crate::core::config_manager::ConfigManager;
use crate::core::statistics::Statistics;
use crate::model::flow_type::FlowType;
use actix::prelude::*;
use actix_web_actors::ws;
use std::time::Duration;
use crate::model::direction::Direction;
use crate::model::time_type::TimeType;
pub struct IPv4FlowWebSocket {
pub flow_type: FlowType,
pub direction: Direction,
pub time_type: TimeType,
pub interval: Option<SpawnHandle>,
}
@ -17,8 +19,9 @@ impl Actor for IPv4FlowWebSocket {
let config = ConfigManager::now_blocking();
let refresh_interval = Duration::from_secs(config.refresh_interval);
let interval = ctx.run_interval(refresh_interval, |actor, ctx| {
let flow_type = actor.flow_type.clone();
let future = async move { Statistics::get_ipv4_flow_data(flow_type).await };
let direction = actor.direction.clone();
let time_type = actor.time_type.clone();
let future = async move { Statistics::get_ipv4_flow_data(direction, time_type).await };
ctx.wait(future.into_actor(actor).map(|flow_data, _, ctx| {
if let Ok(json) = serde_json::to_string(&flow_data) {
ctx.text(json);
@ -50,7 +53,8 @@ impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for IPv4FlowWebSocket
}
pub struct IPv6FlowWebSocket {
pub flow_type: FlowType,
pub direction: Direction,
pub time_type: TimeType,
pub interval: Option<SpawnHandle>,
}
@ -61,8 +65,9 @@ impl Actor for IPv6FlowWebSocket {
let config = ConfigManager::now_blocking();
let refresh_interval = Duration::from_secs(config.refresh_interval);
let interval = ctx.run_interval(refresh_interval, |actor, ctx| {
let flow_type = actor.flow_type.clone();
let future = async move { Statistics::get_ipv6_flow_data(flow_type).await };
let direction = actor.direction.clone();
let time_type = actor.time_type.clone();
let future = async move { Statistics::get_ipv6_flow_data(direction, time_type).await };
ctx.wait(future.into_actor(actor).map(|flow_data, _, ctx| {
if let Ok(json) = serde_json::to_string(&flow_data) {
ctx.text(json);