Add preallocated buffer version of sendmsg

howjmay · howjmay · commit 1eed9f4c501b · 2024-09-22T20:31:04.000+02:00
diff --git a/src/sys/socket/mod.rs b/src/sys/socket/mod.rs
@@ -39,10 +39,20 @@ pub use self::addr::{SockaddrLike, SockaddrStorage};
 pub use self::addr::{AddressFamily, UnixAddr};
 #[cfg(not(solarish))]
 pub use self::addr::{AddressFamily, UnixAddr};
-#[cfg(not(any(solarish, target_os = "haiku", target_os = "hurd", target_os = "redox")))]
+#[cfg(not(any(
+    solarish,
+    target_os = "haiku",
+    target_os = "hurd",
+    target_os = "redox"
+)))]
 #[cfg(feature = "net")]
 pub use self::addr::{LinkAddr, SockaddrIn, SockaddrIn6};
-#[cfg(any(solarish, target_os = "haiku", target_os = "hurd", target_os = "redox"))]
+#[cfg(any(
+    solarish,
+    target_os = "haiku",
+    target_os = "hurd",
+    target_os = "redox"
+))]
 #[cfg(feature = "net")]
 pub use self::addr::{SockaddrIn, SockaddrIn6};
 
@@ -794,17 +804,17 @@ pub enum ControlMessageOwned {
     #[cfg_attr(docsrs, doc(cfg(feature = "net")))]
     Ipv6HopLimit(i32),
 
-    /// Retrieve the DSCP (ToS) header field of the incoming IPv4 packet. 
+    /// Retrieve the DSCP (ToS) header field of the incoming IPv4 packet.
     #[cfg(any(linux_android, target_os = "freebsd"))]
     #[cfg(feature = "net")]
     #[cfg_attr(docsrs, doc(cfg(feature = "net")))]
     Ipv4Tos(u8),
 
-    /// Retrieve the DSCP (Traffic Class) header field of the incoming IPv6 packet. 
+    /// Retrieve the DSCP (Traffic Class) header field of the incoming IPv6 packet.
     #[cfg(any(linux_android, target_os = "freebsd"))]
     #[cfg(feature = "net")]
     #[cfg_attr(docsrs, doc(cfg(feature = "net")))]
-    Ipv6TClass(i32), 
+    Ipv6TClass(i32),
 
     /// UDP Generic Receive Offload (GRO) allows receiving multiple UDP
     /// packets from a single sender.
@@ -1577,7 +1587,7 @@ impl<'a> ControlMessage<'a> {
 /// by ancillary data. Optionally direct the message at the given address,
 /// as with sendto.
 ///
-/// Allocates if cmsgs is nonempty.
+/// Allocates if cmsgs is nonempty, use [`sendmsg_prealloc()`] if you want to use a pre-allocated buffer.
 ///
 /// # Examples
 /// When not directing to any specific address, use `()` for the generic type
@@ -1631,6 +1641,29 @@ pub fn sendmsg<S>(fd: RawFd, iov: &[IoSlice<'_>], cmsgs: &[ControlMessage],
 }
 
 
+/// `sendmsg_prealloc()` is the same as [`sendmsg()`] but it accepts a preallocated
+/// `cmsg` buffer vector.
+///
+/// Send data in scatter-gather vectors to a socket, possibly accompanied
+/// by ancillary data. Optionally direct the message at the given address,
+/// as with sendto.
+pub fn sendmsg_prealloc<S>(fd: RawFd, iov: &[IoSlice<'_>], cmsgs: &[ControlMessage],
+               flags: MsgFlags, addr: Option<&S>, cmsg_buffer: &mut Vec<u8>) -> Result<usize>
+    where S: SockaddrLike
+{
+
+    if cmsg_buffer.len() < cmsgs.len() {
+        return Err(Errno::ENOBUFS);
+    }
+
+    let mhdr = pack_mhdr_to_send(&mut cmsg_buffer[..], iov, cmsgs, addr);
+
+    let ret = unsafe { libc::sendmsg(fd, &mhdr, flags.bits()) };
+
+    Errno::result(ret).map(|r| r as usize)
+}
+
+
 /// An extension of `sendmsg` that allows the caller to transmit multiple
 /// messages on a socket using a single system call. This has performance
 /// benefits for some applications.
@@ -2456,4 +2489,3 @@ pub fn shutdown(df: RawFd, how: Shutdown) -> Result<()> {
         Errno::result(shutdown(df, how)).map(drop)
     }
 }
-
diff --git a/test/sys/test_socket.rs b/test/sys/test_socket.rs
@@ -1204,7 +1204,6 @@ pub fn test_sendmsg_ipv4packetinfo() {
     }
 
     let cmsg = [ControlMessage::Ipv4PacketInfo(&pi)];
-
     sendmsg(
         sock.as_raw_fd(),
         &iov,
@@ -1215,6 +1214,59 @@ pub fn test_sendmsg_ipv4packetinfo() {
     .expect("sendmsg");
 }
 
+#[cfg(any(target_os = "linux", apple_targets, target_os = "netbsd"))]
+#[test]
+pub fn test_sendmsg_prealloc_ipv4packetinfo() {
+    use cfg_if::cfg_if;
+    use nix::sys::socket::{
+        bind, sendmsg_prealloc, socket, AddressFamily, ControlMessage,
+        MsgFlags, SockFlag, SockType, SockaddrIn,
+    };
+    use std::io::IoSlice;
+
+    let sock = socket(
+        AddressFamily::Inet,
+        SockType::Datagram,
+        SockFlag::empty(),
+        None,
+    )
+    .expect("socket failed");
+
+    let sock_addr = SockaddrIn::new(127, 0, 0, 1, 4000);
+
+    bind(sock.as_raw_fd(), &sock_addr).expect("bind failed");
+
+    let slice = [1u8, 2, 3, 4, 5, 6, 7, 8];
+    let iov = [IoSlice::new(&slice)];
+
+    cfg_if! {
+        if #[cfg(target_os = "netbsd")] {
+            let pi = libc::in_pktinfo {
+                ipi_ifindex: 0, /* Unspecified interface */
+                ipi_addr: libc::in_addr { s_addr: 0 },
+            };
+        } else {
+            let pi = libc::in_pktinfo {
+                ipi_ifindex: 0, /* Unspecified interface */
+                ipi_addr: libc::in_addr { s_addr: 0 },
+                ipi_spec_dst: sock_addr.as_ref().sin_addr,
+            };
+        }
+    }
+
+    let cmsg = [ControlMessage::Ipv4PacketInfo(&pi)];
+    let mut cmsg_buffer = vec![0 as u8; 32]; // same size as sendmsg_ipv4packetinfo
+    sendmsg_prealloc(
+        sock.as_raw_fd(),
+        &iov,
+        &cmsg,
+        MsgFlags::empty(),
+        Some(&sock_addr),
+        &mut cmsg_buffer,
+    )
+    .expect("sendmsg_prealloc");
+}
+
 // Verify `ControlMessage::Ipv6PacketInfo` for `sendmsg`.
 // This creates a (udp) socket bound to ip6-localhost, then sends a message to
 // itself but uses Ipv6PacketInfo to force the source address to be
