Skip to content

Commit

Permalink
Add a binding for pcap_loop (#335)
Browse files Browse the repository at this point in the history
* Add a binding for pcap_loop

* Add a test that covers panic in pcap_loop

* Change to mock tests and add a count argument

* Check return code from pcap_loop
  • Loading branch information
saethlin authored Mar 15, 2024
1 parent 587b263 commit d1494d0
Show file tree
Hide file tree
Showing 3 changed files with 239 additions and 4 deletions.
24 changes: 24 additions & 0 deletions examples/loop.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
fn main() {
// get the default Device
let device = pcap::Device::lookup()
.expect("device lookup failed")
.expect("no device available");
println!("Using device {}", device.name);

// Setup Capture
let mut cap = pcap::Capture::from_device(device)
.unwrap()
.immediate_mode(true)
.open()
.unwrap();

let mut count = 0;
cap.for_each(None, |packet| {
println!("Got {:?}", packet.header);
count += 1;
if count > 100 {
panic!("ow");
}
})
.unwrap();
}
205 changes: 205 additions & 0 deletions src/capture/activated/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,11 @@ pub mod iterator;
pub mod offline;

use std::{
any::Any,
convert::TryInto,
ffi::CString,
fmt, mem,
panic::{catch_unwind, resume_unwind, AssertUnwindSafe},
path::Path,
ptr::{self, NonNull},
slice,
Expand Down Expand Up @@ -199,6 +202,39 @@ impl<T: Activated + ?Sized> Capture<T> {
PacketIter::new(self, codec)
}

pub fn for_each<F>(&mut self, count: Option<usize>, handler: F) -> Result<(), Error>
where
F: FnMut(Packet),
{
let cnt = match count {
// Actually passing 0 down to pcap_loop would mean read forever.
// We interpret it as "read nothing", so we just succeed immediately.
Some(0) => return Ok(()),
Some(cnt) => cnt
.try_into()
.expect("count of packets to read cannot exceed c_int::MAX"),
None => -1,
};

let mut handler = Handler {
func: AssertUnwindSafe(handler),
panic_payload: None,
handle: self.handle,
};
let return_code = unsafe {
raw::pcap_loop(
self.handle.as_ptr(),
cnt,
Handler::<F>::callback,
&mut handler as *mut Handler<AssertUnwindSafe<F>> as *mut u8,
)
};
if let Some(e) = handler.panic_payload {
resume_unwind(e);
}
self.check_err(return_code == 0)
}

/// Compiles the string into a filter program using `pcap_compile`.
pub fn compile(&self, program: &str, optimize: bool) -> Result<BpfProgram, Error> {
let program = CString::new(program)?;
Expand Down Expand Up @@ -241,6 +277,45 @@ impl<T: Activated + ?Sized> Capture<T> {
}
}

// Handler and its associated function let us create an extern "C" fn which dispatches to a normal
// Rust FnMut, which may be a closure with a captured environment. The *only* purpose of this
// generic parameter is to ensure that in Capture::pcap_loop that we pass the right function
// pointer and the right data pointer to pcap_loop.
struct Handler<F> {
func: F,
panic_payload: Option<Box<dyn Any + Send>>,
handle: NonNull<raw::pcap_t>,
}

impl<F> Handler<F>
where
F: FnMut(Packet),
{
extern "C" fn callback(
slf: *mut libc::c_uchar,
header: *const raw::pcap_pkthdr,
packet: *const libc::c_uchar,
) {
unsafe {
let packet = Packet::new(
&*(header as *const PacketHeader),
slice::from_raw_parts(packet, (*header).caplen as _),
);

let slf = slf as *mut Self;
let func = &mut (*slf).func;
let mut func = AssertUnwindSafe(func);
// If our handler function panics, we need to prevent it from unwinding across the
// FFI boundary. If the handler panics we catch the unwind here, break out of
// pcap_loop, and resume the unwind outside.
if let Err(e) = catch_unwind(move || func(packet)) {
(*slf).panic_payload = Some(e);
raw::pcap_breakloop((*slf).handle.as_ptr());
}
}
}
}

impl<T: Activated> From<Capture<T>> for Capture<dyn Activated> {
fn from(cap: Capture<T>) -> Capture<dyn Activated> {
unsafe { mem::transmute(cap) }
Expand Down Expand Up @@ -920,4 +995,134 @@ mod tests {
});
assert_eq!(format!("{}", instr), "1 2 3 4");
}

#[test]
fn read_packet_via_pcap_loop() {
let _m = RAWMTX.lock();

let mut value: isize = 777;
let pcap = as_pcap_t(&mut value);

let test_capture = test_capture::<Active>(pcap);
let mut capture: Capture<dyn Activated> = test_capture.capture.into();

let ctx = raw::pcap_loop_context();
ctx.expect()
.withf_st(move |arg1, cnt, _, _| *arg1 == pcap && *cnt == -1)
.return_once_st(move |_, _, func, data| {
let header = raw::pcap_pkthdr {
ts: libc::timeval {
tv_sec: 0,
tv_usec: 0,
},
caplen: 0,
len: 0,
};
let packet_data = &[];
func(data, &header, packet_data.as_ptr());
0
});

let mut packets = 0;
capture
.for_each(None, |_| {
packets += 1;
})
.unwrap();
assert_eq!(packets, 1);
}

#[test]
#[should_panic = "panic in callback"]
fn panic_in_pcap_loop() {
let _m = RAWMTX.lock();

let mut value: isize = 777;
let pcap = as_pcap_t(&mut value);

let test_capture = test_capture::<Active>(pcap);
let mut capture: Capture<dyn Activated> = test_capture.capture.into();

let ctx = raw::pcap_loop_context();
ctx.expect()
.withf_st(move |arg1, cnt, _, _| *arg1 == pcap && *cnt == -1)
.return_once_st(move |_, _, func, data| {
let header = raw::pcap_pkthdr {
ts: libc::timeval {
tv_sec: 0,
tv_usec: 0,
},
caplen: 0,
len: 0,
};
let packet_data = &[];
func(data, &header, packet_data.as_ptr());
0
});

let ctx = raw::pcap_breakloop_context();
ctx.expect()
.withf_st(move |arg1| *arg1 == pcap)
.return_once_st(move |_| {});

capture
.for_each(None, |_| panic!("panic in callback"))
.unwrap();
}

#[test]
fn for_each_with_count() {
let _m = RAWMTX.lock();

let mut value: isize = 777;
let pcap = as_pcap_t(&mut value);

let test_capture = test_capture::<Active>(pcap);
let mut capture: Capture<dyn Activated> = test_capture.capture.into();

let ctx = raw::pcap_loop_context();
ctx.expect()
.withf_st(move |arg1, cnt, _, _| *arg1 == pcap && *cnt == 2)
.return_once_st(move |_, _, func, data| {
let header = raw::pcap_pkthdr {
ts: libc::timeval {
tv_sec: 0,
tv_usec: 0,
},
caplen: 0,
len: 0,
};
let packet_data = &[];
func(data, &header, packet_data.as_ptr());
func(data, &header, packet_data.as_ptr());
0
});

let mut packets = 0;
capture
.for_each(Some(2), |_| {
packets += 1;
})
.unwrap();
assert_eq!(packets, 2);
}

#[test]
fn for_each_with_count_0() {
let _m = RAWMTX.lock();

let mut value: isize = 777;
let pcap = as_pcap_t(&mut value);

let test_capture = test_capture::<Active>(pcap);
let mut capture: Capture<dyn Activated> = test_capture.capture.into();

let mut packets = 0;
capture
.for_each(Some(0), |_| {
packets += 1;
})
.unwrap();
assert_eq!(packets, 0);
}
}
14 changes: 10 additions & 4 deletions src/raw.rs
Original file line number Diff line number Diff line change
Expand Up @@ -100,8 +100,10 @@ pub struct pcap_send_queue {
pub buffer: *mut c_char,
}

// This is not Option<fn>, pcap functions do not check if the handler is null so it is wrong to
// pass them Option::<fn>::None.
pub type pcap_handler =
Option<extern "C" fn(arg1: *mut c_uchar, arg2: *const pcap_pkthdr, arg3: *const c_uchar) -> ()>;
extern "C" fn(arg1: *mut c_uchar, arg2: *const pcap_pkthdr, arg3: *const c_uchar) -> ();

#[cfg_attr(test, automock)]
pub mod ffi {
Expand All @@ -124,8 +126,12 @@ pub mod ffi {
pub fn pcap_open_offline(arg1: *const c_char, arg2: *mut c_char) -> *mut pcap_t;
pub fn pcap_fopen_offline(arg1: *mut FILE, arg2: *mut c_char) -> *mut pcap_t;
pub fn pcap_close(arg1: *mut pcap_t);
// pub fn pcap_loop(arg1: *mut pcap_t, arg2: c_int,
// arg3: pcap_handler, arg4: *mut c_uchar) -> c_int;
pub fn pcap_loop(
arg1: *mut pcap_t,
arg2: c_int,
arg3: pcap_handler,
arg4: *mut c_uchar,
) -> c_int;
// pub fn pcap_dispatch(arg1: *mut pcap_t, arg2: c_int, arg3: pcap_handler,
// arg4: *mut c_uchar)-> c_int;
// pub fn pcap_next(arg1: *mut pcap_t, arg2: *mut pcap_pkthdr) -> *const c_uchar;
Expand All @@ -134,7 +140,7 @@ pub mod ffi {
arg2: *mut *mut pcap_pkthdr,
arg3: *mut *const c_uchar,
) -> c_int;
// pub fn pcap_breakloop(arg1: *mut pcap_t);
pub fn pcap_breakloop(arg1: *mut pcap_t);
pub fn pcap_stats(arg1: *mut pcap_t, arg2: *mut pcap_stat) -> c_int;
pub fn pcap_setfilter(arg1: *mut pcap_t, arg2: *mut bpf_program) -> c_int;
pub fn pcap_setdirection(arg1: *mut pcap_t, arg2: pcap_direction_t) -> c_int;
Expand Down

0 comments on commit d1494d0

Please sign in to comment.