Reduce unsafety

Author: Gleb Pomykalov

src/sys/socket/mod.rs

@@ -774,17 +774,13 @@ impl<'a> ControlMessage<'a> {
 pub fn sendmsg(fd: RawFd, iov: &[IoVec<&[u8]>], cmsgs: &[ControlMessage],
                flags: MsgFlags, addr: Option<&SockAddr>) -> Result<usize>
 {
-    let mut mhdr = mem::MaybeUninit::<msghdr>::zeroed();
-
     let capacity = cmsgs.iter().map(|c| c.space()).sum();
 
     // First size the buffer needed to hold the cmsgs.  It must be zeroed,
     // because subsequent code will not clear the padding bytes.
     let mut cmsg_buffer = vec![0u8; capacity];
 
-    pack_mhdr_to_send(mhdr.as_mut_ptr(), &mut cmsg_buffer[..], &iov, &cmsgs, addr);
-
-    let mhdr = unsafe { mhdr.assume_init() };
+    let mhdr = pack_mhdr_to_send(&mut cmsg_buffer[..], &iov, &cmsgs, addr);
 
     let ret = unsafe { libc::sendmsg(fd, &mhdr, flags.bits()) };
 
@@ -836,8 +832,7 @@ pub struct SendMmsgData<'a, I, C>
 ))]
 pub fn sendmmsg<'a, I, C>(
     fd: RawFd,
-    data: impl std::iter::IntoIterator<Item=&'a SendMmsgData<'a, I, C>,
-        IntoIter=impl ExactSizeIterator + Iterator<Item=&'a SendMmsgData<'a, I, C>>>,
+    data: impl std::iter::IntoIterator<Item=&'a SendMmsgData<'a, I, C>>,
     flags: MsgFlags
 ) -> Result<Vec<usize>>
     where
@@ -846,31 +841,29 @@ pub fn sendmmsg<'a, I, C>(
 {
     let iter = data.into_iter();
 
-    let num_messages = iter.len();
+    let size_hint = iter.size_hint();
+    let reserve_items = size_hint.1.unwrap_or(size_hint.0);
 
-    let mut output = Vec::<libc::mmsghdr>::with_capacity(num_messages);
-    unsafe {
-        output.set_len(num_messages);
-    }
+    let mut output = Vec::<libc::mmsghdr>::with_capacity(reserve_items);
 
     let mut cmsgs_buffer = vec![0u8; 0];
 
-    iter.enumerate().for_each(|(i, d)| {
-        let element = &mut output[i];
-
+    for d in iter {
         let cmsgs_start = cmsgs_buffer.len();
         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);
 
-        pack_mhdr_to_send(
-            &mut element.msg_hdr,
-            &mut cmsgs_buffer[cmsgs_start..],
-            &d.iov,
-            &d.cmsgs,
-            d.addr.as_ref()
-        );
-    });
+        output.push(libc::mmsghdr {
+            msg_hdr: pack_mhdr_to_send(
+                &mut cmsgs_buffer[cmsgs_start..],
+                &d.iov,
+                &d.cmsgs,
+                d.addr.as_ref()
+            ),
+            msg_len: 0,
+        });
+    };
 
     let ret = unsafe { libc::sendmmsg(fd, output.as_mut_ptr(), output.len() as _, flags.bits() as _) };
 
@@ -947,25 +940,27 @@ pub fn recvmmsg<'a, I>(
     let num_messages = iter.len();
 
     let mut output: Vec<libc::mmsghdr> = Vec::with_capacity(num_messages);
-    let mut address: Vec<sockaddr_storage> = Vec::with_capacity(num_messages);
 
-    unsafe {
-        output.set_len(num_messages);
-        address.set_len(num_messages);
-    }
+    // Addresses should be pre-allocated and never change the address during building
+    // of the input data for `recvmmsg`
+    let mut addresses: Vec<sockaddr_storage> = vec![unsafe { mem::zeroed() }; num_messages];
 
     let results: Vec<_> = iter.enumerate().map(|(i, d)| {
-        let element = &mut output[i];
-
-        let msg_controllen = unsafe {
+        let (msg_controllen, mhdr) = unsafe {
             pack_mhdr_to_receive(
-                &mut element.msg_hdr,
                 d.iov.as_ref(),
                 &mut d.cmsg_buffer,
-                &mut address[i]
+                &mut addresses[i],
             )
         };
 
+        output.push(
+            libc::mmsghdr {
+                msg_hdr: mhdr,
+                msg_len: 0,
+            }
+        );
+
         (msg_controllen as usize, &mut d.cmsg_buffer)
     }).collect();
 
@@ -981,15 +976,15 @@ pub fn recvmmsg<'a, I>(
 
     Ok(output
         .into_iter()
-        .zip(address.into_iter())
+        .zip(addresses.into_iter())
         .zip(results.into_iter())
-        .map(|((mmsghdr, mut address), (msg_controllen, cmsg_buffer))| {
+        .map(|((mmsghdr, address), (msg_controllen, cmsg_buffer))| {
             unsafe {
                 read_mhdr(
                     mmsghdr.msg_hdr,
                     r as isize,
                     msg_controllen,
-                    &mut address,
+                    address,
                     cmsg_buffer
                 )
             }
@@ -1001,7 +996,7 @@ unsafe fn read_mhdr<'a, 'b>(
     mhdr: msghdr,
     r: isize,
     msg_controllen: usize,
-    address: *mut sockaddr_storage,
+    address: sockaddr_storage,
     cmsg_buffer: &'a mut Option<&'b mut Vec<u8>>
 ) -> RecvMsg<'b> {
     let cmsghdr = {
@@ -1020,7 +1015,7 @@ unsafe fn read_mhdr<'a, 'b>(
     };
 
     let address = sockaddr_storage_to_addr(
-        &*address ,
+        &address ,
          mhdr.msg_namelen as usize
     ).ok();
 
@@ -1034,42 +1029,46 @@ unsafe fn read_mhdr<'a, 'b>(
 }
 
 unsafe fn pack_mhdr_to_receive<'a, I>(
-    out: *mut msghdr,
     iov: I,
     cmsg_buffer: &mut Option<&mut Vec<u8>>,
     address: *mut sockaddr_storage,
-) -> usize
+) -> (usize, msghdr)
     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 c_void;
-    (*out).msg_namelen = mem::size_of::<sockaddr_storage>() as socklen_t;
-    (*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;
+    let mhdr = {
+        // Musl's msghdr has private fields, so this is the only way to
+        // initialize it.
+        let mut mhdr = mem::MaybeUninit::<msghdr>::zeroed();
+        let p = mhdr.as_mut_ptr();
+        (*p).msg_name = address as *mut c_void;
+        (*p).msg_namelen = mem::size_of::<sockaddr_storage>() as socklen_t;
+        (*p).msg_iov = iov.as_ref().as_ptr() as *mut iovec;
+        (*p).msg_iovlen = iov.as_ref().len() as _;
+        (*p).msg_control = msg_control as *mut c_void;
+        (*p).msg_controllen = msg_controllen as _;
+        (*p).msg_flags = 0;
+        mhdr.assume_init()
+    };
 
-    msg_controllen
+    (msg_controllen, mhdr)
 }
 
-
 fn pack_mhdr_to_send<'a, I, C>(
-    out: *mut msghdr,
     cmsg_buffer: &mut [u8],
     iov: I,
     cmsgs: C,
     addr: Option<&SockAddr>
-)
+) -> msghdr
     where
         I: AsRef<[IoVec<&'a [u8]>]>,
         C: AsRef<[ControlMessage<'a>]>
 {
-    let cmsg_capacity = cmsg_buffer.len();
+    let capacity = cmsg_buffer.len();
 
     // Next encode the sending address, if provided
     let (name, namelen) = match addr {
@@ -1081,38 +1080,43 @@ fn pack_mhdr_to_send<'a, I, C>(
     };
 
     // The message header must be initialized before the individual cmsgs.
-    let cmsg_ptr = if cmsg_capacity > 0 {
+    let cmsg_ptr = if capacity > 0 {
         cmsg_buffer.as_ptr() as *mut c_void
     } else {
         ptr::null_mut()
     };
 
-    // Musl's msghdr has private fields, so this is the only way to
-    // initialize it.
-    unsafe {
-        (*out).msg_name = name as *mut _;
-        (*out).msg_namelen = namelen;
+    let mhdr = unsafe {
+        // Musl's msghdr has private fields, so this is the only way to
+        // initialize it.
+        let mut mhdr = mem::MaybeUninit::<msghdr>::zeroed();
+        let p = mhdr.as_mut_ptr();
+        (*p).msg_name = name as *mut _;
+        (*p).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;
-    }
+        (*p).msg_iov = iov.as_ref().as_ptr() as *mut _;
+        (*p).msg_iovlen = iov.as_ref().len() as _;
+        (*p).msg_control = cmsg_ptr;
+        (*p).msg_controllen = capacity as _;
+        (*p).msg_flags = 0;
+        mhdr.assume_init()
+    };
 
     // 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 = unsafe { CMSG_FIRSTHDR(out) };
+    let mut pmhdr: *mut cmsghdr = unsafe { CMSG_FIRSTHDR(&mhdr as *const msghdr) };
     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
-        unsafe { cmsg.encode_into(pmhdr); }
+        unsafe { cmsg.encode_into(pmhdr) };
         // Safe because mhdr is valid
-        pmhdr = unsafe { CMSG_NXTHDR(out, pmhdr) };
+        pmhdr = unsafe { CMSG_NXTHDR(&mhdr as *const msghdr, pmhdr) };
     }
+
+    mhdr
 }
 
 /// Receive message in scatter-gather vectors from a socket, and
@@ -1133,20 +1137,17 @@ pub fn recvmsg<'a>(fd: RawFd, iov: &[IoVec<&mut [u8]>],
                    mut cmsg_buffer: Option<&'a mut Vec<u8>>,
                    flags: MsgFlags) -> Result<RecvMsg<'a>>
 {
-    let mut out = mem::MaybeUninit::<msghdr>::zeroed();
     let mut address = mem::MaybeUninit::uninit();
 
-    let msg_controllen = unsafe {
-        pack_mhdr_to_receive(out.as_mut_ptr(), &iov, &mut cmsg_buffer, address.as_mut_ptr())
+    let (msg_controllen, mut mhdr) = unsafe {
+        pack_mhdr_to_receive(&iov, &mut cmsg_buffer, address.as_mut_ptr())
     };
 
-    let mut mhdr = unsafe { out.assume_init() };
-
     let ret = unsafe { libc::recvmsg(fd, &mut mhdr, flags.bits()) };
 
     let r = Errno::result(ret)?;
 
-    Ok(unsafe { read_mhdr(mhdr, r, msg_controllen, address.as_mut_ptr(), &mut cmsg_buffer) })
+    Ok(unsafe { read_mhdr(mhdr, r, msg_controllen, address.assume_init(), &mut cmsg_buffer) })
 }