diff --git a/examples/loop.rs b/examples/loop.rs new file mode 100644 index 000000000..f338c76db --- /dev/null +++ b/examples/loop.rs @@ -0,0 +1,23 @@ +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(|packet| { + println!("Got {:?}", packet.header); + count += 1; + if count > 100 { + panic!("ow"); + } + }); +} diff --git a/src/lib.rs b/src/lib.rs index 72e9383d6..bd599424d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -59,6 +59,7 @@ //! ``` use bitflags::bitflags; +use std::any::Any; use std::borrow::Borrow; use std::convert::TryFrom; use std::ffi::{self, CStr, CString}; @@ -69,6 +70,7 @@ use std::net::IpAddr; use std::ops::Deref; #[cfg(not(windows))] use std::os::unix::io::{AsRawFd, RawFd}; +use std::panic::{catch_unwind, resume_unwind, AssertUnwindSafe}; use std::path::Path; use std::ptr::{self, NonNull}; use std::slice; @@ -835,7 +837,68 @@ impl From> for 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 Capture { + pub fn for_each(&mut self, handler: F) + where + F: FnMut(Packet), + { + let mut handler = Handler { + func: AssertUnwindSafe(handler), + panic_payload: None, + handle: self.handle, + }; + unsafe { + raw::pcap_loop( + self.handle.as_ptr(), + -1, + Handler::::callback, + ptr::addr_of_mut!(handler).cast(), + ) + }; + if let Some(e) = handler.panic_payload { + resume_unwind(e); + } + } + fn new_raw(path: Option<&str>, func: F) -> Result, Error> where F: FnOnce(*const libc::c_char, *mut libc::c_char) -> *mut raw::pcap_t, diff --git a/src/raw.rs b/src/raw.rs index 6188184ae..6d400d11f 100644 --- a/src/raw.rs +++ b/src/raw.rs @@ -99,8 +99,9 @@ pub struct pcap_send_queue { pub buffer: *mut c_char, } +// This is not Option, pcap functions do not check if the handler is null. pub type pcap_handler = - Option ()>; + extern "C" fn(user: *mut c_uchar, h: *const pcap_pkthdr, bytes: *const c_uchar) -> (); extern "C" { // [OBSOLETE] pub fn pcap_lookupdev(arg1: *mut c_char) -> *mut c_char; @@ -119,8 +120,12 @@ extern "C" { 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; @@ -129,7 +134,7 @@ extern "C" { 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; diff --git a/tests/lib.rs b/tests/lib.rs index 186f28370..3629fd5b0 100644 --- a/tests/lib.rs +++ b/tests/lib.rs @@ -31,6 +31,16 @@ fn read_packet_with_full_data() { assert_eq!(capture.next_packet().unwrap().len(), 98); } +#[test] +fn read_packet_via_pcap_loop() { + let mut packets = 0; + let mut capture = capture_from_test_file("packet_snaplen_65535.pcap"); + capture.pcap_loop(|_| { + packets += 1; + }); + assert_eq!(packets, 1); +} + #[test] fn read_packet_with_truncated_data() { let mut capture = capture_from_test_file("packet_snaplen_20.pcap");