mirror of https://github.com/rust-bpf/rust-bcc
105 lines
3.1 KiB
Rust
105 lines
3.1 KiB
Rust
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);
|
|
}
|
|
}
|