mirror of https://github.com/rust-bpf/rust-bcc
168 lines
4.2 KiB
Rust
168 lines
4.2 KiB
Rust
use bcc::BccError;
|
|
use bcc::{Kprobe, BPF};
|
|
use clap::{App, Arg};
|
|
|
|
use core::sync::atomic::{AtomicBool, Ordering};
|
|
use std::net::Ipv4Addr;
|
|
use std::net::Ipv6Addr;
|
|
use std::ptr;
|
|
use std::sync::Arc;
|
|
|
|
// A simple tool for tracing tcp retransmits.
|
|
//
|
|
// Based on: https://github.com/iovisor/bcc/blob/master/tools/tcpretrans.py
|
|
|
|
#[repr(C)]
|
|
struct ipv4_data_t {
|
|
pid: u32,
|
|
ip: u64,
|
|
saddr: u32,
|
|
daddr: u32,
|
|
lport: u16,
|
|
dport: u16,
|
|
state: u64,
|
|
type_: u64,
|
|
}
|
|
|
|
#[repr(C)]
|
|
struct ipv6_data_t {
|
|
pid: u32,
|
|
ip: u64,
|
|
saddr: u128,
|
|
daddr: u128,
|
|
lport: u16,
|
|
dport: u16,
|
|
state: u64,
|
|
type_: u64,
|
|
}
|
|
|
|
fn do_main(runnable: Arc<AtomicBool>) -> Result<(), BccError> {
|
|
let matches = App::new("biosnoop")
|
|
.arg(
|
|
Arg::with_name("duration")
|
|
.long("duration")
|
|
.value_name("Seconds")
|
|
.help("The total duration to run this tool")
|
|
.takes_value(true),
|
|
)
|
|
.get_matches();
|
|
|
|
let duration: Option<std::time::Duration> = matches
|
|
.value_of("duration")
|
|
.map(|v| std::time::Duration::new(v.parse().expect("Invalid argument for duration"), 0));
|
|
|
|
let code = include_str!("tcpretrans.c");
|
|
// compile the above BPF code!
|
|
let mut bpf = BPF::new(&code)?;
|
|
Kprobe::new()
|
|
.handler("trace_retransmit")
|
|
.function("tcp_retransmit_skb")
|
|
.attach(&mut bpf)?;
|
|
|
|
let table = bpf.table("ipv4_events")?;
|
|
bpf.init_perf_map(table, print_ipv4_event)?;
|
|
let table = bpf.table("ipv6_events")?;
|
|
bpf.init_perf_map(table, print_ipv6_event)?;
|
|
|
|
println!(
|
|
"{:<-8} {:<-6} {:<-2} {:<-20} {:>-1} {:<-20} {:<-4}",
|
|
"TIME", "PID", "IP", "LADDR:LPORT", "T", "RADDR:RPORT", "STATE"
|
|
);
|
|
let start = std::time::Instant::now();
|
|
while runnable.load(Ordering::SeqCst) {
|
|
bpf.perf_map_poll(200);
|
|
if let Some(d) = duration {
|
|
if std::time::Instant::now() - start >= d {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn get_datetime() -> String {
|
|
let now = time::OffsetDateTime::now_utc();
|
|
let date = now.date();
|
|
let time = now.time();
|
|
format!(
|
|
"{:04}-{:02}-{:02}T{:02}:{:02}:{:02}.{:09}Z",
|
|
date.year(),
|
|
date.month() as u8,
|
|
date.day(),
|
|
time.hour(),
|
|
time.minute(),
|
|
time.second(),
|
|
time.nanosecond(),
|
|
)
|
|
}
|
|
|
|
fn print_ipv4_event() -> Box<dyn FnMut(&[u8]) + Send> {
|
|
Box::new(|x| {
|
|
let event = parse_ipv4_struct(x);
|
|
println!(
|
|
"{:<-8} {:<-6} {:<-2} {:<-20} {:->1}> {:<-20} {:>-4}",
|
|
get_datetime(),
|
|
event.pid,
|
|
event.ip,
|
|
format!(
|
|
"{}:{}",
|
|
Ipv4Addr::from(u32::from_be(event.saddr)),
|
|
event.lport
|
|
),
|
|
event.type_,
|
|
format!(
|
|
"{}:{}",
|
|
Ipv4Addr::from(u32::from_be(event.daddr)),
|
|
event.dport
|
|
),
|
|
event.state,
|
|
);
|
|
})
|
|
}
|
|
|
|
fn print_ipv6_event() -> Box<dyn FnMut(&[u8]) + Send> {
|
|
Box::new(|x| {
|
|
let event = parse_ipv6_struct(x);
|
|
println!(
|
|
"{:<-8} {:<-6} {:<-2} {:<-20} {:->1}> {:<-20} {:>-4}",
|
|
get_datetime(),
|
|
event.pid,
|
|
event.ip,
|
|
format!(
|
|
"{}:{}",
|
|
Ipv6Addr::from(u128::from_be(event.saddr)),
|
|
event.lport
|
|
),
|
|
event.type_,
|
|
format!(
|
|
"{}:{}",
|
|
Ipv6Addr::from(u128::from_be(event.daddr)),
|
|
event.dport
|
|
),
|
|
event.state,
|
|
);
|
|
})
|
|
}
|
|
|
|
fn parse_ipv4_struct(x: &[u8]) -> ipv4_data_t {
|
|
unsafe { ptr::read_unaligned(x.as_ptr() as *const ipv4_data_t) }
|
|
}
|
|
|
|
fn parse_ipv6_struct(x: &[u8]) -> ipv6_data_t {
|
|
unsafe { ptr::read_unaligned(x.as_ptr() as *const ipv6_data_t) }
|
|
}
|
|
|
|
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(x) = do_main(runnable) {
|
|
eprintln!("Error: {}", x);
|
|
std::process::exit(1);
|
|
}
|
|
}
|