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
This commit is contained in:
nutcas3
2026-04-25 03:51:23 +03:00
parent 5467e22d22
commit f90882b4b2
5 changed files with 276 additions and 209 deletions
+24 -99
View File
@@ -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<Self> {
// 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<String> = 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<String> = 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<u32> {
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<u32> = 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)
}
+115
View File
@@ -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<HealthResponse> {
Json(HealthResponse {
status: "healthy".to_string(),
service: "packet-gateway".to_string(),
timestamp: chrono::Utc::now().to_rfc3339(),
})
}
pub async fn liveness_handler() -> Json<HealthResponse> {
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<Arc<redis::Client>>,
State(ebpf_attached): State<Arc<AtomicBool>>,
) -> Result<Json<ReadinessResponse>, 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<Arc<EbpfManager>>,
) -> Json<MetricsResponse> {
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,
})
}
+15
View File
@@ -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<u32> {
let ip_addr: Ipv4Addr = ip_str.parse()
.context("Invalid IP address format")?;
Ok(u32::from(ip_addr))
}
+2 -110
View File
@@ -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<HealthResponse> {
Json(HealthResponse {
status: "healthy".to_string(),
service: "packet-gateway".to_string(),
timestamp: chrono::Utc::now().to_rfc3339(),
})
}
async fn liveness_handler() -> Json<HealthResponse> {
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<Arc<EbpfManager>>,
) -> Json<MetricsResponse> {
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<Arc<redis::Client>>,
State(ebpf_attached): State<Arc<AtomicBool>>,
) -> Result<Json<ReadinessResponse>, 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
+120
View File
@@ -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<F>(
redis_conn: &mut redis::aio::MultiplexedConnection,
update_credit: F,
) -> Result<()>
where
F: Fn(u32, i64) -> Result<()>,
{
let keys: Vec<String> = 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<F>(
redis_conn: &mut redis::aio::MultiplexedConnection,
block_user: F,
) -> Result<()>
where
F: Fn(u32) -> Result<()>,
{
let blocked_keys: Vec<String> = 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(())
}
}