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

            
9
use futures::{Sink, ready};
10

            
11
#[cfg(target_os = "linux")]
12
use {crate::capture::selectable::SelectableCapture, tokio::io::unix::AsyncFd};
13

            
14
#[cfg(not(target_os = "linux"))]
15
use tokio::task::coop;
16

            
17
use crate::{
18
    Error,
19
    capture::{Active, Capture},
20
};
21

            
22
impl Capture<Active> {
23
    /// Returns this capture as a [`futures::Sink`] for sending packets.
24
    ///
25
    /// ```no_run
26
    /// # use futures::SinkExt;
27
    /// # use pcap::{Active, Capture};
28
    /// # async fn doc(capture: Capture<Active>) -> Result<(), pcap::Error> {
29
    /// let mut sink = capture.sink()?;
30
    /// sink.send(vec![0u8; 64]).await?;
31
    /// # Ok(())
32
    /// # }
33
    /// ```
34
    ///
35
    /// # Errors
36
    ///
37
    /// If this capture is set to be blocking, an error will be returned. On Linux, where the sink
38
    /// waits for the interface, an error is also returned if the network device does not support
39
    /// `select()`.
40
4
    pub fn sink<C: AsRef<[u8]>>(self) -> Result<PacketSink<C>, Error> {
41
4
        if !self.is_nonblock() {
42
2
            return Err(Error::NonNonBlock);
43
2
        }
44
2
        PacketSink::new(self)
45
4
    }
46
}
47

            
48
/// Implement Sink for async use of pcap
49
///
50
/// The packet given to `start_send` is held until a later poll can send it, so it is sent without
51
/// being copied. Only one packet is held at a time. A packet that fails to send is dropped rather
52
/// than retried, as libpcap does not tell us how much of it made it onto the wire.
53
///
54
/// Closing the sink flushes it but does not close the capture, which happens when the
55
/// [`PacketSink`] is dropped. Once closed, the sink takes no more packets and reports
56
/// `Error::IoError` with a kind of `BrokenPipe` instead.
57
///
58
/// # Warning
59
///
60
/// Only on Linux does the capture report when the interface is ready for another packet.
61
/// Elsewhere the packet is sent from within the poll, and a full transmit queue comes back as an
62
/// error instead of pausing the sink until there is room.
63
pub struct PacketSink<C> {
64
    #[cfg(target_os = "linux")]
65
    inner: AsyncFd<SelectableCapture<Active>>,
66
    #[cfg(not(target_os = "linux"))]
67
    capture: Capture<Active>,
68
    #[cfg(not(target_os = "linux"))]
69
    sent_since_yield: u32,
70
    packet: Option<C>,
71
    closed: bool,
72
}
73

            
74
#[cfg(target_os = "linux")]
75
impl<C> PacketSink<C> {
76
10
    pub(crate) fn new(capture: Capture<Active>) -> Result<Self, Error> {
77
10
        let capture = SelectableCapture::new(capture)?;
78
        Ok(PacketSink {
79
10
            inner: AsyncFd::with_interest(capture, tokio::io::Interest::WRITABLE)?,
80
10
            packet: None,
81
            closed: false,
82
        })
83
10
    }
84

            
85
    /// Returns a mutable reference to the inner [`Capture`].
86
    ///
87
    /// The caller must ensure the capture will not be set to be blocking.
88
2
    pub fn capture_mut(&mut self) -> &mut Capture<Active> {
89
2
        self.inner.get_mut().get_inner_mut()
90
2
    }
91

            
92
20
    fn poll_send(&mut self, cx: &mut task::Context<'_>) -> Poll<Result<(), Error>>
93
20
    where
94
20
        C: AsRef<[u8]>,
95
    {
96
20
        let Self { inner, packet, .. } = self;
97

            
98
        loop {
99
22
            let buf = match &*packet {
100
14
                Some(buf) => buf.as_ref(),
101
8
                None => return Poll::Ready(Ok(())),
102
            };
103

            
104
14
            let mut guard = ready!(inner.poll_write_ready_mut(cx))?;
105
            // A busy device goes through the io::Result, so that try_io knows to wait for the
106
            // next readiness event. A real error goes through the inner Result untouched.
107
8
            let result = guard.try_io(|inner| {
108
8
                match inner.get_mut().get_inner_mut().sendpacket_nonblock(buf) {
109
4
                    Ok(()) => Ok(Ok(())),
110
2
                    Err(e @ Error::IoError(io::ErrorKind::WouldBlock)) => {
111
2
                        Err(io::Error::new(io::ErrorKind::WouldBlock, e))
112
                    }
113
2
                    Err(e) => Ok(Err(e)),
114
                }
115
8
            });
116

            
117
8
            match result {
118
6
                Ok(result) => {
119
6
                    *packet = None;
120
6
                    return Poll::Ready(result?);
121
                }
122
2
                Err(_would_block) => continue,
123
            }
124
        }
125
20
    }
126
}
127

            
128
/// How many packets the sink sends before it yields to the executor.
129
///
130
/// This count and tokio's task budget are redundant rather than complementary: either alone is
131
/// enough to make the sink yield, and whichever runs out first is the one that does. The value
132
/// is the budget tokio gives a task, so a sink driven on tokio yields at much the same points
133
/// it otherwise would, but correctness does not depend on the two agreeing. Tokio keeps its own
134
/// number private and is free to change it; the smaller of the two then takes effect.
135
#[cfg(not(target_os = "linux"))]
136
const SENDS_BETWEEN_YIELDS: u32 = 128;
137

            
138
#[cfg(not(target_os = "linux"))]
139
impl<C> PacketSink<C> {
140
    pub(crate) fn new(capture: Capture<Active>) -> Result<Self, Error> {
141
        Ok(PacketSink {
142
            capture,
143
            sent_since_yield: 0,
144
            packet: None,
145
            closed: false,
146
        })
147
    }
148

            
149
    /// Returns a mutable reference to the inner [`Capture`].
150
    ///
151
    /// The caller must ensure the capture will not be set to be blocking.
152
    pub fn capture_mut(&mut self) -> &mut Capture<Active> {
153
        &mut self.capture
154
    }
155

            
156
    fn poll_send(&mut self, cx: &mut task::Context<'_>) -> Poll<Result<(), Error>>
157
    where
158
        C: AsRef<[u8]>,
159
    {
160
        let buf = match &self.packet {
161
            Some(packet) => packet.as_ref(),
162
            None => return Poll::Ready(Ok(())),
163
        };
164

            
165
        // Sending here never waits for the interface, so a sink that is kept fed would never
166
        // return Pending and the task it runs in would never let the executor poll anything
167
        // else. Two things stop that. The count applies to whoever is driving the sink, which off
168
        // Linux can be any executor, as it holds nothing of tokio's. Spending the task's budget
169
        // as well holds a task that also does tokio I/O to one budget between yields rather
170
        // than one for each source.
171
        if self.sent_since_yield == SENDS_BETWEEN_YIELDS {
172
            self.sent_since_yield = 0;
173
            cx.waker().wake_by_ref();
174
            return Poll::Pending;
175
        }
176
        let coop = ready!(coop::poll_proceed(cx));
177
        self.sent_since_yield += 1;
178

            
179
        let result = self.capture.sendpacket(buf);
180
        coop.made_progress();
181
        self.packet = None;
182
        Poll::Ready(result)
183
    }
184
}
185

            
186
impl<C> Unpin for PacketSink<C> {}
187

            
188
impl<C: AsRef<[u8]>> Sink<C> for PacketSink<C> {
189
    type Error = Error;
190

            
191
6
    fn poll_ready(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Result<(), Error>> {
192
6
        let sink = Pin::into_inner(self);
193
6
        if sink.closed {
194
2
            return Poll::Ready(Err(Error::IoError(io::ErrorKind::BrokenPipe)));
195
4
        }
196
4
        sink.poll_send(cx)
197
6
    }
198

            
199
8
    fn start_send(self: Pin<&mut Self>, item: C) -> Result<(), Error> {
200
8
        let sink = Pin::into_inner(self);
201
8
        if sink.closed {
202
            return Err(Error::IoError(io::ErrorKind::BrokenPipe));
203
8
        }
204
8
        if sink.packet.is_some() {
205
            // poll_ready did not report the sink ready, so the packet it was given last is
206
            // still waiting. Keep it: dropping it here would lose a packet the caller has
207
            // already been told the sink took.
208
2
            return Err(Error::IoError(io::ErrorKind::WouldBlock));
209
6
        }
210
6
        sink.packet = Some(item);
211
6
        Ok(())
212
8
    }
213

            
214
10
    fn poll_flush(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Result<(), Error>> {
215
10
        Pin::into_inner(self).poll_send(cx)
216
10
    }
217

            
218
4
    fn poll_close(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Result<(), Error>> {
219
4
        let sink = Pin::into_inner(self);
220
4
        ready!(sink.poll_send(cx))?;
221
4
        sink.closed = true;
222
4
        Poll::Ready(Ok(()))
223
4
    }
224
}
225

            
226
#[cfg(test)]
227
mod tests {
228
    use futures::SinkExt;
229

            
230
    use crate::{
231
        capture::testmod::test_capture,
232
        raw::{
233
            mock_ffi::*,
234
            testmod::{RAWMTX, as_pcap_t, geterr_expect},
235
        },
236
    };
237

            
238
    use super::*;
239

            
240
    #[test]
241
    fn test_sink_error() {
242
        let _m = RAWMTX.lock();
243

            
244
        let mut dummy: isize = 777;
245
        let pcap = as_pcap_t(&mut dummy);
246

            
247
        let test_capture = test_capture::<Active>(pcap);
248
        let capture = test_capture.capture;
249
        assert!(!capture.is_nonblock());
250

            
251
        let result = capture.sink::<Vec<u8>>();
252
        assert!(result.is_err());
253
    }
254

            
255
    #[cfg(target_os = "linux")]
256
    mod linux {
257
        use std::os::unix::io::RawFd;
258

            
259
        use crate::{capture::selectable::testmod::FdPair, raw};
260

            
261
        use super::*;
262

            
263
        // The caller holds on to the TestCapture, which owns the pcap_close expectation that fires
264
        // when the sink is dropped.
265
        fn test_sink(
266
            pcap: *mut raw::pcap_t,
267
            capture: Capture<Active>,
268
            fd: RawFd,
269
        ) -> PacketSink<Vec<u8>> {
270
            let ctx = raw::pcap_get_selectable_fd_context();
271
            ctx.expect()
272
                .withf_st(move |arg1| *arg1 == pcap)
273
                .return_once(move |_| fd);
274

            
275
            PacketSink::new(capture).unwrap()
276
        }
277

            
278
        #[tokio::test]
279
        async fn test_sink_ok() {
280
            let _m = RAWMTX.lock();
281

            
282
            let mut dummy: isize = 777;
283
            let pcap = as_pcap_t(&mut dummy);
284
            let fds = FdPair::new();
285
            let fd = fds.0[0];
286

            
287
            let test_capture = test_capture::<Active>(pcap);
288

            
289
            let ctx = raw::pcap_setnonblock_context();
290
            ctx.expect()
291
                .withf_st(move |arg1, arg2, _| (*arg1 == pcap) && (*arg2 == 1))
292
                .return_once(|_, _, _| 0);
293

            
294
            let capture = test_capture.capture.setnonblock().unwrap();
295

            
296
            let ctx = raw::pcap_get_selectable_fd_context();
297
            ctx.expect()
298
                .withf_st(move |arg1| *arg1 == pcap)
299
                .return_once(move |_| fd);
300

            
301
            let mut sink = capture.sink::<Vec<u8>>().unwrap();
302
            assert!(sink.capture_mut().is_nonblock());
303

            
304
            // Closing flushes the sink. The capture stays open until the sink is dropped.
305
            sink.close().await.unwrap();
306
        }
307

            
308
        #[tokio::test]
309
        async fn test_sink_sends() {
310
            let _m = RAWMTX.lock();
311

            
312
            let mut dummy: isize = 777;
313
            let pcap = as_pcap_t(&mut dummy);
314
            let fds = FdPair::new();
315

            
316
            let test_capture = test_capture::<Active>(pcap);
317
            let mut sink = test_sink(pcap, test_capture.capture, fds.0[0]);
318

            
319
            let ctx = pcap_sendpacket_context();
320
            ctx.expect()
321
                .withf_st(move |arg1, _, arg3| (*arg1 == pcap) && (*arg3 == 4))
322
                .return_once(|_, _, _| 0);
323

            
324
            sink.send(vec![1, 2, 3, 4]).await.unwrap();
325
            assert!(sink.packet.is_none());
326

            
327
            let ctx = pcap_sendpacket_context();
328
            ctx.checkpoint();
329
            ctx.expect()
330
                .withf_st(move |arg1, _, _| *arg1 == pcap)
331
                .return_once(|_, _, _| {
332
                    errno::set_errno(errno::Errno(libc::EINVAL));
333
                    -1
334
                });
335

            
336
            let _err = geterr_expect(pcap);
337

            
338
            let result = sink.send(vec![1, 2, 3, 4]).await;
339
            assert!(matches!(result, Err(Error::PcapError(_))));
340

            
341
            // The failed packet is dropped, so the sink can be used again.
342
            assert!(sink.packet.is_none());
343
        }
344

            
345
        #[tokio::test]
346
        async fn test_sink_backpressure() {
347
            let _m = RAWMTX.lock();
348

            
349
            let mut dummy: isize = 777;
350
            let pcap = as_pcap_t(&mut dummy);
351
            let fds = FdPair::new();
352

            
353
            let test_capture = test_capture::<Active>(pcap);
354
            let mut sink = test_sink(pcap, test_capture.capture, fds.0[0]);
355

            
356
            let ctx = pcap_sendpacket_context();
357
            ctx.expect()
358
                .withf_st(move |arg1, _, _| *arg1 == pcap)
359
                .returning(|_, _, _| {
360
                    errno::set_errno(errno::Errno(libc::EAGAIN));
361
                    -1
362
                });
363

            
364
            // Wait until the socketpair is reported writable, so the poll below gets past the
365
            // readiness check and asks libpcap to send. The guard is dropped without clearing,
366
            // which leaves the readiness in place.
367
            drop(sink.inner.writable().await.unwrap());
368

            
369
            sink.packet = Some(vec![1, 2, 3]);
370

            
371
            // A busy device is not an error. The sink waits and the packet stays queued.
372
            let poll = futures::future::poll_fn(|cx| Poll::Ready(sink.poll_send(cx))).await;
373
            assert!(poll.is_pending());
374
            assert_eq!(sink.packet, Some(vec![1, 2, 3]));
375
        }
376

            
377
        #[tokio::test]
378
        async fn test_sink_closed() {
379
            let _m = RAWMTX.lock();
380

            
381
            let mut dummy: isize = 777;
382
            let pcap = as_pcap_t(&mut dummy);
383
            let fds = FdPair::new();
384

            
385
            let test_capture = test_capture::<Active>(pcap);
386
            let mut sink = test_sink(pcap, test_capture.capture, fds.0[0]);
387

            
388
            sink.close().await.unwrap();
389

            
390
            // There is no pcap_sendpacket expectation, so a send here would fail the test.
391
            assert_eq!(
392
                sink.send(vec![1, 2, 3, 4]).await,
393
                Err(Error::IoError(io::ErrorKind::BrokenPipe))
394
            );
395
            assert!(sink.packet.is_none());
396
        }
397

            
398
        #[tokio::test]
399
        async fn test_sink_start_send_twice() {
400
            let _m = RAWMTX.lock();
401

            
402
            let mut dummy: isize = 777;
403
            let pcap = as_pcap_t(&mut dummy);
404
            let fds = FdPair::new();
405

            
406
            let test_capture = test_capture::<Active>(pcap);
407
            let mut sink = test_sink(pcap, test_capture.capture, fds.0[0]);
408

            
409
            Pin::new(&mut sink).start_send(vec![1, 2, 3, 4]).unwrap();
410

            
411
            // Nothing has sent the first packet yet, so the second one is refused rather than
412
            // put in its place.
413
            assert_eq!(
414
                Pin::new(&mut sink).start_send(vec![5, 6, 7]),
415
                Err(Error::IoError(io::ErrorKind::WouldBlock))
416
            );
417
            assert_eq!(sink.packet, Some(vec![1, 2, 3, 4]));
418

            
419
            let ctx = pcap_sendpacket_context();
420
            ctx.expect()
421
                .withf_st(move |arg1, _, arg3| (*arg1 == pcap) && (*arg3 == 4))
422
                .return_once(|_, _, _| 0);
423

            
424
            sink.flush().await.unwrap();
425
            assert!(sink.packet.is_none());
426
        }
427
    }
428

            
429
    #[cfg(not(target_os = "linux"))]
430
    #[tokio::test]
431
    async fn test_sink_ok() {
432
        let _m = RAWMTX.lock();
433

            
434
        let mut dummy: isize = 777;
435
        let pcap = as_pcap_t(&mut dummy);
436

            
437
        let test_capture = test_capture::<Active>(pcap);
438

            
439
        let ctx = pcap_setnonblock_context();
440
        ctx.expect()
441
            .withf_st(move |arg1, arg2, _| (*arg1 == pcap) && (*arg2 == 1))
442
            .return_once(|_, _, _| 0);
443

            
444
        let capture = test_capture.capture.setnonblock().unwrap();
445

            
446
        let mut sink = capture.sink::<Vec<u8>>().unwrap();
447
        assert!(sink.capture_mut().is_nonblock());
448

            
449
        // Closing flushes the sink. The capture stays open until the sink is dropped.
450
        sink.close().await.unwrap();
451
    }
452

            
453
    #[cfg(not(target_os = "linux"))]
454
    #[tokio::test]
455
    async fn test_sink_sends() {
456
        let _m = RAWMTX.lock();
457

            
458
        let mut dummy: isize = 777;
459
        let pcap = as_pcap_t(&mut dummy);
460

            
461
        let test_capture = test_capture::<Active>(pcap);
462
        let mut sink = PacketSink::new(test_capture.capture).unwrap();
463

            
464
        let ctx = pcap_sendpacket_context();
465
        ctx.expect()
466
            .withf_st(move |arg1, _, arg3| (*arg1 == pcap) && (*arg3 == 4))
467
            .return_once(|_, _, _| 0);
468

            
469
        sink.send(vec![1, 2, 3, 4]).await.unwrap();
470
        assert!(sink.packet.is_none());
471

            
472
        let ctx = pcap_sendpacket_context();
473
        ctx.checkpoint();
474
        ctx.expect()
475
            .withf_st(move |arg1, _, _| *arg1 == pcap)
476
            .return_once(|_, _, _| -1);
477

            
478
        let _err = geterr_expect(pcap);
479

            
480
        let result = sink.send(vec![1, 2, 3, 4]).await;
481
        assert!(matches!(result, Err(Error::PcapError(_))));
482
        assert!(sink.packet.is_none());
483
    }
484

            
485
    #[cfg(not(target_os = "linux"))]
486
    #[tokio::test]
487
    async fn test_sink_closed() {
488
        let _m = RAWMTX.lock();
489

            
490
        let mut dummy: isize = 777;
491
        let pcap = as_pcap_t(&mut dummy);
492

            
493
        let test_capture = test_capture::<Active>(pcap);
494
        let mut sink = PacketSink::new(test_capture.capture).unwrap();
495

            
496
        sink.close().await.unwrap();
497

            
498
        // There is no pcap_sendpacket expectation, so a send here would fail the test.
499
        assert_eq!(
500
            sink.send(vec![1, 2, 3, 4]).await,
501
            Err(Error::IoError(io::ErrorKind::BrokenPipe))
502
        );
503
        assert!(sink.packet.is_none());
504
    }
505

            
506
    #[cfg(not(target_os = "linux"))]
507
    #[tokio::test]
508
    async fn test_sink_start_send_twice() {
509
        let _m = RAWMTX.lock();
510

            
511
        let mut dummy: isize = 777;
512
        let pcap = as_pcap_t(&mut dummy);
513

            
514
        let test_capture = test_capture::<Active>(pcap);
515
        let mut sink = PacketSink::new(test_capture.capture).unwrap();
516

            
517
        Pin::new(&mut sink).start_send(vec![1, 2, 3, 4]).unwrap();
518

            
519
        // Nothing has sent the first packet yet, so the second one is refused rather than put
520
        // in its place.
521
        assert_eq!(
522
            Pin::new(&mut sink).start_send(vec![5, 6, 7]),
523
            Err(Error::IoError(io::ErrorKind::WouldBlock))
524
        );
525
        assert_eq!(sink.packet, Some(vec![1, 2, 3, 4]));
526

            
527
        let ctx = pcap_sendpacket_context();
528
        ctx.expect()
529
            .withf_st(move |arg1, _, arg3| (*arg1 == pcap) && (*arg3 == 4))
530
            .return_once(|_, _, _| 0);
531

            
532
        sink.flush().await.unwrap();
533
        assert!(sink.packet.is_none());
534
    }
535

            
536
    // Deliberately not a tokio test. The sink is a futures::Sink and off Linux it holds nothing
537
    // of tokio's, so it has to give up its turn whatever executor is driving it.
538
    #[cfg(not(target_os = "linux"))]
539
    #[test]
540
    fn test_sink_yields() {
541
        let _m = RAWMTX.lock();
542

            
543
        let mut dummy: isize = 777;
544
        let pcap = as_pcap_t(&mut dummy);
545

            
546
        let test_capture = test_capture::<Active>(pcap);
547
        let mut sink = PacketSink::new(test_capture.capture).unwrap();
548

            
549
        let ctx = pcap_sendpacket_context();
550
        ctx.expect()
551
            .withf_st(move |arg1, _, _| *arg1 == pcap)
552
            .returning(|_, _, _| 0);
553

            
554
        let waker = futures::task::noop_waker();
555
        let mut cx = task::Context::from_waker(&waker);
556

            
557
        // Keep the sink fed and poll it until it asks to be polled again later. One that never
558
        // does keeps the thread to itself for as long as there are packets to send.
559
        let mut sends = 0;
560
        loop {
561
            sink.packet = Some(vec![1, 2, 3, 4]);
562
            if sink.poll_send(&mut cx).is_pending() {
563
                break;
564
            }
565
            sends += 1;
566
            assert!(sends < 10_000, "the sink never gave up its turn");
567
        }
568
        assert_eq!(sends, SENDS_BETWEEN_YIELDS);
569

            
570
        // The packet the sink yielded on is still there for the poll after it.
571
        assert_eq!(sink.packet, Some(vec![1, 2, 3, 4]));
572
    }
573

            
574
    // Under tokio the sink spends the task's budget too, so a task that has already spent some
575
    // of it elsewhere gets its turn back sooner than the count alone would give it.
576
    #[cfg(not(target_os = "linux"))]
577
    #[tokio::test]
578
    async fn test_sink_yields_on_task_budget() {
579
        let _m = RAWMTX.lock();
580

            
581
        let mut dummy: isize = 777;
582
        let pcap = as_pcap_t(&mut dummy);
583

            
584
        let test_capture = test_capture::<Active>(pcap);
585
        let mut sink = PacketSink::new(test_capture.capture).unwrap();
586

            
587
        let ctx = pcap_sendpacket_context();
588
        ctx.expect()
589
            .withf_st(move |arg1, _, _| *arg1 == pcap)
590
            .returning(|_, _, _| 0);
591

            
592
        // Spend half the budget on something that is not the sink.
593
        for _ in 0..SENDS_BETWEEN_YIELDS / 2 {
594
            coop::consume_budget().await;
595
        }
596

            
597
        let mut sends = 0;
598
        loop {
599
            sink.packet = Some(vec![1, 2, 3, 4]);
600
            if futures::future::poll_fn(|cx| Poll::Ready(sink.poll_send(cx)))
601
                .await
602
                .is_pending()
603
            {
604
                break;
605
            }
606
            sends += 1;
607
            assert!(sends < 10_000, "the sink never gave up its turn");
608
        }
609
        assert!(
610
            sends < SENDS_BETWEEN_YIELDS,
611
            "the sink sent {sends} before yielding, so it kept a budget of its own"
612
        );
613
    }
614
}