tokio_fast_udp/
socket.rs

1// FastUdpSocket and FastUdpSocketBuilder for Unix (AsyncFd + libc syscalls).
2
3// ---------------------------------------------------------------------------
4// Unix implementation
5// ---------------------------------------------------------------------------
6#[cfg(unix)]
7mod unix_socket {
8    use std::{
9        io,
10        net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV6},
11        os::fd::{AsRawFd, FromRawFd, IntoRawFd, OwnedFd},
12    };
13
14    // -- platform alias --------------------------------------------------------
15    #[cfg(target_os = "linux")]
16    use linux_inner as plat;
17    use tokio::io::{Interest, unix::AsyncFd};
18    #[cfg(all(unix, not(target_os = "linux")))]
19    use unix_inner as plat;
20
21    use crate::{
22        capability::Capabilities,
23        ecn::Ecn,
24        item::{ReceiveItem, SendItem},
25    };
26
27    // -- FastUdpSocket ---------------------------------------------------------
28    pub struct FastUdpSocket {
29        fd: AsyncFd<OwnedFd>,
30        caps: Capabilities,
31    }
32
33    impl FastUdpSocket {
34        pub fn build(addr: SocketAddr) -> FastUdpSocketBuilder {
35            FastUdpSocketBuilder {
36                addr,
37                disable_gso: false,
38                disable_gro: false,
39                disable_sendmmsg: false,
40                disable_ecn: false,
41                max_batch: 64,
42                enable_reuse_port: false,
43            }
44        }
45    }
46
47    impl FastUdpSocket {
48        pub async fn send(&self, item: SendItem<'_>) -> io::Result<()> {
49            let items = [item];
50            self.send_many(&items).await?;
51            Ok(())
52        }
53
54        pub async fn send_many(&self, items: &[SendItem<'_>]) -> io::Result<usize> {
55            if items.is_empty() {
56                return Ok(0);
57            }
58            let caps = &self.caps;
59            let limit = items.len().min(caps.max_batch);
60            let items = &items[..limit];
61            let mut state = plat::SendState::default();
62            self.fd
63                .async_io(Interest::WRITABLE, |inner| {
64                    plat::send_batch(inner.as_raw_fd(), items, caps, &mut state)
65                })
66                .await
67        }
68
69        pub async fn receive(&self, item: &mut ReceiveItem<'_>) -> io::Result<()> {
70            self.receive_many(std::slice::from_mut(item)).await?;
71            Ok(())
72        }
73
74        pub async fn receive_many(&self, items: &mut [ReceiveItem<'_>]) -> io::Result<usize> {
75            if items.is_empty() {
76                return Ok(0);
77            }
78            let caps = &self.caps;
79            let limit = items.len().min(caps.max_batch);
80            let items = &mut items[..limit];
81            self.fd
82                .async_io(Interest::READABLE, |inner| plat::recv_batch(inner.as_raw_fd(), items, caps))
83                .await
84        }
85
86        pub fn capabilities(&self) -> &Capabilities {
87            &self.caps
88        }
89
90        pub fn local_addr(&self) -> io::Result<SocketAddr> {
91            socket_local_addr(self.fd.get_ref())
92        }
93    }
94
95    // -- FastUdpSocketBuilder --------------------------------------------------
96    pub struct FastUdpSocketBuilder {
97        addr: SocketAddr,
98        disable_gso: bool,
99        disable_gro: bool,
100        disable_sendmmsg: bool,
101        disable_ecn: bool,
102        max_batch: usize,
103        enable_reuse_port: bool,
104    }
105
106    impl FastUdpSocketBuilder {
107        pub fn gso(mut self, enable: bool) -> Self {
108            self.disable_gso = !enable;
109            self
110        }
111
112        pub fn gro(mut self, enable: bool) -> Self {
113            self.disable_gro = !enable;
114            self
115        }
116
117        pub fn sendmmsg(mut self, enable: bool) -> Self {
118            self.disable_sendmmsg = !enable;
119            self
120        }
121
122        pub fn ecn(mut self, enable: bool) -> Self {
123            self.disable_ecn = !enable;
124            self
125        }
126
127        pub fn reuse_port(mut self, enable: bool) -> Self {
128            self.enable_reuse_port = enable;
129            self
130        }
131
132        pub fn max_batch_size(mut self, n: usize) -> Self {
133            self.max_batch = n;
134            self
135        }
136
137        pub fn bind(self) -> io::Result<FastUdpSocket> {
138            let is_ipv6 = self.addr.is_ipv6();
139
140            let fd: OwnedFd = if self.enable_reuse_port {
141                create_socket_with_reuse_port(self.addr)?
142            } else {
143                let std_socket = std::net::UdpSocket::bind(self.addr)?;
144                std_socket.set_nonblocking(true)?;
145
146                let raw_fd = std_socket.into_raw_fd();
147                unsafe { OwnedFd::from_raw_fd(raw_fd) }
148            };
149
150            let config = BuildConfig {
151                disable_gso: self.disable_gso,
152                disable_gro: self.disable_gro,
153                disable_sendmmsg: self.disable_sendmmsg,
154                disable_ecn: self.disable_ecn,
155                max_batch: self.max_batch,
156            };
157
158            let caps = plat::setup_and_probe(&fd, &config, is_ipv6)?;
159            // SAFETY: `fd` is an owned, open UDP socket descriptor. `AsyncFd` takes
160            // ownership of it and holds it for its entire lifetime, so the descriptor
161            // stays valid and unchanged until the `AsyncFd` is dropped.
162            let async_fd = unsafe { AsyncFd::register(fd) }?;
163
164            Ok(FastUdpSocket {
165                fd: async_fd,
166                caps,
167            })
168        }
169    }
170
171    // -- shared syscall helpers (used by the platform inner modules via super::) --
172
173    /// Size of the ancillary-data buffer used for one datagram.
174    ///
175    /// The largest control-message set this crate produces is a GSO segment size
176    /// plus an ECN codepoint, which fits comfortably.
177    const CMSG_BUF_SIZE: usize = 128;
178
179    /// Alignment-safe outgoing ancillary-data buffer.
180    #[repr(align(8))]
181    struct CmsgBuf {
182        buf: [u8; CMSG_BUF_SIZE],
183        len: usize,
184    }
185
186    impl CmsgBuf {
187        fn new() -> Self {
188            Self {
189                buf: [0u8; CMSG_BUF_SIZE],
190                len: 0,
191            }
192        }
193
194        /// Append one control message, returning `false` if it does not fit.
195        #[cfg(any(target_os = "linux", target_vendor = "apple"))]
196        fn push(&mut self, level: libc::c_int, ty: libc::c_int, data: *const u8, data_len: usize) -> bool {
197            let space = unsafe { libc::CMSG_SPACE(data_len as u32) as usize };
198            if self.len + space > self.buf.len() {
199                return false;
200            }
201
202            // Control messages are laid out sequentially, each occupying
203            // `CMSG_SPACE(payload_len)` bytes starting at the previous message's
204            // end. The buffer is aligned via `#[repr(align(8))]`.
205            let header = unsafe { self.buf.as_mut_ptr().add(self.len) as *mut libc::cmsghdr };
206            unsafe {
207                (*header).cmsg_len = libc::CMSG_LEN(data_len as u32) as _;
208                (*header).cmsg_level = level;
209                (*header).cmsg_type = ty;
210                std::ptr::copy_nonoverlapping(data, libc::CMSG_DATA(header), data_len);
211            }
212            self.len += space;
213            true
214        }
215
216        /// Pointer to hand to `sendmsg`/`sendmmsg` as the control buffer.
217        ///
218        /// Returns a null pointer when no control message was appended. POSIX
219        /// requires `msg_control` to be null when `msg_controllen` is zero, and
220        /// XNU's `sendit()` rejects a non-null `msg_control` whose
221        /// `msg_controllen` is smaller than `sizeof(struct cmsghdr)` with
222        /// `EINVAL`, so a non-null pointer to an empty buffer makes every
223        /// ancillary-data-free send fail on macOS/iOS.
224        fn as_ptr(&self) -> *const libc::c_void {
225            if self.len == 0 {
226                std::ptr::null()
227            } else {
228                self.buf.as_ptr() as *const libc::c_void
229            }
230        }
231
232        fn len(&self) -> usize {
233            self.len
234        }
235    }
236
237    /// Alignment-safe buffer for received ancillary data.
238    #[repr(align(8))]
239    struct RecvCmsgBuf([u8; CMSG_BUF_SIZE]);
240
241    impl RecvCmsgBuf {
242        fn new() -> Self {
243            RecvCmsgBuf([0u8; CMSG_BUF_SIZE])
244        }
245    }
246
247    /// Reject a `segment_size` of zero up front: it is meaningless and would
248    /// otherwise make `slice::chunks(0)` panic inside the send path.
249    fn reject_zero_segment_size(items: &[SendItem<'_>]) -> io::Result<()> {
250        if items.iter().any(|item| item.segment_size == Some(0)) {
251            return Err(io::Error::new(
252                io::ErrorKind::InvalidInput,
253                "segment_size must be greater than zero",
254            ));
255        }
256        Ok(())
257    }
258
259    /// Set a boolean socket option (an `int` read as `0`/`1`).
260    fn set_bool_opt(fd: std::os::fd::RawFd, level: libc::c_int, name: libc::c_int) -> io::Result<()> {
261        let enable: libc::c_int = 1;
262        let ret = unsafe {
263            libc::setsockopt(
264                fd,
265                level,
266                name,
267                &enable as *const libc::c_int as *const libc::c_void,
268                std::mem::size_of::<libc::c_int>() as libc::socklen_t,
269            )
270        };
271        if ret < 0 {
272            return Err(io::Error::last_os_error());
273        }
274        Ok(())
275    }
276
277    /// Read an `int`-valued socket option.
278    #[cfg(target_os = "linux")]
279    fn get_int_opt(fd: std::os::fd::RawFd, level: libc::c_int, name: libc::c_int) -> io::Result<libc::c_int> {
280        let mut value: libc::c_int = 0;
281        let mut len = std::mem::size_of::<libc::c_int>() as libc::socklen_t;
282        let ret =
283            unsafe { libc::getsockopt(fd, level, name, &mut value as *mut libc::c_int as *mut libc::c_void, &mut len) };
284        if ret < 0 {
285            return Err(io::Error::last_os_error());
286        }
287        Ok(value)
288    }
289
290    fn create_socket_with_reuse_port(addr: SocketAddr) -> io::Result<OwnedFd> {
291        let domain = if addr.is_ipv6() { libc::AF_INET6 } else { libc::AF_INET };
292
293        // `SOCK_NONBLOCK` / `SOCK_CLOEXEC` are not accepted by `socket()` on
294        // every Unix (notably macOS), so create a plain socket and apply the
295        // flags portably via `fcntl`.
296        let raw_fd = unsafe { libc::socket(domain, libc::SOCK_DGRAM, 0) };
297        if raw_fd < 0 {
298            return Err(io::Error::last_os_error());
299        }
300        let fd = unsafe { OwnedFd::from_raw_fd(raw_fd) };
301
302        set_cloexec(&fd)?;
303        set_nonblocking(&fd)?;
304        set_bool_opt(fd.as_raw_fd(), libc::SOL_SOCKET, libc::SO_REUSEPORT)?;
305
306        let mut sock_addr: libc::sockaddr_storage = unsafe { std::mem::zeroed() };
307        let addr_len = fill_sockaddr(&addr, &mut sock_addr);
308        let ret = unsafe { libc::bind(fd.as_raw_fd(), &sock_addr as *const _ as *const libc::sockaddr, addr_len) };
309        if ret < 0 {
310            return Err(io::Error::last_os_error());
311        }
312
313        Ok(fd)
314    }
315
316    /// Set `FD_CLOEXEC` on `fd` using the portable `fcntl` interface.
317    fn set_cloexec(fd: &OwnedFd) -> io::Result<()> {
318        let ret = unsafe { libc::fcntl(fd.as_raw_fd(), libc::F_SETFD, libc::FD_CLOEXEC) };
319        if ret < 0 {
320            return Err(io::Error::last_os_error());
321        }
322        Ok(())
323    }
324
325    /// Put `fd` into non-blocking mode using the portable `fcntl` interface.
326    fn set_nonblocking(fd: &OwnedFd) -> io::Result<()> {
327        let flags = unsafe { libc::fcntl(fd.as_raw_fd(), libc::F_GETFL) };
328        if flags < 0 {
329            return Err(io::Error::last_os_error());
330        }
331        let ret = unsafe { libc::fcntl(fd.as_raw_fd(), libc::F_SETFL, flags | libc::O_NONBLOCK) };
332        if ret < 0 {
333            return Err(io::Error::last_os_error());
334        }
335        Ok(())
336    }
337
338    #[cfg(any(target_os = "linux", target_vendor = "apple"))]
339    fn ecn_to_tos(ecn: Ecn) -> u8 {
340        ecn.to_tos_bits()
341    }
342
343    /// Fill a `sockaddr_storage` from a `SocketAddr`, returning its length.
344    fn fill_sockaddr(addr: &SocketAddr, ss: &mut libc::sockaddr_storage) -> libc::socklen_t {
345        unsafe {
346            std::ptr::write_bytes(ss as *mut _ as *mut u8, 0, std::mem::size_of::<libc::sockaddr_storage>());
347        }
348        match addr {
349            SocketAddr::V4(v4) => {
350                let sin = ss as *mut _ as *mut libc::sockaddr_in;
351                unsafe {
352                    (*sin).sin_family = libc::AF_INET as libc::sa_family_t;
353                    (*sin).sin_port = v4.port().to_be();
354                    (*sin).sin_addr.s_addr = u32::from_be_bytes(v4.ip().octets()).to_be();
355                }
356                // BSD-style sockets carry their own length field.
357                #[cfg(any(
358                    target_vendor = "apple",
359                    target_os = "freebsd",
360                    target_os = "dragonfly",
361                    target_os = "netbsd",
362                    target_os = "openbsd"
363                ))]
364                unsafe {
365                    (*sin).sin_len = std::mem::size_of::<libc::sockaddr_in>() as u8;
366                }
367                std::mem::size_of::<libc::sockaddr_in>() as libc::socklen_t
368            }
369            SocketAddr::V6(v6) => {
370                let sin6 = ss as *mut _ as *mut libc::sockaddr_in6;
371                unsafe {
372                    (*sin6).sin6_family = libc::AF_INET6 as libc::sa_family_t;
373                    (*sin6).sin6_port = v6.port().to_be();
374                    (*sin6).sin6_addr.s6_addr = v6.ip().octets();
375                    (*sin6).sin6_flowinfo = v6.flowinfo().to_be();
376                    (*sin6).sin6_scope_id = v6.scope_id();
377                }
378                #[cfg(any(
379                    target_vendor = "apple",
380                    target_os = "freebsd",
381                    target_os = "dragonfly",
382                    target_os = "netbsd",
383                    target_os = "openbsd"
384                ))]
385                unsafe {
386                    (*sin6).sin6_len = std::mem::size_of::<libc::sockaddr_in6>() as u8;
387                }
388                std::mem::size_of::<libc::sockaddr_in6>() as libc::socklen_t
389            }
390        }
391    }
392
393    /// Convert a `sockaddr_storage` into a `SocketAddr`.
394    ///
395    /// Addresses of an unknown family fall back to `0.0.0.0:0`.
396    fn storage_to_addr(ss: &libc::sockaddr_storage) -> SocketAddr {
397        match ss.ss_family as libc::c_int {
398            libc::AF_INET => {
399                let sin = unsafe { &*(ss as *const _ as *const libc::sockaddr_in) };
400                let ip = Ipv4Addr::from(sin.sin_addr.s_addr.to_ne_bytes());
401                SocketAddr::new(IpAddr::V4(ip), u16::from_be(sin.sin_port))
402            }
403            libc::AF_INET6 => {
404                let sin6 = unsafe { &*(ss as *const _ as *const libc::sockaddr_in6) };
405                let ip = Ipv6Addr::from(sin6.sin6_addr.s6_addr);
406                SocketAddr::V6(SocketAddrV6::new(
407                    ip,
408                    u16::from_be(sin6.sin6_port),
409                    u32::from_be(sin6.sin6_flowinfo),
410                    sin6.sin6_scope_id,
411                ))
412            }
413            _ => SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0),
414        }
415    }
416
417    /// Send a single datagram to `destination`, with optional ancillary data.
418    fn sendmsg_raw(fd: std::os::fd::RawFd, data: &[u8], destination: &SocketAddr, cmsgs: &CmsgBuf) -> io::Result<()> {
419        // `sendmsg` only reads the payload, so the cast from a shared slice to a
420        // mutable pointer is sound.
421        let mut iov = libc::iovec {
422            iov_base: data.as_ptr() as *mut libc::c_void,
423            iov_len: data.len(),
424        };
425
426        let mut addr: libc::sockaddr_storage = unsafe { std::mem::zeroed() };
427        let addr_len = fill_sockaddr(destination, &mut addr);
428
429        let mut mhdr: libc::msghdr = unsafe { std::mem::zeroed() };
430        mhdr.msg_name = &mut addr as *mut _ as *mut libc::c_void;
431        mhdr.msg_namelen = addr_len;
432        mhdr.msg_iov = &mut iov as *mut libc::iovec;
433        mhdr.msg_iovlen = 1;
434        mhdr.msg_control = cmsgs.as_ptr() as *mut libc::c_void;
435        mhdr.msg_controllen = cmsgs.len() as _;
436
437        let ret = unsafe { libc::sendmsg(fd, &mhdr, 0) };
438        if ret < 0 {
439            return Err(io::Error::last_os_error());
440        }
441        Ok(())
442    }
443
444    /// Receive one datagram into `buf`.
445    ///
446    /// Returns the number of bytes received, the source address, and the decoded
447    /// ancillary data: the GRO segment size and the ECN codepoint, if present.
448    fn recvmsg_raw(
449        fd: std::os::fd::RawFd,
450        buf: &mut [u8],
451        cmsg_buf: &mut RecvCmsgBuf,
452        ecn_enabled: bool,
453    ) -> io::Result<(usize, SocketAddr, Option<u32>, Option<Ecn>)> {
454        let mut iov = libc::iovec {
455            iov_base: buf.as_mut_ptr() as *mut libc::c_void,
456            iov_len: buf.len(),
457        };
458        let mut addr: libc::sockaddr_storage = unsafe { std::mem::zeroed() };
459
460        let mut mhdr: libc::msghdr = unsafe { std::mem::zeroed() };
461        mhdr.msg_name = &mut addr as *mut _ as *mut libc::c_void;
462        mhdr.msg_namelen = std::mem::size_of::<libc::sockaddr_storage>() as _;
463        mhdr.msg_iov = &mut iov as *mut libc::iovec;
464        mhdr.msg_iovlen = 1;
465        mhdr.msg_control = cmsg_buf.0.as_mut_ptr() as *mut libc::c_void;
466        mhdr.msg_controllen = cmsg_buf.0.len() as _;
467
468        let ret = unsafe { libc::recvmsg(fd, &mut mhdr, 0) };
469        if ret < 0 {
470            return Err(io::Error::last_os_error());
471        }
472
473        let (gro_seg, ecn) = decode_cmsgs(&mhdr, ecn_enabled);
474        Ok((ret as usize, storage_to_addr(&addr), gro_seg, ecn))
475    }
476
477    /// Decode the GRO/ECN ancillary data attached to a received datagram.
478    ///
479    /// GRO segment sizes and ECN codepoints are read straight from the kernel's
480    /// `cmsghdr` layout, since the levels/types and payload widths differ across
481    /// platforms (and across Apple OS versions for IPv6 traffic class).
482    fn decode_cmsgs(mhdr: &libc::msghdr, ecn_enabled: bool) -> (Option<u32>, Option<Ecn>) {
483        #[cfg_attr(not(target_os = "linux"), allow(unused_mut))]
484        let mut gro_seg = None;
485        let mut ecn = None;
486
487        let mut cmsg = unsafe { libc::CMSG_FIRSTHDR(mhdr) };
488        while !cmsg.is_null() {
489            let level = unsafe { (*cmsg).cmsg_level };
490            let ty = unsafe { (*cmsg).cmsg_type };
491            let data = unsafe { libc::CMSG_DATA(cmsg) };
492            let header_len = unsafe { libc::CMSG_LEN(0) } as usize;
493            let payload_len = (unsafe { (*cmsg).cmsg_len } as usize).saturating_sub(header_len);
494
495            // UDP Generic Receive Offload (Linux): the segment size is an `int`.
496            #[cfg(target_os = "linux")]
497            if level == libc::SOL_UDP && ty == libc::UDP_GRO && gro_seg.is_none() && payload_len >= 4 {
498                gro_seg = Some(unsafe { std::ptr::read_unaligned(data as *const libc::c_int) } as u32);
499            }
500
501            if ecn_enabled && ecn.is_none() {
502                // IPv4: `IP_TOS` on Linux, `IP_RECVTOS` on the BSDs and Apple,
503                // reported as a single byte.
504                if level == libc::IPPROTO_IP && (ty == libc::IP_TOS || ty == libc::IP_RECVTOS) && payload_len >= 1 {
505                    ecn = Some(Ecn::from_tos_bits(unsafe { std::ptr::read_unaligned(data as *const u8) }));
506                }
507                // IPv6: `IPV6_TCLASS` is an `int`, but Apple has historically
508                // reported it as a single byte (Apple bug 48761855).
509                if level == libc::IPPROTO_IPV6 && ty == libc::IPV6_TCLASS {
510                    if payload_len >= std::mem::size_of::<libc::c_int>() {
511                        let tc = unsafe { std::ptr::read_unaligned(data as *const libc::c_int) };
512                        ecn = Some(Ecn::from_tos_bits(tc as u8));
513                    } else if payload_len >= 1 {
514                        ecn = Some(Ecn::from_tos_bits(unsafe { std::ptr::read_unaligned(data as *const u8) }));
515                    }
516                }
517            }
518
519            cmsg = unsafe { libc::CMSG_NXTHDR(mhdr, cmsg) };
520        }
521
522        (gro_seg, ecn)
523    }
524
525    /// Append the UDP GSO segment-size control message.
526    #[cfg(target_os = "linux")]
527    fn push_gso_cmsg(cmsgs: &mut CmsgBuf, segment_size: u16) -> bool {
528        cmsgs.push(
529            libc::SOL_UDP,
530            libc::UDP_SEGMENT,
531            &segment_size as *const u16 as *const u8,
532            std::mem::size_of::<u16>(),
533        )
534    }
535
536    /// Append the outgoing ECN codepoint control message.
537    #[cfg(any(target_os = "linux", target_vendor = "apple"))]
538    fn push_ecn_cmsg(cmsgs: &mut CmsgBuf, ecn: Ecn, is_v6: bool) -> bool {
539        if is_v6 {
540            let value = ecn_to_tos(ecn) as libc::c_int;
541            cmsgs.push(
542                libc::IPPROTO_IPV6,
543                libc::IPV6_TCLASS,
544                &value as *const libc::c_int as *const u8,
545                std::mem::size_of::<libc::c_int>(),
546            )
547        } else if cfg!(target_vendor = "apple") {
548            // Apple expects an `int`-valued `IP_TOS`.
549            let value = ecn_to_tos(ecn) as libc::c_int;
550            cmsgs.push(
551                libc::IPPROTO_IP,
552                libc::IP_TOS,
553                &value as *const libc::c_int as *const u8,
554                std::mem::size_of::<libc::c_int>(),
555            )
556        } else {
557            // Linux requires exactly one byte for `IP_TOS`.
558            let value = ecn_to_tos(ecn);
559            cmsgs.push(libc::IPPROTO_IP, libc::IP_TOS, &value as *const u8, std::mem::size_of::<u8>())
560        }
561    }
562
563    fn socket_local_addr(fd: &OwnedFd) -> io::Result<SocketAddr> {
564        let mut addr: libc::sockaddr_storage = unsafe { std::mem::zeroed() };
565        let mut len = std::mem::size_of::<libc::sockaddr_storage>() as libc::socklen_t;
566        let ret = unsafe { libc::getsockname(fd.as_raw_fd(), &mut addr as *mut _ as *mut libc::sockaddr, &mut len) };
567        if ret < 0 {
568            return Err(io::Error::last_os_error());
569        }
570        Ok(storage_to_addr(&addr))
571    }
572
573    // Fields for Linux-only offload features are still populated by the
574    // builder on every platform, but only the Linux path reads them.
575    #[cfg_attr(not(target_os = "linux"), allow(dead_code))]
576    pub(crate) struct BuildConfig {
577        pub disable_gso: bool,
578        pub disable_gro: bool,
579        pub disable_sendmmsg: bool,
580        pub disable_ecn: bool,
581        pub max_batch: usize,
582    }
583
584    // -- Linux fast path ----------------------------------------------------
585    #[cfg(target_os = "linux")]
586    mod linux_inner {
587        use std::{
588            io,
589            os::fd::{AsRawFd, OwnedFd},
590        };
591
592        use super::{
593            BuildConfig, CmsgBuf, RecvCmsgBuf, fill_sockaddr, get_int_opt, push_ecn_cmsg, push_gso_cmsg, recvmsg_raw,
594            reject_zero_segment_size, sendmsg_raw, set_bool_opt,
595        };
596        use crate::{
597            capability::Capabilities,
598            item::{ReceiveItem, SendItem},
599        };
600
601        /// Per-call state threaded through [`send_batch`] so it survives the
602        /// retries performed by `AsyncFd::async_io` (GSO `EINVAL` back-off and
603        /// manual-chunking progress).
604        #[derive(Default)]
605        pub struct SendState {
606            pub gso_einval: bool,
607            pub chunk_offset: usize,
608        }
609
610        fn send_one(fd: std::os::fd::RawFd, item: &SendItem, ecn_enabled: bool, gso_enabled: bool) -> io::Result<()> {
611            let mut cmsgs = CmsgBuf::new();
612
613            if gso_enabled
614                && let Some(segment_size) = item.segment_size
615                && !push_gso_cmsg(&mut cmsgs, segment_size)
616            {
617                return Err(io::Error::other("control message buffer overflow"));
618            }
619
620            if ecn_enabled
621                && let Some(ecn) = item.ecn
622                && !push_ecn_cmsg(&mut cmsgs, ecn, item.destination.is_ipv6())
623            {
624                return Err(io::Error::other("control message buffer overflow"));
625            }
626
627            sendmsg_raw(fd, item.data, &item.destination, &cmsgs)
628        }
629
630        const SENDMMSG_MAX: usize = 64;
631
632        fn sendmmsg_batch(fd: std::os::fd::RawFd, items: &[SendItem], caps: &Capabilities) -> io::Result<usize> {
633            let mut total_sent = 0;
634
635            while total_sent < items.len() {
636                let chunk_len = (items.len() - total_sent).min(SENDMMSG_MAX);
637                let chunk = &items[total_sent..total_sent + chunk_len];
638
639                let mut addrs: [libc::sockaddr_storage; SENDMMSG_MAX] = unsafe { std::mem::zeroed() };
640                let mut iovs: [libc::iovec; SENDMMSG_MAX] = unsafe { std::mem::zeroed() };
641                let mut cmsg_bufs: [CmsgBuf; SENDMMSG_MAX] = std::array::from_fn(|_| CmsgBuf::new());
642                let mut msgs: [libc::mmsghdr; SENDMMSG_MAX] = unsafe { std::mem::zeroed() };
643
644                for i in 0..chunk_len {
645                    let item = &chunk[i];
646                    let addr_len = fill_sockaddr(&item.destination, &mut addrs[i]);
647
648                    if caps.ecn
649                        && let Some(ecn) = item.ecn
650                        && !push_ecn_cmsg(&mut cmsg_bufs[i], ecn, item.destination.is_ipv6())
651                    {
652                        return Err(io::Error::other("control message buffer overflow"));
653                    }
654
655                    iovs[i].iov_base = item.data.as_ptr() as *mut _;
656                    iovs[i].iov_len = item.data.len();
657
658                    let mhdr = &mut msgs[i].msg_hdr;
659                    mhdr.msg_name = &mut addrs[i] as *mut _ as *mut _;
660                    mhdr.msg_namelen = addr_len;
661                    mhdr.msg_iov = &mut iovs[i] as *mut _ as *mut _;
662                    mhdr.msg_iovlen = 1;
663                    mhdr.msg_control = cmsg_bufs[i].as_ptr() as *mut _;
664                    mhdr.msg_controllen = cmsg_bufs[i].len() as _;
665                    mhdr.msg_flags = 0;
666                }
667
668                let ret = unsafe { libc::sendmmsg(fd, msgs.as_mut_ptr(), chunk_len as u32, 0) };
669                if ret < 0 {
670                    let e = io::Error::last_os_error();
671                    if total_sent > 0 && e.kind() == io::ErrorKind::WouldBlock {
672                        return Ok(total_sent);
673                    }
674                    return Err(e);
675                }
676                let sent = ret as usize;
677                total_sent += sent;
678                if sent < chunk_len {
679                    break;
680                }
681            }
682
683            Ok(total_sent)
684        }
685
686        pub fn send_batch(
687            fd: std::os::fd::RawFd,
688            items: &[SendItem<'_>],
689            caps: &Capabilities,
690            state: &mut SendState,
691        ) -> io::Result<usize> {
692            if items.is_empty() {
693                return Ok(0);
694            }
695
696            reject_zero_segment_size(items)?;
697
698            // Single-item GSO fast path — only attempted once
699            if items.len() == 1 && items[0].segment_size.is_some() && caps.gso && !state.gso_einval {
700                match send_one(fd, &items[0], caps.ecn, true) {
701                    Ok(()) => {
702                        state.chunk_offset = 0;
703                        return Ok(1);
704                    }
705                    Err(e) if e.raw_os_error() == Some(libc::EINVAL) => {
706                        state.gso_einval = true;
707                        // Fall through — GSO rejected by this interface,
708                        // manually chunk below
709                    }
710                    Err(e) if e.kind() == io::ErrorKind::WouldBlock => return Err(e),
711                    Err(e) => return Err(e),
712                }
713            }
714
715            // Multi-item non-segmented batch via sendmmsg
716            if caps.sendmmsg && items.len() > 1 && !items.iter().any(|i| i.segment_size.is_some()) {
717                return sendmmsg_batch(fd, items, caps);
718            }
719
720            let mut count = 0usize;
721            for item in items {
722                if let Some(seg) = item.segment_size {
723                    // Manual chunking — GSO unavailable or rejected
724                    let seg = seg as usize;
725                    let data_slice = &item.data[state.chunk_offset..];
726                    for chunk in data_slice.chunks(seg) {
727                        let mut chunk_item = SendItem::new(item.destination, chunk);
728                        if let Some(ecn) = item.ecn {
729                            chunk_item = chunk_item.ecn(ecn);
730                        }
731                        match send_one(fd, &chunk_item, caps.ecn, false) {
732                            Ok(()) => state.chunk_offset += chunk.len(),
733                            Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
734                                if count == 0 && state.chunk_offset == 0 {
735                                    return Err(e);
736                                }
737                                return Ok(count);
738                            }
739                            Err(e) => return Err(e),
740                        }
741                    }
742                    state.chunk_offset = 0;
743                    count += 1;
744                } else {
745                    match send_one(fd, item, caps.ecn, false) {
746                        Ok(()) => count += 1,
747                        Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
748                            if count == 0 {
749                                return Err(e);
750                            }
751                            return Ok(count);
752                        }
753                        Err(e) => return Err(e),
754                    }
755                }
756            }
757            Ok(count)
758        }
759
760        fn recv_one(fd: std::os::fd::RawFd, item: &mut ReceiveItem<'_>, ecn_enabled: bool) -> io::Result<()> {
761            let mut cmsg_buf = RecvCmsgBuf::new();
762            let (len, source, _gro_seg, ecn) = recvmsg_raw(fd, item.buf, &mut cmsg_buf, ecn_enabled)?;
763
764            item.len = len;
765            item.source = source;
766            item.ecn = ecn;
767
768            Ok(())
769        }
770
771        fn recv_gro(fd: std::os::fd::RawFd, items: &mut [ReceiveItem<'_>], ecn_enabled: bool) -> io::Result<usize> {
772            let mut gro_buf = [0u8; 65536];
773            let mut cmsg_buf = RecvCmsgBuf::new();
774
775            let (total, source, seg_size, ecn_val) = recvmsg_raw(fd, &mut gro_buf, &mut cmsg_buf, ecn_enabled)?;
776
777            if seg_size.is_none() {
778                if items.is_empty() {
779                    return Ok(0);
780                }
781                let copy_len = total.min(items[0].buf.len());
782                items[0].buf[..copy_len].copy_from_slice(&gro_buf[..copy_len]);
783                items[0].len = copy_len;
784                items[0].source = source;
785                items[0].ecn = ecn_val;
786                return Ok(1);
787            }
788
789            let seg = seg_size.unwrap() as usize;
790            let mut offset = 0;
791            let mut count = 0;
792            for item in items.iter_mut() {
793                if offset >= total {
794                    break;
795                }
796                let seg_end = (offset + seg).min(total);
797                let copy_len = (seg_end - offset).min(item.buf.len());
798                item.buf[..copy_len].copy_from_slice(&gro_buf[offset..offset + copy_len]);
799                item.len = copy_len;
800                item.source = source;
801                item.ecn = ecn_val;
802                count += 1;
803                offset = seg_end;
804            }
805            Ok(count)
806        }
807
808        pub fn recv_batch(
809            fd: std::os::fd::RawFd,
810            items: &mut [ReceiveItem<'_>],
811            caps: &Capabilities,
812        ) -> io::Result<usize> {
813            if items.is_empty() {
814                return Ok(0);
815            }
816
817            if caps.gro {
818                return recv_gro(fd, items, caps.ecn);
819            }
820
821            let mut count = 0;
822            for item in items.iter_mut() {
823                match recv_one(fd, item, caps.ecn) {
824                    Ok(()) => count += 1,
825                    Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
826                        if count == 0 {
827                            return Err(e);
828                        }
829                        return Ok(count);
830                    }
831                    Err(e) => return Err(e),
832                }
833            }
834            Ok(count)
835        }
836
837        fn probe_gso(fd: &OwnedFd) -> bool {
838            get_int_opt(fd.as_raw_fd(), libc::SOL_UDP, libc::UDP_SEGMENT).is_ok()
839        }
840
841        fn probe_gro(fd: &OwnedFd) -> bool {
842            set_bool_opt(fd.as_raw_fd(), libc::IPPROTO_UDP, libc::UDP_GRO).is_ok()
843        }
844
845        fn enable_ecn_recv(fd: &OwnedFd, is_ipv6: bool) -> bool {
846            let (level, name) = if is_ipv6 {
847                (libc::IPPROTO_IPV6, libc::IPV6_RECVTCLASS)
848            } else {
849                (libc::IPPROTO_IP, libc::IP_RECVTOS)
850            };
851            set_bool_opt(fd.as_raw_fd(), level, name).is_ok()
852        }
853
854        pub fn setup_and_probe(fd: &OwnedFd, config: &BuildConfig, is_ipv6: bool) -> io::Result<Capabilities> {
855            set_bool_opt(fd.as_raw_fd(), libc::SOL_SOCKET, libc::SO_REUSEADDR)?;
856
857            let gso = if config.disable_gso { false } else { probe_gso(fd) };
858            let gro = if config.disable_gro { false } else { probe_gro(fd) };
859            let sendmmsg = !config.disable_sendmmsg;
860            let ecn = if config.disable_ecn {
861                false
862            } else {
863                enable_ecn_recv(fd, is_ipv6)
864            };
865
866            Ok(Capabilities {
867                gso,
868                gro,
869                sendmmsg,
870                ecn,
871                max_batch: config.max_batch,
872            })
873        }
874    }
875
876    // -- Generic Unix fallback -------------------------------------------------
877    #[cfg(all(unix, not(target_os = "linux")))]
878    mod unix_inner {
879        use std::{
880            io,
881            os::fd::{AsRawFd, OwnedFd},
882        };
883
884        #[cfg(target_vendor = "apple")]
885        use super::push_ecn_cmsg;
886        use super::{
887            BuildConfig, CmsgBuf, RecvCmsgBuf, recvmsg_raw, reject_zero_segment_size, sendmsg_raw, set_bool_opt,
888        };
889        use crate::{
890            capability::Capabilities,
891            item::{ReceiveItem, SendItem},
892        };
893
894        /// Per-call state threaded through [`send_batch`] so that partially
895        /// completed user-space segmentation survives the retries performed by
896        /// `AsyncFd::async_io`.
897        #[derive(Default)]
898        pub struct SendState {
899            pub chunk_offset: usize,
900        }
901
902        fn send_one(fd: std::os::fd::RawFd, item: &SendItem, ecn_enabled: bool) -> io::Result<()> {
903            #[cfg_attr(not(target_vendor = "apple"), allow(unused_mut))]
904            let mut cmsgs = CmsgBuf::new();
905
906            // ECN is only implemented on Apple targets among non-Linux Unixes.
907            #[cfg(target_vendor = "apple")]
908            if ecn_enabled
909                && let Some(ecn) = item.ecn
910                && !push_ecn_cmsg(&mut cmsgs, ecn, item.destination.is_ipv6())
911            {
912                return Err(io::Error::other("control message buffer overflow"));
913            }
914            #[cfg(not(target_vendor = "apple"))]
915            let _ = ecn_enabled;
916
917            sendmsg_raw(fd, item.data, &item.destination, &cmsgs)
918        }
919
920        pub fn send_batch(
921            fd: std::os::fd::RawFd,
922            items: &[SendItem<'_>],
923            caps: &Capabilities,
924            state: &mut SendState,
925        ) -> io::Result<usize> {
926            if items.is_empty() {
927                return Ok(0);
928            }
929
930            reject_zero_segment_size(items)?;
931
932            let mut count = 0;
933            for item in items {
934                if let Some(seg) = item.segment_size {
935                    // No GSO on this platform: emulate it by splitting the
936                    // payload into one datagram per segment in user space.
937                    let seg = seg as usize;
938                    let data_slice = &item.data[state.chunk_offset..];
939                    for chunk in data_slice.chunks(seg) {
940                        let mut chunk_item = SendItem::new(item.destination, chunk);
941                        if let Some(ecn) = item.ecn {
942                            chunk_item = chunk_item.ecn(ecn);
943                        }
944                        match send_one(fd, &chunk_item, caps.ecn) {
945                            Ok(()) => state.chunk_offset += chunk.len(),
946                            Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
947                                if count == 0 && state.chunk_offset == 0 {
948                                    return Err(e);
949                                }
950                                return Ok(count);
951                            }
952                            Err(e) => return Err(e),
953                        }
954                    }
955                    state.chunk_offset = 0;
956                    count += 1;
957                } else {
958                    match send_one(fd, item, caps.ecn) {
959                        Ok(()) => count += 1,
960                        Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
961                            if count == 0 {
962                                return Err(e);
963                            }
964                            return Ok(count);
965                        }
966                        Err(e) => return Err(e),
967                    }
968                }
969            }
970            Ok(count)
971        }
972
973        fn recv_one(fd: std::os::fd::RawFd, item: &mut ReceiveItem<'_>, ecn_enabled: bool) -> io::Result<()> {
974            let mut cmsg_buf = RecvCmsgBuf::new();
975            let (len, source, _gro_seg, ecn) = recvmsg_raw(fd, item.buf, &mut cmsg_buf, ecn_enabled)?;
976
977            item.len = len;
978            item.source = source;
979            item.ecn = ecn;
980
981            Ok(())
982        }
983
984        pub fn recv_batch(
985            fd: std::os::fd::RawFd,
986            items: &mut [ReceiveItem<'_>],
987            caps: &Capabilities,
988        ) -> io::Result<usize> {
989            if items.is_empty() {
990                return Ok(0);
991            }
992
993            let mut count = 0;
994            for item in items.iter_mut() {
995                match recv_one(fd, item, caps.ecn) {
996                    Ok(()) => count += 1,
997                    Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
998                        if count == 0 {
999                            return Err(e);
1000                        }
1001                        return Ok(count);
1002                    }
1003                    Err(e) => return Err(e),
1004                }
1005            }
1006            Ok(count)
1007        }
1008
1009        /// Enable reception of ECN codepoints.
1010        ///
1011        /// An IPv4 socket reports the TOS byte via `IP_RECVTOS`; an IPv6 socket
1012        /// reports the traffic class via `IPV6_RECVTCLASS`, which also covers
1013        /// IPv4-mapped traffic. Setting `IP_RECVTOS` on a dual-stack IPv6 socket
1014        /// fails on macOS, so only the option matching the socket family is set.
1015        #[cfg(target_vendor = "apple")]
1016        fn enable_ecn_recv(fd: &OwnedFd, is_ipv6: bool) -> bool {
1017            let (level, name) = if is_ipv6 {
1018                (libc::IPPROTO_IPV6, libc::IPV6_RECVTCLASS)
1019            } else {
1020                (libc::IPPROTO_IP, libc::IP_RECVTOS)
1021            };
1022            set_bool_opt(fd.as_raw_fd(), level, name).is_ok()
1023        }
1024
1025        /// ECN is only implemented on Apple targets among non-Linux Unixes.
1026        #[cfg(not(target_vendor = "apple"))]
1027        fn enable_ecn_recv(_fd: &OwnedFd, _is_ipv6: bool) -> bool {
1028            false
1029        }
1030
1031        pub fn setup_and_probe(fd: &OwnedFd, config: &BuildConfig, is_ipv6: bool) -> io::Result<Capabilities> {
1032            set_bool_opt(fd.as_raw_fd(), libc::SOL_SOCKET, libc::SO_REUSEADDR)?;
1033
1034            let ecn = !config.disable_ecn && enable_ecn_recv(fd, is_ipv6);
1035
1036            Ok(Capabilities {
1037                gso: false,
1038                gro: false,
1039                sendmmsg: false,
1040                ecn,
1041                max_batch: config.max_batch,
1042            })
1043        }
1044    }
1045}
1046
1047// -- re-exports ------------------------------------------------------------
1048#[cfg(unix)]
1049pub use unix_socket::{FastUdpSocket, FastUdpSocketBuilder};
1050
1051// This crate is Unix-only: it is built on `AsyncFd` and libc syscalls.
1052#[cfg(not(unix))]
1053compile_error!("tokio_fast_udp only supports Unix targets");
tokio_fast_udp/socket.rs

tokio_fast_udp/
socket.rs
