1use std::io;
32use std::net::{SocketAddr, UdpSocket};
33
34#[derive(Debug, Clone, Copy, PartialEq, Eq)]
36pub enum DgramBackend {
37 Udp,
39 IoUring,
41 Wire,
45 Demux,
49}
50
51#[allow(clippy::large_enum_variant)]
56enum Inner {
57 Udp(UdpSocket),
58 #[cfg(target_os = "linux")]
59 IoUring(linux_iou::IoUringDgram),
60 #[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
61 Wire(wire_backend::WireDgram),
62 Demux(DemuxDgram),
65}
66
67pub type DemuxQueue =
72 std::sync::Arc<std::sync::Mutex<std::collections::VecDeque<(Vec<u8>, SocketAddr, Option<i128>)>>>;
73
74pub fn new_demux_queue() -> DemuxQueue {
76 std::sync::Arc::new(std::sync::Mutex::new(std::collections::VecDeque::new()))
77}
78
79struct DemuxDgram {
85 real: std::sync::Arc<UdpSocket>,
86 queue: DemuxQueue,
87 peer: std::sync::Mutex<Option<SocketAddr>>,
90 sent: Option<std::sync::Arc<std::sync::atomic::AtomicU64>>,
93}
94
95impl DemuxDgram {
96 fn recv_with_kts(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr, Option<i128>)> {
97 match self.queue.lock().unwrap().pop_front() {
98 Some((data, from, kts)) => {
99 let n = data.len().min(buf.len());
100 buf[..n].copy_from_slice(&data[..n]);
101 Ok((n, from, kts))
102 }
103 None => Err(io::Error::new(io::ErrorKind::WouldBlock, "demux queue empty")),
106 }
107 }
108
109 fn connect(&self, addr: SocketAddr) {
110 *self.peer.lock().unwrap() = Some(addr);
111 }
112
113 fn count_fwd(&self, buf: &[u8]) {
114 if let Some(c) = &self.sent
118 && matches!(buf.first(), Some(&(1 | 10 | 11)))
119 {
120 c.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
121 }
122 }
123
124 fn send_to(&self, buf: &[u8], addr: SocketAddr) -> io::Result<usize> {
125 self.count_fwd(buf);
126 self.real.send_to(buf, addr)
127 }
128
129 fn send(&self, buf: &[u8]) -> io::Result<usize> {
130 self.count_fwd(buf);
131 match *self.peer.lock().unwrap() {
132 Some(p) => self.real.send_to(buf, p),
133 None => Err(io::Error::new(io::ErrorKind::NotConnected, "demux send before connect")),
134 }
135 }
136}
137
138pub struct DgramSock {
141 inner: Inner,
142}
143
144impl DgramSock {
145 pub fn wrap(sock: UdpSocket) -> Self {
149 enable_rx_timestamp(&sock);
150 let forced = std::env::var("SUBETHA_DGRAM").ok();
151 if forced.as_deref() == Some("udp") {
152 return Self { inner: Inner::Udp(sock) };
153 }
154
155 #[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
165 {
166 let wire_if = std::env::var("SUBETHA_WIRE_IFNAME").ok();
167 let force_wire = forced.as_deref() == Some("wire");
168 let auto_wire = forced.is_none()
169 && wire_if.as_deref().map(wire_link_fast_enough).unwrap_or(false);
170 if force_wire || auto_wire {
171 let port = sock.local_addr().map(|a| a.port()).unwrap_or(0);
172 match wire_backend::WireDgram::from_env(port) {
173 Ok(w) => {
174 if force_wire
175 && !wire_if.as_deref().map(wire_link_fast_enough).unwrap_or(false)
176 {
177 eprintln!(
178 "SUBETHA_DGRAM=wire forced on a link below the {:.0} Gbit/s \
179 gate; the kernel-bypass path is net-negative below line rate",
180 wire_min_link_bps() as f64 / 1e9
181 );
182 }
183 return Self { inner: Inner::Wire(w) };
184 }
185 Err(e) => eprintln!(
186 "Wire backend requested but unavailable ({e}); falling back"
187 ),
188 }
189 }
190 }
191
192 #[cfg(target_os = "linux")]
193 {
194 let force_ring = forced.as_deref() == Some("iouring");
195 match linux_iou::IoUringDgram::new(sock) {
196 Ok(d) => Self { inner: Inner::IoUring(d) },
197 Err((sock, e)) => {
198 if force_ring {
199 eprintln!(
200 "SUBETHA_DGRAM=iouring requested but the ring is unavailable \
201 ({e}); using plain UDP"
202 );
203 }
204 Self { inner: Inner::Udp(sock) }
205 }
206 }
207 }
208
209 #[cfg(not(target_os = "linux"))]
210 {
211 drop(forced);
212 Self { inner: Inner::Udp(sock) }
213 }
214 }
215
216 pub fn demux(real: std::sync::Arc<UdpSocket>, queue: DemuxQueue) -> Self {
222 Self::demux_inner(real, queue, None)
223 }
224
225 pub fn demux_counted(
228 real: std::sync::Arc<UdpSocket>,
229 queue: DemuxQueue,
230 sent: std::sync::Arc<std::sync::atomic::AtomicU64>,
231 ) -> Self {
232 Self::demux_inner(real, queue, Some(sent))
233 }
234
235 fn demux_inner(
236 real: std::sync::Arc<UdpSocket>,
237 queue: DemuxQueue,
238 sent: Option<std::sync::Arc<std::sync::atomic::AtomicU64>>,
239 ) -> Self {
240 Self {
241 inner: Inner::Demux(DemuxDgram {
242 real,
243 queue,
244 peer: std::sync::Mutex::new(None),
245 sent,
246 }),
247 }
248 }
249
250 pub fn from_udp(sock: UdpSocket) -> Self {
258 Self { inner: Inner::Udp(sock) }
259 }
260
261 pub fn as_udp(&self) -> Option<&UdpSocket> {
266 match &self.inner {
267 Inner::Udp(s) => Some(s),
268 _ => None,
269 }
270 }
271
272 pub fn connect(&self, addr: SocketAddr) -> io::Result<()> {
276 match &self.inner {
277 Inner::Udp(s) => s.connect(addr),
278 Inner::Demux(d) => {
279 d.connect(addr);
280 Ok(())
281 }
282 #[allow(unreachable_patterns)]
283 _ => Err(io::Error::new(
284 io::ErrorKind::Unsupported,
285 "connect on a non-Udp/Demux backend",
286 )),
287 }
288 }
289
290 pub fn send(&self, buf: &[u8]) -> io::Result<usize> {
292 match &self.inner {
293 Inner::Udp(s) => s.send(buf),
294 Inner::Demux(d) => d.send(buf),
295 #[allow(unreachable_patterns)]
296 _ => Err(io::Error::new(
297 io::ErrorKind::Unsupported,
298 "send on a non-Udp/Demux backend",
299 )),
300 }
301 }
302
303 pub fn recv(&self, buf: &mut [u8]) -> io::Result<usize> {
306 match &self.inner {
307 Inner::Udp(s) => s.recv(buf),
308 Inner::Demux(d) => d.recv_with_kts(buf).map(|(n, _, _)| n),
309 #[allow(unreachable_patterns)]
310 _ => Err(io::Error::new(
311 io::ErrorKind::Unsupported,
312 "recv on a non-Udp/Demux backend",
313 )),
314 }
315 }
316
317 pub fn backend(&self) -> DgramBackend {
319 match &self.inner {
320 Inner::Udp(_) => DgramBackend::Udp,
321 #[cfg(target_os = "linux")]
322 Inner::IoUring(_) => DgramBackend::IoUring,
323 #[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
324 Inner::Wire(_) => DgramBackend::Wire,
325 Inner::Demux(_) => DgramBackend::Demux,
326 }
327 }
328
329 pub fn send_to(&self, buf: &[u8], addr: SocketAddr) -> io::Result<usize> {
330 match &self.inner {
331 Inner::Udp(s) => s.send_to(buf, addr),
332 #[cfg(target_os = "linux")]
333 Inner::IoUring(d) => d.send_to(buf, addr),
334 #[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
335 Inner::Wire(d) => d.send_to(buf, addr),
336 Inner::Demux(d) => d.send_to(buf, addr),
337 }
338 }
339
340 #[cfg(target_os = "linux")]
344 fn raw_fd(&self) -> Option<i32> {
345 use std::os::unix::io::AsRawFd;
346 match &self.inner {
347 Inner::Udp(s) => Some(s.as_raw_fd()),
348 Inner::IoUring(d) => Some(d.raw_fd()),
349 #[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
350 Inner::Wire(_) => None,
351 Inner::Demux(_) => None,
352 }
353 }
354
355 pub fn send_gso(&self, batch: &[u8], seg_size: u16, addr: SocketAddr) -> io::Result<()> {
364 #[cfg(target_os = "linux")]
365 if let Some(fd) = self.raw_fd() {
366 return linux_gso_send(fd, batch, seg_size, addr);
367 }
368 let n = (seg_size as usize).max(1);
369 let mut off = 0;
370 while off < batch.len() {
371 let end = (off + n).min(batch.len());
372 self.send_to(&batch[off..end], addr)?;
373 off = end;
374 }
375 Ok(())
376 }
377
378 pub fn recv_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
379 match &self.inner {
380 Inner::Udp(s) => s.recv_from(buf),
381 #[cfg(target_os = "linux")]
382 Inner::IoUring(d) => d.recv_with_kts(buf).map(|(n, a, _)| (n, a)),
383 #[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
384 Inner::Wire(d) => d.recv_with_kts(buf).map(|(n, a, _)| (n, a)),
385 Inner::Demux(d) => d.recv_with_kts(buf).map(|(n, a, _)| (n, a)),
386 }
387 }
388
389 pub fn recv_with_kts(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr, Option<i128>)> {
391 match &self.inner {
392 Inner::Udp(s) => udp_recv_with_kts(s, buf),
393 #[cfg(target_os = "linux")]
394 Inner::IoUring(d) => d.recv_with_kts(buf),
395 #[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
396 Inner::Wire(d) => d.recv_with_kts(buf),
397 Inner::Demux(d) => d.recv_with_kts(buf),
398 }
399 }
400
401 pub fn local_addr(&self) -> io::Result<SocketAddr> {
402 match &self.inner {
403 Inner::Udp(s) => s.local_addr(),
404 #[cfg(target_os = "linux")]
405 Inner::IoUring(d) => d.local_addr(),
406 #[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
407 Inner::Wire(d) => d.local_addr(),
408 Inner::Demux(d) => d.real.local_addr(),
409 }
410 }
411
412 pub fn set_nonblocking(&self, nb: bool) -> io::Result<()> {
413 match &self.inner {
414 Inner::Udp(s) => s.set_nonblocking(nb),
415 #[cfg(target_os = "linux")]
416 Inner::IoUring(d) => d.set_nonblocking(nb),
417 #[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
418 Inner::Wire(d) => d.set_nonblocking(nb),
419 Inner::Demux(_) => Ok(()),
422 }
423 }
424}
425
426#[cfg(target_os = "linux")]
432fn enable_rx_timestamp(sock: &UdpSocket) {
433 use std::os::fd::AsRawFd;
434 let on: libc::c_int = 1;
435 unsafe {
436 libc::setsockopt(
437 sock.as_raw_fd(),
438 libc::SOL_SOCKET,
439 libc::SO_TIMESTAMPNS,
440 &on as *const libc::c_int as *const libc::c_void,
441 std::mem::size_of::<libc::c_int>() as libc::socklen_t,
442 );
443 }
444}
445
446#[cfg(not(target_os = "linux"))]
447fn enable_rx_timestamp(_sock: &UdpSocket) {}
448
449#[cfg(target_os = "linux")]
451pub(crate) fn udp_recv_with_kts(sock: &UdpSocket, buf: &mut [u8]) -> io::Result<(usize, SocketAddr, Option<i128>)> {
452 use std::os::fd::AsRawFd;
453 let mut iov = libc::iovec {
454 iov_base: buf.as_mut_ptr() as *mut libc::c_void,
455 iov_len: buf.len(),
456 };
457 let mut name: libc::sockaddr_storage = unsafe { std::mem::zeroed() };
458 let mut control = [0u8; 64];
459 let mut msg: libc::msghdr = unsafe { std::mem::zeroed() };
460 msg.msg_name = &mut name as *mut _ as *mut libc::c_void;
461 msg.msg_namelen = std::mem::size_of::<libc::sockaddr_storage>() as libc::socklen_t;
462 msg.msg_iov = &mut iov;
463 msg.msg_iovlen = 1;
464 msg.msg_control = control.as_mut_ptr() as *mut libc::c_void;
465 msg.msg_controllen = control.len();
466 let n = unsafe { libc::recvmsg(sock.as_raw_fd(), &mut msg, 0) };
467 if n < 0 {
468 return Err(io::Error::last_os_error());
469 }
470 let kts = parse_timestamp(&msg);
471 let from = unsafe { sockaddr_to_socketaddr(&name) }
472 .ok_or_else(|| io::Error::other("non-IP source"))?;
473 Ok((n as usize, from, kts))
474}
475
476#[cfg(not(target_os = "linux"))]
477pub(crate) fn udp_recv_with_kts(sock: &UdpSocket, buf: &mut [u8]) -> io::Result<(usize, SocketAddr, Option<i128>)> {
478 let (n, from) = sock.recv_from(buf)?;
479 Ok((n, from, None))
480}
481
482#[cfg(target_os = "linux")]
483fn parse_timestamp(msg: &libc::msghdr) -> Option<i128> {
484 unsafe {
485 let mut cmsg = libc::CMSG_FIRSTHDR(msg);
486 while !cmsg.is_null() {
487 if (*cmsg).cmsg_level == libc::SOL_SOCKET && (*cmsg).cmsg_type == libc::SCM_TIMESTAMPNS {
488 let ts = (libc::CMSG_DATA(cmsg) as *const libc::timespec).read_unaligned();
489 return Some(ts.tv_sec as i128 * 1_000_000_000 + ts.tv_nsec as i128);
490 }
491 cmsg = libc::CMSG_NXTHDR(msg, cmsg);
492 }
493 }
494 None
495}
496
497#[cfg(target_os = "linux")]
498unsafe fn sockaddr_to_socketaddr(name: &libc::sockaddr_storage) -> Option<SocketAddr> {
499 use std::net::{Ipv4Addr, Ipv6Addr};
500 match name.ss_family as libc::c_int {
501 libc::AF_INET => {
502 let sin = unsafe { &*(name as *const libc::sockaddr_storage as *const libc::sockaddr_in) };
503 let ip = Ipv4Addr::from(sin.sin_addr.s_addr.to_ne_bytes());
504 Some(SocketAddr::new(ip.into(), u16::from_be(sin.sin_port)))
505 }
506 libc::AF_INET6 => {
507 let sin6 = unsafe { &*(name as *const libc::sockaddr_storage as *const libc::sockaddr_in6) };
508 let ip = Ipv6Addr::from(sin6.sin6_addr.s6_addr);
509 Some(SocketAddr::new(ip.into(), u16::from_be(sin6.sin6_port)))
510 }
511 _ => None,
512 }
513}
514
515#[cfg(target_os = "linux")]
518fn socketaddr_to_sockaddr(addr: SocketAddr) -> (libc::sockaddr_storage, libc::socklen_t) {
519 let mut storage: libc::sockaddr_storage = unsafe { std::mem::zeroed() };
520 match addr {
521 SocketAddr::V4(a) => {
522 let sin = unsafe { &mut *(&mut storage as *mut _ as *mut libc::sockaddr_in) };
523 sin.sin_family = libc::AF_INET as libc::sa_family_t;
524 sin.sin_port = a.port().to_be();
525 sin.sin_addr.s_addr = u32::from_ne_bytes(a.ip().octets());
526 (storage, std::mem::size_of::<libc::sockaddr_in>() as libc::socklen_t)
527 }
528 SocketAddr::V6(a) => {
529 let sin6 = unsafe { &mut *(&mut storage as *mut _ as *mut libc::sockaddr_in6) };
530 sin6.sin6_family = libc::AF_INET6 as libc::sa_family_t;
531 sin6.sin6_port = a.port().to_be();
532 sin6.sin6_addr.s6_addr = a.ip().octets();
533 (storage, std::mem::size_of::<libc::sockaddr_in6>() as libc::socklen_t)
534 }
535 }
536}
537
538#[cfg(target_os = "linux")]
542fn linux_gso_send(fd: i32, batch: &[u8], seg_size: u16, addr: SocketAddr) -> io::Result<()> {
543 const UDP_SEGMENT: libc::c_int = 103;
545 const SOL_UDP: libc::c_int = 17;
546 let (storage, addrlen) = socketaddr_to_sockaddr(addr);
547 let mut iov = libc::iovec {
548 iov_base: batch.as_ptr() as *mut libc::c_void,
549 iov_len: batch.len(),
550 };
551 let cmsg_space = unsafe { libc::CMSG_SPACE(2) } as usize;
552 let mut cbuf = vec![0u8; cmsg_space];
553 let mut msg: libc::msghdr = unsafe { std::mem::zeroed() };
554 msg.msg_name = &storage as *const _ as *mut libc::c_void;
555 msg.msg_namelen = addrlen;
556 msg.msg_iov = &mut iov;
557 msg.msg_iovlen = 1;
558 msg.msg_control = cbuf.as_mut_ptr() as *mut libc::c_void;
559 msg.msg_controllen = cmsg_space as _;
560 unsafe {
563 let c = libc::CMSG_FIRSTHDR(&msg);
564 (*c).cmsg_level = SOL_UDP;
565 (*c).cmsg_type = UDP_SEGMENT;
566 (*c).cmsg_len = libc::CMSG_LEN(2) as _;
567 std::ptr::copy_nonoverlapping(&seg_size as *const u16 as *const u8, libc::CMSG_DATA(c), 2);
568 loop {
569 let r = libc::sendmsg(fd, &msg, 0);
570 if r >= 0 {
571 return Ok(());
572 }
573 let e = io::Error::last_os_error();
574 if e.kind() == io::ErrorKind::Interrupted {
575 continue;
576 }
577 return Err(e);
578 }
579 }
580}
581
582#[cfg(target_os = "linux")]
587mod linux_iou {
588 use super::{parse_timestamp, sockaddr_to_socketaddr, socketaddr_to_sockaddr};
589 use std::cell::RefCell;
590 use std::collections::VecDeque;
591 use std::io;
592 use std::net::{SocketAddr, UdpSocket};
593 use std::os::fd::AsRawFd;
594
595 use io_uring::{opcode, types, IoUring};
596
597 const RECV_DEPTH: usize = 32;
598 const SEND_DEPTH: usize = 32;
599 const FRAME_CAP: usize = 2048;
600 const CONTROL_CAP: usize = 64;
601 const SEND_TAG: u64 = 1 << 32; struct RecvCtx {
608 addr: libc::sockaddr_storage,
609 iov: libc::iovec,
610 control: [u8; CONTROL_CAP],
611 msghdr: libc::msghdr,
612 buf: [u8; FRAME_CAP],
613 }
614
615 impl RecvCtx {
616 fn boxed() -> Box<Self> {
617 let mut b: Box<Self> = Box::new(unsafe { std::mem::zeroed() });
618 b.refresh();
619 b
620 }
621 fn refresh(&mut self) {
622 self.iov.iov_base = self.buf.as_mut_ptr() as *mut libc::c_void;
623 self.iov.iov_len = FRAME_CAP;
624 self.msghdr.msg_name = std::ptr::addr_of_mut!(self.addr) as *mut libc::c_void;
625 self.msghdr.msg_namelen = std::mem::size_of::<libc::sockaddr_storage>() as u32;
626 self.msghdr.msg_iov = std::ptr::addr_of_mut!(self.iov);
627 self.msghdr.msg_iovlen = 1;
628 self.msghdr.msg_control = self.control.as_mut_ptr() as *mut libc::c_void;
629 self.msghdr.msg_controllen = CONTROL_CAP;
630 }
631 }
632
633 struct SendCtx {
634 addr: libc::sockaddr_storage,
635 iov: libc::iovec,
636 msghdr: libc::msghdr,
637 buf: [u8; FRAME_CAP],
638 }
639
640 struct State {
641 ring: IoUring,
642 #[allow(clippy::vec_box)]
647 recv: Vec<Box<RecvCtx>>,
648 #[allow(clippy::vec_box)]
649 send: Vec<Box<SendCtx>>,
650 ready: VecDeque<usize>,
651 ready_len: Vec<usize>,
652 free_send: Vec<usize>,
653 }
654
655 pub struct IoUringDgram {
656 sock: UdpSocket,
657 fd: i32,
658 st: RefCell<State>,
659 }
660
661 unsafe impl Send for IoUringDgram {}
662
663 impl IoUringDgram {
664 pub fn new(sock: UdpSocket) -> Result<Self, (UdpSocket, io::Error)> {
665 let entries = ((RECV_DEPTH + SEND_DEPTH) * 2).next_power_of_two() as u32;
666 let ring = match IoUring::new(entries) {
667 Ok(r) => r,
668 Err(e) => return Err((sock, e)),
669 };
670 let fd = sock.as_raw_fd();
671 let recv: Vec<Box<RecvCtx>> = (0..RECV_DEPTH).map(|_| RecvCtx::boxed()).collect();
672 let send: Vec<Box<SendCtx>> =
673 (0..SEND_DEPTH).map(|_| Box::new(unsafe { std::mem::zeroed() })).collect();
674 let st = State {
675 ring,
676 recv,
677 send,
678 ready: VecDeque::new(),
679 ready_len: vec![0usize; RECV_DEPTH],
680 free_send: (0..SEND_DEPTH).collect(),
681 };
682 let me = Self { sock, fd, st: RefCell::new(st) };
683 if let Err(e) = me.submit_all_recv() {
684 return Err((me.sock, e));
685 }
686 Ok(me)
687 }
688
689 pub fn local_addr(&self) -> io::Result<SocketAddr> {
690 self.sock.local_addr()
691 }
692
693 pub fn raw_fd(&self) -> i32 {
697 self.fd
698 }
699
700 pub fn set_nonblocking(&self, nb: bool) -> io::Result<()> {
701 self.sock.set_nonblocking(nb)
702 }
703
704 fn submit_all_recv(&self) -> io::Result<()> {
705 let mut st = self.st.borrow_mut();
706 for i in 0..st.recv.len() {
707 self.push_recv(&mut st, i)?;
708 }
709 st.ring.submit()?;
710 Ok(())
711 }
712
713 fn push_recv(&self, st: &mut State, idx: usize) -> io::Result<()> {
714 st.recv[idx].refresh();
715 let msg: *mut libc::msghdr = std::ptr::addr_of_mut!(st.recv[idx].msghdr);
716 let e = opcode::RecvMsg::new(types::Fd(self.fd), msg)
717 .build()
718 .user_data(idx as u64);
719 unsafe {
722 st.ring
723 .submission()
724 .push(&e)
725 .map_err(|_| io::Error::other("io_uring SQ full (recv)"))?;
726 }
727 Ok(())
728 }
729
730 fn reap(&self, st: &mut State) {
731 let mut completed: Vec<(u64, i32)> = Vec::new();
732 for cqe in st.ring.completion() {
733 completed.push((cqe.user_data(), cqe.result()));
734 }
735 for (ud, res) in completed {
736 if ud >= SEND_TAG {
737 st.free_send.push((ud - SEND_TAG) as usize);
738 } else {
739 let idx = ud as usize;
740 if res >= 0 {
741 st.ready_len[idx] = res as usize;
742 st.ready.push_back(idx);
743 } else {
744 self.push_recv(st, idx).ok();
745 }
746 }
747 }
748 }
749
750 pub fn recv_with_kts(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr, Option<i128>)> {
751 let mut st = self.st.borrow_mut();
752 self.reap(&mut st);
753 if st.ready.is_empty() {
754 st.ring.submit()?;
755 self.reap(&mut st);
756 if st.ready.is_empty() {
757 return Err(io::Error::from(io::ErrorKind::WouldBlock));
758 }
759 }
760 let idx = st.ready.pop_front().unwrap();
761 let n = st.ready_len[idx];
762 let copy = n.min(buf.len());
763 let from;
764 let kts;
765 {
766 let ctx = &st.recv[idx];
767 buf[..copy].copy_from_slice(&ctx.buf[..copy]);
768 from = unsafe { sockaddr_to_socketaddr(&ctx.addr) }
769 .ok_or_else(|| io::Error::other("non-IP source"))?;
770 kts = parse_timestamp(&ctx.msghdr);
771 }
772 self.push_recv(&mut st, idx)?;
773 st.ring.submit()?;
774 Ok((copy, from, kts))
775 }
776
777 pub fn send_to(&self, data: &[u8], addr: SocketAddr) -> io::Result<usize> {
778 if data.len() > FRAME_CAP {
779 return Err(io::Error::other("datagram exceeds frame cap"));
780 }
781 let mut st = self.st.borrow_mut();
782 self.reap(&mut st);
783 if st.free_send.is_empty() {
784 st.ring.submit()?;
785 self.reap(&mut st);
786 if st.free_send.is_empty() {
787 return self.sock.send_to(data, addr);
788 }
789 }
790 let idx = st.free_send.pop().unwrap();
791 let (sa, sa_len) = socketaddr_to_sockaddr(addr);
792 {
793 let ctx = &mut st.send[idx];
794 ctx.buf[..data.len()].copy_from_slice(data);
795 ctx.addr = sa;
796 ctx.iov.iov_base = ctx.buf.as_mut_ptr() as *mut libc::c_void;
797 ctx.iov.iov_len = data.len();
798 ctx.msghdr.msg_name = std::ptr::addr_of_mut!(ctx.addr) as *mut libc::c_void;
799 ctx.msghdr.msg_namelen = sa_len;
800 ctx.msghdr.msg_iov = std::ptr::addr_of_mut!(ctx.iov);
801 ctx.msghdr.msg_iovlen = 1;
802 }
803 let msg: *const libc::msghdr = std::ptr::addr_of!(st.send[idx].msghdr);
804 let e = opcode::SendMsg::new(types::Fd(self.fd), msg)
805 .build()
806 .user_data(SEND_TAG | idx as u64);
807 unsafe {
810 st.ring
811 .submission()
812 .push(&e)
813 .map_err(|_| io::Error::other("io_uring SQ full (send)"))?;
814 }
815 st.ring.submit()?;
816 Ok(data.len())
817 }
818 }
819}
820
821#[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
844pub fn link_speed_bps(ifname: &str) -> Option<u64> {
845 let phys = ifname.strip_prefix("netmap:").unwrap_or(ifname);
848 if phys.starts_with("vale") || phys.contains(':') || phys.contains('{') {
849 return None;
850 }
851 #[cfg(target_os = "linux")]
852 {
853 let mbps: i64 = std::fs::read_to_string(format!("/sys/class/net/{phys}/speed"))
855 .ok()?
856 .trim()
857 .parse()
858 .ok()?;
859 (mbps > 0).then_some(mbps as u64 * 1_000_000)
860 }
861 #[cfg(any(target_os = "freebsd", target_os = "macos"))]
862 {
863 link_speed_baudrate(phys)
864 }
865}
866
867#[cfg(all(feature = "wire-locale", any(target_os = "freebsd", target_os = "macos")))]
870fn link_speed_baudrate(ifname: &str) -> Option<u64> {
871 use std::ffi::CStr;
872 let mut ifap: *mut libc::ifaddrs = std::ptr::null_mut();
873 if unsafe { libc::getifaddrs(&mut ifap) } != 0 {
874 return None;
875 }
876 let mut speed = None;
877 let mut cur = ifap;
878 while !cur.is_null() {
879 let ifa = unsafe { &*cur };
883 if !ifa.ifa_name.is_null() && !ifa.ifa_addr.is_null() && !ifa.ifa_data.is_null() {
884 let name = unsafe { CStr::from_ptr(ifa.ifa_name) }.to_string_lossy();
885 let family = unsafe { (*ifa.ifa_addr).sa_family } as i32;
886 if name == ifname && family == libc::AF_LINK {
887 let baud =
888 unsafe { (*(ifa.ifa_data as *const libc::if_data)).ifi_baudrate };
889 if baud > 0 {
890 #[cfg(target_os = "freebsd")]
892 {
893 speed = Some(baud);
894 }
895 #[cfg(target_os = "macos")]
896 {
897 speed = Some(u64::from(baud));
898 }
899 }
900 }
901 }
902 cur = ifa.ifa_next;
903 }
904 unsafe { libc::freeifaddrs(ifap) };
905 speed
906}
907
908#[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
914fn wire_min_link_bps() -> u64 {
915 std::env::var("SUBETHA_WIRE_MIN_GBPS")
916 .ok()
917 .and_then(|s| s.trim().parse::<f64>().ok())
918 .map(|g| (g * 1e9) as u64)
919 .unwrap_or(10_000_000_000)
920}
921
922#[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
927pub fn wire_gate_admits(ifname: &str) -> bool {
928 wire_link_fast_enough(ifname)
929}
930
931#[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
932fn wire_link_fast_enough(ifname: &str) -> bool {
933 let threshold = wire_min_link_bps();
934 if threshold == 0 {
935 return true;
936 }
937 link_speed_bps(ifname).map(|bps| bps >= threshold).unwrap_or(false)
938}
939
940#[cfg(all(feature = "wire-locale", any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
941mod wire_backend {
942 use super::*;
943 use std::cell::RefCell;
944 use std::net::{IpAddr, Ipv4Addr};
945
946 use crate::locale_wire::WireSocket;
947
948 const ETH_HDR: usize = 14;
949 const IP_HDR: usize = 20;
950 const UDP_HDR: usize = 8;
951 const HDRS: usize = ETH_HDR + IP_HDR + UDP_HDR; fn ipv4_csum(h: &[u8]) -> u16 {
954 let mut sum = 0u32;
955 let mut i = 0;
956 while i + 1 < h.len() {
957 sum += u16::from_be_bytes([h[i], h[i + 1]]) as u32;
958 i += 2;
959 }
960 while (sum >> 16) != 0 {
961 sum = (sum & 0xffff) + (sum >> 16);
962 }
963 !(sum as u16)
964 }
965
966 #[allow(clippy::too_many_arguments)]
969 fn build_frame(
970 dst_mac: [u8; 6],
971 src_mac: [u8; 6],
972 src_ip: [u8; 4],
973 dst_ip: [u8; 4],
974 src_port: u16,
975 dst_port: u16,
976 payload: &[u8],
977 out: &mut Vec<u8>,
978 ) {
979 out.clear();
980 out.resize(HDRS + payload.len(), 0);
981 out[0..6].copy_from_slice(&dst_mac);
982 out[6..12].copy_from_slice(&src_mac);
983 out[12..14].copy_from_slice(&0x0800u16.to_be_bytes());
984 out[14] = 0x45;
985 out[16..18].copy_from_slice(&((IP_HDR + UDP_HDR + payload.len()) as u16).to_be_bytes());
986 out[22] = 64;
987 out[23] = 17;
988 out[26..30].copy_from_slice(&src_ip);
989 out[30..34].copy_from_slice(&dst_ip);
990 let c = ipv4_csum(&out[14..34]);
991 out[24..26].copy_from_slice(&c.to_be_bytes());
992 out[34..36].copy_from_slice(&src_port.to_be_bytes());
993 out[36..38].copy_from_slice(&dst_port.to_be_bytes());
994 out[38..40].copy_from_slice(&((UDP_HDR + payload.len()) as u16).to_be_bytes());
995 out[42..].copy_from_slice(payload);
996 }
997
998 fn parse_frame(f: &[u8]) -> Option<([u8; 4], u16, u16, usize, usize)> {
1001 if f.len() < HDRS || f[12..14] != 0x0800u16.to_be_bytes() {
1002 return None;
1003 }
1004 let ihl = (f[14] & 0x0f) as usize * 4;
1005 if ihl < IP_HDR || f[23] != 17 {
1006 return None;
1007 }
1008 let l4 = ETH_HDR + ihl;
1009 if f.len() < l4 + UDP_HDR {
1010 return None;
1011 }
1012 let src_ip = [f[26], f[27], f[28], f[29]];
1013 let src_port = u16::from_be_bytes([f[l4], f[l4 + 1]]);
1014 let dst_port = u16::from_be_bytes([f[l4 + 2], f[l4 + 3]]);
1015 let udp_len = u16::from_be_bytes([f[l4 + 4], f[l4 + 5]]) as usize;
1016 let pstart = l4 + UDP_HDR;
1017 let plen = udp_len.saturating_sub(UDP_HDR).min(f.len() - pstart);
1018 Some((src_ip, src_port, dst_port, pstart, plen))
1019 }
1020
1021 fn parse_mac(s: &str) -> Option<[u8; 6]> {
1022 let mut mac = [0u8; 6];
1023 let parts: Vec<&str> = s.split(':').collect();
1024 if parts.len() != 6 {
1025 return None;
1026 }
1027 for (i, p) in parts.iter().enumerate() {
1028 mac[i] = u8::from_str_radix(p, 16).ok()?;
1029 }
1030 Some(mac)
1031 }
1032
1033 fn parse_ipv4(s: &str) -> Option<[u8; 4]> {
1034 s.parse::<Ipv4Addr>().ok().map(|a| a.octets())
1035 }
1036
1037 pub struct WireDgram {
1043 wire: RefCell<WireSocket>,
1044 scratch: RefCell<Vec<u8>>,
1045 local_ip: [u8; 4],
1046 local_mac: [u8; 6],
1047 peer_mac: [u8; 6],
1048 local_port: u16,
1049 }
1050
1051 impl WireDgram {
1052 pub fn from_env(local_port: u16) -> io::Result<Self> {
1053 let getv = |k: &str| -> io::Result<String> {
1054 std::env::var(k).map_err(|_| io::Error::other(format!("{k} unset")))
1055 };
1056 let ifname = getv("SUBETHA_WIRE_IFNAME")?;
1057 let local_ip = parse_ipv4(&getv("SUBETHA_WIRE_LOCAL_IP")?)
1058 .ok_or_else(|| io::Error::other("bad SUBETHA_WIRE_LOCAL_IP"))?;
1059 let local_mac = parse_mac(&getv("SUBETHA_WIRE_LOCAL_MAC")?)
1060 .ok_or_else(|| io::Error::other("bad SUBETHA_WIRE_LOCAL_MAC"))?;
1061 let peer_mac = parse_mac(&getv("SUBETHA_WIRE_PEER_MAC")?)
1062 .ok_or_else(|| io::Error::other("bad SUBETHA_WIRE_PEER_MAC"))?;
1063 let wire = WireSocket::bind(&ifname, 0)?;
1064 Ok(Self {
1065 wire: RefCell::new(wire),
1066 scratch: RefCell::new(Vec::with_capacity(HDRS + 2048)),
1067 local_ip,
1068 local_mac,
1069 peer_mac,
1070 local_port,
1071 })
1072 }
1073
1074 pub fn local_addr(&self) -> io::Result<SocketAddr> {
1075 Ok(SocketAddr::new(IpAddr::V4(Ipv4Addr::from(self.local_ip)), self.local_port))
1076 }
1077
1078 pub fn set_nonblocking(&self, _nb: bool) -> io::Result<()> {
1079 Ok(())
1081 }
1082
1083 pub fn send_to(&self, buf: &[u8], addr: SocketAddr) -> io::Result<usize> {
1084 let dst_ip = match addr.ip() {
1085 IpAddr::V4(v) => v.octets(),
1086 IpAddr::V6(_) => return Err(io::Error::other("wire backend is IPv4-only")),
1087 };
1088 let mut scratch = self.scratch.borrow_mut();
1089 build_frame(
1090 self.peer_mac,
1091 self.local_mac,
1092 self.local_ip,
1093 dst_ip,
1094 self.local_port,
1095 addr.port(),
1096 buf,
1097 &mut scratch,
1098 );
1099 self.wire.borrow_mut().send_frame(&scratch)?;
1100 Ok(buf.len())
1101 }
1102
1103 pub fn recv_with_kts(
1104 &self,
1105 buf: &mut [u8],
1106 ) -> io::Result<(usize, SocketAddr, Option<i128>)> {
1107 let mut fb = [0u8; 2048];
1108 let mut wire = self.wire.borrow_mut();
1109 loop {
1110 let n = wire.recv_frame(&mut fb, 0)?;
1112 if n == 0 {
1113 return Err(io::Error::from(io::ErrorKind::WouldBlock));
1114 }
1115 if let Some((src_ip, src_port, dst_port, pstart, plen)) = parse_frame(&fb[..n])
1118 && dst_port == self.local_port
1119 {
1120 let copy = plen.min(buf.len());
1121 buf[..copy].copy_from_slice(&fb[pstart..pstart + copy]);
1122 let from = SocketAddr::new(IpAddr::V4(Ipv4Addr::from(src_ip)), src_port);
1123 return Ok((copy, from, None));
1124 }
1125 }
1127 }
1128 }
1129}