用 Rust + eBPF + XDP 实现一个简易但实用的 DDoS 防护网络监控 Agent

网卡硬件 → DMA → 驱动 ring buffer → 分配 sk_buff → 内核 L2/L3/L4 处理 → netfilter → socket → 用户态

物理网卡 → [XDP eBPF 程序]
              ↓
   解析 Eth → IP → UDP → 计数/判断
              ↓
         BPF Map（PerCpuArray / HashMap / RingBuf）
              ↓
   Rust 用户态 Agent（aya）
   → 实时统计 / 告警 / 热更新阈值 / 下发规则

cargo install cargo-generate cargo-xtask bpf-linker

cargo generate --name rust-xdp-ddos-agent \
  https://github.com/aya-rs/aya-template

cd rust-xdp-ddos-agent
# 选择 program_type = xdp

#![no_std]
#![no_main]

use aya_ebpf::{bindings::xdp_action, macros::xdp, programs::XdpContext};
use aya_ebpf::maps::PerCpuArray;
use aya_log_ebpf::{info};
use network_types::{
    eth::{EtherType, EthHdr},
    ip::{IpProto, Ipv4Hdr},
    udp::UdpHdr,
};

#[derive(Copy, Clone)]
#[repr(C)]
pub struct Counter {
    pub udp_packets: u64,
    pub dropped: u64,
}

#[map]
static mut GLOBAL_COUNTER: PerCpuArray<Counter> = PerCpuArray::with_max_entries(1, 0);

const UDP_FLOOD_THRESHOLD: u64 = 5000;  // 生产建议 10k–100k+ 根据网卡能力调整

#[inline(always)]
unsafe fn ptr_at<T>(ctx: &XdpContext, offset: usize) -> Result<*const T, ()> {
    let ptr = ctx.data() + offset;
    if ptr + core::mem::size_of::<T>() > ctx.data_end() {
        return Err(());
    }
    Ok(ptr as *const T)
}

#[xdp]
pub fn xdp_ddos_guard(ctx: XdpContext) -> u32 {
    if let Ok(action) = try_xdp_ddos_guard(&ctx) {
        action
    } else {
        xdp_action::XDP_ABORTED
    }
}

fn try_xdp_ddos_guard(ctx: &XdpContext) -> Result<u32, ()> {
    let eth = unsafe { *ptr_at::<EthHdr>(ctx, 0)? };
    if eth.ether_type != EtherType::Ipv4.to_be() {
        return Ok(xdp_action::XDP_PASS);
    }

    let ip = unsafe { *ptr_at::<Ipv4Hdr>(ctx, EthHdr::LEN)? };
    if ip.proto != IpProto::Udp.to_be() {
        return Ok(xdp_action::XDP_PASS);
    }

    // 示例：放行 DNS 53 端口（可扩展白名单）
    let udp_offset = EthHdr::LEN + (ip.ihl() as usize * 4);
    let udp = unsafe { *ptr_at::<UdpHdr>(ctx, udp_offset)? };
    if u16::from_be(udp.dest) == 53 {
        return Ok(xdp_action::XDP_PASS);
    }

    let counter = unsafe { GLOBAL_COUNTER.get_ptr_mut(0).ok_or(())? };
    unsafe {
        (*counter).udp_packets = (*counter).udp_packets.wrapping_add(1);

        if (*counter).udp_packets > UDP_FLOOD_THRESHOLD {
            (*counter).dropped = (*counter).dropped.wrapping_add(1);
            info!(ctx, "UDP flood! > {} pps, dropping", UDP_FLOOD_THRESHOLD);
            return Ok(xdp_action::XDP_DROP);
        }
    }

    Ok(xdp_action::XDP_PASS)
}

use anyhow::Context;
use aya::{Ebpf, maps::PerCpuArray, programs::{Xdp, XdpFlags}};
use aya_log::EbpfLogger;
use clap::Parser;
use std::time::Duration;
use tokio::time::sleep;

#[derive(Parser, Debug)]
struct Opt {
    #[clap(short, long, default_value = "eth0")]
    iface: String,
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let opt = Opt::parse();

    // 加载并附加程序
    let mut bpf = Ebpf::load_file("target/bpfel-unknown-none/debug/rust-xdp-ddos-agent")?;
    EbpfLogger::init(&mut bpf)?;

    let prog: &mut Xdp = bpf.program_mut("xdp_ddos_guard")?.try_into()?;
    prog.load()?;
    prog.attach(&opt.iface, XdpFlags::default())?;

    let mut counter: PerCpuArray<_, Counter> = bpf.take_map("GLOBAL_COUNTER")?.try_into()?;

    println!("🚀 XDP DDoS Agent 已启动 | 接口: {} | 阈值: >{} UDP pps → DROP", opt.iface, 5000);

    loop {
        let stats: Vec<Counter> = counter.iter().collect();
        let udp_total: u64 = stats.iter().map(|c| c.udp_packets).sum();
        let dropped: u64 = stats.iter().map(|c| c.dropped).sum();

        if udp_total > 3000 {
            println!("[ALERT] 疑似 UDP flood！当前 ≈ {} pkt/s，已丢包 {} 个", udp_total, dropped);
        } else {
            println!("正常 → UDP: {} pkt/s", udp_total);
        }

        // 每秒重置计数（生产建议用 ringbuf + 用户态滑动窗口）
        counter.set(0, Counter { udp_packets: 0, dropped }, 0)?;

        sleep(Duration::from_secs(1)).await;
    }
}

cargo xtask build-ebpf
cargo build --release
sudo RUST_LOG=info target/release/rust-xdp-ddos-agent --iface eth0