simoom_sys/
udp.rs

1//! Batched UDP I/O with per-datagram metadata: `recvmmsg`/`sendmmsg` on Linux with GRO/GSO
2//! (`UDP_GRO`/`UDP_SEGMENT`), ECN and packet-info control messages; `WSARecvMsg`/`WSASendMsg`
3//! on Windows with URO/USO, ECN and packet info; plain `recvmsg`/`sendmsg` on macOS and BSD
4//! (no control messages there yet).
5
6use crate::net::{decode_addr, encode_addr, SockAddrBuf, Socket, UdpOptions};
7use crate::os::sys;
8#[cfg(not(windows))]
9use core::ffi::c_void;
10use std::io;
11#[cfg(any(target_os = "linux", windows))]
12use std::net::Ipv6Addr;
13use std::net::{IpAddr, Ipv4Addr, SocketAddr};
14
15/// What the OS reported about one received datagram (or one GRO batch of datagrams).
16#[derive(Clone, Copy, Debug, PartialEq, Eq)]
17pub struct RecvMeta {
18    /// The sender.
19    pub peer: SocketAddr,
20    /// Bytes received into the buffer.
21    pub len: usize,
22    /// With GRO, the buffer holds several datagrams of this size (the last may be shorter).
23    /// Equal to `len` when the buffer holds one datagram.
24    pub segment_size: usize,
25    /// The ECN bits (0..=3) when ECN reporting is on, else 0.
26    pub ecn: u8,
27    /// The local address the datagram was sent to, when packet info is on.
28    pub local_ip: Option<IpAddr>,
29    /// The datagram did not fit the buffer and was cut.
30    pub truncated: bool,
31}
32
33impl RecvMeta {
34    fn empty() -> RecvMeta {
35        RecvMeta {
36            peer: SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0),
37            len: 0,
38            segment_size: 0,
39            ecn: 0,
40            local_ip: None,
41            truncated: false,
42        }
43    }
44}
45
46impl Default for RecvMeta {
47    fn default() -> Self {
48        RecvMeta::empty()
49    }
50}
51
52/// One datagram (or, with `segment_size`, a GSO batch of datagrams) to send.
53#[derive(Clone, Copy, Debug)]
54pub struct Transmit<'a> {
55    /// The destination.
56    pub peer: SocketAddr,
57    /// The payload. With `segment_size` set it is cut into datagrams of that size by the
58    /// kernel (Linux GSO, Windows USO); elsewhere this module cuts it itself.
59    pub data: &'a [u8],
60    /// Segment size for GSO, or 0 for one datagram.
61    pub segment_size: usize,
62    /// ECN bits to set (0..=3).
63    pub ecn: u8,
64    /// Source address to send from (multi-homed hosts), when packet info is on.
65    pub src_ip: Option<IpAddr>,
66}
67
68/// How many datagrams one batched call handles at most.
69pub const BATCH: usize = 32;
70
71#[cfg(any(target_os = "linux", windows))]
72const CONTROL: usize = 128;
73
74#[cfg(any(target_os = "linux", windows))]
75#[repr(C, align(8))]
76#[derive(Clone, Copy)]
77struct Control([u8; CONTROL]);
78
79/// Applies the receive options of `opts` to a fresh UDP socket.
80pub(crate) fn configure(s: &Socket, v6: bool, opts: &UdpOptions) -> io::Result<()> {
81    #[cfg(target_os = "linux")]
82    {
83        if opts.gro {
84            // UDP_GRO is 5.0+; ignore failure on older kernels.
85            let _ = s.setsockopt_int(sys::IPPROTO_UDP, sys::UDP_GRO, 1);
86        }
87        if opts.pktinfo {
88            s.setsockopt_int(sys::IPPROTO_IP, sys::IP_PKTINFO, 1).ok();
89            if v6 {
90                s.setsockopt_int(sys::IPPROTO_IPV6, sys::IPV6_RECVPKTINFO, 1)?;
91            }
92        }
93        if opts.ecn {
94            s.setsockopt_int(sys::IPPROTO_IP, sys::IP_RECVTOS, 1).ok();
95            if v6 {
96                s.setsockopt_int(sys::IPPROTO_IPV6, sys::IPV6_RECVTCLASS, 1)?;
97            }
98        }
99        if opts.dont_fragment {
100            s.setsockopt_int(sys::IPPROTO_IP, sys::IP_MTU_DISCOVER, sys::IP_PMTUDISC_PROBE).ok();
101            if v6 {
102                s.setsockopt_int(sys::IPPROTO_IPV6, sys::IPV6_MTU_DISCOVER, sys::IP_PMTUDISC_PROBE)?;
103            }
104        }
105    }
106    #[cfg(windows)]
107    {
108        // A UDP socket on Windows reports an ICMP "port unreachable" for an earlier send as an
109        // error on the next receive (WSAECONNRESET); a server must not see that.
110        let mut off: u32 = 0;
111        let mut ret = 0u32;
112        // SAFETY: SIO_UDP_CONNRESET takes a BOOL input; no output, no overlapped.
113        unsafe {
114            sys::WSAIoctl(s.raw(), sys::SIO_UDP_CONNRESET, (&mut off as *mut u32).cast(), 4, core::ptr::null_mut(), 0, &mut ret, core::ptr::null_mut(), core::ptr::null_mut());
115        }
116        if opts.gro {
117            let _ = s.setsockopt_int(sys::IPPROTO_UDP, sys::UDP_RECV_MAX_COALESCED_SIZE, 65527);
118        }
119        if opts.pktinfo {
120            s.setsockopt_int(sys::IPPROTO_IP, sys::IP_PKTINFO, 1).ok();
121            if v6 {
122                s.setsockopt_int(sys::IPPROTO_IPV6, sys::IPV6_PKTINFO, 1)?;
123            }
124        }
125        if opts.ecn {
126            s.setsockopt_int(sys::IPPROTO_IP, sys::IP_ECN, 1).ok();
127            if v6 {
128                s.setsockopt_int(sys::IPPROTO_IPV6, sys::IPV6_ECN, 1).ok();
129            }
130        }
131        if opts.dont_fragment {
132            s.setsockopt_int(sys::IPPROTO_IP, sys::IP_DONTFRAGMENT, 1).ok();
133            if v6 {
134                s.setsockopt_int(sys::IPPROTO_IPV6, sys::IPV6_DONTFRAG, 1).ok();
135            }
136        }
137    }
138    #[cfg(not(any(target_os = "linux", windows)))]
139    {
140        let _ = (s, v6, opts);
141    }
142    Ok(())
143}
144
145// ---------------------------------------------------------------------------------------------
146// Control messages (Linux and Windows share the 64-bit layout: a 16-byte header, 8-byte align)
147// ---------------------------------------------------------------------------------------------
148
149#[cfg(any(target_os = "linux", windows))]
150mod cmsg {
151    pub const HDR: usize = 16;
152    pub const fn align(n: usize) -> usize {
153        (n + 7) & !7
154    }
155    pub const fn space(n: usize) -> usize {
156        HDR + align(n)
157    }
158
159    /// Walks the control buffer and calls `f(level, type, data)` for each message.
160    pub fn parse(buf: &[u8], mut f: impl FnMut(i32, i32, &[u8])) {
161        let mut off = 0;
162        while off + HDR <= buf.len() {
163            let len = usize::from_ne_bytes(buf[off..off + 8].try_into().unwrap());
164            let level = i32::from_ne_bytes(buf[off + 8..off + 12].try_into().unwrap());
165            let ty = i32::from_ne_bytes(buf[off + 12..off + 16].try_into().unwrap());
166            if len < HDR || off + len > buf.len() {
167                break;
168            }
169            f(level, ty, &buf[off + HDR..off + len]);
170            off += align(len);
171        }
172    }
173
174    /// Appends one control message; returns the new used length.
175    pub fn push(buf: &mut [u8], used: usize, level: i32, ty: i32, data: &[u8]) -> usize {
176        let total = space(data.len());
177        if used + total > buf.len() {
178            return used;
179        }
180        let b = &mut buf[used..used + total];
181        b.fill(0);
182        b[0..8].copy_from_slice(&(HDR + data.len()).to_ne_bytes());
183        b[8..12].copy_from_slice(&level.to_ne_bytes());
184        b[12..16].copy_from_slice(&ty.to_ne_bytes());
185        b[HDR..HDR + data.len()].copy_from_slice(data);
186        used + total
187    }
188}
189
190#[cfg(target_os = "linux")]
191fn apply_cmsg(meta: &mut RecvMeta, level: i32, ty: i32, data: &[u8]) {
192    let int = |d: &[u8]| if d.len() >= 4 { i32::from_ne_bytes(d[..4].try_into().unwrap()) } else { d.first().copied().unwrap_or(0) as i32 };
193    match (level, ty) {
194        (sys::IPPROTO_UDP, sys::UDP_GRO) => meta.segment_size = int(data) as usize,
195        (sys::IPPROTO_IP, sys::IP_TOS) => meta.ecn = (data.first().copied().unwrap_or(0)) & 3,
196        (sys::IPPROTO_IPV6, sys::IPV6_TCLASS) => meta.ecn = (int(data) & 3) as u8,
197        (sys::IPPROTO_IP, sys::IP_PKTINFO) if data.len() >= 12 => {
198            meta.local_ip = Some(IpAddr::V4(Ipv4Addr::new(data[8], data[9], data[10], data[11])));
199        }
200        (sys::IPPROTO_IPV6, sys::IPV6_PKTINFO) if data.len() >= 16 => {
201            let mut a = [0u8; 16];
202            a.copy_from_slice(&data[..16]);
203            meta.local_ip = Some(IpAddr::V6(Ipv6Addr::from(a)));
204        }
205        _ => {}
206    }
207}
208
209#[cfg(windows)]
210fn apply_cmsg(meta: &mut RecvMeta, level: i32, ty: i32, data: &[u8]) {
211    let int = |d: &[u8]| if d.len() >= 4 { i32::from_ne_bytes(d[..4].try_into().unwrap()) } else { 0 };
212    match (level, ty) {
213        (sys::IPPROTO_UDP, sys::UDP_COALESCED_INFO) => meta.segment_size = int(data) as usize,
214        (sys::IPPROTO_IP, sys::IP_ECN) | (sys::IPPROTO_IPV6, sys::IPV6_ECN) => meta.ecn = (int(data) & 3) as u8,
215        (sys::IPPROTO_IP, sys::IP_PKTINFO) if data.len() >= 4 => {
216            meta.local_ip = Some(IpAddr::V4(Ipv4Addr::new(data[0], data[1], data[2], data[3])));
217        }
218        (sys::IPPROTO_IPV6, sys::IPV6_PKTINFO) if data.len() >= 16 => {
219            let mut a = [0u8; 16];
220            a.copy_from_slice(&data[..16]);
221            meta.local_ip = Some(IpAddr::V6(Ipv6Addr::from(a)));
222        }
223        _ => {}
224    }
225}
226
227/// Builds the control messages for one transmit; returns the used length.
228#[cfg(any(target_os = "linux", windows))]
229fn build_cmsg(t: &Transmit<'_>, ctl: &mut [u8]) -> usize {
230    let mut used = 0;
231    if t.segment_size > 0 && t.segment_size < t.data.len() {
232        #[cfg(target_os = "linux")]
233        {
234            used = cmsg::push(ctl, used, sys::IPPROTO_UDP, sys::UDP_SEGMENT, &(t.segment_size as u16).to_ne_bytes());
235        }
236        #[cfg(windows)]
237        {
238            used = cmsg::push(ctl, used, sys::IPPROTO_UDP, sys::UDP_SEND_MSG_SIZE, &(t.segment_size as u32).to_ne_bytes());
239        }
240    }
241    if t.ecn != 0 {
242        let v = (t.ecn & 3) as i32;
243        #[cfg(target_os = "linux")]
244        let (l4, t4, l6, t6) = (sys::IPPROTO_IP, sys::IP_TOS, sys::IPPROTO_IPV6, sys::IPV6_TCLASS);
245        #[cfg(windows)]
246        let (l4, t4, l6, t6) = (sys::IPPROTO_IP, sys::IP_ECN, sys::IPPROTO_IPV6, sys::IPV6_ECN);
247        used = if t.peer.is_ipv4() { cmsg::push(ctl, used, l4, t4, &v.to_ne_bytes()) } else { cmsg::push(ctl, used, l6, t6, &v.to_ne_bytes()) };
248    }
249    match t.src_ip {
250        #[cfg(target_os = "linux")]
251        Some(IpAddr::V4(a)) => {
252            let mut d = [0u8; 12];
253            d[4..8].copy_from_slice(&a.octets());
254            used = cmsg::push(ctl, used, sys::IPPROTO_IP, sys::IP_PKTINFO, &d);
255        }
256        #[cfg(windows)]
257        Some(IpAddr::V4(a)) => {
258            let mut d = [0u8; 8];
259            d[..4].copy_from_slice(&a.octets());
260            used = cmsg::push(ctl, used, sys::IPPROTO_IP, sys::IP_PKTINFO, &d);
261        }
262        Some(IpAddr::V6(a)) => {
263            let mut d = [0u8; 20];
264            d[..16].copy_from_slice(&a.octets());
265            used = cmsg::push(ctl, used, sys::IPPROTO_IPV6, sys::IPV6_PKTINFO, &d);
266        }
267        None => {}
268    }
269    used
270}
271
272// ---------------------------------------------------------------------------------------------
273// Windows extension function
274// ---------------------------------------------------------------------------------------------
275
276#[cfg(windows)]
277pub(crate) fn wsa_recvmsg(s: &Socket) -> io::Result<sys::LPFN_WSARECVMSG> {
278    use std::sync::OnceLock;
279    static FN: OnceLock<usize> = OnceLock::new();
280    if let Some(&f) = FN.get() {
281        // SAFETY: the value was produced by WSAIoctl for WSAID_WSARECVMSG below.
282        return Ok(unsafe { core::mem::transmute::<usize, sys::LPFN_WSARECVMSG>(f) });
283    }
284    let mut f: usize = 0;
285    let mut ret = 0u32;
286    let guid = sys::WSAID_WSARECVMSG;
287    // SAFETY: SIO_GET_EXTENSION_FUNCTION_POINTER reads the GUID and writes one pointer.
288    let r = unsafe {
289        sys::WSAIoctl(
290            s.raw(),
291            sys::SIO_GET_EXTENSION_FUNCTION_POINTER,
292            (&guid as *const sys::GUID).cast(),
293            core::mem::size_of::<sys::GUID>() as u32,
294            (&mut f as *mut usize).cast(),
295            core::mem::size_of::<usize>() as u32,
296            &mut ret,
297            core::ptr::null_mut(),
298            core::ptr::null_mut(),
299        )
300    };
301    if r != 0 || f == 0 {
302        return Err(crate::net::last_error());
303    }
304    let _ = FN.set(f);
305    // SAFETY: `f` is the WSARecvMsg entry point returned by Winsock.
306    Ok(unsafe { core::mem::transmute::<usize, sys::LPFN_WSARECVMSG>(f) })
307}
308
309/// Parses the result of a `WSARecvMsg` into `meta` (shared with the IOCP driver).
310#[cfg(windows)]
311pub(crate) fn finish_meta_windows(name: &SockAddrBuf, namelen: i32, control: &[u8], flags: u32, len: usize) -> RecvMeta {
312    let mut meta = RecvMeta::empty();
313    if let Some(p) = decode_addr(name, namelen.max(0) as usize) {
314        meta.peer = p;
315    }
316    meta.len = len;
317    cmsg::parse(control, |l, t, d| apply_cmsg(&mut meta, l, t, d));
318    if meta.segment_size == 0 {
319        meta.segment_size = len;
320    }
321    meta.truncated = flags & sys::MSG_TRUNC != 0;
322    meta
323}
324
325// ---------------------------------------------------------------------------------------------
326// Batched calls
327// ---------------------------------------------------------------------------------------------
328
329impl Socket {
330    /// Receives up to `bufs.len()` (at most [`BATCH`]) datagrams without blocking. Each buffer
331    /// is cleared and filled up to its capacity; `meta[i]` describes `bufs[i]`. Returns how many
332    /// were received, or a `WouldBlock` error when none was waiting.
333    pub fn recv_msgs(&self, bufs: &mut [Vec<u8>], meta: &mut [RecvMeta]) -> io::Result<usize> {
334        let n = bufs.len().min(meta.len()).min(BATCH);
335        if n == 0 {
336            return Ok(0);
337        }
338        #[cfg(target_os = "linux")]
339        {
340            let mut names = [SockAddrBuf::new(); BATCH];
341            let mut ctls = [Control([0; CONTROL]); BATCH];
342            let mut iovs = [sys::iovec { iov_base: core::ptr::null_mut(), iov_len: 0 }; BATCH];
343            // SAFETY: an all-zero mmsghdr is valid.
344            let mut hdrs: [sys::mmsghdr; BATCH] = unsafe { core::mem::zeroed() };
345            for i in 0..n {
346                bufs[i].clear();
347                iovs[i] = sys::iovec { iov_base: bufs[i].as_mut_ptr().cast(), iov_len: bufs[i].capacity() };
348                let h = &mut hdrs[i].msg_hdr;
349                h.msg_name = names[i].as_mut_ptr();
350                h.msg_namelen = 128;
351                h.msg_iov = &mut iovs[i];
352                h.msg_iovlen = 1;
353                h.msg_control = ctls[i].0.as_mut_ptr().cast();
354                h.msg_controllen = CONTROL;
355            }
356            // SAFETY: every header points at live name, iovec and control storage on this stack
357            // frame, and each iovec at a buffer's spare capacity.
358            let got = unsafe { sys::recvmmsg(self.raw(), hdrs.as_mut_ptr(), n as u32, sys::MSG_DONTWAIT, core::ptr::null_mut()) };
359            if got < 0 {
360                return Err(crate::net::last_error());
361            }
362            let got = got as usize;
363            for i in 0..got {
364                let h = &hdrs[i].msg_hdr;
365                let len = (hdrs[i].msg_len as usize).min(bufs[i].capacity());
366                // SAFETY: the kernel wrote `len` bytes into the buffer's capacity.
367                unsafe { bufs[i].set_len(len) };
368                let mut m = RecvMeta::empty();
369                if let Some(p) = decode_addr(&names[i], h.msg_namelen as usize) {
370                    m.peer = p;
371                }
372                m.len = len;
373                cmsg::parse(&ctls[i].0[..h.msg_controllen.min(CONTROL)], |l, t, d| apply_cmsg(&mut m, l, t, d));
374                if m.segment_size == 0 {
375                    m.segment_size = len;
376                }
377                m.truncated = h.msg_flags & sys::MSG_TRUNC != 0;
378                meta[i] = m;
379            }
380            Ok(got)
381        }
382        #[cfg(windows)]
383        {
384            let f = wsa_recvmsg(self)?;
385            let mut got = 0;
386            while got < n {
387                let mut name = SockAddrBuf::new();
388                let mut ctl = Control([0; CONTROL]);
389                bufs[got].clear();
390                let mut wb = sys::WSABUF { len: bufs[got].capacity().min(u32::MAX as usize) as u32, buf: bufs[got].as_mut_ptr() };
391                let mut msg = sys::WSAMSG {
392                    name: name.as_mut_ptr(),
393                    namelen: 128,
394                    lpBuffers: &mut wb,
395                    dwBufferCount: 1,
396                    Control: sys::WSABUF { len: CONTROL as u32, buf: ctl.0.as_mut_ptr() },
397                    dwFlags: 0,
398                };
399                let mut len = 0u32;
400                // SAFETY: `msg` points at live name, buffer and control storage; a null
401                // OVERLAPPED makes this a non-blocking call that completes before returning.
402                let r = unsafe { f(self.raw(), &mut msg, &mut len, core::ptr::null_mut(), core::ptr::null_mut()) };
403                if r != 0 {
404                    let e = crate::net::last_error();
405                    if e.raw_os_error() == Some(sys::WSAEMSGSIZE) {
406                        msg.dwFlags |= sys::MSG_TRUNC;
407                        len = wb.len;
408                    } else if got > 0 && e.kind() == io::ErrorKind::WouldBlock {
409                        break;
410                    } else {
411                        return Err(e);
412                    }
413                }
414                // SAFETY: Winsock wrote `len` bytes into the buffer's capacity.
415                unsafe { bufs[got].set_len((len as usize).min(bufs[got].capacity())) };
416                meta[got] = finish_meta_windows(&name, msg.namelen, &ctl.0[..(msg.Control.len as usize).min(CONTROL)], msg.dwFlags, len as usize);
417                got += 1;
418            }
419            Ok(got)
420        }
421        #[cfg(not(any(target_os = "linux", windows)))]
422        {
423            let mut got = 0;
424            while got < n {
425                let mut name = SockAddrBuf::new();
426                bufs[got].clear();
427                let mut iov = sys::iovec { iov_base: bufs[got].as_mut_ptr().cast(), iov_len: bufs[got].capacity() };
428                // SAFETY: an all-zero msghdr is valid.
429                let mut h: sys::msghdr = unsafe { core::mem::zeroed() };
430                h.msg_name = name.as_mut_ptr();
431                h.msg_namelen = 128;
432                h.msg_iov = &mut iov;
433                h.msg_iovlen = 1;
434                // SAFETY: `h` points at live name and buffer storage.
435                let r = unsafe { sys::recvmsg(self.raw(), &mut h, 0) };
436                if r < 0 {
437                    let e = crate::net::last_error();
438                    if got > 0 && e.kind() == io::ErrorKind::WouldBlock {
439                        break;
440                    }
441                    return Err(e);
442                }
443                let len = (r as usize).min(bufs[got].capacity());
444                // SAFETY: the kernel wrote `len` bytes.
445                unsafe { bufs[got].set_len(len) };
446                let mut m = RecvMeta::empty();
447                if let Some(p) = decode_addr(&name, h.msg_namelen as usize) {
448                    m.peer = p;
449                }
450                m.len = len;
451                m.segment_size = len;
452                meta[got] = m;
453                got += 1;
454            }
455            Ok(got)
456        }
457    }
458
459    /// Sends up to [`BATCH`] transmits without blocking. Returns how many were handed to the
460    /// kernel; a `WouldBlock` error means none was (the socket buffer is full).
461    pub fn send_msgs(&self, msgs: &[Transmit<'_>]) -> io::Result<usize> {
462        let n = msgs.len().min(BATCH);
463        if n == 0 {
464            return Ok(0);
465        }
466        #[cfg(target_os = "linux")]
467        {
468            let mut names = [SockAddrBuf::new(); BATCH];
469            let mut ctls = [Control([0; CONTROL]); BATCH];
470            let mut iovs = [sys::iovec { iov_base: core::ptr::null_mut(), iov_len: 0 }; BATCH];
471            // SAFETY: an all-zero mmsghdr is valid.
472            let mut hdrs: [sys::mmsghdr; BATCH] = unsafe { core::mem::zeroed() };
473            for (i, t) in msgs[..n].iter().enumerate() {
474                let nl = encode_addr(&t.peer, &mut names[i]);
475                let cl = build_cmsg(t, &mut ctls[i].0);
476                iovs[i] = sys::iovec { iov_base: t.data.as_ptr() as *mut c_void, iov_len: t.data.len() };
477                let h = &mut hdrs[i].msg_hdr;
478                h.msg_name = names[i].as_mut_ptr();
479                h.msg_namelen = nl;
480                h.msg_iov = &mut iovs[i];
481                h.msg_iovlen = 1;
482                if cl > 0 {
483                    h.msg_control = ctls[i].0.as_mut_ptr().cast();
484                    h.msg_controllen = cl;
485                }
486            }
487            // SAFETY: every header points at live storage on this frame and at the payloads.
488            let r = unsafe { sys::sendmmsg(self.raw(), hdrs.as_mut_ptr(), n as u32, sys::MSG_DONTWAIT | sys::MSG_NOSIGNAL) };
489            if r < 0 {
490                Err(crate::net::last_error())
491            } else {
492                Ok(r as usize)
493            }
494        }
495        #[cfg(windows)]
496        {
497            for (i, t) in msgs[..n].iter().enumerate() {
498                let mut name = SockAddrBuf::new();
499                let nl = encode_addr(&t.peer, &mut name);
500                let mut ctl = Control([0; CONTROL]);
501                let cl = build_cmsg(t, &mut ctl.0);
502                let mut wb = sys::WSABUF { len: t.data.len() as u32, buf: t.data.as_ptr() as *mut u8 };
503                let msg = sys::WSAMSG {
504                    name: name.as_mut_ptr(),
505                    namelen: nl as i32,
506                    lpBuffers: &mut wb,
507                    dwBufferCount: 1,
508                    Control: sys::WSABUF { len: cl as u32, buf: if cl > 0 { ctl.0.as_mut_ptr() } else { core::ptr::null_mut() } },
509                    dwFlags: 0,
510                };
511                let mut sent = 0u32;
512                // SAFETY: `msg` points at live storage; a null OVERLAPPED makes this a
513                // non-blocking call that completes before returning.
514                let r = unsafe { sys::WSASendMsg(self.raw(), &msg, 0, &mut sent, core::ptr::null_mut(), core::ptr::null_mut()) };
515                if r != 0 {
516                    let e = crate::net::last_error();
517                    if i > 0 && e.kind() == io::ErrorKind::WouldBlock {
518                        return Ok(i);
519                    }
520                    return Err(e);
521                }
522            }
523            Ok(n)
524        }
525        #[cfg(not(any(target_os = "linux", windows)))]
526        {
527            for (i, t) in msgs[..n].iter().enumerate() {
528                let seg = if t.segment_size > 0 { t.segment_size } else { t.data.len().max(1) };
529                for chunk in t.data.chunks(seg) {
530                    let mut name = SockAddrBuf::new();
531                    let nl = encode_addr(&t.peer, &mut name);
532                    let mut iov = sys::iovec { iov_base: chunk.as_ptr() as *mut c_void, iov_len: chunk.len() };
533                    // SAFETY: an all-zero msghdr is valid.
534                    let mut h: sys::msghdr = unsafe { core::mem::zeroed() };
535                    h.msg_name = name.as_mut_ptr();
536                    h.msg_namelen = nl;
537                    h.msg_iov = &mut iov;
538                    h.msg_iovlen = 1;
539                    // SAFETY: `h` points at live storage.
540                    let r = unsafe { sys::sendmsg(self.raw(), &h, sys::MSG_NOSIGNAL) };
541                    if r < 0 {
542                        let e = crate::net::last_error();
543                        if i > 0 && e.kind() == io::ErrorKind::WouldBlock {
544                            return Ok(i);
545                        }
546                        return Err(e);
547                    }
548                }
549            }
550            Ok(n)
551        }
552    }
553}
554
555#[cfg(test)]
556mod tests {
557    use super::*;
558
559    fn recv_wait(s: &Socket, bufs: &mut [Vec<u8>], meta: &mut [RecvMeta], want: usize) -> usize {
560        let mut total = 0;
561        let start = std::time::Instant::now();
562        while total < want && start.elapsed() < std::time::Duration::from_secs(5) {
563            match s.recv_msgs(&mut bufs[total..], &mut meta[total..]) {
564                Ok(n) => total += n,
565                Err(e) if e.kind() == io::ErrorKind::WouldBlock => std::thread::sleep(std::time::Duration::from_millis(1)),
566                Err(e) => panic!("{e}"),
567            }
568        }
569        total
570    }
571
572    #[test]
573    fn batch_send_and_receive_with_metadata() {
574        let opts = UdpOptions { pktinfo: true, ecn: true, ..Default::default() };
575        let a = Socket::udp_bind("127.0.0.1:0".parse().unwrap(), &opts).unwrap();
576        let b = Socket::udp_bind("127.0.0.1:0".parse().unwrap(), &opts).unwrap();
577        let to = a.local_addr().unwrap();
578        let payloads: Vec<Vec<u8>> = (0..5u8).map(|i| vec![i; 100 + i as usize]).collect();
579        let tx: Vec<Transmit<'_>> = payloads.iter().map(|p| Transmit { peer: to, data: p, segment_size: 0, ecn: 2, src_ip: None }).collect();
580        assert_eq!(b.send_msgs(&tx).unwrap(), 5);
581        let mut bufs: Vec<Vec<u8>> = (0..8).map(|_| Vec::with_capacity(2048)).collect();
582        let mut meta = [RecvMeta::default(); 8];
583        assert_eq!(recv_wait(&a, &mut bufs, &mut meta, 5), 5);
584        for i in 0..5 {
585            assert_eq!(bufs[i], payloads[i]);
586            assert_eq!(meta[i].peer, b.local_addr().unwrap());
587            assert_eq!(meta[i].len, 100 + i);
588            assert!(!meta[i].truncated);
589            #[cfg(target_os = "linux")]
590            {
591                assert_eq!(meta[i].ecn, 2, "ECN bits reported");
592                assert_eq!(meta[i].local_ip, Some("127.0.0.1".parse().unwrap()));
593            }
594            #[cfg(windows)]
595            assert_eq!(meta[i].local_ip, Some("127.0.0.1".parse().unwrap()));
596        }
597        assert!(matches!(a.recv_msgs(&mut bufs, &mut meta), Err(e) if e.kind() == io::ErrorKind::WouldBlock));
598    }
599
600    #[test]
601    fn segmented_send_arrives_as_separate_datagrams() {
602        // GSO where the kernel supports it, else cut in user space; without GRO the receiver
603        // sees one datagram per segment either way.
604        let a = Socket::udp_bind("127.0.0.1:0".parse().unwrap(), &UdpOptions::default()).unwrap();
605        let b = Socket::udp_bind("127.0.0.1:0".parse().unwrap(), &UdpOptions::default()).unwrap();
606        let data: Vec<u8> = (0..1000u32).map(|i| i as u8).collect();
607        let t = Transmit { peer: a.local_addr().unwrap(), data: &data, segment_size: 400, ecn: 0, src_ip: None };
608        match b.send_msgs(&[t]) {
609            Ok(1) => {}
610            // Kernels or drivers without segmentation offload refuse the control message.
611            Err(e) if cfg!(any(target_os = "linux", windows)) => {
612                eprintln!("segmentation offload unavailable: {e}");
613                return;
614            }
615            other => panic!("{other:?}"),
616        }
617        let mut bufs: Vec<Vec<u8>> = (0..4).map(|_| Vec::with_capacity(2048)).collect();
618        let mut meta = [RecvMeta::default(); 4];
619        assert_eq!(recv_wait(&a, &mut bufs, &mut meta, 3), 3);
620        assert_eq!(bufs[0], data[..400]);
621        assert_eq!(bufs[1], data[400..800]);
622        assert_eq!(bufs[2], data[800..]);
623    }
624
625    #[cfg(target_os = "linux")]
626    #[test]
627    fn gro_coalesces_or_reports_single_segments() {
628        let a = Socket::udp_bind("127.0.0.1:0".parse().unwrap(), &UdpOptions { gro: true, ..Default::default() }).unwrap();
629        let b = Socket::udp_bind("127.0.0.1:0".parse().unwrap(), &UdpOptions::default()).unwrap();
630        let data = vec![7u8; 1200];
631        let t = Transmit { peer: a.local_addr().unwrap(), data: &data, segment_size: 300, ecn: 0, src_ip: None };
632        assert_eq!(b.send_msgs(&[t]).unwrap(), 1);
633        let mut bufs: Vec<Vec<u8>> = (0..4).map(|_| Vec::with_capacity(65536)).collect();
634        let mut meta = [RecvMeta::default(); 4];
635        let mut total = 0;
636        let start = std::time::Instant::now();
637        let mut got = 0;
638        while total < 1200 && start.elapsed() < std::time::Duration::from_secs(5) {
639            match a.recv_msgs(&mut bufs[got..], &mut meta[got..]) {
640                Ok(n) => {
641                    for m in &meta[got..got + n] {
642                        assert_eq!(m.segment_size, 300);
643                        total += m.len;
644                    }
645                    got += n;
646                }
647                Err(e) if e.kind() == io::ErrorKind::WouldBlock => std::thread::sleep(std::time::Duration::from_millis(1)),
648                Err(e) => panic!("{e}"),
649            }
650        }
651        assert_eq!(total, 1200);
652    }
653
654    #[test]
655    fn ipv6_datagrams() {
656        let Ok(a) = Socket::udp_bind("[::1]:0".parse().unwrap(), &UdpOptions { pktinfo: true, ..Default::default() }) else { return };
657        let b = Socket::udp_bind("[::1]:0".parse().unwrap(), &UdpOptions::default()).unwrap();
658        let t = Transmit { peer: a.local_addr().unwrap(), data: b"six", segment_size: 0, ecn: 0, src_ip: None };
659        assert_eq!(b.send_msgs(&[t]).unwrap(), 1);
660        let mut bufs = vec![Vec::with_capacity(64)];
661        let mut meta = [RecvMeta::default()];
662        assert_eq!(recv_wait(&a, &mut bufs, &mut meta, 1), 1);
663        assert_eq!(bufs[0], b"six");
664        assert_eq!(meta[0].peer, b.local_addr().unwrap());
665        #[cfg(any(target_os = "linux", windows))]
666        assert_eq!(meta[0].local_ip, Some("::1".parse().unwrap()));
667    }
668}
simoom_sys/udp.rs

simoom_sys/
udp.rs
