Improve ergonomics of recvmmsg/sendmmsg

Author: Gleb Pomykalov

src/sys/socket/mod.rs

@@ -772,7 +772,7 @@ pub fn sendmsg(fd: RawFd, iov: &[IoVec<&[u8]>], cmsgs: &[ControlMessage],
     // because subsequent code will not clear the padding bytes.
     let mut cmsg_buffer = vec![0u8; capacity];
 
-    unsafe { send_pack_mhdr(mhdr.as_mut_ptr(), &mut cmsg_buffer[..], iov, cmsgs, addr) };
+    unsafe { send_pack_mhdr(mhdr.as_mut_ptr(), &mut cmsg_buffer[..], &iov, &cmsgs, addr) };
 
     let mhdr = unsafe { mhdr.assume_init() };
 
@@ -783,10 +783,15 @@ pub fn sendmsg(fd: RawFd, iov: &[IoVec<&[u8]>], cmsgs: &[ControlMessage],
 
 #[cfg(any(target_os = "linux", target_os = "freebsd"))]
 #[derive(Debug)]
-pub struct SendMmsgData<'a> {
-    pub iov: &'a [IoVec<&'a [u8]>],
-    pub cmsgs: &'a [ControlMessage<'a>],
-    pub addr: Option<&'a SockAddr>,
+pub struct SendMmsgData<'a, I, C>
+    where
+        I: AsRef<[IoVec<&'a [u8]>]>,
+        C: AsRef<[ControlMessage<'a>]>
+{
+    pub iov: I,
+    pub cmsgs: C,
+    pub addr: Option<SockAddr>,
+    pub _lt: std::marker::PhantomData<&'a I>,
 }
 
 /// An extension of `sendmsg`` that allows the caller to transmit multiple
@@ -804,7 +809,11 @@ pub struct SendMmsgData<'a> {
 /// # References
 /// [sendmmsg(2)](http://man7.org/linux/man-pages/man2/sendmmsg.2.html)
 #[cfg(any(target_os = "linux", target_os = "freebsd"))]
-pub fn sendmmsg<'a>(fd: RawFd, data: impl std::iter::IntoIterator<Item=&'a SendMmsgData<'a>>, flags: MsgFlags) -> Result<(usize, Vec<usize>)>
+pub fn sendmmsg<'a, I, C>(fd: RawFd, data: impl std::iter::IntoIterator<Item=&'a SendMmsgData<'a, I, C>>, flags: MsgFlags)
+    -> Result<(usize, Vec<usize>)>
+    where
+        I: AsRef<[IoVec<&'a [u8]>]> + 'a,
+        C: AsRef<[ControlMessage<'a>]> + 'a,
 {
     let iter = data.into_iter();
 
@@ -823,11 +832,19 @@ pub fn sendmmsg<'a>(fd: RawFd, data: impl std::iter::IntoIterator<Item=&'a SendM
         let element = &mut output[i];
 
         let cmsgs_start = cmsgs_buffer.len();
-        let cmsgs_required_capacity: usize = d.cmsgs.iter().map(|c| c.space()).sum();
+        let cmsgs_required_capacity: usize = d.cmsgs.as_ref().iter().map(|c| c.space()).sum();
         let cmsgs_buffer_need_capacity = cmsgs_start + cmsgs_required_capacity;
         cmsgs_buffer.resize(cmsgs_buffer_need_capacity, 0);
 
-        unsafe { send_pack_mhdr(&mut (*element.as_mut_ptr()).msg_hdr, &mut cmsgs_buffer[cmsgs_start..], d.iov, d.cmsgs, d.addr) };
+        unsafe {
+            send_pack_mhdr(
+                &mut (*element.as_mut_ptr()).msg_hdr,
+                &mut cmsgs_buffer[cmsgs_start..],
+                &d.iov,
+                &d.cmsgs,
+                d.addr.as_ref()
+            )
+        };
     });
 
     let mut initialized_data = unsafe { mem::transmute::<_, Vec<libc::mmsghdr>>(output) };
@@ -848,8 +865,11 @@ pub fn sendmmsg<'a>(fd: RawFd, data: impl std::iter::IntoIterator<Item=&'a SendM
 
 #[cfg(any(target_os = "linux", target_os = "freebsd"))]
 #[derive(Debug)]
-pub struct RecvMmsgData<'a> {
-    pub iov: &'a [IoVec<&'a mut [u8]>],
+pub struct RecvMmsgData<'a, I>
+    where
+        I: AsRef<[IoVec<&'a mut [u8]>]> + 'a,
+{
+    pub iov: I,
     pub cmsg_buffer: Option<&'a mut Vec<u8>>,
 }
 
@@ -876,9 +896,13 @@ pub struct RecvMmsgData<'a> {
 /// # References
 /// [recvmmsg(2)](http://man7.org/linux/man-pages/man2/recvmmsg.2.html)
 #[cfg(any(target_os = "linux", target_os = "freebsd"))]
-pub fn recvmmsg<'a>(fd: RawFd,
-                    data: impl std::iter::IntoIterator<Item=&'a mut RecvMmsgData<'a>>,
-                    flags: MsgFlags, timeout: Option<crate::sys::time::TimeSpec>) -> Result<Vec<RecvMsg<'a>>>
+pub fn recvmmsg<'a, I>(
+    fd: RawFd,
+    data: impl std::iter::IntoIterator<Item=&'a mut RecvMmsgData<'a, I>>,
+    flags: MsgFlags, timeout: Option<crate::sys::time::TimeSpec>
+) -> Result<Vec<RecvMsg<'a>>>
+    where
+        I: AsRef<[IoVec<&'a mut [u8]>]> + 'a,
 {
     let iter = data.into_iter();
 
@@ -899,7 +923,7 @@ pub fn recvmmsg<'a>(fd: RawFd,
         let msg_controllen = unsafe {
             recv_pack_mhdr(
                 &mut (*element.as_mut_ptr()).msg_hdr,
-                d.iov,
+                d.iov.as_ref(),
                 &mut d.cmsg_buffer,
                 &mut address[i]
             )
@@ -974,15 +998,23 @@ unsafe fn recv_read_mhdr<'a, 'b>(
     }
 }
 
-unsafe fn recv_pack_mhdr(out: *mut msghdr, iov: &[IoVec<&mut [u8]>], cmsg_buffer: &mut Option<&mut Vec<u8>>, address: &mut mem::MaybeUninit<sockaddr_storage>) -> usize {
+unsafe fn recv_pack_mhdr<'a, I>(
+    out: *mut msghdr,
+    iov: I,
+    cmsg_buffer: &mut Option<&mut Vec<u8>>,
+    address: &mut mem::MaybeUninit<sockaddr_storage>
+) -> usize
+    where
+        I: AsRef<[IoVec<&'a mut [u8]>]> + 'a,
+{
     let (msg_control, msg_controllen) = cmsg_buffer.as_mut()
         .map(|v| (v.as_mut_ptr(), v.capacity()))
         .unwrap_or((ptr::null_mut(), 0));
 
     (*out).msg_name = address.as_mut_ptr() as *mut c_void;
     (*out).msg_namelen = mem::size_of::<sockaddr_storage>() as socklen_t;
-    (*out).msg_iov = iov.as_ptr() as *mut iovec;
-    (*out).msg_iovlen = iov.len() as _;
+    (*out).msg_iov = iov.as_ref().as_ptr() as *mut iovec;
+    (*out).msg_iovlen = iov.as_ref().len() as _;
     (*out).msg_control = msg_control as *mut c_void;
     (*out).msg_controllen = msg_controllen as _;
     (*out).msg_flags = 0;
@@ -991,7 +1023,17 @@ unsafe fn recv_pack_mhdr(out: *mut msghdr, iov: &[IoVec<&mut [u8]>], cmsg_buffer
 }
 
 
-unsafe fn send_pack_mhdr(out: *mut msghdr, cmsg_buffer: &mut [u8], iov: &[IoVec<&[u8]>], cmsgs: &[ControlMessage], addr: Option<&SockAddr>) {
+unsafe fn send_pack_mhdr<'a, I, C>(
+    out: *mut msghdr,
+    cmsg_buffer: &mut [u8],
+    iov: I,
+    cmsgs: C,
+    addr: Option<&SockAddr>
+)
+    where
+        I: AsRef<[IoVec<&'a [u8]>]>,
+        C: AsRef<[ControlMessage<'a>]>
+{
     let cmsg_capacity = cmsg_buffer.len();
 
     // Next encode the sending address, if provided
@@ -1016,8 +1058,8 @@ unsafe fn send_pack_mhdr(out: *mut msghdr, cmsg_buffer: &mut [u8], iov: &[IoVec<
     (*out).msg_namelen = namelen;
     // transmute iov into a mutable pointer.  sendmsg doesn't really mutate
     // the buffer, but the standard says that it takes a mutable pointer
-    (*out).msg_iov = iov.as_ptr() as *mut _;
-    (*out).msg_iovlen = iov.len() as _;
+    (*out).msg_iov = iov.as_ref().as_ptr() as *mut _;
+    (*out).msg_iovlen = iov.as_ref().len() as _;
     (*out).msg_control = cmsg_ptr;
     (*out).msg_controllen = cmsg_capacity as _;
     (*out).msg_flags = 0;
@@ -1026,7 +1068,7 @@ unsafe fn send_pack_mhdr(out: *mut msghdr, cmsg_buffer: &mut [u8], iov: &[IoVec<
     // CMSG_NEXT_HDR and friends read the msg_control and msg_controllen fields.
     // CMSG_FIRSTHDR is always safe
     let mut pmhdr: *mut cmsghdr = CMSG_FIRSTHDR(out);
-    for cmsg in cmsgs {
+    for cmsg in cmsgs.as_ref() {
         assert_ne!(pmhdr, ptr::null_mut());
         // Safe because we know that pmhdr is valid, and we initialized it with
         // sufficient space
@@ -1057,7 +1099,9 @@ pub fn recvmsg<'a>(fd: RawFd, iov: &'a [IoVec<&'a mut [u8]>],
     let mut out = mem::MaybeUninit::<msghdr>::zeroed();
     let mut address = mem::MaybeUninit::uninit();
 
-    let msg_controllen = unsafe { recv_pack_mhdr(out.as_mut_ptr(), iov, &mut cmsg_buffer, &mut address) };
+    let msg_controllen = unsafe {
+        recv_pack_mhdr(out.as_mut_ptr(), &iov, &mut cmsg_buffer, &mut address)
+    };
 
     let mut mhdr = unsafe { out.assume_init() };
 

test/sys/test_socket.rs

@@ -254,30 +254,33 @@ mod recvfrom {
 
         let from = sendrecv(rsock, ssock, move |s, m, flags| {
             let iov = [IoVec::from_slice(m)];
-            let mut msgs = std::collections::LinkedList::new();
-            msgs.push_back(
+            let mut msgs = Vec::new();
+            msgs.push(
                 SendMmsgData {
                     iov: &iov,
                     cmsgs: &[],
-                    addr: Some(&sock_addr),
+                    addr: Some(sock_addr),
+                    _lt: Default::default(),
                 });
             for _ in 0..15 {
-                msgs.push_back(
+                msgs.push(
                     SendMmsgData {
                         iov: &iov,
                         cmsgs: &[],
-                        addr: Some(&sock_addr2),
+                        addr: Some(sock_addr2),
+                        _lt: Default::default(),
                     }
                 );
             }
-            sendmmsg(s, &msgs, flags).map(move |(sent_messages, sent_bytes)| {
-                assert!(sent_messages >= 1);
-                assert_eq!(sent_bytes.len(), sent_messages);
-                for sent in &sent_bytes {
-                    assert_eq!(*sent, m.len());
-                }
-                sent_messages
-            })
+            sendmmsg(s, msgs.iter(), flags)
+                .map(move |(sent_messages, sent_bytes)| {
+                    assert!(sent_messages >= 1);
+                    assert_eq!(sent_bytes.len(), sent_messages);
+                    for sent in &sent_bytes {
+                        assert_eq!(*sent, m.len());
+                    }
+                    sent_messages
+                })
         });
         // UDP sockets should set the from address
         assert_eq!(AddressFamily::Inet, from.unwrap().family());
@@ -332,7 +335,7 @@ mod recvfrom {
         let res = recvmmsg(rsock, &mut msgs, MsgFlags::empty(), None).expect("recvmmsg");
         assert_eq!(res.len(), 2);
 
-        for RecvMsg { address, bytes, ..} in res.into_iter() {
+        for RecvMsg { address, bytes, .. } in res.into_iter() {
             assert_eq!(AddressFamily::Inet, address.unwrap().family());
             assert_eq!(2, bytes);
         }