Refactor Socket::recv_from

Author: Thomasdezeeuw

src/sockaddr.rs

@@ -39,6 +39,11 @@ impl SockAddr {
         }
     }
 
+    /// Constructs a `SockAddr` from its raw components.
+    pub(crate) const fn from_raw(storage: sockaddr_storage, len: socklen_t) -> SockAddr {
+        SockAddr { storage, len }
+    }
+
     /// Returns this address's family.
     pub fn family(&self) -> sa_family_t {
         self.storage.ss_family

src/socket.rs

@@ -352,19 +352,19 @@ impl Socket {
     /// Receives data from the socket. On success, returns the number of bytes
     /// read and the address from whence the data came.
     pub fn recv_from(&self, buf: &mut [u8]) -> io::Result<(usize, SockAddr)> {
-        self.inner().recv_from(buf, 0)
+        self.recv_from_with_flags(buf, 0)
     }
 
-    /// Identical to [`recv_from`] but allows for specification of arbitrary flags to the underlying
-    /// `recvfrom` call.
+    /// Identical to [`recv_from`] but allows for specification of arbitrary
+    /// flags to the underlying `recvfrom` call.
     ///
-    /// [`recv_from`]: #method.recv_from
+    /// [`recv_from`]: Socket::recv_from
     pub fn recv_from_with_flags(
         &self,
         buf: &mut [u8],
         flags: i32,
     ) -> io::Result<(usize, SockAddr)> {
-        self.inner().recv_from(buf, flags)
+        sys::recv_from(self.inner, buf, flags)
     }
 
     /// Identical to [`recv_from_with_flags`] but reads into a slice of buffers.

src/sys/unix.rs

@@ -390,6 +390,7 @@ pub(crate) fn shutdown(fd: SysSocket, how: Shutdown) -> io::Result<()> {
     };
     syscall!(shutdown(fd, how)).map(|_| ())
 }
+
 pub(crate) fn recv(fd: SysSocket, buf: &mut [u8], flags: c_int) -> io::Result<usize> {
     syscall!(recv(
         fd,
@@ -400,6 +401,30 @@ pub(crate) fn recv(fd: SysSocket, buf: &mut [u8], flags: c_int) -> io::Result<us
     .map(|n| n as usize)
 }
 
+pub(crate) fn recv_from(
+    fd: SysSocket,
+    buf: &mut [u8],
+    flags: c_int,
+) -> io::Result<(usize, SockAddr)> {
+    let mut storage: MaybeUninit<libc::sockaddr_storage> = MaybeUninit::zeroed();
+    let mut addrlen = size_of_val(&storage) as socklen_t;
+    syscall!(recvfrom(
+        fd,
+        buf.as_mut_ptr().cast(),
+        min(buf.len(), MAX_BUF_LEN),
+        flags,
+        storage.as_mut_ptr().cast(),
+        &mut addrlen,
+    ))
+    .map(|n| {
+        // Safety: `recvfrom` 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, addr)
+    })
+}
+
 #[cfg(not(target_os = "redox"))]
 pub(crate) fn recv_vectored(
     fd: SysSocket,
@@ -554,23 +579,7 @@ pub struct Socket {
 
 impl Socket {
     pub fn peek_from(&self, buf: &mut [u8]) -> io::Result<(usize, SockAddr)> {
-        self.recv_from(buf, libc::MSG_PEEK)
-    }
-
-    pub fn recv_from(&self, buf: &mut [u8], flags: c_int) -> io::Result<(usize, SockAddr)> {
-        let mut storage: libc::sockaddr_storage = unsafe { mem::zeroed() };
-        let mut addrlen = mem::size_of_val(&storage) as socklen_t;
-
-        let n = syscall!(recvfrom(
-            self.fd,
-            buf.as_mut_ptr() as *mut c_void,
-            cmp::min(buf.len(), max_len()),
-            flags,
-            &mut storage as *mut _ as *mut _,
-            &mut addrlen,
-        ))?;
-        let addr = unsafe { SockAddr::from_raw_parts(&storage as *const _ as *const _, addrlen) };
-        Ok((n as usize, addr))
+        recv_from(self.fd, buf, libc::MSG_PEEK)
     }
 
     #[cfg(not(target_os = "redox"))]

src/sys/windows.rs

@@ -325,6 +325,42 @@ pub(crate) fn recv_vectored(
     }
 }
 
+pub(crate) fn recv_from(
+    socket: SysSocket,
+    buf: &mut [u8],
+    flags: c_int,
+) -> io::Result<(usize, SockAddr)> {
+    let mut storage: MaybeUninit<SOCKADDR_STORAGE> = MaybeUninit::zeroed();
+    let mut addrlen = size_of_val(&storage) as socklen_t;
+    let res = syscall!(
+        recvfrom(
+            socket,
+            buf.as_mut_ptr().cast(),
+            min(buf.len(), MAX_BUF_LEN) as c_int,
+            flags,
+            storage.as_mut_ptr().cast(),
+            &mut addrlen,
+        ),
+        PartialEq::eq,
+        sock::SOCKET_ERROR
+    );
+    match res {
+        Ok(n) => {
+            // Safety: `recvfrom` 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((n as usize, 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((0, 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();
@@ -379,32 +415,7 @@ pub struct Socket {
 
 impl Socket {
     pub fn peek_from(&self, buf: &mut [u8]) -> io::Result<(usize, SockAddr)> {
-        self.recv_from(buf, MSG_PEEK)
-    }
-
-    pub fn recv_from(&self, buf: &mut [u8], flags: c_int) -> io::Result<(usize, SockAddr)> {
-        unsafe {
-            let mut storage: SOCKADDR_STORAGE = mem::zeroed();
-            let mut addrlen = mem::size_of_val(&storage) as c_int;
-
-            let n = {
-                sock::recvfrom(
-                    self.socket,
-                    buf.as_mut_ptr() as *mut c_char,
-                    clamp(buf.len()),
-                    flags,
-                    &mut storage as *mut _ as *mut _,
-                    &mut addrlen,
-                )
-            };
-            let n = match n {
-                sock::SOCKET_ERROR if sock::WSAGetLastError() == sock::WSAESHUTDOWN as i32 => 0,
-                sock::SOCKET_ERROR => return Err(last_error()),
-                n => n as usize,
-            };
-            let addr = SockAddr::from_raw_parts(&storage as *const _ as *const _, addrlen);
-            Ok((n, addr))
-        }
+        recv_from(self.socket, buf, MSG_PEEK)
     }
 
     pub fn recv_from_vectored(