ya_relay_stack/
device.rs

1use smoltcp::phy;
2use smoltcp::time::Instant;
3use std::cell::RefCell;
4use std::collections::VecDeque;
5use std::rc::Rc;
6
7use crate::metrics::ChannelMetrics;
8
9use ya_relay_util::Payload;
10
11type Pcap = RefCell<Box<dyn phy::PcapSink>>;
12
13/// Network device capable of injecting and extracting packets
14pub struct CaptureDevice {
15    tx_queue: VecDeque<Vec<u8>>,
16    rx_queue: VecDeque<Payload>,
17    medium: phy::Medium,
18    max_transmission_unit: usize,
19    pcap: Option<Pcap>,
20    metrics: Rc<RefCell<ChannelMetrics>>,
21}
22
23impl Default for CaptureDevice {
24    fn default() -> Self {
25        Self {
26            tx_queue: Default::default(),
27            rx_queue: Default::default(),
28            medium: Default::default(),
29            max_transmission_unit: 1280,
30            pcap: Default::default(),
31            metrics: Default::default(),
32        }
33    }
34}
35
36impl CaptureDevice {
37    pub fn tap(mtu: usize) -> Self {
38        Self {
39            max_transmission_unit: mtu,
40            ..Default::default()
41        }
42    }
43
44    pub fn tun(mtu: usize) -> Self {
45        Self {
46            medium: phy::Medium::Ip,
47            max_transmission_unit: mtu,
48            ..Default::default()
49        }
50    }
51
52    pub fn pcap_tap<S>(mtu: usize, mut pcap: S) -> Self
53    where
54        S: phy::PcapSink + 'static,
55    {
56        pcap.global_header(phy::PcapLinkType::Ethernet);
57        let pcap = RefCell::new(Box::new(pcap));
58
59        Self {
60            max_transmission_unit: mtu,
61            pcap: Some(pcap),
62            ..Default::default()
63        }
64    }
65
66    pub fn pcap_tun<S>(mtu: usize, mut pcap: S) -> Self
67    where
68        S: phy::PcapSink + 'static,
69    {
70        pcap.global_header(phy::PcapLinkType::Ip);
71        let pcap = RefCell::new(Box::new(pcap));
72
73        Self {
74            medium: phy::Medium::Ip,
75            max_transmission_unit: mtu,
76            pcap: Some(pcap),
77            ..Default::default()
78        }
79    }
80
81    #[inline]
82    pub fn is_tun(&self) -> bool {
83        self.medium == phy::Medium::Ip
84    }
85
86    #[inline]
87    pub fn metrics(&self) -> ChannelMetrics {
88        self.metrics.borrow().clone()
89    }
90
91    #[inline]
92    pub fn phy_rx(&mut self, data: impl Into<Payload>) {
93        self.rx_queue.push_back(data.into());
94    }
95
96    #[inline]
97    pub fn next_phy_tx(&mut self) -> Option<Vec<u8>> {
98        self.tx_queue.pop_front()
99    }
100}
101
102impl phy::Device for CaptureDevice {
103    type RxToken<'a> = RxToken<'a>
104    where
105        Self: 'a;
106    type TxToken<'a> = TxToken<'a>
107    where
108        Self: 'a;
109
110    fn receive(&mut self, timestamp: Instant) -> Option<(Self::RxToken<'_>, Self::TxToken<'_>)> {
111        let item = self.rx_queue.pop_front();
112        item.map(move |buffer| {
113            let rx = RxToken {
114                buffer,
115                pcap: &self.pcap,
116                metrics: self.metrics.clone(),
117                timestamp,
118            };
119            let tx = TxToken {
120                queue: &mut self.tx_queue,
121                pcap: &self.pcap,
122                metrics: self.metrics.clone(),
123                timestamp,
124            };
125            (rx, tx)
126        })
127    }
128
129    fn transmit(&mut self, timestamp: Instant) -> Option<Self::TxToken<'_>> {
130        Some(TxToken {
131            queue: &mut self.tx_queue,
132            pcap: &self.pcap,
133            metrics: self.metrics.clone(),
134            timestamp,
135        })
136    }
137
138    fn capabilities(&self) -> phy::DeviceCapabilities {
139        let mut caps = phy::DeviceCapabilities::default();
140        caps.max_transmission_unit = self.max_transmission_unit;
141        caps.medium = self.medium;
142        caps
143    }
144}
145
146/// Receipt token
147pub struct RxToken<'a> {
148    buffer: Payload,
149    pcap: &'a Option<Pcap>,
150    metrics: Rc<RefCell<ChannelMetrics>>,
151    timestamp: Instant,
152}
153
154impl<'a> phy::RxToken for RxToken<'a> {
155    fn consume<R, F>(mut self, f: F) -> R
156    where
157        F: FnOnce(&mut [u8]) -> R,
158    {
159        let result = f(self.buffer.as_mut());
160
161        {
162            let mut metrics = self.metrics.borrow_mut();
163            metrics.rx.push(self.buffer.len() as f32);
164        }
165
166        if let Some(pcap) = self.pcap {
167            pcap.borrow_mut()
168                .packet(self.timestamp, self.buffer.as_ref());
169        }
170
171        result
172    }
173}
174
175/// Transmission token
176pub struct TxToken<'a> {
177    queue: &'a mut VecDeque<Vec<u8>>,
178    pcap: &'a Option<Pcap>,
179    metrics: Rc<RefCell<ChannelMetrics>>,
180    timestamp: Instant,
181}
182
183impl<'a> phy::TxToken for TxToken<'a> {
184    fn consume<R, F>(self, len: usize, f: F) -> R
185    where
186        F: FnOnce(&mut [u8]) -> R,
187    {
188        let mut buffer = vec![0; len];
189        buffer.resize(len, 0);
190        let result = f(&mut buffer);
191
192        {
193            let mut metrics = self.metrics.borrow_mut();
194            metrics.tx.push(buffer.len() as f32);
195        }
196
197        if let Some(pcap) = self.pcap {
198            pcap.borrow_mut().packet(self.timestamp, buffer.as_ref());
199        }
200
201        self.queue.push_back(buffer);
202        result
203    }
204}
ya_relay_stack/device.rs

ya_relay_stack/
device.rs
