DaLaw2/NetGuardia
1
0
Fork 0
mirror of https://github.com/DaLaw2/NetGuardia.git synced 2026-08-24 14:10:28 +09:00
Code Issues Packages Projects Releases Wiki Activity

Fix expire key not clear bug

Browse Source
This commit is contained in:
DaLaw2 2024-11-30 21:18:01 +08:00
parent efc3635246
commit c3ce182c6e
7 changed files with 178 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
Cargo.lock generated
View File

@ -1,6 +1,6 @@
# This file is automatically @generated by Cargo.
# It is not intended for manual editing.
version = 3
version = 4
[[package]]
name = "actix"
@ -1186,6 +1186,7 @@ dependencies = [
"net-guardia-common",
"rust-embed",
"serde",
"serde_json",
"thiserror 2.0.3",
"tokio",
"toml",

1
net-guardia/Cargo.toml
View File

@ -23,6 +23,7 @@ actix-cors = "0.7.0"
actix = "0.13.5"
mime_guess = "2.0.5"
actix-web-actors = "4.3.0"
serde_json = "1.0.133"
[build-dependencies]
cargo_metadata = { workspace = true }

65
net-guardia/src/core/monitor.rs
View File

@ -2,6 +2,7 @@ use crate::core::system::System;
use crate::model::flow_type::{IPv4FlowType, IPv6FlowType};
use crate::utils::log_entry::system::SystemEntry;
use aya::maps::{HashMap as AyaHashMap, MapData};
use aya::Pod;
use net_guardia_common::model::flow_stats::FlowStats as EbpfFlowStats;
use net_guardia_common::model::ip_address::{
AddrPortV4 as EbpfAddrPortV4, AddrPortV6 as EbpfAddrPortV6,
@ -14,6 +15,7 @@ use tracing::info;
static MONITOR: OnceLock<RwLock<Monitor>> = OnceLock::new();
pub struct Monitor {
terminate: bool,
src_ipv4_1min: AyaHashMap<MapData, EbpfAddrPortV4, EbpfFlowStats>,
src_ipv4_10min: AyaHashMap<MapData, EbpfAddrPortV4, EbpfFlowStats>,
src_ipv4_1hour: AyaHashMap<MapData, EbpfAddrPortV4, EbpfFlowStats>,
@ -34,6 +36,7 @@ impl Monitor {
let mut system = System::instance_mut().await;
let mut ebpf = &mut system.ebpf;
let monitor = Monitor {
terminate: false,
src_ipv4_1min: AyaHashMap::try_from(ebpf.take_map("SRC_IPV4_1MIN").unwrap())?,
src_ipv4_10min: AyaHashMap::try_from(ebpf.take_map("SRC_IPV4_10MIN").unwrap())?,
src_ipv4_1hour: AyaHashMap::try_from(ebpf.take_map("SRC_IPV4_1HOUR").unwrap())?,
@ -68,6 +71,68 @@ impl Monitor {
once_lock.write().await
}
pub async fn run() {
tokio::spawn(async {
loop {
Monitor::cleanup_expired_flows().await;
tokio::time::sleep(tokio::time::Duration::from_secs(60)).await;
}
});
}
pub async fn terminate() {
let mut monitor = Monitor::instance_mut().await;
monitor.terminate = true;
}
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 = Monitor::instance_mut().await;
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos() as u64;
Monitor::cleanup_map(&mut monitor.src_ipv4_1min, now, ONE_MIN).await;
Monitor::cleanup_map(&mut monitor.dst_ipv4_1min, now, ONE_MIN).await;
Monitor::cleanup_map(&mut monitor.src_ipv4_10min, now, TEN_MIN).await;
Monitor::cleanup_map(&mut monitor.dst_ipv4_10min, now, TEN_MIN).await;
Monitor::cleanup_map(&mut monitor.src_ipv4_1hour, now, ONE_HOUR).await;
Monitor::cleanup_map(&mut monitor.dst_ipv4_1hour, now, ONE_HOUR).await;
Monitor::cleanup_map(&mut monitor.src_ipv6_1min, now, ONE_MIN).await;
Monitor::cleanup_map(&mut monitor.dst_ipv6_1min, now, ONE_MIN).await;
Monitor::cleanup_map(&mut monitor.src_ipv6_10min, now, TEN_MIN).await;
Monitor::cleanup_map(&mut monitor.dst_ipv6_10min, now, TEN_MIN).await;
Monitor::cleanup_map(&mut monitor.src_ipv6_1hour, now, ONE_HOUR).await;
Monitor::cleanup_map(&mut monitor.dst_ipv6_1hour, now, ONE_HOUR).await;
}
async fn cleanup_map<K>(map: &mut AyaHashMap<MapData, K, EbpfFlowStats>, now: u64, window: u64)
where
K: Pod,
{
let expired_keys: Vec<K> = map
.iter()
.filter_map(|result| {
if let Ok((key, stats)) = result {
if now - stats[2] > window {
Some(key)
} else {
None
}
} else {
None
}
})
.collect();
for key in expired_keys {
let _ = map.remove(&key);
}
}
pub async fn get_ipv4_flow_data(
ipv4_flow_type: IPv4FlowType,
) -> StdHashMap<EbpfAddrPortV4, EbpfFlowStats> {

2
net-guardia/src/core/system.rs
View File

@ -53,6 +53,7 @@ impl System {
pub async fn run() -> anyhow::Result<()> {
info!("{}", SystemEntry::Online);
Monitor::run().await;
let config = ConfigManager::now().await;
HttpServer::new(|| {
let cors = actix_cors::Cors::default()
@ -73,6 +74,7 @@ impl System {
pub async fn terminate() -> anyhow::Result<()> {
info!("{}", SystemEntry::Terminating);
Monitor::terminate().await;
info!("{}", SystemEntry::TerminateComplete);
Ok(())
}

7
net-guardia/src/model/flow_type.rs
View File

@ -1,4 +1,7 @@
#[derive(Copy, Clone)]
use std::str::FromStr;
use serde::Deserialize;
#[derive(Deserialize, Debug, Copy, Clone)]
pub enum IPv4FlowType {
SrcIPv4_1Min,
SrcIPv4_10Min,
@ -8,7 +11,7 @@ pub enum IPv4FlowType {
DstIPv4_1Hour,
}
#[derive(Copy, Clone)]
#[derive(Deserialize, Debug, Copy, Clone)]
pub enum IPv6FlowType {
SrcIPv6_1Min,
SrcIPv6_10Min,

37
net-guardia/src/web/api/monitor.rs
View File

@ -1,7 +1,10 @@
use crate::core::monitor::Monitor;
use crate::model::flow_type::{IPv4FlowType, IPv6FlowType};
use actix_web::{get, web, Error, HttpRequest, HttpResponse, Responder, Scope};
use crate::web::utils::map_util::{transform_ipv4_flow_data, transform_ipv6_flow_data};
use actix_web::{get, web, Error, HttpRequest, HttpResponse, Responder, Scope};
use actix_web_actors::ws::start;
use tracing::info;
use crate::web::utils::flow_websocket::{IPv4FlowWebSocket, IPv6FlowWebSocket};
pub fn initialize() -> Scope {
web::scope("/monitor")
@ -17,6 +20,8 @@ pub fn initialize() -> Scope {
.service(get_dst_ipv6_1min)
.service(get_dst_ipv6_10min)
.service(get_dst_ipv6_1hour)
.service(websocket_ipv4)
.service(websocket_ipv6)
}
#[get("/get/src/ipv4/1min")]
@ -103,7 +108,31 @@ async fn get_dst_ipv6_1hour() -> impl Responder {
HttpResponse::Ok().json(web::Json(formated))
}
#[get("/websocket")]
async fn websocket(req: HttpRequest, stream: web::Payload) -> Result<HttpResponse, Error> {
Ok(HttpResponse::Forbidden().finish())
#[get("/websocket/ipv4/{flow_type}")]
async fn websocket_ipv4(
req: HttpRequest,
stream: web::Payload,
path: web::Path<IPv4FlowType>,
) -> Result<HttpResponse, Error> {
info!("Try to connect ipv4 websocket");
let flow_type = path.into_inner();
let websocket = IPv4FlowWebSocket {
flow_type,
interval: None,
};
start(websocket, &req, stream)
}
#[get("/websocket/ipv6/{flow_type}")]
async fn websocket_ipv6(
req: HttpRequest,
stream: web::Payload,
path: web::Path<IPv6FlowType>,
) -> Result<HttpResponse, Error> {
let flow_type = path.into_inner();
let websocket = IPv6FlowWebSocket {
flow_type,
interval: None,
};
start(websocket, &req, stream)
}

73
net-guardia/src/web/utils/flow_websocket.rs
View File

@ -1,7 +1,13 @@
use crate::core::config_manager::ConfigManager;
use crate::core::monitor::Monitor;
use crate::model::flow_type::{IPv4FlowType, IPv6FlowType};
use crate::web::utils::map_util::{transform_ipv4_flow_data, transform_ipv6_flow_data};
use actix::prelude::*;
use actix_web_actors::ws;
use std::time::Duration;
pub struct IPv4FlowWebSocket {
pub flow_type: IPv4FlowType,
pub interval: Option<SpawnHandle>,
}
@ -9,15 +15,46 @@ impl Actor for IPv4FlowWebSocket {
type Context = ws::WebsocketContext<Self>;
fn started(&mut self, ctx: &mut Self::Context) {
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 {
let flow_data = Monitor::get_ipv4_flow_data(flow_type).await;
transform_ipv4_flow_data(flow_data)
};
ctx.wait(future.into_actor(actor).map(|flow_data, _, ctx| {
if let Ok(json) = serde_json::to_string(&flow_data) {
ctx.text(json);
}
}));
});
self.interval = Some(interval);
}
fn stopping(&mut self, ctx: &mut Self::Context) -> Running {
if let Some(interval) = self.interval.take() {
ctx.cancel_future(interval);
}
Running::Stop
}
}
impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for IPv4FlowWebSocket {
fn handle(&mut self, msg: Result<ws::Message, ws::ProtocolError>, ctx: &mut Self::Context) {
match msg {
Ok(ws::Message::Ping(msg)) => ctx.pong(&msg),
Ok(ws::Message::Pong(_)) => (),
Ok(ws::Message::Text(text)) => ctx.text(text),
Ok(ws::Message::Binary(bin)) => ctx.binary(bin),
Ok(ws::Message::Close(reason)) => ctx.close(reason),
_ => (),
}
}
}
pub struct IPv6FlowWebSocket {
pub flow_type: IPv6FlowType,
pub interval: Option<SpawnHandle>,
}
@ -25,10 +62,40 @@ impl Actor for IPv6FlowWebSocket {
type Context = ws::WebsocketContext<Self>;
fn started(&mut self, ctx: &mut Self::Context) {
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 {
let flow_data = Monitor::get_ipv6_flow_data(flow_type).await;
transform_ipv6_flow_data(flow_data)
};
ctx.wait(future.into_actor(actor).map(|flow_data, _, ctx| {
if let Ok(json) = serde_json::to_string(&flow_data) {
ctx.text(json);
}
}));
});
self.interval = Some(interval);
}
fn stopping(&mut self, ctx: &mut Self::Context) -> Running {
if let Some(interval) = self.interval.take() {
ctx.cancel_future(interval);
}
Running::Stop
}
}
}
impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for IPv6FlowWebSocket {
fn handle(&mut self, msg: Result<ws::Message, ws::ProtocolError>, ctx: &mut Self::Context) {
match msg {
Ok(ws::Message::Ping(msg)) => ctx.pong(&msg),
Ok(ws::Message::Pong(_)) => (),
Ok(ws::Message::Text(text)) => ctx.text(text),
Ok(ws::Message::Binary(bin)) => ctx.binary(bin),
Ok(ws::Message::Close(reason)) => ctx.close(reason),
_ => (),
}
}
}