diff --git a/examples/loop.rs b/examples/loop.rs new file mode 100644 index 00000000..9344b29e --- /dev/null +++ b/examples/loop.rs @@ -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(); +} diff --git a/src/capture/activated/mod.rs b/src/capture/activated/mod.rs index 0b08bf88..5a4287c3 100644 --- a/src/capture/activated/mod.rs +++ b/src/capture/activated/mod.rs @@ -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, @@ -199,6 +202,39 @@ impl Capture { PacketIter::new(self, codec) } + pub fn for_each(&mut self, count: Option, 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::::callback, + &mut handler as *mut Handler> 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 { let program = CString::new(program)?; @@ -241,6 +277,45 @@ impl Capture { } } +// 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 { + func: F, + panic_payload: Option>, + handle: NonNull, +} + +impl Handler +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 From> for Capture { fn from(cap: Capture) -> Capture { unsafe { mem::transmute(cap) } @@ -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::(pcap); + let mut capture: Capture = 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::(pcap); + let mut capture: Capture = 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::(pcap); + let mut capture: Capture = 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::(pcap); + let mut capture: Capture = test_capture.capture.into(); + + let mut packets = 0; + capture + .for_each(Some(0), |_| { + packets += 1; + }) + .unwrap(); + assert_eq!(packets, 0); + } } diff --git a/src/raw.rs b/src/raw.rs index 0961ee0d..ce8dca86 100644 --- a/src/raw.rs +++ b/src/raw.rs @@ -100,8 +100,10 @@ pub struct pcap_send_queue { pub buffer: *mut c_char, } +// This is not Option, pcap functions do not check if the handler is null so it is wrong to +// pass them Option::::None. pub type pcap_handler = - Option ()>; + extern "C" fn(arg1: *mut c_uchar, arg2: *const pcap_pkthdr, arg3: *const c_uchar) -> (); #[cfg_attr(test, automock)] pub mod ffi { @@ -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; @@ -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;