Change Socket::recv_vectored to accept uninitialised buffers

Also includes:
* Socket::recv_vectored_with_flags
* Socket::recv_from_vectored
* Socket::recv_from_vectored_with_flags

Author: Thomasdezeeuw

src/lib.rs

@@ -66,7 +66,10 @@
 // Disallow warnings in examples.
 #![doc(test(attr(deny(warnings))))]
 
+use std::fmt;
+use std::mem::MaybeUninit;
 use std::net::SocketAddr;
+use std::ops::{Deref, DerefMut};
 use std::time::Duration;
 
 /// Macro to implement `fmt::Debug` for a type, printing the constant names
@@ -277,6 +280,47 @@ impl RecvFlags {
     }
 }
 
+/// A version of [`IoSliceMut`] that allows the buffer to be uninitialised.
+///
+/// [`IoSliceMut`]: std::io::IoSliceMut
+#[repr(transparent)]
+pub struct MaybeUninitSlice<'a>(sys::MaybeUninitSlice<'a>);
+
+unsafe impl<'a> Send for MaybeUninitSlice<'a> {}
+
+unsafe impl<'a> Sync for MaybeUninitSlice<'a> {}
+
+impl<'a> fmt::Debug for MaybeUninitSlice<'a> {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        fmt::Debug::fmt(self.0.as_slice(), fmt)
+    }
+}
+
+impl<'a> MaybeUninitSlice<'a> {
+    /// Creates a new `MaybeUninitSlice` wrapping a byte slice.
+    ///
+    /// # Panics
+    ///
+    /// Panics on Windows if the slice is larger than 4GB.
+    pub fn new(buf: &'a mut [MaybeUninit<u8>]) -> MaybeUninitSlice<'a> {
+        MaybeUninitSlice(sys::MaybeUninitSlice::new(buf))
+    }
+}
+
+impl<'a> Deref for MaybeUninitSlice<'a> {
+    type Target = [MaybeUninit<u8>];
+
+    fn deref(&self) -> &[MaybeUninit<u8>] {
+        self.0.as_slice()
+    }
+}
+
+impl<'a> DerefMut for MaybeUninitSlice<'a> {
+    fn deref_mut(&mut self) -> &mut [MaybeUninit<u8>] {
+        self.0.as_mut_slice()
+    }
+}
+
 /// Configures a socket's TCP keepalive parameters.
 ///
 /// See [`Socket::set_tcp_keepalive`].

src/socket.rs

@@ -21,7 +21,7 @@ use std::time::Duration;
 use crate::sys::{self, c_int, getsockopt, setsockopt, Bool};
 #[cfg(not(target_os = "redox"))]
 use crate::RecvFlags;
-use crate::{Domain, Protocol, SockAddr, TcpKeepalive, Type};
+use crate::{Domain, MaybeUninitSlice, Protocol, SockAddr, TcpKeepalive, Type};
 
 /// Owned wrapper around a system socket.
 ///
@@ -339,19 +339,42 @@ impl Socket {
     ///
     /// [`recv`]: Socket::recv
     /// [`connect`]: Socket::connect
+    ///
+    /// # Safety
+    ///
+    /// Normally casting a `IoSliceMut` to `MaybeUninitSlice` would be unsound,
+    /// as that allows us to write uninitialised bytes to the buffer. However
+    /// this implementation promises to not write uninitialised bytes to the
+    /// `bufs` and passes it directly to `recvmsg(2)` system call. This promise
+    /// ensures that this function can be called using `bufs` of type `&mut
+    /// [IoSliceMut]`.
+    ///
+    /// Note that the [`io::Read::read_vectored`] implementation calls this
+    /// function with `buf`s of type `&mut [IoSliceMut]`, allowing initialised
+    /// buffers to be used without using `unsafe`.
     #[cfg(not(target_os = "redox"))]
-    pub fn recv_vectored(&self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<(usize, RecvFlags)> {
+    pub fn recv_vectored(
+        &self,
+        bufs: &mut [MaybeUninitSlice<'_>],
+    ) -> io::Result<(usize, RecvFlags)> {
         self.recv_vectored_with_flags(bufs, 0)
     }
 
     /// Identical to [`recv_vectored`] but allows for specification of arbitrary
     /// flags to the underlying `recvmsg`/`WSARecv` call.
     ///
     /// [`recv_vectored`]: Socket::recv_vectored
+    ///
+    /// # Safety
+    ///
+    /// `recv_from_vectored` makes the same safety guarantees regarding `bufs`
+    /// as [`recv_vectored`].
+    ///
+    /// [`recv_vectored`]: Socket::recv_vectored
     #[cfg(not(target_os = "redox"))]
     pub fn recv_vectored_with_flags(
         &self,
-        bufs: &mut [IoSliceMut<'_>],
+        bufs: &mut [MaybeUninitSlice<'_>],
         flags: i32,
     ) -> io::Result<(usize, RecvFlags)> {
         sys::recv_vectored(self.inner, bufs, flags)
@@ -404,10 +427,17 @@ impl Socket {
     /// [`recv_from`] this allows passing multiple buffers.
     ///
     /// [`recv_from`]: Socket::recv_from
+    ///
+    /// # Safety
+    ///
+    /// `recv_from_vectored` makes the same safety guarantees regarding `bufs`
+    /// as [`recv_vectored`].
+    ///
+    /// [`recv_vectored`]: Socket::recv_vectored
     #[cfg(not(target_os = "redox"))]
     pub fn recv_from_vectored(
         &self,
-        bufs: &mut [IoSliceMut<'_>],
+        bufs: &mut [MaybeUninitSlice<'_>],
     ) -> io::Result<(usize, RecvFlags, SockAddr)> {
         self.recv_from_vectored_with_flags(bufs, 0)
     }
@@ -416,10 +446,17 @@ impl Socket {
     /// arbitrary flags to the underlying `recvmsg`/`WSARecvFrom` call.
     ///
     /// [`recv_from_vectored`]: Socket::recv_from_vectored
+    ///
+    /// # Safety
+    ///
+    /// `recv_from_vectored` makes the same safety guarantees regarding `bufs`
+    /// as [`recv_vectored`].
+    ///
+    /// [`recv_vectored`]: Socket::recv_vectored
     #[cfg(not(target_os = "redox"))]
     pub fn recv_from_vectored_with_flags(
         &self,
-        bufs: &mut [IoSliceMut<'_>],
+        bufs: &mut [MaybeUninitSlice<'_>],
         flags: i32,
     ) -> io::Result<(usize, RecvFlags, SockAddr)> {
         sys::recv_from_vectored(self.inner, bufs, flags)
@@ -1292,6 +1329,11 @@ impl Read for Socket {
 
     #[cfg(not(target_os = "redox"))]
     fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
+        // Safety: both `IoSliceMut` and `MaybeUninitSlice` promise to have the
+        // same layout, that of `iovec`/`WSABUF`. Furthermore `recv_vectored`
+        // promises to not write unitialised bytes to the `bufs` and pass it
+        // directly to the `recvmsg` system call, so this is safe.
+        let bufs = unsafe { &mut *(bufs as *mut [IoSliceMut<'_>] as *mut [MaybeUninitSlice<'_>]) };
         self.recv_vectored(bufs).map(|(n, _)| n)
     }
 }
@@ -1305,6 +1347,8 @@ impl<'a> Read for &'a Socket {
 
     #[cfg(not(target_os = "redox"))]
     fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
+        // Safety: see other `Read::read` impl.
+        let bufs = unsafe { &mut *(bufs as *mut [IoSliceMut<'_>] as *mut [MaybeUninitSlice<'_>]) };
         self.recv_vectored(bufs).map(|(n, _)| n)
     }
 }

src/sys/unix.rs

@@ -10,7 +10,8 @@ use std::cmp::min;
 #[cfg(all(feature = "all", target_os = "linux"))]
 use std::ffi::{CStr, CString};
 #[cfg(not(target_os = "redox"))]
-use std::io::{IoSlice, IoSliceMut};
+use std::io::IoSlice;
+use std::marker::PhantomData;
 use std::mem::{self, size_of, MaybeUninit};
 use std::net::Shutdown;
 use std::net::{Ipv4Addr, Ipv6Addr};
@@ -21,10 +22,8 @@ use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd};
 use std::os::unix::net::{UnixDatagram, UnixListener, UnixStream};
 #[cfg(feature = "all")]
 use std::path::Path;
-#[cfg(all(feature = "all", target_os = "linux"))]
-use std::slice;
 use std::time::Duration;
-use std::{io, ptr};
+use std::{io, ptr, slice};
 
 #[cfg(not(target_vendor = "apple"))]
 use libc::ssize_t;
@@ -308,6 +307,32 @@ impl std::fmt::Debug for RecvFlags {
     }
 }
 
+#[repr(transparent)]
+pub struct MaybeUninitSlice<'a> {
+    vec: libc::iovec,
+    _lifetime: PhantomData<&'a mut [MaybeUninit<u8>]>,
+}
+
+impl<'a> MaybeUninitSlice<'a> {
+    pub(crate) fn new(buf: &'a mut [MaybeUninit<u8>]) -> MaybeUninitSlice<'a> {
+        MaybeUninitSlice {
+            vec: libc::iovec {
+                iov_base: buf.as_mut_ptr().cast(),
+                iov_len: buf.len(),
+            },
+            _lifetime: PhantomData,
+        }
+    }
+
+    pub(crate) fn as_slice(&self) -> &[MaybeUninit<u8>] {
+        unsafe { slice::from_raw_parts(self.vec.iov_base.cast(), self.vec.iov_len) }
+    }
+
+    pub(crate) fn as_mut_slice(&mut self) -> &mut [MaybeUninit<u8>] {
+        unsafe { slice::from_raw_parts_mut(self.vec.iov_base.cast(), self.vec.iov_len) }
+    }
+}
+
 /// Unix only API.
 impl SockAddr {
     /// Constructs a `SockAddr` with the family `AF_UNIX` and the provided path.
@@ -453,7 +478,7 @@ pub(crate) fn recv_from(
 #[cfg(not(target_os = "redox"))]
 pub(crate) fn recv_vectored(
     fd: Socket,
-    bufs: &mut [IoSliceMut<'_>],
+    bufs: &mut [crate::MaybeUninitSlice<'_>],
     flags: c_int,
 ) -> io::Result<(usize, RecvFlags)> {
     recvmsg(fd, ptr::null_mut(), bufs, flags).map(|(n, _, recv_flags)| (n, recv_flags))
@@ -462,7 +487,7 @@ pub(crate) fn recv_vectored(
 #[cfg(not(target_os = "redox"))]
 pub(crate) fn recv_from_vectored(
     fd: Socket,
-    bufs: &mut [IoSliceMut<'_>],
+    bufs: &mut [crate::MaybeUninitSlice<'_>],
     flags: c_int,
 ) -> io::Result<(usize, RecvFlags, SockAddr)> {
     // Safety: `recvmsg` initialises the address storage and we set the length
@@ -484,7 +509,7 @@ pub(crate) fn recv_from_vectored(
 fn recvmsg(
     fd: Socket,
     msg_name: *mut sockaddr_storage,
-    bufs: &mut [IoSliceMut<'_>],
+    bufs: &mut [crate::MaybeUninitSlice<'_>],
     flags: c_int,
 ) -> io::Result<(usize, libc::socklen_t, RecvFlags)> {
     let msg_namelen = if msg_name.is_null() {

src/sys/windows.rs

@@ -7,21 +7,24 @@
 // except according to those terms.
 
 use std::cmp::min;
-use std::io::{self, IoSlice, IoSliceMut};
+use std::io::{self, IoSlice};
+use std::marker::PhantomData;
 use std::mem::{self, size_of, MaybeUninit};
 use std::net::{self, Ipv4Addr, Ipv6Addr, Shutdown};
 use std::os::windows::prelude::*;
-use std::ptr;
 use std::sync::Once;
 use std::time::Duration;
+use std::{ptr, slice};
 
 use winapi::ctypes::c_long;
 use winapi::shared::in6addr::*;
 use winapi::shared::inaddr::*;
 use winapi::shared::minwindef::DWORD;
+use winapi::shared::minwindef::ULONG;
 use winapi::shared::mstcpip::{tcp_keepalive, SIO_KEEPALIVE_VALS};
 use winapi::shared::ntdef::HANDLE;
 use winapi::shared::ws2def;
+use winapi::shared::ws2def::WSABUF;
 use winapi::um::handleapi::SetHandleInformation;
 use winapi::um::processthreadsapi::GetCurrentProcessId;
 use winapi::um::winbase::{self, INFINITE};
@@ -144,6 +147,33 @@ impl std::fmt::Debug for RecvFlags {
     }
 }
 
+#[repr(transparent)]
+pub struct MaybeUninitSlice<'a> {
+    vec: WSABUF,
+    _lifetime: PhantomData<&'a mut [MaybeUninit<u8>]>,
+}
+
+impl<'a> MaybeUninitSlice<'a> {
+    pub fn new(buf: &'a mut [MaybeUninit<u8>]) -> MaybeUninitSlice<'a> {
+        assert!(buf.len() <= ULONG::MAX as usize);
+        MaybeUninitSlice {
+            vec: WSABUF {
+                len: buf.len() as ULONG,
+                buf: buf.as_mut_ptr().cast(),
+            },
+            _lifetime: PhantomData,
+        }
+    }
+
+    pub fn as_slice(&self) -> &[MaybeUninit<u8>] {
+        unsafe { slice::from_raw_parts(self.vec.buf.cast(), self.vec.len as usize) }
+    }
+
+    pub fn as_mut_slice(&mut self) -> &mut [MaybeUninit<u8>] {
+        unsafe { slice::from_raw_parts_mut(self.vec.buf.cast(), self.vec.len as usize) }
+    }
+}
+
 fn init() {
     static INIT: Once = Once::new();
 
@@ -293,7 +323,7 @@ pub(crate) fn recv(socket: Socket, buf: &mut [MaybeUninit<u8>], flags: c_int) ->
 
 pub(crate) fn recv_vectored(
     socket: Socket,
-    bufs: &mut [IoSliceMut<'_>],
+    bufs: &mut [crate::MaybeUninitSlice<'_>],
     flags: c_int,
 ) -> io::Result<(usize, RecvFlags)> {
     let mut nread = 0;
@@ -354,7 +384,7 @@ pub(crate) fn recv_from(
 
 pub(crate) fn recv_from_vectored(
     socket: Socket,
-    bufs: &mut [IoSliceMut<'_>],
+    bufs: &mut [crate::MaybeUninitSlice<'_>],
     flags: c_int,
 ) -> io::Result<(usize, RecvFlags, SockAddr)> {
     // Safety: `recvfrom` initialises the `SockAddr` for us.