1use std::io;
11use std::net::SocketAddr;
12use std::os::fd::RawFd;
13use std::time::Instant;
14
15use crate::{sockaddr_to_std, std_to_sockaddr, MAX_DATAGRAM_BATCH, MAX_WIRE_DATAGRAM};
16
17const GSO_CONTROL_PAYLOAD_LEN: u32 = 2;
19const GSO_CONTROL_SPACE: usize = 64;
21
22pub trait Dispatch {
36 fn on_udp(&mut self, src: SocketAddr, dg: &[u8], now_ms: u64) -> DispatchOut<'_>;
40
41 fn on_tun(&mut self, inner: &[u8], now_ms: u64) -> &[EgressDatagram];
46
47 fn tick(&mut self, now_ms: u64) -> Option<&[EgressDatagram]>;
50
51 fn next_timeout_ms(&self, _now_ms: u64) -> Option<u64> {
55 None
56 }
57}
58
59#[derive(Debug, Clone)]
73pub struct EgressDatagram {
74 pub fate: u16,
75 pub dst: SocketAddr,
76 pub bytes: Vec<u8>,
77}
78
79impl AsRef<[u8]> for EgressDatagram {
80 fn as_ref(&self) -> &[u8] {
81 &self.bytes
82 }
83}
84
85pub enum DispatchOut<'a> {
87 None,
89 Tun(&'a [u8]),
91 Udp(&'a [EgressDatagram]),
93 Both(&'a [u8], &'a [EgressDatagram]),
95}
96
97fn set_nonblocking(fd: RawFd) -> io::Result<()> {
101 let rc = unsafe { libc::fcntl(fd, libc::F_SETFL, libc::O_NONBLOCK) };
105 if rc < 0 {
106 Err(io::Error::last_os_error())
107 } else {
108 Ok(())
109 }
110}
111
112#[expect(
122 clippy::too_many_arguments,
123 reason = "batch buffers (the GSO scratch `gso`, and the TUN-offload `vnet_hdr`/`coalescer` \
124 state) are owned by run_poll and threaded through explicitly rather than bundled \
125 into a struct, to keep the hot-path buffers' aliasing/lifetimes obvious"
126)]
127fn drain_udp(
128 udp_fd: RawFd,
129 tun_fd: RawFd,
130 d: &mut impl Dispatch,
131 now_ms: u64,
132 bufs: &mut [[u8; MAX_WIRE_DATAGRAM]; MAX_DATAGRAM_BATCH],
133 lens: &mut [usize; MAX_DATAGRAM_BATCH],
134 srcs: &mut [SocketAddr; MAX_DATAGRAM_BATCH],
135 tx: &mut Vec<EgressDatagram>,
136 gso: &mut GsoScratch,
137 vnet_hdr: bool,
138 coalescer: &mut crate::tun_offload::Coalescer,
139) -> io::Result<()> {
140 loop {
141 let n = recv_mmsg(udp_fd, bufs, lens, srcs)?;
142 if n == 0 {
143 break; }
145 for i in 0..n {
146 let dg = &bufs[i][..lens[i]];
147 match d.on_udp(srcs[i], dg, now_ms) {
148 DispatchOut::None => {}
149 DispatchOut::Tun(inner) => write_to_tun(tun_fd, inner, vnet_hdr, coalescer),
150 DispatchOut::Udp(pkts) => tx.extend(pkts.iter().cloned()),
151 DispatchOut::Both(inner, pkts) => {
152 write_to_tun(tun_fd, inner, vnet_hdr, coalescer);
153 tx.extend(pkts.iter().cloned());
154 }
155 }
156 }
157 flush_tx(udp_fd, tx, gso)?; if n < MAX_DATAGRAM_BATCH {
159 break; }
161 }
162 if vnet_hdr {
163 if let Some(frame) = coalescer.flush() {
164 send_to_tun(tun_fd, frame);
165 }
166 }
167 Ok(())
168}
169
170#[inline]
175fn write_to_tun(
176 tun_fd: RawFd,
177 inner: &[u8],
178 vnet_hdr: bool,
179 coalescer: &mut crate::tun_offload::Coalescer,
180) {
181 if vnet_hdr {
182 if let Some(frame) = coalescer.push(inner) {
183 send_to_tun(tun_fd, frame);
184 }
185 } else {
186 send_to_tun(tun_fd, inner);
187 }
188}
189
190#[expect(
202 clippy::too_many_arguments,
203 reason = "batch buffers (the GSO scratch `gso`, the reusable read buffer `buf`, and the \
204 TUN-offload `vnet_hdr`/split scratch) are owned by run_poll and threaded through \
205 explicitly rather than bundled into a struct, to keep the hot-path buffers' \
206 aliasing/lifetimes obvious"
207)]
208fn drain_tun(
209 tun_fd: RawFd,
210 udp_fd: RawFd,
211 d: &mut impl Dispatch,
212 now_ms: u64,
213 tx: &mut Vec<EgressDatagram>,
214 gso: &mut GsoScratch,
215 vnet_hdr: bool,
216 buf: &mut [u8],
217 split_out: &mut Vec<u8>,
218 split_offs: &mut Vec<(usize, usize)>,
219) -> io::Result<()> {
220 loop {
221 let n = unsafe { libc::read(tun_fd, buf.as_mut_ptr().cast(), buf.len()) };
224 if n < 0 {
225 let e = io::Error::last_os_error();
226 let raw = e.raw_os_error().unwrap_or(0);
227 if raw == libc::EWOULDBLOCK || raw == libc::EAGAIN {
228 break;
229 }
230 if raw == libc::EINTR {
231 continue;
232 }
233 return Err(e);
234 }
235 if n == 0 {
236 break;
237 }
238 let n = usize::try_from(n).expect("non-negative read return fits usize");
239 if vnet_hdr {
240 if crate::tun_offload::split_gro(&buf[..n], split_out, split_offs) {
241 for &(s, l) in split_offs.iter() {
242 tx.extend(d.on_tun(&split_out[s..s + l], now_ms).iter().cloned());
243 if tx.len() >= MAX_DATAGRAM_BATCH {
244 flush_tx(udp_fd, tx, gso)?;
245 }
246 }
247 } else if let Some(body) = buf.get(crate::tun_offload::VNET_HDR_LEN..n) {
248 tx.extend(d.on_tun(body, now_ms).iter().cloned());
252 if tx.len() >= MAX_DATAGRAM_BATCH {
253 flush_tx(udp_fd, tx, gso)?;
254 }
255 }
256 } else {
258 let inner = &buf[..n];
259 tx.extend(d.on_tun(inner, now_ms).iter().cloned());
260 if tx.len() >= MAX_DATAGRAM_BATCH {
261 flush_tx(udp_fd, tx, gso)?;
262 }
263 }
264 }
265 flush_tx(udp_fd, tx, gso)?; Ok(())
267}
268
269struct GsoScratch {
271 enabled: bool,
274 runs: Vec<crate::gso::GsoRun>,
275 payload: Vec<u8>,
276}
277
278impl GsoScratch {
279 fn new() -> Self {
280 Self {
281 enabled: true,
282 runs: Vec::with_capacity(MAX_DATAGRAM_BATCH),
283 payload: Vec::with_capacity(MAX_WIRE_DATAGRAM * crate::gso::MAX_GSO_SEGMENTS_PER_SEND),
284 }
285 }
286}
287
288fn send_run_plain(udp_fd: RawFd, run: &[EgressDatagram]) -> io::Result<()> {
292 let mut sent = 0;
293 while sent < run.len() {
294 let n = send_mmsg(udp_fd, &run[sent..])?;
295 if n == 0 {
296 break;
297 }
298 sent += n;
299 }
300 Ok(())
301}
302
303fn flush_tx(udp_fd: RawFd, tx: &mut Vec<EgressDatagram>, gso: &mut GsoScratch) -> io::Result<()> {
309 if !gso.enabled {
310 let r = send_run_plain(udp_fd, tx);
311 tx.clear();
312 return r;
313 }
314 crate::gso::partition_fate_safe(tx, MAX_DATAGRAM_BATCH, &mut gso.runs);
315 let runs = std::mem::take(&mut gso.runs);
318 let mut outcome = Ok(());
319 for run in &runs {
320 if gso.enabled && run.members.len() >= 2 && run.segment_size > 0 {
321 let dst = tx[run.members[0]].dst;
322 match send_gso_indexed(
323 udp_fd,
324 tx,
325 &run.members,
326 run.segment_size,
327 dst,
328 &mut gso.payload,
329 ) {
330 Ok(true) => {}
331 Ok(false) => {
332 gso.enabled = false; for &i in &run.members {
334 if let Err(e) = send_run_plain(udp_fd, std::slice::from_ref(&tx[i])) {
335 outcome = Err(e);
336 break;
337 }
338 }
339 }
340 Err(e) => outcome = Err(e),
341 }
342 } else {
343 for &i in &run.members {
344 if let Err(e) = send_run_plain(udp_fd, std::slice::from_ref(&tx[i])) {
345 outcome = Err(e);
346 break;
347 }
348 }
349 }
350 if outcome.is_err() {
351 break;
352 }
353 }
354 gso.runs = runs; tx.clear();
356 outcome
357}
358
359#[inline]
364fn send_to_tun(tun_fd: RawFd, buf: &[u8]) {
365 let rc = unsafe { libc::write(tun_fd, buf.as_ptr().cast(), buf.len()) };
371 if rc < 0 {
372 eprintln!("poll: tun write error: {}", io::Error::last_os_error());
373 }
374}
375
376fn send_mmsg(udp_fd: RawFd, datagrams: &[EgressDatagram]) -> io::Result<usize> {
380 if datagrams.is_empty() {
381 return Ok(0);
382 }
383 let count = datagrams.len().min(MAX_DATAGRAM_BATCH);
384
385 let mut storages: [libc::sockaddr_storage; MAX_DATAGRAM_BATCH] = unsafe { std::mem::zeroed() };
390 let mut addrlens = [0 as libc::socklen_t; MAX_DATAGRAM_BATCH];
391 let mut iovecs = [libc::iovec {
392 iov_base: std::ptr::null_mut(),
393 iov_len: 0,
394 }; MAX_DATAGRAM_BATCH];
395 let mut msgs = [crate::empty_mmsghdr(); MAX_DATAGRAM_BATCH];
396
397 for (i, dg) in datagrams[..count].iter().enumerate() {
398 let (storage, addr_len) = std_to_sockaddr(dg.dst);
399 storages[i] = storage;
400 addrlens[i] = addr_len;
401 iovecs[i].iov_base = dg.bytes.as_ptr().cast_mut().cast::<libc::c_void>();
405 iovecs[i].iov_len = dg.bytes.len();
406 msgs[i].msg_hdr.msg_iov = &raw mut iovecs[i];
407 msgs[i].msg_hdr.msg_iovlen = 1;
408 msgs[i].msg_hdr.msg_name = std::ptr::from_mut(&mut storages[i]).cast::<libc::c_void>();
409 msgs[i].msg_hdr.msg_namelen = addrlens[i];
410 }
411
412 let ret = unsafe {
416 libc::sendmmsg(
417 udp_fd,
418 msgs.as_mut_ptr(),
419 u32::try_from(count).expect("count ≤ 64 fits u32"),
420 libc::MSG_NOSIGNAL as crate::MsgFlags,
421 )
422 };
423 if ret < 0 {
424 let e = io::Error::last_os_error();
425 let raw = e.raw_os_error().unwrap_or(0);
426 if raw == libc::EWOULDBLOCK || raw == libc::EAGAIN || raw == libc::ENOBUFS {
428 return Ok(0);
429 }
430 return Err(e);
431 }
432 Ok(usize::try_from(ret).expect("non-negative sendmmsg return fits usize"))
433}
434
435fn send_gso_payload(
441 udp_fd: RawFd,
442 payload: &[u8],
443 segment_size: u16,
444 dst: SocketAddr,
445) -> io::Result<bool> {
446 let (mut storage, addr_len) = std_to_sockaddr(dst);
447 let mut iov = libc::iovec {
448 iov_base: payload.as_ptr().cast_mut().cast::<libc::c_void>(),
451 iov_len: payload.len(),
452 };
453 let mut control = [0u8; GSO_CONTROL_SPACE];
454 let mut msg: libc::msghdr = unsafe { std::mem::zeroed() };
457 msg.msg_name = std::ptr::from_mut(&mut storage).cast::<libc::c_void>();
458 msg.msg_namelen = addr_len;
459 msg.msg_iov = &raw mut iov;
460 msg.msg_iovlen = 1;
461 msg.msg_control = control.as_mut_ptr().cast::<libc::c_void>();
462 let cmsg_space = usize::try_from(unsafe { libc::CMSG_SPACE(GSO_CONTROL_PAYLOAD_LEN) })
464 .expect("cmsg space fits usize");
465 if cmsg_space > control.len() {
466 return Err(io::Error::other("GSO control buffer too small"));
467 }
468 msg.msg_controllen = cmsg_space as crate::CmsgLen;
469
470 unsafe {
473 let cmsg = libc::CMSG_FIRSTHDR(&raw const msg);
474 if cmsg.is_null() {
475 return Err(io::Error::other("missing first cmsg header"));
476 }
477 (*cmsg).cmsg_level = libc::SOL_UDP;
478 (*cmsg).cmsg_type = libc::UDP_SEGMENT;
479 (*cmsg).cmsg_len = usize::try_from(libc::CMSG_LEN(GSO_CONTROL_PAYLOAD_LEN))
480 .expect("cmsg len fits usize") as crate::CmsgLen;
481 let seg_ptr = libc::CMSG_DATA(cmsg).cast::<u16>();
482 *seg_ptr = segment_size;
483 }
484
485 let ret = unsafe { libc::sendmsg(udp_fd, &raw const msg, libc::MSG_NOSIGNAL) };
488 if ret < 0 {
489 let e = io::Error::last_os_error();
490 let raw = e.raw_os_error().unwrap_or(0);
491 if raw == libc::EWOULDBLOCK || raw == libc::EAGAIN || raw == libc::ENOBUFS {
492 return Ok(true); }
494 if raw == libc::EIO || raw == libc::EINVAL {
495 return Ok(false); }
497 return Err(e);
498 }
499 Ok(true)
500}
501
502#[cfg_attr(
504 not(test),
505 expect(
506 dead_code,
507 reason = "run-slice wrapper kept for its own unit test; flush_tx uses \
508 send_gso_indexed to avoid copying the run out of the egress batch"
509 )
510)]
511fn send_gso(
512 udp_fd: RawFd,
513 run: &[EgressDatagram],
514 segment_size: u16,
515 dst: SocketAddr,
516 payload: &mut Vec<u8>,
517) -> io::Result<bool> {
518 payload.clear();
519 for dg in run {
520 payload.extend_from_slice(&dg.bytes);
521 }
522 send_gso_payload(udp_fd, payload, segment_size, dst)
523}
524
525fn send_gso_indexed(
528 udp_fd: RawFd,
529 tx: &[EgressDatagram],
530 indices: &[usize],
531 segment_size: u16,
532 dst: SocketAddr,
533 payload: &mut Vec<u8>,
534) -> io::Result<bool> {
535 payload.clear();
536 for &i in indices {
537 payload.extend_from_slice(&tx[i].bytes);
538 }
539 send_gso_payload(udp_fd, payload, segment_size, dst)
540}
541
542fn recv_mmsg(
547 udp_fd: RawFd,
548 bufs: &mut [[u8; MAX_WIRE_DATAGRAM]],
549 lens: &mut [usize],
550 srcs: &mut [SocketAddr],
551) -> io::Result<usize> {
552 let count = bufs
553 .len()
554 .min(lens.len())
555 .min(srcs.len())
556 .min(MAX_DATAGRAM_BATCH);
557 if count == 0 {
558 return Ok(0);
559 }
560 let mut storages: [libc::sockaddr_storage; MAX_DATAGRAM_BATCH] = unsafe { std::mem::zeroed() };
563 let addrlens = [libc::socklen_t::try_from(std::mem::size_of::<libc::sockaddr_storage>())
564 .expect("size fits socklen_t"); MAX_DATAGRAM_BATCH];
565 let mut iovecs = [libc::iovec {
566 iov_base: std::ptr::null_mut(),
567 iov_len: 0,
568 }; MAX_DATAGRAM_BATCH];
569 let mut msgs = [crate::empty_mmsghdr(); MAX_DATAGRAM_BATCH];
570
571 for i in 0..count {
572 iovecs[i].iov_base = bufs[i].as_mut_ptr().cast::<libc::c_void>();
575 iovecs[i].iov_len = MAX_WIRE_DATAGRAM;
576 msgs[i].msg_hdr.msg_iov = &raw mut iovecs[i];
577 msgs[i].msg_hdr.msg_iovlen = 1;
578 msgs[i].msg_hdr.msg_name = std::ptr::from_mut(&mut storages[i]).cast::<libc::c_void>();
579 msgs[i].msg_hdr.msg_namelen = addrlens[i];
580 }
581
582 let ret = unsafe {
586 libc::recvmmsg(
587 udp_fd,
588 msgs.as_mut_ptr(),
589 u32::try_from(count).expect("count ≤ 64 fits u32"),
590 libc::MSG_DONTWAIT as crate::MsgFlags,
591 std::ptr::null_mut(),
592 )
593 };
594 if ret < 0 {
595 let e = io::Error::last_os_error();
596 let raw = e.raw_os_error().unwrap_or(0);
597 if raw == libc::EWOULDBLOCK || raw == libc::EAGAIN {
598 return Ok(0);
599 }
600 return Err(e);
601 }
602 let received = usize::try_from(ret).expect("non-negative recvmmsg return fits usize");
603 for i in 0..received {
604 lens[i] = usize::try_from(msgs[i].msg_len).expect("msg_len fits usize");
605 srcs[i] = sockaddr_to_std(&storages[i], msgs[i].msg_hdr.msg_namelen)
607 .unwrap_or_else(|_| SocketAddr::from(([0, 0, 0, 0], 0)));
608 }
609 Ok(received)
610}
611
612pub fn run_poll<D: Dispatch>(
629 udp_fd: RawFd,
630 tun_fd: RawFd,
631 vnet_hdr: bool,
632 d: &mut D,
633) -> io::Result<()> {
634 set_nonblocking(udp_fd)?;
635 set_nonblocking(tun_fd)?;
636
637 let epoll_fd = unsafe { libc::epoll_create1(libc::EPOLL_CLOEXEC) };
640 if epoll_fd < 0 {
641 return Err(io::Error::last_os_error());
642 }
643
644 let mut ev_udp = libc::epoll_event {
646 events: libc::EPOLLIN as u32,
647 u64: udp_fd as u64,
648 };
649 let rc = unsafe { libc::epoll_ctl(epoll_fd, libc::EPOLL_CTL_ADD, udp_fd, &raw mut ev_udp) };
653 if rc < 0 {
654 unsafe { libc::close(epoll_fd) };
656 return Err(io::Error::last_os_error());
657 }
658
659 let mut ev_tun = libc::epoll_event {
661 events: libc::EPOLLIN as u32,
662 u64: tun_fd as u64,
663 };
664 let rc = unsafe { libc::epoll_ctl(epoll_fd, libc::EPOLL_CTL_ADD, tun_fd, &raw mut ev_tun) };
667 if rc < 0 {
668 unsafe { libc::close(epoll_fd) };
670 return Err(io::Error::last_os_error());
671 }
672
673 let start = Instant::now();
674 let mut events = [libc::epoll_event { events: 0, u64: 0 }; 4];
677
678 let mut rx_bufs = Box::new([[0u8; MAX_WIRE_DATAGRAM]; MAX_DATAGRAM_BATCH]);
680 let mut rx_lens = [0usize; MAX_DATAGRAM_BATCH];
681 let mut rx_srcs = [SocketAddr::from(([0, 0, 0, 0], 0)); MAX_DATAGRAM_BATCH];
682 let mut tx_batch: Vec<EgressDatagram> = Vec::with_capacity(MAX_DATAGRAM_BATCH);
683 let mut gso = GsoScratch::new();
684
685 let mut coalescer = crate::tun_offload::Coalescer::new();
692 let mut split_out: Vec<u8> = Vec::new();
693 let mut split_offs: Vec<(usize, usize)> = Vec::new();
694 let mut tun_buf = vec![
695 0u8;
696 if vnet_hdr {
697 crate::tun_offload::VNET_HDR_LEN + crate::tun_offload::MAX_GSO_PAYLOAD
698 } else {
699 MAX_WIRE_DATAGRAM
700 }
701 ];
702
703 loop {
704 let now_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
705 let timeout_ms = match d.next_timeout_ms(now_ms) {
706 Some(delay_ms) => {
707 let clamped = i32::try_from(delay_ms).unwrap_or(10);
708 10.min(clamped)
709 }
710 None => 10,
711 };
712
713 let nfds = unsafe {
718 libc::epoll_wait(
719 epoll_fd,
720 events.as_mut_ptr(),
721 libc::c_int::try_from(events.len()).expect("event array len fits c_int"),
722 timeout_ms,
723 )
724 };
725
726 if nfds < 0 {
727 let e = io::Error::last_os_error();
728 if e.raw_os_error() == Some(libc::EINTR) {
729 continue;
731 }
732 unsafe { libc::close(epoll_fd) };
734 return Err(e);
735 }
736
737 let now_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
738
739 for ev in &events[..usize::try_from(nfds).expect("non-negative nfds fits usize")] {
740 let ready_fd = ev.u64 as RawFd;
741 if ready_fd == udp_fd {
742 if let Err(e) = drain_udp(
743 udp_fd,
744 tun_fd,
745 d,
746 now_ms,
747 &mut rx_bufs,
748 &mut rx_lens,
749 &mut rx_srcs,
750 &mut tx_batch,
751 &mut gso,
752 vnet_hdr,
753 &mut coalescer,
754 ) {
755 unsafe { libc::close(epoll_fd) };
757 return Err(e);
758 }
759 } else if ready_fd == tun_fd {
760 if let Err(e) = drain_tun(
761 tun_fd,
762 udp_fd,
763 d,
764 now_ms,
765 &mut tx_batch,
766 &mut gso,
767 vnet_hdr,
768 &mut tun_buf,
769 &mut split_out,
770 &mut split_offs,
771 ) {
772 unsafe { libc::close(epoll_fd) };
774 return Err(e);
775 }
776 }
777 }
778
779 if let Some(pkts) = d.tick(now_ms) {
781 tx_batch.extend(pkts.iter().cloned());
782 if let Err(e) = flush_tx(udp_fd, &mut tx_batch, &mut gso) {
783 unsafe { libc::close(epoll_fd) };
785 return Err(e);
786 }
787 }
788 }
789}
790
791#[cfg(test)]
794mod tests {
795 use super::*;
796 use std::net::UdpSocket;
797 use std::os::fd::AsRawFd;
798
799 struct CountDispatch {
802 received: Vec<Vec<u8>>,
803 call_count: usize,
804 }
805
806 impl CountDispatch {
807 fn new() -> Self {
808 Self {
809 received: Vec::new(),
810 call_count: 0,
811 }
812 }
813 }
814
815 impl Dispatch for CountDispatch {
816 fn on_udp(&mut self, _src: SocketAddr, dg: &[u8], _now_ms: u64) -> DispatchOut<'_> {
817 self.received.push(dg.to_vec());
818 self.call_count += 1;
819 DispatchOut::None
820 }
821
822 fn on_tun(&mut self, _inner: &[u8], _now_ms: u64) -> &[EgressDatagram] {
823 &[]
824 }
825
826 fn tick(&mut self, _now_ms: u64) -> Option<&[EgressDatagram]> {
827 None
828 }
829 }
830
831 #[test]
834 fn drain_udp_delivers_datagram() {
835 let a = UdpSocket::bind("127.0.0.1:0").unwrap();
837 let b = UdpSocket::bind("127.0.0.1:0").unwrap();
838 a.connect(b.local_addr().unwrap()).unwrap();
839 b.connect(a.local_addr().unwrap()).unwrap();
840
841 set_nonblocking(b.as_raw_fd()).unwrap();
843
844 a.send(b"hello from drain_udp test").unwrap();
846
847 let devnull = std::fs::OpenOptions::new()
850 .write(true)
851 .open("/dev/null")
852 .unwrap();
853 let null_fd = devnull.as_raw_fd();
854
855 let mut d = CountDispatch::new();
856 let mut bufs = Box::new([[0u8; MAX_WIRE_DATAGRAM]; MAX_DATAGRAM_BATCH]);
857 let mut lens = [0usize; MAX_DATAGRAM_BATCH];
858 let mut srcs = [std::net::SocketAddr::from(([0, 0, 0, 0], 0)); MAX_DATAGRAM_BATCH];
859 let mut tx = Vec::new();
860 let mut gso = GsoScratch::new();
861 let mut coalescer = crate::tun_offload::Coalescer::new();
862 drain_udp(
863 b.as_raw_fd(),
864 null_fd,
865 &mut d,
866 0,
867 &mut bufs,
868 &mut lens,
869 &mut srcs,
870 &mut tx,
871 &mut gso,
872 false,
873 &mut coalescer,
874 )
875 .unwrap();
876
877 assert_eq!(d.call_count, 1, "on_udp must be called exactly once");
878 assert_eq!(d.received[0], b"hello from drain_udp test");
879 }
880
881 #[test]
883 fn drain_udp_drains_multiple_datagrams() {
884 let a = UdpSocket::bind("127.0.0.1:0").unwrap();
885 let b = UdpSocket::bind("127.0.0.1:0").unwrap();
886 a.connect(b.local_addr().unwrap()).unwrap();
887 b.connect(a.local_addr().unwrap()).unwrap();
888 set_nonblocking(b.as_raw_fd()).unwrap();
889
890 a.send(b"first").unwrap();
891 a.send(b"second").unwrap();
892 a.send(b"third").unwrap();
893
894 let devnull = std::fs::OpenOptions::new()
895 .write(true)
896 .open("/dev/null")
897 .unwrap();
898 let null_fd = devnull.as_raw_fd();
899
900 let mut d = CountDispatch::new();
901 let mut bufs = Box::new([[0u8; MAX_WIRE_DATAGRAM]; MAX_DATAGRAM_BATCH]);
902 let mut lens = [0usize; MAX_DATAGRAM_BATCH];
903 let mut srcs = [std::net::SocketAddr::from(([0, 0, 0, 0], 0)); MAX_DATAGRAM_BATCH];
904 let mut tx = Vec::new();
905 let mut gso = GsoScratch::new();
906 let mut coalescer = crate::tun_offload::Coalescer::new();
907 drain_udp(
908 b.as_raw_fd(),
909 null_fd,
910 &mut d,
911 0,
912 &mut bufs,
913 &mut lens,
914 &mut srcs,
915 &mut tx,
916 &mut gso,
917 false,
918 &mut coalescer,
919 )
920 .unwrap();
921
922 assert_eq!(d.call_count, 3);
923 }
924
925 #[test]
927 fn set_nonblocking_is_idempotent() {
928 let sock = UdpSocket::bind("127.0.0.1:0").unwrap();
929 set_nonblocking(sock.as_raw_fd()).unwrap();
930 set_nonblocking(sock.as_raw_fd()).unwrap(); }
932
933 struct ForwardDispatch {
939 reply: Vec<u8>,
941 scratch: Vec<EgressDatagram>,
944 }
945
946 impl ForwardDispatch {
947 fn new(reply: impl Into<Vec<u8>>) -> Self {
948 Self {
949 reply: reply.into(),
950 scratch: Vec::new(),
951 }
952 }
953 }
954
955 impl Dispatch for ForwardDispatch {
956 fn on_udp(&mut self, src: SocketAddr, _dg: &[u8], _now_ms: u64) -> DispatchOut<'_> {
957 self.scratch = vec![EgressDatagram {
958 fate: 0,
959 dst: src,
960 bytes: self.reply.clone(),
961 }];
962 DispatchOut::Udp(&self.scratch)
963 }
964
965 fn on_tun(&mut self, _inner: &[u8], _now_ms: u64) -> &[EgressDatagram] {
966 &[]
967 }
968
969 fn tick(&mut self, _now_ms: u64) -> Option<&[EgressDatagram]> {
970 None
971 }
972 }
973
974 #[test]
979 fn drain_udp_forwards_dispatch_out_udp_to_peer() {
980 let a = UdpSocket::bind("127.0.0.1:0").unwrap();
982 let b = UdpSocket::bind("127.0.0.1:0").unwrap();
983 a.connect(b.local_addr().unwrap()).unwrap();
984 b.connect(a.local_addr().unwrap()).unwrap();
985
986 set_nonblocking(b.as_raw_fd()).unwrap();
988 set_nonblocking(a.as_raw_fd()).unwrap();
990
991 a.send(b"trigger").unwrap();
993
994 let devnull = std::fs::OpenOptions::new()
996 .write(true)
997 .open("/dev/null")
998 .unwrap();
999 let null_fd = devnull.as_raw_fd();
1000
1001 let mut d = ForwardDispatch::new(b"forwarded".as_slice());
1003 let mut bufs = Box::new([[0u8; MAX_WIRE_DATAGRAM]; MAX_DATAGRAM_BATCH]);
1004 let mut lens = [0usize; MAX_DATAGRAM_BATCH];
1005 let mut srcs = [std::net::SocketAddr::from(([0, 0, 0, 0], 0)); MAX_DATAGRAM_BATCH];
1006 let mut tx = Vec::new();
1007 let mut gso = GsoScratch::new();
1008 let mut coalescer = crate::tun_offload::Coalescer::new();
1009 drain_udp(
1010 b.as_raw_fd(),
1011 null_fd,
1012 &mut d,
1013 0,
1014 &mut bufs,
1015 &mut lens,
1016 &mut srcs,
1017 &mut tx,
1018 &mut gso,
1019 false,
1020 &mut coalescer,
1021 )
1022 .unwrap();
1023
1024 let mut recv_buf = [0u8; 64];
1026 let n = a
1027 .recv(&mut recv_buf)
1028 .expect("forwarded datagram must arrive on peer socket");
1029 assert_eq!(&recv_buf[..n], b"forwarded");
1030 }
1031
1032 #[test]
1033 fn send_mmsg_delivers_each_datagram_to_its_own_dst() {
1034 use std::net::UdpSocket;
1035 let rx_a = UdpSocket::bind("127.0.0.1:0").unwrap();
1037 let rx_b = UdpSocket::bind("127.0.0.1:0").unwrap();
1038 rx_a.set_read_timeout(Some(std::time::Duration::from_secs(2)))
1039 .unwrap();
1040 rx_b.set_read_timeout(Some(std::time::Duration::from_secs(2)))
1041 .unwrap();
1042 let tx = UdpSocket::bind("127.0.0.1:0").unwrap();
1043 let dst_a = rx_a.local_addr().unwrap();
1044 let dst_b = rx_b.local_addr().unwrap();
1045
1046 let batch = [
1047 EgressDatagram {
1048 fate: 0,
1049 dst: dst_a,
1050 bytes: b"to-a-1".to_vec(),
1051 },
1052 EgressDatagram {
1053 fate: 0,
1054 dst: dst_b,
1055 bytes: b"to-b".to_vec(),
1056 },
1057 EgressDatagram {
1058 fate: 0,
1059 dst: dst_a,
1060 bytes: b"to-a-2".to_vec(),
1061 },
1062 ];
1063 let mut sent = 0;
1064 while sent < batch.len() {
1065 sent += send_mmsg(tx.as_raw_fd(), &batch[sent..]).unwrap();
1066 }
1067 assert_eq!(sent, 3);
1068
1069 let mut buf = [0u8; 64];
1070 let (n1, _) = rx_a.recv_from(&mut buf).unwrap();
1072 assert_eq!(&buf[..n1], b"to-a-1");
1073 let (n2, _) = rx_a.recv_from(&mut buf).unwrap();
1074 assert_eq!(&buf[..n2], b"to-a-2");
1075 let (n3, _) = rx_b.recv_from(&mut buf).unwrap();
1077 assert_eq!(&buf[..n3], b"to-b");
1078 }
1079
1080 #[test]
1081 fn recv_mmsg_returns_bytes_and_source_per_datagram() {
1082 use std::net::UdpSocket;
1083 let rx = UdpSocket::bind("127.0.0.1:0").unwrap();
1084 let tx = UdpSocket::bind("127.0.0.1:0").unwrap();
1085 let rx_addr = rx.local_addr().unwrap();
1086 let tx_addr = tx.local_addr().unwrap();
1087 tx.send_to(b"hello", rx_addr).unwrap();
1088 tx.send_to(b"world!!", rx_addr).unwrap();
1089 std::thread::sleep(std::time::Duration::from_millis(50));
1091
1092 let mut bufs = [[0u8; MAX_WIRE_DATAGRAM]; 4];
1093 let mut lens = [0usize; 4];
1094 let mut srcs = [std::net::SocketAddr::from(([0, 0, 0, 0], 0)); 4];
1095 let n = recv_mmsg(rx.as_raw_fd(), &mut bufs, &mut lens, &mut srcs).unwrap();
1096 assert_eq!(n, 2);
1097 assert_eq!(&bufs[0][..lens[0]], b"hello");
1098 assert_eq!(&bufs[1][..lens[1]], b"world!!");
1099 assert_eq!(srcs[0], tx_addr);
1100 assert_eq!(srcs[1], tx_addr);
1101 }
1102
1103 #[test]
1104 fn recv_mmsg_returns_zero_when_nothing_queued() {
1105 use std::net::UdpSocket;
1106 let rx = UdpSocket::bind("127.0.0.1:0").unwrap();
1107 let mut bufs = [[0u8; MAX_WIRE_DATAGRAM]; 4];
1108 let mut lens = [0usize; 4];
1109 let mut srcs = [std::net::SocketAddr::from(([0, 0, 0, 0], 0)); 4];
1110 assert_eq!(
1111 recv_mmsg(rx.as_raw_fd(), &mut bufs, &mut lens, &mut srcs).unwrap(),
1112 0
1113 );
1114 }
1115
1116 #[test]
1117 fn send_gso_delivers_each_segment_when_supported() {
1118 use std::net::UdpSocket;
1119 let rx = UdpSocket::bind("127.0.0.1:0").unwrap();
1120 let rx_addr = rx.local_addr().unwrap();
1121 let tx = UdpSocket::bind("127.0.0.1:0").unwrap();
1122 rx.set_nonblocking(true).unwrap();
1123
1124 let run = vec![
1125 EgressDatagram {
1126 fate: 1,
1127 dst: rx_addr,
1128 bytes: vec![0xAA; 1000],
1129 },
1130 EgressDatagram {
1131 fate: 2,
1132 dst: rx_addr,
1133 bytes: vec![0xBB; 1000],
1134 },
1135 EgressDatagram {
1136 fate: 3,
1137 dst: rx_addr,
1138 bytes: vec![0xCC; 1000],
1139 },
1140 ];
1141 let mut payload = Vec::new();
1142 let accepted =
1143 send_gso(tx.as_raw_fd(), &run, 1000, rx_addr, &mut payload).expect("send_gso");
1144 if !accepted {
1145 return; }
1147 std::thread::sleep(std::time::Duration::from_millis(50));
1148 let mut got = Vec::new();
1149 let mut buf = [0u8; 2048];
1150 while let Ok((n, _)) = rx.recv_from(&mut buf) {
1151 got.push(buf[..n].to_vec());
1152 }
1153 assert_eq!(got.len(), 3, "GSO must segment into 3 separate datagrams");
1154 assert!(got.iter().all(|d| d.len() == 1000));
1155 }
1156
1157 #[test]
1158 fn flush_tx_delivers_all_datagrams_gso_or_fallback() {
1159 use std::net::UdpSocket;
1160 let rx = UdpSocket::bind("127.0.0.1:0").unwrap();
1161 let rx_addr = rx.local_addr().unwrap();
1162 let tx_sock = UdpSocket::bind("127.0.0.1:0").unwrap();
1163 rx.set_nonblocking(true).unwrap();
1164
1165 let mut tx: Vec<EgressDatagram> = [1u16, 2, 3, 1]
1167 .iter()
1168 .map(|&f| EgressDatagram {
1169 fate: f,
1170 dst: rx_addr,
1171 bytes: vec![f as u8; 1000],
1172 })
1173 .collect();
1174 let mut gso = GsoScratch::new();
1175 flush_tx(tx_sock.as_raw_fd(), &mut tx, &mut gso).expect("flush_tx");
1176 assert!(tx.is_empty(), "flush_tx must drain tx");
1177
1178 std::thread::sleep(std::time::Duration::from_millis(50));
1179 let mut count = 0;
1180 let mut buf = [0u8; 2048];
1181 while let Ok((n, _)) = rx.recv_from(&mut buf) {
1182 assert_eq!(n, 1000);
1183 count += 1;
1184 }
1185 assert_eq!(
1186 count, 4,
1187 "all four datagrams must arrive (GSO or plain fallback)"
1188 );
1189 }
1190}