Address several change requests

Author: Gleb Pomykalov

src/sys/socket/mod.rs

@@ -783,7 +783,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) };
+    pack_mhdr(mhdr.as_mut_ptr(), &mut cmsg_buffer[..], &iov, &cmsgs, addr);
 
     let mhdr = unsafe { mhdr.assume_init() };
 
@@ -805,9 +805,9 @@ pub struct SendMmsgData<'a, I, C>
     pub _lt: std::marker::PhantomData<&'a I>,
 }
 
-/// An extension of `sendmsg`` that allows the caller to transmit multiple
+/// 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.). Supported on Linux and FreeBSD
+/// benefits for some applications.).
 ///
 /// Allocations are performed for cmsgs and to build `msghdr` buffer
 ///
@@ -817,11 +817,23 @@ pub struct SendMmsgData<'a, I, C>
 /// * `data`:           Struct that implements `IntoIterator` with `SendMmsgData` items
 /// * `flags`:          Optional flags passed directly to the operating system.
 ///
-/// # 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, I, C>(fd: RawFd, data: impl std::iter::IntoIterator<Item=&'a SendMmsgData<'a, I, C>>, flags: MsgFlags)
-    -> Result<(usize, Vec<usize>)>
+/// Returns tuple, where the first value is number of sent messages, and the second one
+/// it a `Vec` with numbers of bytes sent on each appropriate message.
+///
+///# References
+/// [`sendmmsg`](fn.sendmsg.html)
+#[cfg(any(
+    target_os = "linux",
+    target_os = "android",
+    target_os = "freebsd",
+    target_os = "openbsd",
+    target_os = "netbsd",
+))]
+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,
@@ -848,7 +860,7 @@ pub fn sendmmsg<'a, I, C>(fd: RawFd, data: impl std::iter::IntoIterator<Item=&'a
         cmsgs_buffer.resize(cmsgs_buffer_need_capacity, 0);
 
         unsafe {
-            send_pack_mhdr(
+            pack_mhdr(
                 &mut (*element.as_mut_ptr()).msg_hdr,
                 &mut cmsgs_buffer[cmsgs_start..],
                 &d.iov,
@@ -884,7 +896,7 @@ pub struct RecvMmsgData<'a, I>
     pub cmsg_buffer: Option<&'a mut Vec<u8>>,
 }
 
-/// An extension of recvmsg(2) that allows the caller to receive multiple
+/// An extension of [`recvmsg`] that allows the caller to receive multiple
 /// messages from a socket using a single system call.  (This has
 /// performance benefits for some applications.)
 ///
@@ -905,8 +917,13 @@ pub struct RecvMmsgData<'a, I>
 ///                     [`cmsg_space!`](macro.cmsg_space.html)
 ///
 /// # References
-/// [recvmmsg(2)](http://man7.org/linux/man-pages/man2/recvmmsg.2.html)
-#[cfg(any(target_os = "linux", target_os = "freebsd"))]
+/// [`recvmsg`](fn.recvmsg.html)
+#[cfg(any(
+    target_os = "linux",
+    target_os = "android",
+    target_os = "freebsd",
+    target_os = "netbsd",
+))]
 pub fn recvmmsg<'a, I>(
     fd: RawFd,
     data: impl std::iter::IntoIterator<Item=&'a mut RecvMmsgData<'a, I>>,
@@ -1034,7 +1051,7 @@ unsafe fn recv_pack_mhdr<'a, I>(
 }
 
 
-unsafe fn send_pack_mhdr<'a, I, C>(
+fn pack_mhdr<'a, I, C>(
     out: *mut msghdr,
     cmsg_buffer: &mut [u8],
     iov: I,
@@ -1050,7 +1067,7 @@ unsafe fn send_pack_mhdr<'a, I, C>(
     // Next encode the sending address, if provided
     let (name, namelen) = match addr {
         Some(addr) => {
-            let (x, y) = addr.as_ffi_pair();
+            let (x, y) = unsafe { addr.as_ffi_pair() };
             (x as *const _, y)
         },
         None => (ptr::null(), 0),
@@ -1065,27 +1082,31 @@ unsafe fn send_pack_mhdr<'a, I, C>(
 
     // Musl's msghdr has private fields, so this is the only way to
     // initialize it.
-    (*out).msg_name = name as *mut _;
-    (*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_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;
+    {
+        let mut hdr = unsafe { &mut *out };
+
+        hdr.msg_name = name as *mut _;
+        hdr.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
+        hdr.msg_iov = iov.as_ref().as_ptr() as *mut _;
+        hdr.msg_iovlen = iov.as_ref().len() as _;
+        hdr.msg_control = cmsg_ptr;
+        hdr.msg_controllen = cmsg_capacity as _;
+        hdr.msg_flags = 0;
+    }
 
     // Encode each cmsg.  This must happen after initializing the header because
     // 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);
+    let mut pmhdr: *mut cmsghdr = unsafe { CMSG_FIRSTHDR(out) };
     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
-        cmsg.encode_into(pmhdr);
+        unsafe { cmsg.encode_into(pmhdr); }
         // Safe because mhdr is valid
-        pmhdr = CMSG_NXTHDR(out, pmhdr);
+        pmhdr = unsafe { CMSG_NXTHDR(out, pmhdr) };
     }
 }
 

test/sys/test_socket.rs

@@ -227,7 +227,13 @@ mod recvfrom {
         assert_eq!(AddressFamily::Inet, from.unwrap().family());
     }
 
-    #[cfg(any(target_os = "linux", target_os = "freebsd"))]
+    #[cfg(any(
+        target_os = "linux",
+        target_os = "android",
+        target_os = "freebsd",
+        target_os = "openbsd",
+        target_os = "netbsd",
+    ))]
     #[test]
     pub fn udp_sendmmsg() {
         use nix::sys::uio::IoVec;
@@ -262,7 +268,10 @@ mod recvfrom {
                     addr: Some(sock_addr),
                     _lt: Default::default(),
                 });
-            for _ in 0..15 {
+
+            let batch_size = 15;
+
+            for _ in 0..batch_size {
                 msgs.push(
                     SendMmsgData {
                         iov: &iov,
@@ -286,12 +295,20 @@ mod recvfrom {
         assert_eq!(AddressFamily::Inet, from.unwrap().family());
     }
 
-    #[cfg(any(target_os = "linux", target_os = "freebsd"))]
+    #[cfg(any(
+        target_os = "linux",
+        target_os = "android",
+        target_os = "freebsd",
+        target_os = "netbsd",
+    ))]
     #[test]
     pub fn udp_recvmmsg() {
         use nix::sys::uio::IoVec;
         use nix::sys::socket::{MsgFlags, recvmmsg};
 
+        const NUM_MESSAGES_SENT: usize = 2;
+        const DATA: [u8; 2] = [1,2];
+
         let std_sa = SocketAddr::from_str("127.0.0.1:6798").unwrap();
         let inet_addr = InetAddr::from_std(&std_sa);
         let sock_addr = SockAddr::new_inet(inet_addr);
@@ -309,38 +326,39 @@ mod recvfrom {
             None,
         ).expect("send socket failed");
 
+
         let send_thread = thread::spawn(move || {
-            for _ in 0..2 {
-                sendto(ssock, &b"12"[..], &sock_addr, MsgFlags::empty()).unwrap();
+            for _ in 0..NUM_MESSAGES_SENT {
+                sendto(ssock, &DATA[..], &sock_addr, MsgFlags::empty()).unwrap();
             }
         });
 
         let mut msgs = std::collections::LinkedList::new();
-        let mut rec_buf1 = [0u8; 32];
-        let mut rec_buf2 = [0u8; 32];
-        let iov1 = [IoVec::from_mut_slice(&mut rec_buf1[..])];
-        let iov2 = [IoVec::from_mut_slice(&mut rec_buf2[..])];
-        msgs.push_back(
-            RecvMmsgData {
-                iov: &iov1,
-                cmsg_buffer: None,
-            }
-        );
-        msgs.push_back(
-            RecvMmsgData {
-                iov: &iov2,
+
+        // Buffers to receive exactly `NUM_MESSAGES_SENT` messages
+        let mut receive_buffers = [[0u8; 32]; NUM_MESSAGES_SENT];
+        let iovs: Vec<_> = receive_buffers.iter_mut().map(|buf| {
+            [IoVec::from_mut_slice(&mut buf[..])]
+        }).collect();
+
+        for iov in &iovs {
+            msgs.push_back(RecvMmsgData {
+                iov: iov,
                 cmsg_buffer: None,
-            }
-        );
+            })
+        };
+
         let res = recvmmsg(rsock, &mut msgs, MsgFlags::empty(), None).expect("recvmmsg");
-        assert_eq!(res.len(), 2);
+        assert_eq!(res.len(), DATA.len());
 
         for RecvMsg { address, bytes, .. } in res.into_iter() {
             assert_eq!(AddressFamily::Inet, address.unwrap().family());
-            assert_eq!(2, bytes);
+            assert_eq!(DATA.len(), bytes);
+        }
+
+        for buf in &receive_buffers {
+            assert_eq!(&buf[..DATA.len()], DATA);
         }
-        assert_eq!(&rec_buf1[..2], b"12");
-        assert_eq!(&rec_buf2[..2], b"12");
 
         send_thread.join().unwrap();
     }