fips_core/upper/
tun_write.rs

1use std::collections::VecDeque;
2use std::sync::{Arc, Condvar, Mutex, MutexGuard, mpsc};
3
4use crate::transport::PacketBuffer;
5
6#[derive(Debug)]
7struct TunWriteQueue {
8    state: Mutex<TunWriteState>,
9    ready: Condvar,
10}
11
12#[derive(Debug)]
13struct TunWriteState {
14    priority: VecDeque<TunWritePacket>,
15    senders: usize,
16    receiver_alive: bool,
17}
18
19#[derive(Debug)]
20pub(crate) enum TunWritePacket {
21    Vec(Vec<u8>),
22    Pooled(PacketBuffer),
23}
24
25impl TunWritePacket {
26    pub(crate) fn as_slice(&self) -> &[u8] {
27        match self {
28            Self::Vec(packet) => packet,
29            Self::Pooled(packet) => packet.as_slice(),
30        }
31    }
32
33    #[cfg(any(test, target_os = "linux", target_os = "macos", windows))]
34    pub(crate) fn as_mut_slice(&mut self) -> &mut [u8] {
35        match self {
36            Self::Vec(packet) => packet,
37            Self::Pooled(packet) => packet.as_mut_slice(),
38        }
39    }
40
41    #[cfg(any(test, target_os = "linux", target_os = "macos", windows))]
42    pub(crate) fn len(&self) -> usize {
43        self.as_slice().len()
44    }
45}
46
47impl From<Vec<u8>> for TunWritePacket {
48    fn from(packet: Vec<u8>) -> Self {
49        Self::Vec(packet)
50    }
51}
52
53impl From<PacketBuffer> for TunWritePacket {
54    fn from(packet: PacketBuffer) -> Self {
55        Self::Pooled(packet)
56    }
57}
58
59/// Channel sender for packets to be written to TUN.
60#[derive(Debug)]
61pub struct TunTx {
62    queue: Arc<TunWriteQueue>,
63}
64
65/// Channel receiver consumed by the blocking TUN writer.
66#[derive(Debug)]
67pub(crate) struct TunRx {
68    queue: Arc<TunWriteQueue>,
69}
70
71impl Clone for TunTx {
72    fn clone(&self) -> Self {
73        {
74            let mut state = self.queue.lock();
75            state.senders = state.senders.saturating_add(1);
76        }
77        Self {
78            queue: Arc::clone(&self.queue),
79        }
80    }
81}
82
83impl Drop for TunTx {
84    fn drop(&mut self) {
85        let should_notify = {
86            let mut state = self.queue.lock();
87            state.senders = state.senders.saturating_sub(1);
88            state.senders == 0
89        };
90        if should_notify {
91            self.queue.ready.notify_all();
92        }
93    }
94}
95
96impl TunTx {
97    /// Queue a priority packet for TUN delivery.
98    pub fn send(&self, packet: Vec<u8>) -> Result<(), mpsc::SendError<Vec<u8>>> {
99        let mut state = self.queue.lock();
100        if !state.receiver_alive {
101            return Err(mpsc::SendError(packet));
102        }
103
104        state.priority.push_back(TunWritePacket::Vec(packet));
105        drop(state);
106        self.queue.ready.notify_one();
107        Ok(())
108    }
109
110    /// Queue a batch of packets for TUN delivery. Returns the number of
111    /// packets dropped because the receiver is closed.
112    pub(crate) fn send_batch<I, P>(&self, packets: I) -> usize
113    where
114        I: IntoIterator<Item = P>,
115        P: Into<TunWritePacket>,
116    {
117        let packets = packets.into_iter();
118        let mut state = self.queue.lock();
119        if !state.receiver_alive {
120            return packets.count();
121        }
122
123        let mut sent = 0;
124        for packet in packets {
125            state.priority.push_back(packet.into());
126            sent += 1;
127        }
128        drop(state);
129        if sent > 0 {
130            self.queue.ready.notify_one();
131        }
132        0
133    }
134}
135
136impl TunRx {
137    pub(crate) fn recv(&self) -> Option<TunWritePacket> {
138        let mut state = self.queue.lock();
139        loop {
140            if let Some(packet) = state.priority.pop_front() {
141                return Some(packet);
142            }
143            if state.senders == 0 {
144                state.receiver_alive = false;
145                return None;
146            }
147            state = self
148                .queue
149                .ready
150                .wait(state)
151                .unwrap_or_else(|poisoned| poisoned.into_inner());
152        }
153    }
154
155    #[cfg(any(test, target_os = "linux"))]
156    pub(crate) fn try_recv_packet(&self) -> Result<TunWritePacket, mpsc::TryRecvError> {
157        let mut state = self.queue.lock();
158        if let Some(packet) = state.priority.pop_front() {
159            return Ok(packet);
160        }
161        if state.senders == 0 {
162            state.receiver_alive = false;
163            Err(mpsc::TryRecvError::Disconnected)
164        } else {
165            Err(mpsc::TryRecvError::Empty)
166        }
167    }
168}
169
170impl Drop for TunRx {
171    fn drop(&mut self) {
172        let mut state = self.queue.lock();
173        state.receiver_alive = false;
174        state.priority.clear();
175        drop(state);
176        self.queue.ready.notify_all();
177    }
178}
179
180impl TunWriteQueue {
181    fn lock(&self) -> MutexGuard<'_, TunWriteState> {
182        self.state
183            .lock()
184            .unwrap_or_else(|poisoned| poisoned.into_inner())
185    }
186}
187
188pub(crate) fn write_channel() -> (TunTx, TunRx) {
189    let queue = Arc::new(TunWriteQueue {
190        state: Mutex::new(TunWriteState {
191            priority: VecDeque::new(),
192            senders: 1,
193            receiver_alive: true,
194        }),
195        ready: Condvar::new(),
196    });
197    (
198        TunTx {
199            queue: Arc::clone(&queue),
200        },
201        TunRx { queue },
202    )
203}
fips_core/upper/tun_write.rs

fips_core/upper/
tun_write.rs
