use bcc::XDPMode::{XDP_FLAGS_HW_MODE, XDP_FLAGS_SKB_MODE};
use bcc::{BPFBuilder, BccError, XDP};
use byteorder::{LittleEndian, ReadBytesExt};
use clap::{App, Arg};
use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::io::Cursor;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::thread::sleep;
use std::time::Duration;
fn do_main(runnable: Arc<AtomicBool>) -> Result<(), BccError> {
let matches = App::new("xdpdropcount")
.about("Drop incoming packets on XDP layer and count for which protocol type")
.arg(
Arg::with_name("device")
.long("device")
.short("d")
.help("Device name to attach the XDP program")
.default_value("eth0")
.takes_value(true),
)
.arg(
Arg::with_name("hardware-offload-mode")
.long("hw-offload-mode")
.short("h")
.help("Run the XDP program in hardware offload mode XDP_FLAGS_HW_MODE")
.takes_value(true),
)
.get_matches();
let hw_offload_mode_enabled = matches.is_present("hw-offload-mode");
let mode = if hw_offload_mode_enabled {
XDP_FLAGS_HW_MODE
} else {
XDP_FLAGS_SKB_MODE
};
let device = matches
.value_of("device")
.expect("safe since `device` has a default value");
let code = include_str!("xdp_drop_count.c");
let cflags = &["-w", "-DRETURNCODE=XDP_DROP", "-DCTXTYPE=xdp_md"];
let builder = {
let builder = BPFBuilder::new(code)?.cflags(cflags)?;
if hw_offload_mode_enabled {
builder.device(device)?
} else {
builder
}
};
let mut bpf = builder.build()?;
XDP::new()
.handler("xdp_prog1")
.device(device)
.mode(mode)
.attach(&mut bpf)?;
let table = bpf.table("dropcnt")?;
println!("Printing drops per IP protocol-number, hit CTRL+C to stop");
let mut state = HashMap::new();
while runnable.load(Ordering::SeqCst) {
for entry in table.iter() {
let protocol = Cursor::new(entry.key).read_u32::<LittleEndian>().unwrap();
let current_count = Cursor::new(entry.value).read_u64::<LittleEndian>().unwrap();
let delta = match state.entry(protocol) {
Entry::Vacant(entry) => {
entry.insert(current_count);
current_count
}
Entry::Occupied(mut entry) => {
let delta = current_count - *entry.get();
entry.insert(current_count);
delta
}
};
println!("{:?} => {:?}pkts/s", protocol, delta);
}
sleep(Duration::from_secs(1));
}
Ok(())
}
fn main() {
let runnable = Arc::new(AtomicBool::new(true));
let r = runnable.clone();
ctrlc::set_handler(move || {
r.store(false, Ordering::SeqCst);
})
.expect("Failed to set handler for SIGINT / SIGTERM");
if let Err(err) = do_main(runnable) {
eprintln!("Error: {}", err);
std::process::exit(1);
}
}