Add sendmmsg syscall

Author: Hasan6979

src/lib.rs

@@ -172,6 +172,21 @@ macro_rules! man_links {
     };
 }
 
+macro_rules! define_mmsg_if_supported {
+    ($($item:item)*) => {
+        $(
+            #[cfg(all(
+                feature = "all",
+                any(
+                    target_os = "linux",
+                    target_os = "android",
+                )
+            ))]
+            $item
+        )*
+    };
+}
+
 mod sockaddr;
 mod socket;
 mod sockref;
@@ -697,3 +712,107 @@ impl<'name, 'bufs, 'control> fmt::Debug for MsgHdrMut<'name, 'bufs, 'control> {
         "MsgHdrMut".fmt(fmt)
     }
 }
+
+define_mmsg_if_supported! {
+    /// Wraps `mmsghdr` on Unix for a `sendmmsg(2)` system call.
+    ///
+    /// Also see [`MsgHdr`] for the variant used by `sendmsg(2)`.
+    #[repr(transparent)]
+    pub struct MMsgHdr<'addr, 'bufs, 'control> {
+        inner: sys::mmsghdr,
+        #[allow(clippy::type_complexity)]
+        _lifetimes: PhantomData<(&'addr SockAddr, &'bufs IoSlice<'bufs>, &'control [u8])>,
+    }
+
+    impl<'addr, 'bufs, 'control> MMsgHdr<'addr, 'bufs, 'control> {
+        /// Create a new `MMsgHdr` from `MsgHdr` and with the `msg_len` set to zero.
+        pub fn new(msg: MsgHdr<'_, '_, '_>) -> Self {
+            Self {
+                inner: sys::mmsghdr {
+                    msg_hdr: msg.inner,
+                    msg_len: 0,
+                },
+                _lifetimes: PhantomData,
+            }
+        }
+
+        /// Number of bytes transmitted.
+        pub fn transmitted_bytes(&self) -> u32 {
+            self.inner.msg_len
+        }
+    }
+
+    impl<'addr, 'bufs, 'control> From<MsgHdr<'addr, 'bufs, 'control>>
+        for MMsgHdr<'addr, 'bufs, 'control>
+    {
+        fn from(value: MsgHdr<'_, '_, '_>) -> Self {
+            Self::new(value)
+        }
+    }
+
+    impl<'addr, 'bufs, 'control> fmt::Debug for MMsgHdr<'addr, 'bufs, 'control> {
+        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+            f.write_str("MMsgHdr({})", self.len())
+        }
+    }
+
+    /// Wraps `mmsghdr` on Unix for a `recvmmsg(2)` system call.
+    ///
+    /// Also see [`MsgHdrMut`] for the variant used by `recvmsg(2)`.
+    #[repr(transparent)]
+    pub struct MMsgHdrMut<'addr, 'bufs, 'control> {
+        inner: sys::mmsghdr,
+        #[allow(clippy::type_complexity)]
+        _lifetimes: PhantomData<(
+            &'addr mut SockAddr,
+            &'bufs mut MaybeUninitSlice<'bufs>,
+            &'control mut [u8],
+        )>,
+    }
+
+    impl<'addr, 'bufs, 'control> MMsgHdrMut<'addr, 'bufs, 'control> {
+        /// Create a new `MMsgHdrMut` from `MsgHdrMut` and with the `msg_len` set to zero.
+        pub fn new(msg: MsgHdrMut<'_, '_, '_>) -> Self {
+            Self {
+                inner: sys::mmsghdr {
+                    msg_hdr: msg.inner,
+                    msg_len: 0,
+                },
+                _lifetimes: PhantomData,
+            }
+        }
+
+        /// Number of received bytes.
+        pub fn transmitted_bytes(&self) -> u32 {
+            self.inner.msg_len
+        }
+
+        /// Returns the flags of the message.
+        pub fn flags(&self) -> RecvFlags {
+            sys::msghdr_flags(&self.inner.msg_hdr)
+        }
+
+        /// Gets the length of the control buffer.
+        ///
+        /// Can be used to determine how much, if any, of the control buffer was filled by `recvmsg`.
+        ///
+        /// Corresponds to `msg_controllen` on Unix and `Control.len` on Windows.
+        pub fn control_len(&self) -> usize {
+            sys::msghdr_control_len(&self.inner.msg_hdr)
+        }
+    }
+
+    impl<'addr, 'bufs, 'control> From<MsgHdrMut<'addr, 'bufs, 'control>>
+        for MMsgHdrMut<'addr, 'bufs, 'control>
+    {
+        fn from(value: MsgHdrMut<'_, '_, '_>) -> Self {
+            Self::new(value)
+        }
+    }
+
+    impl<'addr, 'bufs, 'control> fmt::Debug for MMsgHdrMut<'addr, 'bufs, 'control> {
+        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+            f.write_str("MMsgHdrMut({})", self.len())
+        }
+    }
+}

src/socket.rs

@@ -27,6 +27,10 @@ use crate::{Domain, Protocol, SockAddr, TcpKeepalive, Type};
 #[cfg(not(target_os = "redox"))]
 use crate::{MaybeUninitSlice, MsgHdr, RecvFlags};
 
+define_mmsg_if_supported! {
+    use crate::{MMsgHdr, MMsgHdrMut};
+}
+
 /// Owned wrapper around a system socket.
 ///
 /// This type simply wraps an instance of a file descriptor (`c_int`) on Unix
@@ -2196,6 +2200,25 @@ impl Socket {
     pub fn original_dst_ipv6(&self) -> io::Result<SockAddr> {
         sys::original_dst_ipv6(self.as_raw())
     }
+
+    define_mmsg_if_supported! {
+        /// Receive multiple messages on the socket using a single system call.
+        #[doc = man_links!(unix: recvmmsg(2))]
+        pub fn recvmmsg(
+            &self,
+            msgvec: &mut [MMsgHdrMut<'_, '_, '_>],
+            flags: c_int,
+            timeout: Option<Duration>,
+        ) -> io::Result<usize> {
+            sys::recvmmsg(self.as_raw(), msgvec, flags, timeout)
+        }
+
+        /// Send multiple messages on the socket using a single system call.
+        #[doc = man_links!(unix: sendmmsg(2))]
+        pub fn sendmmsg(&self, msgvec: &mut [MMsgHdr<'_, '_, '_>], flags: c_int) -> io::Result<usize> {
+            sys::sendmmsg(self.as_raw(), msgvec, flags)
+        }
+    }
 }
 
 impl Read for Socket {

src/sys/unix.rs

@@ -80,6 +80,13 @@ use crate::{Domain, Protocol, SockAddr, TcpKeepalive, Type};
 #[cfg(not(target_os = "redox"))]
 use crate::{MsgHdr, MsgHdrMut, RecvFlags};
 
+define_mmsg_if_supported! {
+    use crate::{MMsgHdr, MMsgHdrMut};
+
+    // Used in `MMsgHdr`.
+    pub(crate) use libc::mmsghdr;
+}
+
 pub(crate) use libc::c_int;
 
 // Used in `Domain`.
@@ -1076,6 +1083,37 @@ pub(crate) fn recvmsg(
     syscall!(recvmsg(fd, &mut msg.inner, flags)).map(|n| n as usize)
 }
 
+define_mmsg_if_supported! {
+    pub(crate) fn recvmmsg(
+        fd: Socket,
+        msgvec: &mut [MMsgHdrMut<'_, '_, '_>],
+        flags: c_int,
+        timeout: Option<Duration>,
+    ) -> io::Result<usize> {
+        if cfg!(target_env = "musl") {
+            debug_assert!(flags >= 0, "socket flags must be non-negative");
+        }
+
+        let mut timeout = timeout.map(into_timespec);
+        let timeout_ptr = timeout
+            .as_mut()
+            .map(|t| t as *mut _)
+            .unwrap_or(ptr::null_mut());
+
+        syscall!(recvmmsg(
+            fd,
+            // SAFETY: `MMsgHdrMut` is `#[repr(transparent)]` and wraps a `libc::mmsghdr`
+            msgvec.as_mut_ptr() as *mut mmsghdr,
+            msgvec.len() as _,
+            // On glibc: `c_int` to `c_int` (no change).
+            // On musl: `c_int` to `c_uint`.
+            flags as _,
+            timeout_ptr
+        ))
+        .map(|n| n as usize)
+    }
+}
+
 pub(crate) fn send(fd: Socket, buf: &[u8], flags: c_int) -> io::Result<usize> {
     syscall!(send(
         fd,
@@ -1120,6 +1158,29 @@ pub(crate) fn sendmsg(fd: Socket, msg: &MsgHdr<'_, '_, '_>, flags: c_int) -> io:
     syscall!(sendmsg(fd, &msg.inner, flags)).map(|n| n as usize)
 }
 
+define_mmsg_if_supported! {
+    pub(crate) fn sendmmsg(
+        fd: Socket,
+        msgvec: &mut [MMsgHdr<'_, '_, '_>],
+        flags: c_int,
+    ) -> io::Result<usize> {
+        if cfg!(target_env = "musl") {
+            debug_assert!(flags >= 0, "socket flags must be non-negative");
+        }
+
+        syscall!(sendmmsg(
+            fd,
+            // SAFETY: `MMsgHdr` is `#[repr(transparent)]` and wraps a `libc::mmsghdr`
+            msgvec.as_mut_ptr() as *mut mmsghdr,
+            msgvec.len() as _,
+            // On glibc: `c_int` to `c_int` (no change).
+            // On musl: `c_int` to `c_uint`.
+            flags as _
+        ))
+        .map(|n| n as usize)
+    }
+}
+
 /// Wrapper around `getsockopt` to deal with platform specific timeouts.
 pub(crate) fn timeout_opt(fd: Socket, opt: c_int, val: c_int) -> io::Result<Option<Duration>> {
     unsafe { getsockopt(fd, opt, val).map(from_timeval) }
@@ -1161,6 +1222,26 @@ fn into_timeval(duration: Option<Duration>) -> libc::timeval {
     }
 }
 
+#[allow(dead_code)] // Used by `recvmmsg()`.
+fn into_timespec(duration: Duration) -> libc::timespec {
+    // https://github.com/rust-lang/libc/issues/1848
+    #[cfg_attr(target_env = "musl", allow(deprecated))]
+    libc::timespec {
+        tv_sec: min(duration.as_secs(), libc::time_t::MAX as u64) as libc::time_t,
+        #[cfg(any(
+            all(target_arch = "x86_64", target_pointer_width = "32"),
+            target_pointer_width = "64"
+        ))]
+        tv_nsec: duration.subsec_nanos() as i64,
+
+        #[cfg(not(any(
+            all(target_arch = "x86_64", target_pointer_width = "32"),
+            target_pointer_width = "64"
+        )))]
+        tv_nsec: duration.subsec_nanos().clamp(0, i32::MAX as u32) as i32,
+    }
+}
+
 #[cfg(all(
     feature = "all",
     not(any(target_os = "haiku", target_os = "openbsd", target_os = "vita"))

tests/socket.rs

@@ -769,6 +769,32 @@ fn sendmsg() {
     assert_eq!(received, DATA.len());
 }
 
+#[test]
+#[cfg(any(target_os = "linux", target_os = "android"))]
+fn sendmmsg() {
+    let (socket_a, socket_b) = udp_pair_unconnected();
+
+    const DATA: &[u8] = b"Hello, World!";
+
+    let bufs = &[IoSlice::new(DATA)];
+    let addr_b = socket_b.local_addr().unwrap();
+    let batched_msgs = Vec::new();
+    for _ in 0..5 {
+        batched_msgs.push(socket2::MMsgHdr::new(
+            socket2::MsgHdr::new().with_addr(&addr_b).with_buffers(bufs),
+        ));
+    }
+
+    let sent = socket_a.sendmmsg(batched_msgs.as_mut_slice(), 0).unwrap();
+
+    let mut sent_data = 0;
+    // Calculate transmitted length
+    for i in 0..5 {
+        sent_data += batched_msgs[i].transmitted_bytes()
+    }
+    assert!(sent_data == 5 * DATA.len());
+}
+
 #[test]
 #[cfg(not(any(target_os = "redox", target_os = "vita")))]
 fn recv_vectored_truncated() {