use tokio::sync::mpsc;
#[derive(Debug, Clone)]
pub struct TunOutboundTx {
bulk: mpsc::Sender<QueuedTunOutboundPacket>,
}
#[derive(Debug)]
pub struct TunOutboundRx {
bulk: mpsc::Receiver<QueuedTunOutboundPacket>,
bulk_closed: bool,
}
#[derive(Debug)]
struct QueuedTunOutboundPacket {
packet: Vec<u8>,
}
impl QueuedTunOutboundPacket {
fn new(packet: Vec<u8>) -> Self {
Self { packet }
}
fn into_packet(self) -> Vec<u8> {
self.packet
}
}
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub(crate) enum TunOutboundAdmission {
Enqueued,
BulkDropped,
}
pub(crate) fn tun_outbound_channel(capacity: usize) -> (TunOutboundTx, TunOutboundRx) {
let capacity = capacity.max(1);
let (bulk_tx, bulk_rx) = mpsc::channel(capacity);
(
TunOutboundTx { bulk: bulk_tx },
TunOutboundRx {
bulk: bulk_rx,
bulk_closed: false,
},
)
}
impl TunOutboundTx {
pub async fn send(&self, packet: Vec<u8>) -> Result<(), mpsc::error::SendError<Vec<u8>>> {
self.send_queued(QueuedTunOutboundPacket::new(packet)).await
}
async fn send_queued(
&self,
queued: QueuedTunOutboundPacket,
) -> Result<(), mpsc::error::SendError<Vec<u8>>> {
self.bulk
.send(queued)
.await
.map_err(|error| mpsc::error::SendError(error.0.into_packet()))
}
pub fn blocking_send(&self, packet: Vec<u8>) -> Result<(), mpsc::error::SendError<Vec<u8>>> {
self.blocking_send_queued(QueuedTunOutboundPacket::new(packet))
}
fn blocking_send_queued(
&self,
queued: QueuedTunOutboundPacket,
) -> Result<(), mpsc::error::SendError<Vec<u8>>> {
self.bulk
.blocking_send(queued)
.map_err(|error| mpsc::error::SendError(error.0.into_packet()))
}
pub fn try_send(&self, packet: Vec<u8>) -> Result<(), mpsc::error::TrySendError<Vec<u8>>> {
self.try_send_queued(QueuedTunOutboundPacket::new(packet))
}
fn try_send_queued(
&self,
queued: QueuedTunOutboundPacket,
) -> Result<(), mpsc::error::TrySendError<Vec<u8>>> {
self.bulk
.try_send(queued)
.map_err(map_queued_try_send_error)
}
pub(crate) fn admit_from_tun_reader(
&self,
packet: Vec<u8>,
) -> Result<TunOutboundAdmission, mpsc::error::SendError<Vec<u8>>> {
let queued = QueuedTunOutboundPacket::new(packet);
match self.bulk.try_send(queued) {
Ok(()) => Ok(TunOutboundAdmission::Enqueued),
Err(mpsc::error::TrySendError::Full(queued)) => {
crate::perf_profile::record_event(
crate::perf_profile::Event::PendingTunPacketDropped,
);
tracing::debug!(
len = queued.packet.len(),
"Dropping TUN outbound packet because admission queue is full"
);
Ok(TunOutboundAdmission::BulkDropped)
}
Err(mpsc::error::TrySendError::Closed(queued)) => {
Err(mpsc::error::SendError(queued.into_packet()))
}
}
}
}
impl TunOutboundRx {
pub(crate) async fn recv(&mut self) -> Option<Vec<u8>> {
match self.bulk.recv().await {
Some(packet) => Some(packet.into_packet()),
None => {
self.bulk_closed = true;
None
}
}
}
pub(crate) fn try_recv(&mut self) -> Result<Vec<u8>, mpsc::error::TryRecvError> {
match self.bulk.try_recv() {
Ok(packet) => Ok(packet.into_packet()),
Err(mpsc::error::TryRecvError::Empty) => Err(mpsc::error::TryRecvError::Empty),
Err(mpsc::error::TryRecvError::Disconnected) => {
self.bulk_closed = true;
Err(mpsc::error::TryRecvError::Disconnected)
}
}
}
}
fn map_queued_try_send_error(
error: mpsc::error::TrySendError<QueuedTunOutboundPacket>,
) -> mpsc::error::TrySendError<Vec<u8>> {
match error {
mpsc::error::TrySendError::Full(packet) => {
mpsc::error::TrySendError::Full(packet.into_packet())
}
mpsc::error::TrySendError::Closed(packet) => {
mpsc::error::TrySendError::Closed(packet.into_packet())
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn ipv4_packet(proto: u8, body_len: usize) -> Vec<u8> {
let total_len = 20 + body_len;
let mut packet = vec![0u8; total_len];
packet[0] = 0x45;
packet[2..4].copy_from_slice(&(total_len as u16).to_be_bytes());
packet[9] = proto;
packet
}
fn ipv4_tcp_bulk_packet() -> Vec<u8> {
let mut packet = ipv4_packet(6, 20 + 300);
let tcp_offset = 20;
packet[tcp_offset + 12] = 5 << 4;
packet[tcp_offset + 13] = 0x10;
packet
}
fn ipv4_icmp_packet() -> Vec<u8> {
ipv4_packet(1, 8)
}
#[tokio::test]
async fn tun_outbound_recv_preserves_app_packet_order() {
let (tx, mut rx) = tun_outbound_channel(4);
let bulk = ipv4_tcp_bulk_packet();
let icmp = ipv4_icmp_packet();
tx.try_send(bulk.clone())
.expect("first app packet should enqueue");
tx.try_send(icmp.clone())
.expect("second app packet should enqueue");
assert_eq!(rx.recv().await, Some(bulk));
assert_eq!(rx.recv().await, Some(icmp));
}
#[test]
fn tun_outbound_capacity_applies_to_icmp_app_payload() {
let (tx, mut rx) = tun_outbound_channel(1);
let first_bulk = ipv4_tcp_bulk_packet();
let icmp = ipv4_icmp_packet();
tx.try_send(first_bulk.clone())
.expect("first app packet should fit");
assert!(
tx.try_send(icmp).is_err(),
"ICMP app payload should share bounded app packet capacity"
);
assert_eq!(rx.try_recv(), Ok(first_bulk));
}
#[test]
fn tun_reader_admission_sheds_icmp_app_payload_when_full() {
let (tx, mut rx) = tun_outbound_channel(1);
let first_bulk = ipv4_tcp_bulk_packet();
let icmp = ipv4_icmp_packet();
assert!(matches!(
tx.admit_from_tun_reader(first_bulk.clone()),
Ok(TunOutboundAdmission::Enqueued)
));
assert!(matches!(
tx.admit_from_tun_reader(icmp),
Ok(TunOutboundAdmission::BulkDropped)
));
assert_eq!(rx.try_recv(), Ok(first_bulk));
}
}