use std::io;
use std::net::SocketAddr;
use std::os::fd::RawFd;
use std::time::Instant;
use crate::{sockaddr_to_std, std_to_sockaddr, MAX_DATAGRAM_BATCH, MAX_WIRE_DATAGRAM};
const GSO_CONTROL_PAYLOAD_LEN: u32 = 2;
const GSO_CONTROL_SPACE: usize = 64;
pub trait Dispatch {
fn on_udp(&mut self, src: SocketAddr, dg: &[u8], now_ms: u64) -> DispatchOut<'_>;
fn on_tun(&mut self, inner: &[u8], now_ms: u64) -> &[EgressDatagram];
fn tick(&mut self, now_ms: u64) -> Option<&[EgressDatagram]>;
fn next_timeout_ms(&self, _now_ms: u64) -> Option<u64> {
None
}
}
#[derive(Debug, Clone)]
pub struct EgressDatagram {
pub fate: u16,
pub dst: SocketAddr,
pub bytes: Vec<u8>,
}
impl AsRef<[u8]> for EgressDatagram {
fn as_ref(&self) -> &[u8] {
&self.bytes
}
}
pub enum DispatchOut<'a> {
None,
Tun(&'a [u8]),
Udp(&'a [EgressDatagram]),
Both(&'a [u8], &'a [EgressDatagram]),
}
fn set_nonblocking(fd: RawFd) -> io::Result<()> {
let rc = unsafe { libc::fcntl(fd, libc::F_SETFL, libc::O_NONBLOCK) };
if rc < 0 {
Err(io::Error::last_os_error())
} else {
Ok(())
}
}
#[expect(
clippy::too_many_arguments,
reason = "batch buffers (the GSO scratch `gso`, and the TUN-offload `vnet_hdr`/`coalescer` \
state) are owned by run_poll and threaded through explicitly rather than bundled \
into a struct, to keep the hot-path buffers' aliasing/lifetimes obvious"
)]
fn drain_udp(
udp_fd: RawFd,
tun_fd: RawFd,
d: &mut impl Dispatch,
now_ms: u64,
bufs: &mut [[u8; MAX_WIRE_DATAGRAM]; MAX_DATAGRAM_BATCH],
lens: &mut [usize; MAX_DATAGRAM_BATCH],
srcs: &mut [SocketAddr; MAX_DATAGRAM_BATCH],
tx: &mut Vec<EgressDatagram>,
gso: &mut GsoScratch,
vnet_hdr: bool,
coalescer: &mut crate::tun_offload::Coalescer,
) -> io::Result<()> {
loop {
let n = recv_mmsg(udp_fd, bufs, lens, srcs)?;
if n == 0 {
break; }
for i in 0..n {
let dg = &bufs[i][..lens[i]];
match d.on_udp(srcs[i], dg, now_ms) {
DispatchOut::None => {}
DispatchOut::Tun(inner) => write_to_tun(tun_fd, inner, vnet_hdr, coalescer),
DispatchOut::Udp(pkts) => tx.extend(pkts.iter().cloned()),
DispatchOut::Both(inner, pkts) => {
write_to_tun(tun_fd, inner, vnet_hdr, coalescer);
tx.extend(pkts.iter().cloned());
}
}
}
flush_tx(udp_fd, tx, gso)?; if n < MAX_DATAGRAM_BATCH {
break; }
}
if vnet_hdr {
if let Some(frame) = coalescer.flush() {
send_to_tun(tun_fd, frame);
}
}
Ok(())
}
#[inline]
fn write_to_tun(
tun_fd: RawFd,
inner: &[u8],
vnet_hdr: bool,
coalescer: &mut crate::tun_offload::Coalescer,
) {
if vnet_hdr {
if let Some(frame) = coalescer.push(inner) {
send_to_tun(tun_fd, frame);
}
} else {
send_to_tun(tun_fd, inner);
}
}
#[expect(
clippy::too_many_arguments,
reason = "batch buffers (the GSO scratch `gso`, the reusable read buffer `buf`, and the \
TUN-offload `vnet_hdr`/split scratch) are owned by run_poll and threaded through \
explicitly rather than bundled into a struct, to keep the hot-path buffers' \
aliasing/lifetimes obvious"
)]
fn drain_tun(
tun_fd: RawFd,
udp_fd: RawFd,
d: &mut impl Dispatch,
now_ms: u64,
tx: &mut Vec<EgressDatagram>,
gso: &mut GsoScratch,
vnet_hdr: bool,
buf: &mut [u8],
split_out: &mut Vec<u8>,
split_offs: &mut Vec<(usize, usize)>,
) -> io::Result<()> {
loop {
let n = unsafe { libc::read(tun_fd, buf.as_mut_ptr().cast(), buf.len()) };
if n < 0 {
let e = io::Error::last_os_error();
let raw = e.raw_os_error().unwrap_or(0);
if raw == libc::EWOULDBLOCK || raw == libc::EAGAIN {
break;
}
if raw == libc::EINTR {
continue;
}
return Err(e);
}
if n == 0 {
break;
}
let n = usize::try_from(n).expect("non-negative read return fits usize");
if vnet_hdr {
if crate::tun_offload::split_gro(&buf[..n], split_out, split_offs) {
for &(s, l) in split_offs.iter() {
tx.extend(d.on_tun(&split_out[s..s + l], now_ms).iter().cloned());
if tx.len() >= MAX_DATAGRAM_BATCH {
flush_tx(udp_fd, tx, gso)?;
}
}
} else if let Some(body) = buf.get(crate::tun_offload::VNET_HDR_LEN..n) {
tx.extend(d.on_tun(body, now_ms).iter().cloned());
if tx.len() >= MAX_DATAGRAM_BATCH {
flush_tx(udp_fd, tx, gso)?;
}
}
} else {
let inner = &buf[..n];
tx.extend(d.on_tun(inner, now_ms).iter().cloned());
if tx.len() >= MAX_DATAGRAM_BATCH {
flush_tx(udp_fd, tx, gso)?;
}
}
}
flush_tx(udp_fd, tx, gso)?; Ok(())
}
struct GsoScratch {
enabled: bool,
runs: Vec<crate::gso::GsoRun>,
payload: Vec<u8>,
}
impl GsoScratch {
fn new() -> Self {
Self {
enabled: true,
runs: Vec::with_capacity(MAX_DATAGRAM_BATCH),
payload: Vec::with_capacity(MAX_WIRE_DATAGRAM * crate::gso::MAX_GSO_SEGMENTS_PER_SEND),
}
}
}
fn send_run_plain(udp_fd: RawFd, run: &[EgressDatagram]) -> io::Result<()> {
let mut sent = 0;
while sent < run.len() {
let n = send_mmsg(udp_fd, &run[sent..])?;
if n == 0 {
break;
}
sent += n;
}
Ok(())
}
fn flush_tx(udp_fd: RawFd, tx: &mut Vec<EgressDatagram>, gso: &mut GsoScratch) -> io::Result<()> {
if !gso.enabled {
let r = send_run_plain(udp_fd, tx);
tx.clear();
return r;
}
crate::gso::partition_fate_safe(tx, MAX_DATAGRAM_BATCH, &mut gso.runs);
let runs = std::mem::take(&mut gso.runs);
let mut outcome = Ok(());
for run in &runs {
if gso.enabled && run.members.len() >= 2 && run.segment_size > 0 {
let dst = tx[run.members[0]].dst;
match send_gso_indexed(
udp_fd,
tx,
&run.members,
run.segment_size,
dst,
&mut gso.payload,
) {
Ok(true) => {}
Ok(false) => {
gso.enabled = false; for &i in &run.members {
if let Err(e) = send_run_plain(udp_fd, std::slice::from_ref(&tx[i])) {
outcome = Err(e);
break;
}
}
}
Err(e) => outcome = Err(e),
}
} else {
for &i in &run.members {
if let Err(e) = send_run_plain(udp_fd, std::slice::from_ref(&tx[i])) {
outcome = Err(e);
break;
}
}
}
if outcome.is_err() {
break;
}
}
gso.runs = runs; tx.clear();
outcome
}
#[inline]
fn send_to_tun(tun_fd: RawFd, buf: &[u8]) {
let rc = unsafe { libc::write(tun_fd, buf.as_ptr().cast(), buf.len()) };
if rc < 0 {
eprintln!("poll: tun write error: {}", io::Error::last_os_error());
}
}
fn send_mmsg(udp_fd: RawFd, datagrams: &[EgressDatagram]) -> io::Result<usize> {
if datagrams.is_empty() {
return Ok(0);
}
let count = datagrams.len().min(MAX_DATAGRAM_BATCH);
let mut storages: [libc::sockaddr_storage; MAX_DATAGRAM_BATCH] = unsafe { std::mem::zeroed() };
let mut addrlens = [0 as libc::socklen_t; MAX_DATAGRAM_BATCH];
let mut iovecs = [libc::iovec {
iov_base: std::ptr::null_mut(),
iov_len: 0,
}; MAX_DATAGRAM_BATCH];
let mut msgs = [crate::empty_mmsghdr(); MAX_DATAGRAM_BATCH];
for (i, dg) in datagrams[..count].iter().enumerate() {
let (storage, addr_len) = std_to_sockaddr(dg.dst);
storages[i] = storage;
addrlens[i] = addr_len;
iovecs[i].iov_base = dg.bytes.as_ptr().cast_mut().cast::<libc::c_void>();
iovecs[i].iov_len = dg.bytes.len();
msgs[i].msg_hdr.msg_iov = &raw mut iovecs[i];
msgs[i].msg_hdr.msg_iovlen = 1;
msgs[i].msg_hdr.msg_name = std::ptr::from_mut(&mut storages[i]).cast::<libc::c_void>();
msgs[i].msg_hdr.msg_namelen = addrlens[i];
}
let ret = unsafe {
libc::sendmmsg(
udp_fd,
msgs.as_mut_ptr(),
u32::try_from(count).expect("count ≤ 64 fits u32"),
libc::MSG_NOSIGNAL as crate::MsgFlags,
)
};
if ret < 0 {
let e = io::Error::last_os_error();
let raw = e.raw_os_error().unwrap_or(0);
if raw == libc::EWOULDBLOCK || raw == libc::EAGAIN || raw == libc::ENOBUFS {
return Ok(0);
}
return Err(e);
}
Ok(usize::try_from(ret).expect("non-negative sendmmsg return fits usize"))
}
fn send_gso_payload(
udp_fd: RawFd,
payload: &[u8],
segment_size: u16,
dst: SocketAddr,
) -> io::Result<bool> {
let (mut storage, addr_len) = std_to_sockaddr(dst);
let mut iov = libc::iovec {
iov_base: payload.as_ptr().cast_mut().cast::<libc::c_void>(),
iov_len: payload.len(),
};
let mut control = [0u8; GSO_CONTROL_SPACE];
let mut msg: libc::msghdr = unsafe { std::mem::zeroed() };
msg.msg_name = std::ptr::from_mut(&mut storage).cast::<libc::c_void>();
msg.msg_namelen = addr_len;
msg.msg_iov = &raw mut iov;
msg.msg_iovlen = 1;
msg.msg_control = control.as_mut_ptr().cast::<libc::c_void>();
let cmsg_space = usize::try_from(unsafe { libc::CMSG_SPACE(GSO_CONTROL_PAYLOAD_LEN) })
.expect("cmsg space fits usize");
if cmsg_space > control.len() {
return Err(io::Error::other("GSO control buffer too small"));
}
msg.msg_controllen = cmsg_space as crate::CmsgLen;
unsafe {
let cmsg = libc::CMSG_FIRSTHDR(&raw const msg);
if cmsg.is_null() {
return Err(io::Error::other("missing first cmsg header"));
}
(*cmsg).cmsg_level = libc::SOL_UDP;
(*cmsg).cmsg_type = libc::UDP_SEGMENT;
(*cmsg).cmsg_len = usize::try_from(libc::CMSG_LEN(GSO_CONTROL_PAYLOAD_LEN))
.expect("cmsg len fits usize") as crate::CmsgLen;
let seg_ptr = libc::CMSG_DATA(cmsg).cast::<u16>();
*seg_ptr = segment_size;
}
let ret = unsafe { libc::sendmsg(udp_fd, &raw const msg, libc::MSG_NOSIGNAL) };
if ret < 0 {
let e = io::Error::last_os_error();
let raw = e.raw_os_error().unwrap_or(0);
if raw == libc::EWOULDBLOCK || raw == libc::EAGAIN || raw == libc::ENOBUFS {
return Ok(true); }
if raw == libc::EIO || raw == libc::EINVAL {
return Ok(false); }
return Err(e);
}
Ok(true)
}
#[cfg_attr(
not(test),
expect(
dead_code,
reason = "run-slice wrapper kept for its own unit test; flush_tx uses \
send_gso_indexed to avoid copying the run out of the egress batch"
)
)]
fn send_gso(
udp_fd: RawFd,
run: &[EgressDatagram],
segment_size: u16,
dst: SocketAddr,
payload: &mut Vec<u8>,
) -> io::Result<bool> {
payload.clear();
for dg in run {
payload.extend_from_slice(&dg.bytes);
}
send_gso_payload(udp_fd, payload, segment_size, dst)
}
fn send_gso_indexed(
udp_fd: RawFd,
tx: &[EgressDatagram],
indices: &[usize],
segment_size: u16,
dst: SocketAddr,
payload: &mut Vec<u8>,
) -> io::Result<bool> {
payload.clear();
for &i in indices {
payload.extend_from_slice(&tx[i].bytes);
}
send_gso_payload(udp_fd, payload, segment_size, dst)
}
fn recv_mmsg(
udp_fd: RawFd,
bufs: &mut [[u8; MAX_WIRE_DATAGRAM]],
lens: &mut [usize],
srcs: &mut [SocketAddr],
) -> io::Result<usize> {
let count = bufs
.len()
.min(lens.len())
.min(srcs.len())
.min(MAX_DATAGRAM_BATCH);
if count == 0 {
return Ok(0);
}
let mut storages: [libc::sockaddr_storage; MAX_DATAGRAM_BATCH] = unsafe { std::mem::zeroed() };
let addrlens = [libc::socklen_t::try_from(std::mem::size_of::<libc::sockaddr_storage>())
.expect("size fits socklen_t"); MAX_DATAGRAM_BATCH];
let mut iovecs = [libc::iovec {
iov_base: std::ptr::null_mut(),
iov_len: 0,
}; MAX_DATAGRAM_BATCH];
let mut msgs = [crate::empty_mmsghdr(); MAX_DATAGRAM_BATCH];
for i in 0..count {
iovecs[i].iov_base = bufs[i].as_mut_ptr().cast::<libc::c_void>();
iovecs[i].iov_len = MAX_WIRE_DATAGRAM;
msgs[i].msg_hdr.msg_iov = &raw mut iovecs[i];
msgs[i].msg_hdr.msg_iovlen = 1;
msgs[i].msg_hdr.msg_name = std::ptr::from_mut(&mut storages[i]).cast::<libc::c_void>();
msgs[i].msg_hdr.msg_namelen = addrlens[i];
}
let ret = unsafe {
libc::recvmmsg(
udp_fd,
msgs.as_mut_ptr(),
u32::try_from(count).expect("count ≤ 64 fits u32"),
libc::MSG_DONTWAIT as crate::MsgFlags,
std::ptr::null_mut(),
)
};
if ret < 0 {
let e = io::Error::last_os_error();
let raw = e.raw_os_error().unwrap_or(0);
if raw == libc::EWOULDBLOCK || raw == libc::EAGAIN {
return Ok(0);
}
return Err(e);
}
let received = usize::try_from(ret).expect("non-negative recvmmsg return fits usize");
for i in 0..received {
lens[i] = usize::try_from(msgs[i].msg_len).expect("msg_len fits usize");
srcs[i] = sockaddr_to_std(&storages[i], msgs[i].msg_hdr.msg_namelen)
.unwrap_or_else(|_| SocketAddr::from(([0, 0, 0, 0], 0)));
}
Ok(received)
}
pub fn run_poll<D: Dispatch>(
udp_fd: RawFd,
tun_fd: RawFd,
vnet_hdr: bool,
d: &mut D,
) -> io::Result<()> {
set_nonblocking(udp_fd)?;
set_nonblocking(tun_fd)?;
let epoll_fd = unsafe { libc::epoll_create1(libc::EPOLL_CLOEXEC) };
if epoll_fd < 0 {
return Err(io::Error::last_os_error());
}
let mut ev_udp = libc::epoll_event {
events: libc::EPOLLIN as u32,
u64: udp_fd as u64,
};
let rc = unsafe { libc::epoll_ctl(epoll_fd, libc::EPOLL_CTL_ADD, udp_fd, &raw mut ev_udp) };
if rc < 0 {
unsafe { libc::close(epoll_fd) };
return Err(io::Error::last_os_error());
}
let mut ev_tun = libc::epoll_event {
events: libc::EPOLLIN as u32,
u64: tun_fd as u64,
};
let rc = unsafe { libc::epoll_ctl(epoll_fd, libc::EPOLL_CTL_ADD, tun_fd, &raw mut ev_tun) };
if rc < 0 {
unsafe { libc::close(epoll_fd) };
return Err(io::Error::last_os_error());
}
let start = Instant::now();
let mut events = [libc::epoll_event { events: 0, u64: 0 }; 4];
let mut rx_bufs = Box::new([[0u8; MAX_WIRE_DATAGRAM]; MAX_DATAGRAM_BATCH]);
let mut rx_lens = [0usize; MAX_DATAGRAM_BATCH];
let mut rx_srcs = [SocketAddr::from(([0, 0, 0, 0], 0)); MAX_DATAGRAM_BATCH];
let mut tx_batch: Vec<EgressDatagram> = Vec::with_capacity(MAX_DATAGRAM_BATCH);
let mut gso = GsoScratch::new();
let mut coalescer = crate::tun_offload::Coalescer::new();
let mut split_out: Vec<u8> = Vec::new();
let mut split_offs: Vec<(usize, usize)> = Vec::new();
let mut tun_buf = vec![
0u8;
if vnet_hdr {
crate::tun_offload::VNET_HDR_LEN + crate::tun_offload::MAX_GSO_PAYLOAD
} else {
MAX_WIRE_DATAGRAM
}
];
loop {
let now_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
let timeout_ms = match d.next_timeout_ms(now_ms) {
Some(delay_ms) => {
let clamped = i32::try_from(delay_ms).unwrap_or(10);
10.min(clamped)
}
None => 10,
};
let nfds = unsafe {
libc::epoll_wait(
epoll_fd,
events.as_mut_ptr(),
libc::c_int::try_from(events.len()).expect("event array len fits c_int"),
timeout_ms,
)
};
if nfds < 0 {
let e = io::Error::last_os_error();
if e.raw_os_error() == Some(libc::EINTR) {
continue;
}
unsafe { libc::close(epoll_fd) };
return Err(e);
}
let now_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
for ev in &events[..usize::try_from(nfds).expect("non-negative nfds fits usize")] {
let ready_fd = ev.u64 as RawFd;
if ready_fd == udp_fd {
if let Err(e) = drain_udp(
udp_fd,
tun_fd,
d,
now_ms,
&mut rx_bufs,
&mut rx_lens,
&mut rx_srcs,
&mut tx_batch,
&mut gso,
vnet_hdr,
&mut coalescer,
) {
unsafe { libc::close(epoll_fd) };
return Err(e);
}
} else if ready_fd == tun_fd {
if let Err(e) = drain_tun(
tun_fd,
udp_fd,
d,
now_ms,
&mut tx_batch,
&mut gso,
vnet_hdr,
&mut tun_buf,
&mut split_out,
&mut split_offs,
) {
unsafe { libc::close(epoll_fd) };
return Err(e);
}
}
}
if let Some(pkts) = d.tick(now_ms) {
tx_batch.extend(pkts.iter().cloned());
if let Err(e) = flush_tx(udp_fd, &mut tx_batch, &mut gso) {
unsafe { libc::close(epoll_fd) };
return Err(e);
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::net::UdpSocket;
use std::os::fd::AsRawFd;
struct CountDispatch {
received: Vec<Vec<u8>>,
call_count: usize,
}
impl CountDispatch {
fn new() -> Self {
Self {
received: Vec::new(),
call_count: 0,
}
}
}
impl Dispatch for CountDispatch {
fn on_udp(&mut self, _src: SocketAddr, dg: &[u8], _now_ms: u64) -> DispatchOut<'_> {
self.received.push(dg.to_vec());
self.call_count += 1;
DispatchOut::None
}
fn on_tun(&mut self, _inner: &[u8], _now_ms: u64) -> &[EgressDatagram] {
&[]
}
fn tick(&mut self, _now_ms: u64) -> Option<&[EgressDatagram]> {
None
}
}
#[test]
fn drain_udp_delivers_datagram() {
let a = UdpSocket::bind("127.0.0.1:0").unwrap();
let b = UdpSocket::bind("127.0.0.1:0").unwrap();
a.connect(b.local_addr().unwrap()).unwrap();
b.connect(a.local_addr().unwrap()).unwrap();
set_nonblocking(b.as_raw_fd()).unwrap();
a.send(b"hello from drain_udp test").unwrap();
let devnull = std::fs::OpenOptions::new()
.write(true)
.open("/dev/null")
.unwrap();
let null_fd = devnull.as_raw_fd();
let mut d = CountDispatch::new();
let mut bufs = Box::new([[0u8; MAX_WIRE_DATAGRAM]; MAX_DATAGRAM_BATCH]);
let mut lens = [0usize; MAX_DATAGRAM_BATCH];
let mut srcs = [std::net::SocketAddr::from(([0, 0, 0, 0], 0)); MAX_DATAGRAM_BATCH];
let mut tx = Vec::new();
let mut gso = GsoScratch::new();
let mut coalescer = crate::tun_offload::Coalescer::new();
drain_udp(
b.as_raw_fd(),
null_fd,
&mut d,
0,
&mut bufs,
&mut lens,
&mut srcs,
&mut tx,
&mut gso,
false,
&mut coalescer,
)
.unwrap();
assert_eq!(d.call_count, 1, "on_udp must be called exactly once");
assert_eq!(d.received[0], b"hello from drain_udp test");
}
#[test]
fn drain_udp_drains_multiple_datagrams() {
let a = UdpSocket::bind("127.0.0.1:0").unwrap();
let b = UdpSocket::bind("127.0.0.1:0").unwrap();
a.connect(b.local_addr().unwrap()).unwrap();
b.connect(a.local_addr().unwrap()).unwrap();
set_nonblocking(b.as_raw_fd()).unwrap();
a.send(b"first").unwrap();
a.send(b"second").unwrap();
a.send(b"third").unwrap();
let devnull = std::fs::OpenOptions::new()
.write(true)
.open("/dev/null")
.unwrap();
let null_fd = devnull.as_raw_fd();
let mut d = CountDispatch::new();
let mut bufs = Box::new([[0u8; MAX_WIRE_DATAGRAM]; MAX_DATAGRAM_BATCH]);
let mut lens = [0usize; MAX_DATAGRAM_BATCH];
let mut srcs = [std::net::SocketAddr::from(([0, 0, 0, 0], 0)); MAX_DATAGRAM_BATCH];
let mut tx = Vec::new();
let mut gso = GsoScratch::new();
let mut coalescer = crate::tun_offload::Coalescer::new();
drain_udp(
b.as_raw_fd(),
null_fd,
&mut d,
0,
&mut bufs,
&mut lens,
&mut srcs,
&mut tx,
&mut gso,
false,
&mut coalescer,
)
.unwrap();
assert_eq!(d.call_count, 3);
}
#[test]
fn set_nonblocking_is_idempotent() {
let sock = UdpSocket::bind("127.0.0.1:0").unwrap();
set_nonblocking(sock.as_raw_fd()).unwrap();
set_nonblocking(sock.as_raw_fd()).unwrap(); }
struct ForwardDispatch {
reply: Vec<u8>,
scratch: Vec<EgressDatagram>,
}
impl ForwardDispatch {
fn new(reply: impl Into<Vec<u8>>) -> Self {
Self {
reply: reply.into(),
scratch: Vec::new(),
}
}
}
impl Dispatch for ForwardDispatch {
fn on_udp(&mut self, src: SocketAddr, _dg: &[u8], _now_ms: u64) -> DispatchOut<'_> {
self.scratch = vec![EgressDatagram {
fate: 0,
dst: src,
bytes: self.reply.clone(),
}];
DispatchOut::Udp(&self.scratch)
}
fn on_tun(&mut self, _inner: &[u8], _now_ms: u64) -> &[EgressDatagram] {
&[]
}
fn tick(&mut self, _now_ms: u64) -> Option<&[EgressDatagram]> {
None
}
}
#[test]
fn drain_udp_forwards_dispatch_out_udp_to_peer() {
let a = UdpSocket::bind("127.0.0.1:0").unwrap();
let b = UdpSocket::bind("127.0.0.1:0").unwrap();
a.connect(b.local_addr().unwrap()).unwrap();
b.connect(a.local_addr().unwrap()).unwrap();
set_nonblocking(b.as_raw_fd()).unwrap();
set_nonblocking(a.as_raw_fd()).unwrap();
a.send(b"trigger").unwrap();
let devnull = std::fs::OpenOptions::new()
.write(true)
.open("/dev/null")
.unwrap();
let null_fd = devnull.as_raw_fd();
let mut d = ForwardDispatch::new(b"forwarded".as_slice());
let mut bufs = Box::new([[0u8; MAX_WIRE_DATAGRAM]; MAX_DATAGRAM_BATCH]);
let mut lens = [0usize; MAX_DATAGRAM_BATCH];
let mut srcs = [std::net::SocketAddr::from(([0, 0, 0, 0], 0)); MAX_DATAGRAM_BATCH];
let mut tx = Vec::new();
let mut gso = GsoScratch::new();
let mut coalescer = crate::tun_offload::Coalescer::new();
drain_udp(
b.as_raw_fd(),
null_fd,
&mut d,
0,
&mut bufs,
&mut lens,
&mut srcs,
&mut tx,
&mut gso,
false,
&mut coalescer,
)
.unwrap();
let mut recv_buf = [0u8; 64];
let n = a
.recv(&mut recv_buf)
.expect("forwarded datagram must arrive on peer socket");
assert_eq!(&recv_buf[..n], b"forwarded");
}
#[test]
fn send_mmsg_delivers_each_datagram_to_its_own_dst() {
use std::net::UdpSocket;
let rx_a = UdpSocket::bind("127.0.0.1:0").unwrap();
let rx_b = UdpSocket::bind("127.0.0.1:0").unwrap();
rx_a.set_read_timeout(Some(std::time::Duration::from_secs(2)))
.unwrap();
rx_b.set_read_timeout(Some(std::time::Duration::from_secs(2)))
.unwrap();
let tx = UdpSocket::bind("127.0.0.1:0").unwrap();
let dst_a = rx_a.local_addr().unwrap();
let dst_b = rx_b.local_addr().unwrap();
let batch = [
EgressDatagram {
fate: 0,
dst: dst_a,
bytes: b"to-a-1".to_vec(),
},
EgressDatagram {
fate: 0,
dst: dst_b,
bytes: b"to-b".to_vec(),
},
EgressDatagram {
fate: 0,
dst: dst_a,
bytes: b"to-a-2".to_vec(),
},
];
let mut sent = 0;
while sent < batch.len() {
sent += send_mmsg(tx.as_raw_fd(), &batch[sent..]).unwrap();
}
assert_eq!(sent, 3);
let mut buf = [0u8; 64];
let (n1, _) = rx_a.recv_from(&mut buf).unwrap();
assert_eq!(&buf[..n1], b"to-a-1");
let (n2, _) = rx_a.recv_from(&mut buf).unwrap();
assert_eq!(&buf[..n2], b"to-a-2");
let (n3, _) = rx_b.recv_from(&mut buf).unwrap();
assert_eq!(&buf[..n3], b"to-b");
}
#[test]
fn recv_mmsg_returns_bytes_and_source_per_datagram() {
use std::net::UdpSocket;
let rx = UdpSocket::bind("127.0.0.1:0").unwrap();
let tx = UdpSocket::bind("127.0.0.1:0").unwrap();
let rx_addr = rx.local_addr().unwrap();
let tx_addr = tx.local_addr().unwrap();
tx.send_to(b"hello", rx_addr).unwrap();
tx.send_to(b"world!!", rx_addr).unwrap();
std::thread::sleep(std::time::Duration::from_millis(50));
let mut bufs = [[0u8; MAX_WIRE_DATAGRAM]; 4];
let mut lens = [0usize; 4];
let mut srcs = [std::net::SocketAddr::from(([0, 0, 0, 0], 0)); 4];
let n = recv_mmsg(rx.as_raw_fd(), &mut bufs, &mut lens, &mut srcs).unwrap();
assert_eq!(n, 2);
assert_eq!(&bufs[0][..lens[0]], b"hello");
assert_eq!(&bufs[1][..lens[1]], b"world!!");
assert_eq!(srcs[0], tx_addr);
assert_eq!(srcs[1], tx_addr);
}
#[test]
fn recv_mmsg_returns_zero_when_nothing_queued() {
use std::net::UdpSocket;
let rx = UdpSocket::bind("127.0.0.1:0").unwrap();
let mut bufs = [[0u8; MAX_WIRE_DATAGRAM]; 4];
let mut lens = [0usize; 4];
let mut srcs = [std::net::SocketAddr::from(([0, 0, 0, 0], 0)); 4];
assert_eq!(
recv_mmsg(rx.as_raw_fd(), &mut bufs, &mut lens, &mut srcs).unwrap(),
0
);
}
#[test]
fn send_gso_delivers_each_segment_when_supported() {
use std::net::UdpSocket;
let rx = UdpSocket::bind("127.0.0.1:0").unwrap();
let rx_addr = rx.local_addr().unwrap();
let tx = UdpSocket::bind("127.0.0.1:0").unwrap();
rx.set_nonblocking(true).unwrap();
let run = vec![
EgressDatagram {
fate: 1,
dst: rx_addr,
bytes: vec![0xAA; 1000],
},
EgressDatagram {
fate: 2,
dst: rx_addr,
bytes: vec![0xBB; 1000],
},
EgressDatagram {
fate: 3,
dst: rx_addr,
bytes: vec![0xCC; 1000],
},
];
let mut payload = Vec::new();
let accepted =
send_gso(tx.as_raw_fd(), &run, 1000, rx_addr, &mut payload).expect("send_gso");
if !accepted {
return; }
std::thread::sleep(std::time::Duration::from_millis(50));
let mut got = Vec::new();
let mut buf = [0u8; 2048];
while let Ok((n, _)) = rx.recv_from(&mut buf) {
got.push(buf[..n].to_vec());
}
assert_eq!(got.len(), 3, "GSO must segment into 3 separate datagrams");
assert!(got.iter().all(|d| d.len() == 1000));
}
#[test]
fn flush_tx_delivers_all_datagrams_gso_or_fallback() {
use std::net::UdpSocket;
let rx = UdpSocket::bind("127.0.0.1:0").unwrap();
let rx_addr = rx.local_addr().unwrap();
let tx_sock = UdpSocket::bind("127.0.0.1:0").unwrap();
rx.set_nonblocking(true).unwrap();
let mut tx: Vec<EgressDatagram> = [1u16, 2, 3, 1]
.iter()
.map(|&f| EgressDatagram {
fate: f,
dst: rx_addr,
bytes: vec![f as u8; 1000],
})
.collect();
let mut gso = GsoScratch::new();
flush_tx(tx_sock.as_raw_fd(), &mut tx, &mut gso).expect("flush_tx");
assert!(tx.is_empty(), "flush_tx must drain tx");
std::thread::sleep(std::time::Duration::from_millis(50));
let mut count = 0;
let mut buf = [0u8; 2048];
while let Ok((n, _)) = rx.recv_from(&mut buf) {
assert_eq!(n, 1000);
count += 1;
}
assert_eq!(
count, 4,
"all four datagrams must arrive (GSO or plain fallback)"
);
}
}