From f90882b4b2b5b72092d9a48c3342ed606a256db4 Mon Sep 17 00:00:00 2001 From: nutcas3 Date: Sat, 25 Apr 2026 03:51:23 +0300 Subject: [PATCH] refactor: Extract health checks and utility modules from packet-gateway ebpf module - Extract health check handlers to new health.rs module (health, liveness, readiness, metrics endpoints) - Extract IP conversion utilities to new ip_utils.rs module (u32_to_ipv4_string, ipv4_string_to_u32) - Extract Redis synchronization logic to new redis_sync.rs module (RedisSyncer with sync_packet_stats, sync_credit_info, sync_credits_from_redis, sync_blocked_from_redis, sync_batch_entry methods) - Refactor EbpfManager --- apps/packet-gateway/src/ebpf.rs | 123 +++++--------------------- apps/packet-gateway/src/health.rs | 115 ++++++++++++++++++++++++ apps/packet-gateway/src/ip_utils.rs | 15 ++++ apps/packet-gateway/src/main.rs | 112 +---------------------- apps/packet-gateway/src/redis_sync.rs | 120 +++++++++++++++++++++++++ 5 files changed, 276 insertions(+), 209 deletions(-) create mode 100644 apps/packet-gateway/src/health.rs create mode 100644 apps/packet-gateway/src/ip_utils.rs create mode 100644 apps/packet-gateway/src/redis_sync.rs diff --git a/apps/packet-gateway/src/ebpf.rs b/apps/packet-gateway/src/ebpf.rs index 604d227..69aea54 100644 --- a/apps/packet-gateway/src/ebpf.rs +++ b/apps/packet-gateway/src/ebpf.rs @@ -1,12 +1,14 @@ use aya::{programs::Xdp, Bpf, maps::HashMap}; use aya_log::BpfLogger; use anyhow::{Context, Result}; -use std::net::Ipv4Addr; -use tokio::time::{interval, Duration}; use tracing::{info, warn, error, debug}; -use std::collections::HashMap as StdHashMap; use redis::AsyncCommands; +mod ip_utils; +mod redis_sync; +use ip_utils::{u32_to_ipv4_string, ipv4_string_to_u32}; +use redis_sync::RedisSyncer; + pub struct EbpfManager { bpf: Bpf, interface: String, @@ -35,10 +37,8 @@ pub struct CreditInfo { impl EbpfManager { pub async fn new(interface: String) -> Result { - // Load the compiled eBPF program let mut bpf = Bpf::load(include_bytes_aligned!("packet_filter"))?; - // Initialize eBPF logging if let Err(e) = BpfLogger::init(&mut bpf) { warn!("Failed to initialize eBPF logger: {}", e); } @@ -49,16 +49,13 @@ impl EbpfManager { } pub fn attach(&mut self) -> Result<()> { - // Get the XDP program let program: &mut Xdp = self.bpf.program_mut("packet_filter") .ok_or_else(|| anyhow::anyhow!("Program 'packet_filter' not found"))? .try_into()?; - // Load the program program.load()?; info!("Loaded XDP program into kernel"); - // Attach to the specified interface program.attach(&self.interface, aya::programs::XdpFlags::default()) .context("Failed to attach XDP program")?; @@ -73,7 +70,6 @@ impl EbpfManager { let mut stats = Vec::new(); - // Iterate through all entries in the packet_stats map for (ip, bytes) in packet_stats_map.iter()? { stats.push(PacketStats { ip: *ip, bytes: *bytes }); } @@ -88,7 +84,6 @@ impl EbpfManager { let mut credits = Vec::new(); - // Iterate through all entries in the user_credits map for (ip, credit) in user_credits_map.iter()? { credits.push(CreditInfo { ip: *ip, credit: *credit }); } @@ -120,22 +115,14 @@ impl EbpfManager { } pub async fn sync_to_redis(&self, redis_conn: &mut redis::aio::MultiplexedConnection) -> Result<()> { - // Sync packet stats let stats = self.get_packet_stats()?; - for stat in stats { - let ip_str = self.u32_to_ipv4_string(stat.ip); - let key = format!("packet_stats:{}", ip_str); - let _: () = redis_conn.set(&key, stat.bytes).await - .context("Failed to sync packet stats to Redis")?; - } + let stats_vec: Vec<(u32, u64)> = stats.iter().map(|s| (s.ip, s.bytes)).collect(); + RedisSyncer::sync_packet_stats(&stats_vec, redis_conn).await?; - // Sync credit info - let credits = self.get_credit_info()?; - for credit in credits { - let ip_str = self.u32_to_ipv4_string(credit.ip); - let key = format!("user_credit:{}", ip_str); - let _: () = redis_conn.set(&key, credit.credit).await - .context("Failed to sync credit info to Redis")?; + let credits = self.get_credit_info(); + if let Ok(credits) = credits { + let credits_vec: Vec<(u32, i64)> = credits.iter().map(|c| (c.ip, c.credit)).collect(); + RedisSyncer::sync_credit_info(&credits_vec, redis_conn).await?; } debug!("Synced eBPF maps to Redis"); @@ -143,60 +130,23 @@ impl EbpfManager { } pub async fn sync_from_redis(&self, redis_conn: &mut redis::aio::MultiplexedConnection) -> Result<()> { - // Get all credit keys from Redis - let keys: Vec = redis_conn.keys("user_credit:*").await - .context("Failed to get credit keys from Redis")?; + RedisSyncer::sync_credits_from_redis(redis_conn, |ip, credit| { + self.update_user_credit(ip, credit) + }).await?; - for key in keys { - if let Some(ip_part) = key.strip_prefix("user_credit:") { - let credit: i64 = redis_conn.get(&key).await - .context("Failed to get credit from Redis")?; - - if let Ok(ip) = self.ipv4_string_to_u32(ip_part) { - self.update_user_credit(ip, credit)?; - debug!("Synced credit from Redis for IP {}: {}", ip_part, credit); - } else { - warn!("Invalid IP format in Redis key: {}", key); - } - } - } - - // Get blocked users from Redis - let blocked_keys: Vec = redis_conn.keys("blocked_user:*").await - .context("Failed to get blocked user keys from Redis")?; - - for key in blocked_keys { - if let Some(ip_part) = key.strip_prefix("blocked_user:") { - if let Ok(ip) = self.ipv4_string_to_u32(ip_part) { - self.block_user(ip, true)?; - debug!("Synced block status from Redis for IP: {}", ip_part); - } else { - warn!("Invalid IP format in Redis block key: {}", key); - } - } - } + RedisSyncer::sync_blocked_from_redis(redis_conn, |ip| { + self.block_user(ip, true) + }).await?; debug!("Synced data from Redis to eBPF maps"); Ok(()) } - fn u32_to_ipv4_string(&self, ip: u32) -> String { - let ip_addr = Ipv4Addr::from(ip); - ip_addr.to_string() - } - - fn ipv4_string_to_u32(&self, ip_str: &str) -> Result { - let ip_addr: Ipv4Addr = ip_str.parse() - .context("Invalid IP address format")?; - Ok(u32::from(ip_addr)) - } - pub fn trigger_sync(&self) -> Result<()> { let sync_control_map: &HashMap<_, u32, u64> = self.bpf.map("sync_control") .ok_or_else(|| anyhow::anyhow!("sync_control map not found"))? .try_into()?; - // Set sync flag to trigger eBPF sync function let key = 0u32; let sync_flag = 1u64; sync_control_map.insert(&key, &sync_flag, 0)?; @@ -212,7 +162,6 @@ impl EbpfManager { let mut entries = Vec::new(); - // Read from sync buffer (index 0 for now, could expand for batching) let key = 0u32; if let Some(entry) = sync_buffer_map.get(&key, 0)? { if entry.valid == 1 { @@ -224,38 +173,19 @@ impl EbpfManager { } pub fn sync_batch_to_redis(&self, redis_conn: &mut redis::aio::MultiplexedConnection) -> Result<()> { - // Trigger eBPF sync function self.trigger_sync()?; - - // Give eBPF program a moment to process std::thread::sleep(std::time::Duration::from_millis(10)); - // Read batched data from sync buffer let entries = self.read_sync_buffer()?; - // Sync each entry to Redis for entry in entries { - let ip_str = self.u32_to_ipv4_string(entry.ip); - - // Sync packet stats - let stats_key = format!("packet_stats:{}", ip_str); - let _: () = redis_conn.set(&stats_key, entry.bytes).await - .context("Failed to sync packet stats to Redis")?; - - // Sync credit info - let credit_key = format!("user_credit:{}", ip_str); - let _: () = redis_conn.set(&credit_key, entry.credit).await - .context("Failed to sync credit info to Redis")?; - - // Sync block status - if entry.blocked == 1 { - let block_key = format!("blocked_user:{}", ip_str); - let _: () = redis_conn.set(&block_key, 1u8).await - .context("Failed to sync block status to Redis")?; - } - - debug!("Batch synced IP {}: {} bytes, credit {}, blocked {}", - ip_str, entry.bytes, entry.credit, entry.blocked); + RedisSyncer::sync_batch_entry( + entry.ip, + entry.bytes, + entry.credit, + entry.blocked, + redis_conn, + ).await?; } info!("Batch synced {} entries to Redis", entries.len()); @@ -263,7 +193,6 @@ impl EbpfManager { } pub fn cleanup_maps(&self) -> Result<()> { - // Clear packet stats map if let Ok(packet_stats_map) = self.bpf.map("packet_stats").and_then(|m| m.try_into()) { let map: &HashMap<_, u32, u64> = packet_stats_map; let keys: Vec = map.iter()?.map(|(k, _)| *k).collect(); @@ -272,7 +201,6 @@ impl EbpfManager { } } - // Clear sync control if let Ok(sync_control_map) = self.bpf.map("sync_control").and_then(|m| m.try_into()) { let map: &HashMap<_, u32, u64> = sync_control_map; let key = 0u32; @@ -285,9 +213,6 @@ impl EbpfManager { } } -// Helper function to include compiled eBPF bytecode fn include_bytes_aligned(s: &str) -> &'static [u8] { - // Include the compiled eBPF bytecode - // The actual path will be generated by the build process aya::include_bytes_aligned!(s) } diff --git a/apps/packet-gateway/src/health.rs b/apps/packet-gateway/src/health.rs new file mode 100644 index 0000000..9d3d9b6 --- /dev/null +++ b/apps/packet-gateway/src/health.rs @@ -0,0 +1,115 @@ +use axum::{ + Json, + extract::State, +}; +use redis::AsyncCommands; +use serde::Serialize; +use std::sync::atomic::AtomicBool; +use std::sync::Arc; +use crate::ebpf::EbpfManager; + +#[derive(Serialize)] +pub struct HealthResponse { + pub status: String, + pub service: String, + pub timestamp: String, +} + +#[derive(Serialize)] +pub struct ReadinessResponse { + pub status: String, + pub service: String, + pub timestamp: String, + pub checks: serde_json::Value, +} + +#[derive(Serialize)] +pub struct MetricsResponse { + pub status: String, + pub service: String, + pub timestamp: String, + pub tracked_ips: usize, + pub total_bytes_processed: u64, + pub low_credit_users: usize, +} + +pub async fn health_handler() -> Json { + Json(HealthResponse { + status: "healthy".to_string(), + service: "packet-gateway".to_string(), + timestamp: chrono::Utc::now().to_rfc3339(), + }) +} + +pub async fn liveness_handler() -> Json { + Json(HealthResponse { + status: "alive".to_string(), + service: "packet-gateway".to_string(), + timestamp: chrono::Utc::now().to_rfc3339(), + }) +} + +pub async fn readiness_handler( + State(redis_client): State>, + State(ebpf_attached): State>, +) -> Result, axum::http::StatusCode> { + let mut checks = serde_json::Map::new(); + + let redis_ok = match redis_client.get_multiplexed_async_connection().await { + Ok(mut conn) => { + match redis::cmd("PING").query_async::<_, String>(&mut conn).await { + Ok(_) => { + checks.insert("redis".to_string(), serde_json::Value::String("ok".to_string())); + true + } + Err(e) => { + checks.insert("redis".to_string(), serde_json::Value::String(format!("error: {}", e))); + false + } + } + } + Err(e) => { + checks.insert("redis".to_string(), serde_json::Value::String(format!("connection error: {}", e))); + false + } + }; + + let ebpf_ok = ebpf_attached.load(std::sync::atomic::Ordering::Relaxed); + checks.insert("ebpf".to_string(), serde_json::Value::String(if ebpf_ok { "ok".to_string() } else { "not attached".to_string() })); + + if redis_ok && ebpf_ok { + Ok(Json(ReadinessResponse { + status: "ready".to_string(), + service: "packet-gateway".to_string(), + timestamp: chrono::Utc::now().to_rfc3339(), + checks: serde_json::Value::Object(checks), + })) + } else { + Err(axum::http::StatusCode::SERVICE_UNAVAILABLE) + } +} + +pub async fn metrics_handler( + State(ebpf_manager): State>, +) -> Json { + let tracked_ips = ebpf_manager.get_packet_stats() + .map(|stats| stats.len()) + .unwrap_or(0); + + let total_bytes_processed = ebpf_manager.get_packet_stats() + .map(|stats| stats.iter().map(|s| s.bytes).sum()) + .unwrap_or(0); + + let low_credit_users = ebpf_manager.get_credit_info() + .map(|credits| credits.iter().filter(|c| c.credit < 1000).count()) + .unwrap_or(0); + + Json(MetricsResponse { + status: "ok".to_string(), + service: "packet-gateway".to_string(), + timestamp: chrono::Utc::now().to_rfc3339(), + tracked_ips, + total_bytes_processed, + low_credit_users, + }) +} diff --git a/apps/packet-gateway/src/ip_utils.rs b/apps/packet-gateway/src/ip_utils.rs new file mode 100644 index 0000000..2064bf2 --- /dev/null +++ b/apps/packet-gateway/src/ip_utils.rs @@ -0,0 +1,15 @@ +use anyhow::{Context, Result}; +use std::net::Ipv4Addr; + +/// Converts a u32 IP address to an IPv4 string +pub fn u32_to_ipv4_string(ip: u32) -> String { + let ip_addr = Ipv4Addr::from(ip); + ip_addr.to_string() +} + +/// Converts an IPv4 string to a u32 IP address +pub fn ipv4_string_to_u32(ip_str: &str) -> Result { + let ip_addr: Ipv4Addr = ip_str.parse() + .context("Invalid IP address format")?; + Ok(u32::from(ip_addr)) +} diff --git a/apps/packet-gateway/src/main.rs b/apps/packet-gateway/src/main.rs index b878c40..f41ba79 100644 --- a/apps/packet-gateway/src/main.rs +++ b/apps/packet-gateway/src/main.rs @@ -6,14 +6,14 @@ use tracing::{info, warn, error, debug}; use axum::{ routing::get, Router, - Json, }; -use serde::Serialize; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; mod ebpf; +mod health; use ebpf::EbpfManager; +use health::{health_handler, liveness_handler, readiness_handler, metrics_handler}; #[derive(Parser, Debug)] #[command(author, version, about, long_about = None)] @@ -35,114 +35,6 @@ struct Args { health_port: u16, } -#[derive(Serialize)] -struct HealthResponse { - status: String, - service: String, - timestamp: String, -} - -#[derive(Serialize)] -struct ReadinessResponse { - status: String, - service: String, - timestamp: String, - checks: serde_json::Value, -} - -async fn health_handler() -> Json { - Json(HealthResponse { - status: "healthy".to_string(), - service: "packet-gateway".to_string(), - timestamp: chrono::Utc::now().to_rfc3339(), - }) -} - -async fn liveness_handler() -> Json { - Json(HealthResponse { - status: "alive".to_string(), - service: "packet-gateway".to_string(), - timestamp: chrono::Utc::now().to_rfc3339(), - }) -} - -#[derive(Serialize)] -struct MetricsResponse { - status: String, - service: String, - timestamp: String, - tracked_ips: usize, - total_bytes_processed: u64, - low_credit_users: usize, -} - -async fn metrics_handler( - State(ebpf_manager): State>, -) -> Json { - let tracked_ips = ebpf_manager.get_packet_stats() - .map(|stats| stats.len()) - .unwrap_or(0); - - let total_bytes_processed = ebpf_manager.get_packet_stats() - .map(|stats| stats.iter().map(|s| s.bytes).sum()) - .unwrap_or(0); - - let low_credit_users = ebpf_manager.get_credit_info() - .map(|credits| credits.iter().filter(|c| c.credit < 1000).count()) - .unwrap_or(0); - - Json(MetricsResponse { - status: "ok".to_string(), - service: "packet-gateway".to_string(), - timestamp: chrono::Utc::now().to_rfc3339(), - tracked_ips, - total_bytes_processed, - low_credit_users, - }) -} - -async fn readiness_handler( - State(redis_client): State>, - State(ebpf_attached): State>, -) -> Result, axum::http::StatusCode> { - let mut checks = serde_json::Map::new(); - - // Check Redis connectivity - let redis_ok = match redis_client.get_multiplexed_async_connection().await { - Ok(mut conn) => { - match redis::cmd("PING").query_async::<_, String>(&mut conn).await { - Ok(_) => { - checks.insert("redis".to_string(), serde_json::Value::String("ok".to_string())); - true - } - Err(e) => { - checks.insert("redis".to_string(), serde_json::Value::String(format!("error: {}", e))); - false - } - } - } - Err(e) => { - checks.insert("redis".to_string(), serde_json::Value::String(format!("connection error: {}", e))); - false - } - }; - - // Check eBPF attachment - let ebpf_ok = ebpf_attached.load(Ordering::Relaxed); - checks.insert("ebpf".to_string(), serde_json::Value::String(if ebpf_ok { "ok".to_string() } else { "not attached".to_string() })); - - if redis_ok && ebpf_ok { - Ok(Json(ReadinessResponse { - status: "ready".to_string(), - service: "packet-gateway".to_string(), - timestamp: chrono::Utc::now().to_rfc3339(), - checks: serde_json::Value::Object(checks), - })) - } else { - Err(axum::http::StatusCode::SERVICE_UNAVAILABLE) - } -} - #[tokio::main] async fn main() -> Result<()> { // Initialize tracing diff --git a/apps/packet-gateway/src/redis_sync.rs b/apps/packet-gateway/src/redis_sync.rs new file mode 100644 index 0000000..70e3a5c --- /dev/null +++ b/apps/packet-gateway/src/redis_sync.rs @@ -0,0 +1,120 @@ +use anyhow::{Context, Result}; +use redis::AsyncCommands; +use tracing::{debug, warn}; +use super::ip_utils::{u32_to_ipv4_string, ipv4_string_to_u32}; + +/// Handles synchronization between eBPF maps and Redis +pub struct RedisSyncer; + +impl RedisSyncer { + /// Syncs packet stats from eBPF to Redis + pub async fn sync_packet_stats( + stats: &[(u32, u64)], + redis_conn: &mut redis::aio::MultiplexedConnection, + ) -> Result<()> { + for (ip, bytes) in stats { + let ip_str = u32_to_ipv4_string(*ip); + let key = format!("packet_stats:{}", ip_str); + let _: () = redis_conn.set(&key, bytes).await + .context("Failed to sync packet stats to Redis")?; + } + Ok(()) + } + + /// Syncs credit info from eBPF to Redis + pub async fn sync_credit_info( + credits: &[(u32, i64)], + redis_conn: &mut redis::aio::MultiplexedConnection, + ) -> Result<()> { + for (ip, credit) in credits { + let ip_str = u32_to_ipv4_string(*ip); + let key = format!("user_credit:{}", ip_str); + let _: () = redis_conn.set(&key, credit).await + .context("Failed to sync credit info to Redis")?; + } + Ok(()) + } + + /// Syncs credit data from Redis to eBPF + pub async fn sync_credits_from_redis( + redis_conn: &mut redis::aio::MultiplexedConnection, + update_credit: F, + ) -> Result<()> + where + F: Fn(u32, i64) -> Result<()>, + { + let keys: Vec = redis_conn.keys("user_credit:*").await + .context("Failed to get credit keys from Redis")?; + + for key in keys { + if let Some(ip_part) = key.strip_prefix("user_credit:") { + let credit: i64 = redis_conn.get(&key).await + .context("Failed to get credit from Redis")?; + + if let Ok(ip) = ipv4_string_to_u32(ip_part) { + update_credit(ip, credit)?; + debug!("Synced credit from Redis for IP {}: {}", ip_part, credit); + } else { + warn!("Invalid IP format in Redis key: {}", key); + } + } + } + Ok(()) + } + + /// Syncs blocked users from Redis to eBPF + pub async fn sync_blocked_from_redis( + redis_conn: &mut redis::aio::MultiplexedConnection, + block_user: F, + ) -> Result<()> + where + F: Fn(u32) -> Result<()>, + { + let blocked_keys: Vec = redis_conn.keys("blocked_user:*").await + .context("Failed to get blocked user keys from Redis")?; + + for key in blocked_keys { + if let Some(ip_part) = key.strip_prefix("blocked_user:") { + if let Ok(ip) = ipv4_string_to_u32(ip_part) { + block_user(ip)?; + debug!("Synced block status from Redis for IP: {}", ip_part); + } else { + warn!("Invalid IP format in Redis block key: {}", key); + } + } + } + Ok(()) + } + + /// Syncs a batch entry to Redis + pub async fn sync_batch_entry( + ip: u32, + bytes: u64, + credit: i64, + blocked: u8, + redis_conn: &mut redis::aio::MultiplexedConnection, + ) -> Result<()> { + let ip_str = u32_to_ipv4_string(ip); + + // Sync packet stats + let stats_key = format!("packet_stats:{}", ip_str); + let _: () = redis_conn.set(&stats_key, bytes).await + .context("Failed to sync packet stats to Redis")?; + + // Sync credit info + let credit_key = format!("user_credit:{}", ip_str); + let _: () = redis_conn.set(&credit_key, credit).await + .context("Failed to sync credit info to Redis")?; + + // Sync block status + if blocked == 1 { + let block_key = format!("blocked_user:{}", ip_str); + let _: () = redis_conn.set(&block_key, 1u8).await + .context("Failed to sync block status to Redis")?; + } + + debug!("Batch synced IP {}: {} bytes, credit {}, blocked {}", + ip_str, bytes, credit, blocked); + Ok(()) + } +}