forked from rust-pcap/pcap
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge remote-tracking branch 'origin/main'
- Loading branch information
Showing
5 changed files
with
162 additions
and
36 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,124 @@ | ||
| //! Support for asynchronous packet iteration. | ||
| //! | ||
| //! See [`Capture::stream`](super::Capture::stream). | ||
| use super::Activated; | ||
| use super::Capture; | ||
| use super::Error; | ||
| use crate::PacketCodec; | ||
| use futures::ready; | ||
| use futures::FutureExt; | ||
| use std::marker::Unpin; | ||
| use std::pin::Pin; | ||
| use std::task::{self, Poll}; | ||
| use tokio::task::JoinHandle; | ||
| use windows_sys::Win32::Foundation::HANDLE; | ||
| use windows_sys::Win32::System::Threading::WaitForSingleObject; | ||
|
|
||
| /// Implement Stream for async use of pcap | ||
| pub struct PacketStream<T: Activated + ?Sized, C> { | ||
| event_handle: EventHandle, | ||
| capture: Capture<T>, | ||
| codec: C, | ||
| } | ||
|
|
||
| impl<T: Activated + ?Sized, C> PacketStream<T, C> { | ||
| pub(crate) fn new(capture: Capture<T>, codec: C) -> Result<Self, Error> { | ||
| Ok(Self { | ||
| event_handle: EventHandle::new(&capture), | ||
| capture, | ||
| codec, | ||
| }) | ||
| } | ||
|
|
||
| /// Returns a mutable reference to the inner [`Capture`]. | ||
| /// | ||
| /// The caller must ensure the capture will not be set to be | ||
| /// blocking. | ||
| pub fn capture_mut(&mut self) -> &mut Capture<T> { | ||
| &mut self.capture | ||
| } | ||
| } | ||
|
|
||
| impl<T: Activated + ?Sized, C> Unpin for PacketStream<T, C> {} | ||
|
|
||
| impl<T: Activated + ?Sized, C: PacketCodec> futures::Stream for PacketStream<T, C> { | ||
| type Item = Result<C::Item, Error>; | ||
|
|
||
| fn poll_next(self: Pin<&mut Self>, cx: &mut task::Context) -> Poll<Option<Self::Item>> { | ||
| let stream = Pin::into_inner(self); | ||
| let codec = &mut stream.codec; | ||
|
|
||
| loop { | ||
| ready!(stream.event_handle.poll_ready(cx)); | ||
|
|
||
| let res = match stream.capture.next_packet() { | ||
| Ok(p) => Ok(codec.decode(p)), | ||
| Err(Error::TimeoutExpired) => { | ||
| stream.event_handle.clear_ready(); | ||
| continue; | ||
| } | ||
| Err(e) => Err(e), | ||
| }; | ||
| return Poll::Ready(Some(res)); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| /// A wrapper around a HANDLE that can be used to call `WaitForSingleObject` | ||
| /// from an asynchronous context. Once the call to `WaitForSingleObject` | ||
| /// completes, the handle is considered ready and will keep returning `Ready` | ||
| /// until it's reset. | ||
| struct EventHandle { | ||
| handle: HANDLE, | ||
| state: EventHandleState, | ||
| } | ||
|
|
||
| enum EventHandleState { | ||
| /// We haven't started waiting for an event yet. | ||
| Init, | ||
| /// We're currently waiting for an event. | ||
| Polling(JoinHandle<()>), | ||
| /// We waited for an event. | ||
| Ready, | ||
| } | ||
|
|
||
| impl EventHandle { | ||
| pub fn new<T: Activated + ?Sized>(capture: &Capture<T>) -> Self { | ||
| Self { | ||
| handle: unsafe { | ||
| // SAFETY: PacketStream stores the handle before the capture, | ||
| // so the handle will be dropped before the capture. | ||
| capture.get_event() | ||
| }, | ||
| state: EventHandleState::Init, | ||
| } | ||
| } | ||
|
|
||
| pub fn poll_ready(&mut self, cx: &mut task::Context) -> Poll<()> { | ||
| loop { | ||
| match self.state { | ||
| EventHandleState::Init => { | ||
| let handle = self.handle; | ||
| self.state = | ||
| EventHandleState::Polling(tokio::task::spawn_blocking(move || { | ||
| const INFINITE: u32 = !0; | ||
| unsafe { | ||
| WaitForSingleObject(handle, INFINITE); | ||
| } | ||
| })); | ||
| } | ||
| EventHandleState::Polling(ref mut join_handle) => { | ||
| let _ = ready!(join_handle.poll_unpin(cx)); | ||
| self.state = EventHandleState::Ready; | ||
| } | ||
| EventHandleState::Ready => return Poll::Ready(()), | ||
| } | ||
| } | ||
| } | ||
|
|
||
| /// Reset the internal state. This will trigger a call to | ||
| /// `WaitForSingleObject` the next time `poll_ready` is called. | ||
| pub fn clear_ready(&mut self) { | ||
| self.state = EventHandleState::Init; | ||
| } | ||
| } |