硬核！用Rust写eBPF程序抓取TCP包，性能甩传统方案几条街

my-ebpf-filter/
├── Cargo.toml
├── src/
│   └── main.rs        # 用户态程序
└── xdp_filter/        # eBPF程序（用C或Rust写都可以）
    ├── Cargo.toml
    └── src/
        └── lib.rs     # 内核态eBPF程序

use aya_bpf::{
    macros::{xdp, map},
    maps::HashMap,
    programs::XdpContext,
};

#[map(name = "PORT_FILTER")]
static mut PORT_FILTER: HashMap<u16, u8> = HashMap::with_max_entries(1024, 0);

// 需要过滤的目标端口列表
#[map(name = "TARGET_PORTS")]
static mut TARGET_PORTS: HashMap<u16, u8> = HashMap::with_max_entries(16, 0);

// 解析以太网帧头
#[inline]
fn parse_ethernet(ctx: &XdpContext) -> Option<(*const u8, u16)> {
    let data = ctx.data();
    let data_end = ctx.data_end();

    let eth_len = 14;
    if data + eth_len > data_end {
        return None;
    }

    let eth_ptr = data as *const u8;
    Some((eth_ptr, eth_len))
}

// 解析IP头并检查是否是TCP
#[inline]
fn parse_ip_tcp(ctx: &XdpContext, eth_ptr: *const u8) -> Option<(u32, u16, u16)> {
    let data = ctx.data();
    let data_end = ctx.data_end();

    // IP头通常在偏移14的位置
    let ip_offset = 14;
    let ip_ptr = (eth_ptr as usize + ip_offset) as *const u8;

    if ip_ptr.add(20) > data_end as *const u8 {
        return None;
    }

    let version_ihl = unsafe { *ip_ptr };
    let version = version_ihl >> 4;
    let ihl = (version_ihl & 0x0F) * 4;  // IP头长度

    if version != 4 {
        return None;
    }

    // 检查协议类型（TCP = 6）
    let protocol = unsafe { *(ip_ptr.add(9)) };
    if protocol != 6 {
        return None;  // 不是TCP，直接放行
    }

    // 提取源IP和端口
    let src_ip = unsafe { *(ip_ptr.add(12) as *const u32) };
    let src_port = unsafe { *(ip_ptr.add(20) as *const u16) };
    let dst_port = unsafe { *(ip_ptr.add(22) as *const u16) };

    Some((src_ip, src_port, dst_port))
}

#[xdp(name = "tcp_filter")]
pub fn tcp_filter(ctx: XdpContext) -> u32 {
    match try_tcp_filter(ctx) {
        Ok(_) => aya_bpf::programs::XdpAction::Pass as u32,
        Err(_) => aya_bpf::programs::XdpAction::Drop as u32,
    }
}

fn try_tcp_filter(ctx: XdpContext) -> Result<(), ()> {
    let (eth_ptr, _) = parse_ethernet(&ctx).ok_or(())?;
    let (_, src_port, dst_port) = parse_ip_tcp(&ctx, eth_ptr).ok_or(())?;

    // 检查目标端口是否在我们关注的列表中
    // 这里的逻辑是：只过滤特定端口的包
    unsafe {
        if TARGET_PORTS.get(&dst_port).is_none() {
            // 不在目标列表中，直接放行
            return Ok(());
        }

        // 可以在这里做更复杂的过滤逻辑
        // 比如：根据源IP白名单过滤
        // 如果需要丢弃，直接返回 Err(())
    }

    // 放行符合条件的包到用户态
    Ok(())
}

use aya::programs::Xdp, BpfLoader;
use aya::maps::HashMap;
use std::net::Ipv4Addr;
use tokio;

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    // 1. 加载eBPF程序
    let mut bpf = BpfLoader::new()
        .load("my_filter.o")?;

    // 2. 获取XDP程序
    let program: &mut Xdp = bpf.program_mut("tcp_filter")unwrap();
    program.load()?;

    // 3. 附加到网卡
    let iface = "eth0";
    program.attach(iface, 0)?;
    println!("eBPF filter loaded on {}", iface);

    // 4. 设置要过滤的目标端口
    let mut target_ports: HashMap<_, u16, u8> =
        bpf.map_mut("TARGET_PORTS")?;

    target_ports.insert(80, 1)?;  // HTTP
    target_ports.insert(443, 1)?; // HTTPS
    target_ports.insert(8080, 1)?; // 自定义服务

    println!("Target ports configured: 80, 443, 8080");

    // 5. 通过perf event接收数据包（如果需要完整包内容）
    // 这里可以用ring buffer把包送到用户态
    let mut perf_buffer = aya::maps::PerfBuffer::builder(
        bpf.map("PACKETS")?,
        128 // buffer size
    );

    perf_buffer.set_loss_cb(|count| {
        eprintln!("Lost {} packets", count);
    });

    let perf_buffer = perf_buffer.build()?;

    // 6. 处理数据包
    loop {
        perf_buffer.read_events(|_, event| {
            // event包含从内核传来的数据包信息
            println!("Received packet: {:?}", event);
        }).await?;
    }
}

[package]
name = "my-ebpf-filter"
version = "0.1.0"
edition = "2021"

[dependencies]
aya = { version = "0.11", features = ["async_tokio"] }
anyhow = "1.0"
tokio = { version = "1.0", features = ["full"] }

[lib]
name = "xdp_filter"
path = "xdp_filter/src/lib.rs"
crate-type = ["cdylib"]