From c3ce182c6ebd89bed4af439eae495cbd252ecb08 Mon Sep 17 00:00:00 2001 From: DaLaw2 Date: Sat, 30 Nov 2024 21:18:01 +0800 Subject: [PATCH] Fix expire key not clear bug --- Cargo.lock | 3 +- net-guardia/Cargo.toml | 1 + net-guardia/src/core/monitor.rs | 65 ++++++++++++++++++ net-guardia/src/core/system.rs | 2 + net-guardia/src/model/flow_type.rs | 7 +- net-guardia/src/web/api/monitor.rs | 37 +++++++++-- net-guardia/src/web/utils/flow_websocket.rs | 73 ++++++++++++++++++++- 7 files changed, 178 insertions(+), 10 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b33dac2..494de79 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -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", diff --git a/net-guardia/Cargo.toml b/net-guardia/Cargo.toml index 3f893cb..c1667d6 100644 --- a/net-guardia/Cargo.toml +++ b/net-guardia/Cargo.toml @@ -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 } diff --git a/net-guardia/src/core/monitor.rs b/net-guardia/src/core/monitor.rs index 73caf86..2e4cbaf 100644 --- a/net-guardia/src/core/monitor.rs +++ b/net-guardia/src/core/monitor.rs @@ -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> = OnceLock::new(); pub struct Monitor { + terminate: bool, src_ipv4_1min: AyaHashMap, src_ipv4_10min: AyaHashMap, src_ipv4_1hour: AyaHashMap, @@ -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(map: &mut AyaHashMap, now: u64, window: u64) + where + K: Pod, + { + let expired_keys: Vec = 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 { diff --git a/net-guardia/src/core/system.rs b/net-guardia/src/core/system.rs index 589fe5b..f4bc9d4 100644 --- a/net-guardia/src/core/system.rs +++ b/net-guardia/src/core/system.rs @@ -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(()) } diff --git a/net-guardia/src/model/flow_type.rs b/net-guardia/src/model/flow_type.rs index fbb0d5f..4c2bcf9 100644 --- a/net-guardia/src/model/flow_type.rs +++ b/net-guardia/src/model/flow_type.rs @@ -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, diff --git a/net-guardia/src/web/api/monitor.rs b/net-guardia/src/web/api/monitor.rs index 1cb762d..13375a2 100644 --- a/net-guardia/src/web/api/monitor.rs +++ b/net-guardia/src/web/api/monitor.rs @@ -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 { - Ok(HttpResponse::Forbidden().finish()) +#[get("/websocket/ipv4/{flow_type}")] +async fn websocket_ipv4( + req: HttpRequest, + stream: web::Payload, + path: web::Path, +) -> Result { + 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, +) -> Result { + let flow_type = path.into_inner(); + let websocket = IPv6FlowWebSocket { + flow_type, + interval: None, + }; + start(websocket, &req, stream) } diff --git a/net-guardia/src/web/utils/flow_websocket.rs b/net-guardia/src/web/utils/flow_websocket.rs index f7b625c..96ff929 100644 --- a/net-guardia/src/web/utils/flow_websocket.rs +++ b/net-guardia/src/web/utils/flow_websocket.rs @@ -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, } @@ -9,15 +15,46 @@ impl Actor for IPv4FlowWebSocket { type Context = ws::WebsocketContext; 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> for IPv4FlowWebSocket { + fn handle(&mut self, msg: Result, 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, } @@ -25,10 +62,40 @@ impl Actor for IPv6FlowWebSocket { type Context = ws::WebsocketContext; 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 } -} \ No newline at end of file +} + +impl StreamHandler> for IPv6FlowWebSocket { + fn handle(&mut self, msg: Result, 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), + _ => (), + } + } +}