1
//! Support for asynchronous packet iteration.
2
//!
3
//! See [`Capture::stream`](super::Capture::stream).
4
use std::io;
5
use std::marker::Unpin;
6
use std::pin::Pin;
7
use std::task::{self, Poll};
8

            
9
use futures::ready;
10
use tokio::io::unix::AsyncFd;
11

            
12
use crate::{
13
    Error,
14
    capture::{Activated, Capture, selectable::SelectableCapture},
15
    codec::PacketCodec,
16
};
17

            
18
/// Implement Stream for async use of pcap
19
pub struct PacketStream<T: Activated + ?Sized, C> {
20
    inner: AsyncFd<SelectableCapture<T>>,
21
    codec: C,
22
}
23

            
24
impl<T: Activated + ?Sized, C> PacketStream<T, C> {
25
6
    pub(crate) fn new(capture: Capture<T>, codec: C) -> Result<Self, Error> {
26
6
        let capture = SelectableCapture::new(capture)?;
27
        Ok(PacketStream {
28
6
            inner: AsyncFd::with_interest(capture, tokio::io::Interest::READABLE)?,
29
6
            codec,
30
        })
31
6
    }
32

            
33
    /// Returns a mutable reference to the inner [`Capture`].
34
    ///
35
    /// The caller must ensure the capture will not be set to be blocking.
36
2
    pub fn capture_mut(&mut self) -> &mut Capture<T> {
37
2
        self.inner.get_mut().get_inner_mut()
38
2
    }
39
}
40

            
41
impl<T: Activated + ?Sized, C> Unpin for PacketStream<T, C> {}
42

            
43
impl<T: Activated + ?Sized, C: PacketCodec> futures::Stream for PacketStream<T, C> {
44
    type Item = Result<C::Item, Error>;
45

            
46
8
    fn poll_next(self: Pin<&mut Self>, cx: &mut task::Context) -> Poll<Option<Self::Item>> {
47
8
        let stream = Pin::into_inner(self);
48
8
        let codec = &mut stream.codec;
49

            
50
        loop {
51
10
            let mut guard = ready!(stream.inner.poll_read_ready_mut(cx))?;
52
6
            match guard.try_io(
53
6
                |inner| match inner.get_mut().get_inner_mut().next_packet() {
54
2
                    Ok(p) => Ok(Ok(codec.decode(p))),
55
2
                    Err(e @ Error::TimeoutExpired) => {
56
2
                        Err(io::Error::new(io::ErrorKind::WouldBlock, e))
57
                    }
58
2
                    Err(e) => Ok(Err(e)),
59
6
                },
60
            ) {
61
4
                Ok(result) => {
62
4
                    return Poll::Ready(Some(result?));
63
                }
64
2
                Err(_would_block) => continue,
65
            }
66
        }
67
8
    }
68
}
69

            
70
#[cfg(test)]
71
mod tests {
72
    use std::os::unix::io::RawFd;
73

            
74
    use futures::{Stream, StreamExt};
75

            
76
    use crate::{
77
        capture::{
78
            Active,
79
            activated::testmod::{PACKET, next_ex_expect},
80
            selectable::testmod::FdPair,
81
            testmod::test_capture,
82
        },
83
        codec::testmod::Codec,
84
        raw::{
85
            self,
86
            testmod::{RAWMTX, as_pcap_t, geterr_expect},
87
        },
88
    };
89

            
90
    use super::*;
91

            
92
    // The caller holds on to the TestCapture, which owns the pcap_close expectation that fires
93
    // when the stream is dropped.
94
    fn test_stream(
95
        pcap: *mut raw::pcap_t,
96
        capture: Capture<Active>,
97
        fd: RawFd,
98
    ) -> PacketStream<Active, Codec> {
99
        let ctx = raw::pcap_get_selectable_fd_context();
100
        ctx.expect()
101
            .withf_st(move |arg1| *arg1 == pcap)
102
            .return_once(move |_| fd);
103

            
104
        PacketStream::new(capture, Codec).unwrap()
105
    }
106

            
107
    #[tokio::test]
108
    async fn test_stream_ok() {
109
        let _m = RAWMTX.lock();
110

            
111
        let mut dummy: isize = 777;
112
        let pcap = as_pcap_t(&mut dummy);
113
        let fds = FdPair::new();
114
        let fd = fds.0[0];
115

            
116
        let test_capture = test_capture::<Active>(pcap);
117

            
118
        let ctx = raw::pcap_setnonblock_context();
119
        ctx.expect()
120
            .withf_st(move |arg1, arg2, _| (*arg1 == pcap) && (*arg2 == 1))
121
            .return_once(|_, _, _| 0);
122

            
123
        let capture = test_capture.capture.setnonblock().unwrap();
124

            
125
        let ctx = raw::pcap_get_selectable_fd_context();
126
        ctx.expect()
127
            .withf_st(move |arg1| *arg1 == pcap)
128
            .return_once(move |_| fd);
129

            
130
        let mut stream = capture.stream(Codec).unwrap();
131
        assert!(stream.capture_mut().is_nonblock());
132
    }
133

            
134
    #[tokio::test]
135
    async fn test_stream_reads() {
136
        let _m = RAWMTX.lock();
137

            
138
        let mut dummy: isize = 777;
139
        let pcap = as_pcap_t(&mut dummy);
140
        let fds = FdPair::new();
141

            
142
        let test_capture = test_capture::<Active>(pcap);
143
        let mut stream = test_stream(pcap, test_capture.capture, fds.0[0]);
144

            
145
        fds.make_readable();
146

            
147
        let _nxt = next_ex_expect(pcap);
148

            
149
        let next_packet = stream.next().await.unwrap().unwrap();
150
        assert_eq!(next_packet.header, *PACKET.header);
151
        assert_eq!(*next_packet.data, *PACKET.data);
152

            
153
        let ctx = raw::pcap_next_ex_context();
154
        ctx.checkpoint();
155
        ctx.expect()
156
            .withf_st(move |arg1, _, _| *arg1 == pcap)
157
            .return_once(|_, _, _| -1);
158

            
159
        let _err = geterr_expect(pcap);
160

            
161
        // What libpcap failed with arrives as an item. It does not end the stream.
162
        let result = stream.next().await.unwrap();
163
        assert!(matches!(result, Err(Error::PcapError(_))));
164
    }
165

            
166
    #[tokio::test]
167
    async fn test_stream_timeout() {
168
        let _m = RAWMTX.lock();
169

            
170
        let mut dummy: isize = 777;
171
        let pcap = as_pcap_t(&mut dummy);
172
        let fds = FdPair::new();
173

            
174
        let test_capture = test_capture::<Active>(pcap);
175
        let mut stream = test_stream(pcap, test_capture.capture, fds.0[0]);
176

            
177
        fds.make_readable();
178

            
179
        let ctx = raw::pcap_next_ex_context();
180
        ctx.expect()
181
            .withf_st(move |arg1, _, _| *arg1 == pcap)
182
            .return_once(|_, _, _| 0);
183

            
184
        // Wait until the socketpair is reported readable, so the poll below gets past the
185
        // readiness check and asks libpcap for a packet. The guard is dropped without clearing,
186
        // which leaves the readiness in place.
187
        drop(stream.inner.readable().await.unwrap());
188

            
189
        // A timeout is not the end of the stream. It waits for the capture to be readable again.
190
        let poll =
191
            futures::future::poll_fn(|cx| Poll::Ready(Pin::new(&mut stream).poll_next(cx))).await;
192
        assert!(poll.is_pending());
193
    }
194
}