pcap: impl stream for PacketCapture

1 files changed, 57 insertions, 1 deletions

diff --git a/src/packet/capture.rs b/src/packet/capture.rs
index 645a303..c2157ca 100644
--- a/src/packet/capture.rs
+++ b/src/packet/capture.rs
@@ -1,5 +1,11 @@
+use std::{
+    pin::Pin,
+    task::{Context, Poll},
+};
+
 use crate::packet::packet::Packet;
-use pcap::Capture;
+use futures::Stream;
+use pcap::{Capture, Error, Packet as PacketPcap};
 
 pub struct PacketCapture {
     capture: Capture<pcap::Active>,
@@ -29,4 +35,54 @@ impl PacketCapture {
             println!("{:?}", device);
         }
     }
+
+    /// Returns a stream base on PacketCapture.
+    pub fn stream<C: PacketCodec>(self, codec: C) -> PacketStream<C> {
+        PacketStream::new(self.capture, codec)
+    }
+}
+
+//from https://github.com/rust-pcap/pcap
+pub struct PacketStream<C: PacketCodec> {
+    pub capture: Capture<pcap::Active>,
+    pub codec: C,
+}
+
+pub trait PacketCodec {
+    type Output: Unpin;
+    fn decode(&mut self, packet: PacketPcap) -> Self::Output;
+}
+
+impl<C: PacketCodec> PacketStream<C> {
+    pub fn new(capture: Capture<pcap::Active>, codec: C) -> Self {
+        Self { capture, codec }
+    }
+}
+
+impl<C: PacketCodec> Unpin for PacketStream<C> {}
+
+impl<C: PacketCodec> Stream for PacketStream<C> {
+    type Item = Result<C::Output, Error>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
+        let me = self.get_mut();
+        loop {
+            match me.capture.next_packet() {
+                Ok(packet) => {
+                    let output = me.codec.decode(packet);
+                    return Poll::Ready(Some(Ok(output)));
+                }
+                Err(e) => return Poll::Ready(Some(Err(e))),
+            }
+        }
+    }
+}
+
+pub struct BoxCodec;
+
+impl PacketCodec for BoxCodec {
+    type Output = Vec<u8>;
+    fn decode(&mut self, packet: PacketPcap) -> Self::Output {
+        packet.data.to_vec()
+    }
 }