Refactor Socket::recv_from_vectored

And add Socket::recv_from_vectored_with_flags.

Author: Thomasdezeeuw

src/socket.rs

@@ -310,7 +310,7 @@ impl Socket {
     }
 
     /// Receives data on the socket from the remote address to which it is
-    /// connected. Unlike [`recv`] that allows passing multiple buffers.
+    /// connected. Unlike [`recv`] this allows passing multiple buffers.
     ///
     /// The [`connect`] method will connect this socket to a remote address.
     /// This method might fail if the socket is not connected.
@@ -367,19 +367,30 @@ impl Socket {
         sys::recv_from(self.inner, buf, flags)
     }
 
-    /// Identical to [`recv_from_with_flags`] but reads into a slice of buffers.
+    /// Receives data from the socket. Returns the amount of bytes read, the
+    /// [`RecvFlags`] and the remote address from the data is coming. Unlike
+    /// [`recv_from`] this allows passing multiple buffers.
     ///
-    /// In addition to the number of bytes read, this function returns the flags for the received message.
-    /// See [`RecvFlags`] for more information about the flags.
-    ///
-    /// [`recv_from_with_flags`]: #method.recv_from_with_flags
+    /// [`recv_from`]: Socket::recv_from
     #[cfg(not(target_os = "redox"))]
     pub fn recv_from_vectored(
         &self,
         bufs: &mut [IoSliceMut<'_>],
+    ) -> io::Result<(usize, RecvFlags, SockAddr)> {
+        self.recv_from_vectored_with_flags(bufs, 0)
+    }
+
+    /// Identical to [`recv_from_vectored`] but allows for specification of
+    /// arbitrary flags to the underlying `recvmsg`/`WSARecvFrom` call.
+    ///
+    /// [`recv_from_vectored`]: Socket::recv_from_vectored
+    #[cfg(not(target_os = "redox"))]
+    pub fn recv_from_vectored_with_flags(
+        &self,
+        bufs: &mut [IoSliceMut<'_>],
         flags: i32,
     ) -> io::Result<(usize, RecvFlags, SockAddr)> {
-        self.inner().recv_from_vectored(bufs, flags)
+        sys::recv_from_vectored(self.inner, bufs, flags)
     }
 
     /// Receives data from the socket, without removing it from the queue.

src/sys/unix.rs

@@ -431,16 +431,48 @@ pub(crate) fn recv_vectored(
     bufs: &mut [IoSliceMut<'_>],
     flags: c_int,
 ) -> io::Result<(usize, RecvFlags)> {
+    recvmsg(fd, ptr::null_mut(), bufs, flags).map(|(n, _, recv_flags)| (n, recv_flags))
+}
+
+#[cfg(not(target_os = "redox"))]
+pub fn recv_from_vectored(
+    fd: SysSocket,
+    bufs: &mut [IoSliceMut<'_>],
+    flags: c_int,
+) -> io::Result<(usize, RecvFlags, SockAddr)> {
+    let mut storage: MaybeUninit<libc::sockaddr_storage> = MaybeUninit::zeroed();
+    recvmsg(fd, storage.as_mut_ptr(), bufs, flags).map(|(n, addrlen, recv_flags)| {
+        // Safety: `recvmsg` wrote an address of `addrlen` bytes for us. The
+        // remaining bytes are initialised to zero (which is valid for
+        // `sockaddr_storage`).
+        let addr = SockAddr::from_raw(unsafe { storage.assume_init() }, addrlen);
+        (n as usize, recv_flags, addr)
+    })
+}
+
+/// Returns the (bytes received, sending address len, `RecvFlags`).
+fn recvmsg(
+    fd: SysSocket,
+    msg_name: *mut sockaddr_storage,
+    bufs: &mut [IoSliceMut<'_>],
+    flags: c_int,
+) -> io::Result<(usize, libc::socklen_t, RecvFlags)> {
+    let msg_namelen = if msg_name.is_null() {
+        0
+    } else {
+        size_of::<libc::sockaddr_storage>() as libc::socklen_t
+    };
     let mut msg = libc::msghdr {
-        msg_name: ptr::null_mut(),
-        msg_namelen: 0,
+        msg_name: msg_name.cast(),
+        msg_namelen,
         msg_iov: bufs.as_mut_ptr().cast(),
         msg_iovlen: min(bufs.len(), IovLen::MAX as usize) as IovLen,
         msg_control: ptr::null_mut(),
         msg_controllen: 0,
         msg_flags: 0,
     };
-    syscall!(recvmsg(fd, &mut msg as *mut _, flags)).map(|n| (n as usize, RecvFlags(msg.msg_flags)))
+    syscall!(recvmsg(fd, &mut msg as *mut _, flags))
+        .map(|n| (n as usize, msg.msg_namelen, RecvFlags(msg.msg_flags)))
 }
 
 /// Unix only API.
@@ -578,29 +610,6 @@ pub struct Socket {
 }
 
 impl Socket {
-    #[cfg(not(target_os = "redox"))]
-    pub fn recv_from_vectored(
-        &self,
-        bufs: &mut [IoSliceMut<'_>],
-        flags: c_int,
-    ) -> io::Result<(usize, RecvFlags, SockAddr)> {
-        let mut storage: libc::sockaddr_storage = unsafe { mem::zeroed() };
-        let mut msg = libc::msghdr {
-            msg_name: &mut storage as *mut libc::sockaddr_storage as *mut c_void,
-            msg_namelen: mem::size_of_val(&storage) as socklen_t,
-            msg_iov: bufs.as_mut_ptr().cast(),
-            msg_iovlen: bufs.len().min(IovLen::MAX as usize) as IovLen,
-            msg_control: std::ptr::null_mut(),
-            msg_controllen: 0,
-            msg_flags: 0,
-        };
-
-        let n = syscall!(recvmsg(self.fd, &mut msg as *mut _, flags))?;
-        let addr =
-            unsafe { SockAddr::from_raw_parts(&storage as *const _ as *const _, msg.msg_namelen) };
-        Ok((n as usize, RecvFlags(msg.msg_flags), addr))
-    }
-
     pub fn send(&self, buf: &[u8], flags: c_int) -> io::Result<usize> {
         let n = syscall!(send(
             self.fd,

src/sys/windows.rs

@@ -361,6 +361,52 @@ pub(crate) fn recv_from(
     }
 }
 
+pub(crate) fn recv_from_vectored(
+    socket: SysSocket,
+    bufs: &mut [IoSliceMut<'_>],
+    flags: c_int,
+) -> io::Result<(usize, RecvFlags, SockAddr)> {
+    let mut storage: MaybeUninit<SOCKADDR_STORAGE> = MaybeUninit::zeroed();
+    let mut addrlen = size_of_val(&storage) as socklen_t;
+    let mut nread = 0;
+    let mut flags = flags as DWORD;
+    let res = syscall!(
+        WSARecvFrom(
+            socket,
+            bufs.as_mut_ptr().cast(),
+            min(bufs.len(), DWORD::max_value() as usize) as DWORD,
+            &mut nread,
+            &mut flags,
+            storage.as_mut_ptr().cast(),
+            &mut addrlen,
+            ptr::null_mut(),
+            None,
+        ),
+        PartialEq::eq,
+        sock::SOCKET_ERROR
+    );
+    match res {
+        Ok(_) => {
+            // Safety: `WSARecvFrom` wrote an address of `addrlen` bytes for us.
+            // The remaining bytes are initialised to zero (which is valid for
+            // `sockaddr_storage`).
+            let addr = SockAddr::from_raw(unsafe { storage.assume_init() }, addrlen);
+            Ok((nread as usize, RecvFlags(0), addr))
+        }
+        Err(ref err) if err.raw_os_error() == Some(sock::WSAESHUTDOWN as i32) => {
+            // Safety: see above.
+            let addr = SockAddr::from_raw(unsafe { storage.assume_init() }, addrlen);
+            Ok((nread as usize, RecvFlags(0), addr))
+        }
+        Err(ref err) if err.raw_os_error() == Some(sock::WSAEMSGSIZE as i32) => {
+            // Safety: see above.
+            let addr = SockAddr::from_raw(unsafe { storage.assume_init() }, addrlen);
+            Ok((nread as usize, RecvFlags(MSG_TRUNC), addr))
+        }
+        Err(err) => Err(err),
+    }
+}
+
 /// Caller must ensure `T` is the correct type for `opt` and `val`.
 unsafe fn getsockopt<T>(socket: SysSocket, opt: c_int, val: c_int) -> io::Result<T> {
     let mut payload: MaybeUninit<T> = MaybeUninit::uninit();
@@ -414,45 +460,6 @@ pub struct Socket {
 }
 
 impl Socket {
-    pub fn recv_from_vectored(
-        &self,
-        bufs: &mut [IoSliceMut<'_>],
-        flags: c_int,
-    ) -> io::Result<(usize, RecvFlags, SockAddr)> {
-        let mut nread = 0;
-        let mut flags = flags as DWORD;
-        let mut storage: SOCKADDR_STORAGE = unsafe { mem::zeroed() };
-        let mut addrlen = mem::size_of_val(&storage) as c_int;
-        let ret = unsafe {
-            sock::WSARecvFrom(
-                self.socket,
-                bufs.as_mut_ptr() as *mut WSABUF,
-                bufs.len().min(DWORD::MAX as usize) as DWORD,
-                &mut nread,
-                &mut flags,
-                &mut storage as *mut SOCKADDR_STORAGE as *mut SOCKADDR,
-                &mut addrlen,
-                ptr::null_mut(),
-                None,
-            )
-        };
-
-        let flags;
-        if ret == 0 {
-            flags = RecvFlags(0);
-        } else {
-            let error = last_error();
-            if error.raw_os_error() == Some(sock::WSAEMSGSIZE) {
-                flags = RecvFlags(MSG_TRUNC)
-            } else {
-                return Err(error);
-            }
-        }
-
-        let addr = unsafe { SockAddr::from_raw_parts(&storage as *const _ as *const _, addrlen) };
-        Ok((nread as usize, flags, addr))
-    }
-
     pub fn send(&self, buf: &[u8], flags: c_int) -> io::Result<usize> {
         unsafe {
             let n = {

src/tests.rs

@@ -228,15 +228,12 @@ fn send_from_recv_to_vectored() {
     let mut men = [0u8; 3];
     let mut swear = [0u8; 5];
     let (received, flags, addr) = socket_b
-        .recv_from_vectored(
-            &mut [
-                IoSliceMut::new(&mut surgeon),
-                IoSliceMut::new(&mut has),
-                IoSliceMut::new(&mut men),
-                IoSliceMut::new(&mut swear),
-            ],
-            0,
-        )
+        .recv_from_vectored(&mut [
+            IoSliceMut::new(&mut surgeon),
+            IoSliceMut::new(&mut has),
+            IoSliceMut::new(&mut men),
+            IoSliceMut::new(&mut swear),
+        ])
         .unwrap();
 
     assert_eq!(received, 18);
@@ -291,7 +288,7 @@ fn recv_from_vectored_truncated() {
     let mut buffer = [0u8; 24];
 
     let (received, flags, addr) = socket_b
-        .recv_from_vectored(&mut [IoSliceMut::new(&mut buffer)], 0)
+        .recv_from_vectored(&mut [IoSliceMut::new(&mut buffer)])
         .unwrap();
     assert_eq!(received, 24);
     assert_eq!(flags.is_truncated(), true);