bach/environment/net/pcap/
queue.rs

1use crate::{
2    environment::net::{ip::Packet, pcap::Writer},
3    queue::{CloseError, PopError, PushError, Pushable, Queue},
4};
5use std::{io, task::Context};
6
7#[derive(Clone, Copy, Debug, PartialEq, Eq)]
8pub enum Direction {
9    Push,
10    Pop,
11}
12
13pub trait QueueExt
14where
15    Self: Sized + Queue<Packet>,
16{
17    fn pcap<O: io::Write>(self, out: Writer<O>, direction: Direction) -> PcapQueue<Self, O> {
18        PcapQueue::new(self, out, direction)
19    }
20
21    fn pcap_push<O: io::Write>(self, out: Writer<O>) -> PcapQueue<Self, O> {
22        self.pcap(out, Direction::Push)
23    }
24
25    fn pcap_pop<O: io::Write>(self, out: Writer<O>) -> PcapQueue<Self, O> {
26        self.pcap(out, Direction::Pop)
27    }
28}
29
30impl<Q: Queue<Packet>> QueueExt for Q {}
31
32pub struct PcapQueue<Q, O>
33where
34    Q: Queue<Packet>,
35    O: io::Write,
36{
37    queue: Q,
38    pcap: Writer<O>,
39    direction: Direction,
40}
41
42impl<Q, O> PcapQueue<Q, O>
43where
44    Q: Queue<Packet>,
45    O: io::Write,
46{
47    pub fn new(queue: Q, out: Writer<O>, direction: Direction) -> Self {
48        PcapQueue {
49            queue,
50            pcap: out,
51            direction,
52        }
53    }
54}
55
56impl<Q, O> Queue<Packet> for PcapQueue<Q, O>
57where
58    Q: Queue<Packet>,
59    O: io::Write,
60{
61    fn push_lazy(&mut self, value: &mut dyn Pushable<Packet>) -> Result<Option<Packet>, PushError> {
62        if self.direction == Direction::Pop {
63            return self.queue.push_lazy(value);
64        }
65        let mut value = PushablePacket {
66            inner: value,
67            writer: &mut self.pcap,
68        };
69        self.queue.push_lazy(&mut value)
70    }
71
72    fn push_with_notify(
73        &mut self,
74        value: &mut dyn Pushable<Packet>,
75        cx: &mut Context,
76    ) -> Result<Option<Packet>, PushError> {
77        if self.direction == Direction::Pop {
78            return self.queue.push_with_notify(value, cx);
79        }
80        let mut value = PushablePacket {
81            inner: value,
82            writer: &mut self.pcap,
83        };
84        self.queue.push_with_notify(&mut value, cx)
85    }
86
87    fn pop(&mut self) -> Result<Packet, PopError> {
88        let mut packet = self.queue.pop()?;
89        if self.direction == Direction::Pop {
90            self.pcap
91                .write_packet(&mut packet)
92                .expect("failed to write pcap");
93        }
94        Ok(packet)
95    }
96
97    fn pop_with_notify(&mut self, cx: &mut Context) -> Result<Packet, PopError> {
98        let mut packet = self.queue.pop_with_notify(cx)?;
99        if self.direction == Direction::Pop {
100            self.pcap
101                .write_packet(&mut packet)
102                .expect("failed to write pcap");
103        }
104        Ok(packet)
105    }
106
107    fn close(&mut self) -> Result<(), CloseError> {
108        self.queue.close()
109    }
110
111    fn is_closed(&self) -> bool {
112        self.queue.is_closed()
113    }
114
115    fn is_empty(&self) -> bool {
116        self.queue.is_empty()
117    }
118
119    fn is_full(&self) -> bool {
120        self.queue.is_full()
121    }
122
123    fn len(&self) -> usize {
124        self.queue.len()
125    }
126
127    fn capacity(&self) -> Option<usize> {
128        self.queue.capacity()
129    }
130}
131
132struct PushablePacket<'a, O> {
133    inner: &'a mut dyn Pushable<Packet>,
134    writer: &'a mut Writer<O>,
135}
136
137impl<O> Pushable<Packet> for PushablePacket<'_, O>
138where
139    O: io::Write,
140{
141    fn produce(&mut self) -> Packet {
142        let mut packet = self.inner.produce();
143
144        self.writer
145            .write_packet(&mut packet)
146            .expect("failed to write pcap");
147
148        packet
149    }
150}