Refactor Socket::try_clone

Author: Thomasdezeeuw

src/socket.rs

@@ -223,14 +223,19 @@ impl Socket {
 
     /// Creates a new independently owned handle to the underlying socket.
     ///
-    /// The returned `TcpStream` is a reference to the same stream that this
-    /// object references. Both handles will read and write the same stream of
-    /// data, and options set on one stream will be propagated to the other
-    /// stream.
+    /// # Notes
+    ///
+    /// On Unix this uses `F_DUPFD_CLOEXEC` and thus sets the `FD_CLOEXEC` on
+    /// the returned socket.
+    ///
+    /// On Windows this uses `WSA_FLAG_NO_HANDLE_INHERIT` setting inheriting to
+    /// false.
+    ///
+    /// On Windows this can **not** be used function cannot be used on a
+    /// QOS-enabled socket, see
+    /// https://docs.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-wsaduplicatesocketw.
     pub fn try_clone(&self) -> io::Result<Socket> {
-        self.inner()
-            .try_clone()
-            .map(|s| Socket { inner: s.inner() })
+        sys::try_clone(self.inner).map(|inner| Socket { inner })
     }
 
     /// Get the value of the `SO_ERROR` option on this socket.

src/sys/unix.rs

@@ -20,7 +20,6 @@ use std::os::unix::net::{UnixDatagram, UnixListener, UnixStream};
 use std::path::Path;
 #[cfg(feature = "all")]
 use std::ptr;
-use std::sync::atomic::{AtomicBool, Ordering};
 use std::time::Duration;
 use std::{cmp, fmt, io};
 
@@ -346,6 +345,10 @@ pub(crate) fn getpeername(fd: SysSocket) -> io::Result<SockAddr> {
         .map(|_| unsafe { SockAddr::from_raw_parts(&storage as *const _ as *const _, len) })
 }
 
+pub(crate) fn try_clone(fd: SysSocket) -> io::Result<SysSocket> {
+    syscall!(fcntl(fd, libc::F_DUPFD_CLOEXEC, 0))
+}
+
 /// Unix only API.
 impl crate::Socket {
     /// Accept a new incoming connection from this listener.
@@ -408,6 +411,7 @@ impl crate::Socket {
     }
 }
 
+/// Add `flag` to the current set flags of `F_GETFD`.
 fn fcntl_add(fd: SysSocket, flag: c_int) -> io::Result<()> {
     let previous = syscall!(fcntl(fd, libc::F_GETFD))?;
     let new = previous | flag;
@@ -441,35 +445,6 @@ pub struct Socket {
 }
 
 impl Socket {
-    pub fn try_clone(&self) -> io::Result<Socket> {
-        // implementation lifted from libstd
-        #[cfg(any(target_os = "android", target_os = "haiku"))]
-        use libc::F_DUPFD as F_DUPFD_CLOEXEC;
-        #[cfg(not(any(target_os = "android", target_os = "haiku")))]
-        use libc::F_DUPFD_CLOEXEC;
-
-        static CLOEXEC_FAILED: AtomicBool = AtomicBool::new(false);
-        if !CLOEXEC_FAILED.load(Ordering::Relaxed) {
-            match syscall!(fcntl(self.fd, F_DUPFD_CLOEXEC, 0)) {
-                Ok(fd) => {
-                    let fd = unsafe { Socket::from_raw_fd(fd) };
-                    if cfg!(target_os = "linux") {
-                        set_cloexec(fd.as_raw_fd())?;
-                    }
-                    return Ok(fd);
-                }
-                Err(ref e) if e.raw_os_error() == Some(libc::EINVAL) => {
-                    CLOEXEC_FAILED.store(true, Ordering::Relaxed);
-                }
-                Err(e) => return Err(e),
-            }
-        }
-        let fd = syscall!(fcntl(self.fd, libc::F_DUPFD, 0))?;
-        let fd = unsafe { Socket::from_raw_fd(fd) };
-        set_cloexec(fd.as_raw_fd())?;
-        Ok(fd)
-    }
-
     pub fn take_error(&self) -> io::Result<Option<io::Error>> {
         unsafe {
             let raw: c_int = self.getsockopt(libc::SOL_SOCKET, libc::SO_ERROR)?;
@@ -1191,15 +1166,6 @@ fn max_len() -> usize {
     }
 }
 
-fn set_cloexec(fd: c_int) -> io::Result<()> {
-    let previous = syscall!(fcntl(fd, libc::F_GETFD))?;
-    let new = previous | libc::FD_CLOEXEC;
-    if new != previous {
-        syscall!(fcntl(fd, libc::F_SETFD, new))?;
-    }
-    Ok(())
-}
-
 fn dur2timeval(dur: Option<Duration>) -> io::Result<libc::timeval> {
     match dur {
         Some(dur) => {

src/sys/windows.rs

@@ -10,7 +10,7 @@ use std::cmp;
 use std::fmt;
 use std::io;
 use std::io::{IoSlice, IoSliceMut, Read, Write};
-use std::mem::{self, size_of_val};
+use std::mem::{self, size_of_val, MaybeUninit};
 use std::net::Shutdown;
 use std::net::{self, Ipv4Addr, Ipv6Addr};
 use std::os::windows::prelude::*;
@@ -36,13 +36,11 @@ use winapi::um::winsock2 as sock;
 
 use crate::{RecvFlags, SockAddr};
 
-const HANDLE_FLAG_INHERIT: DWORD = 0x00000001;
 const MSG_PEEK: c_int = 0x2;
 const SD_BOTH: c_int = 2;
 const SD_RECEIVE: c_int = 0;
 const SD_SEND: c_int = 1;
 const SIO_KEEPALIVE_VALS: DWORD = 0x98000004;
-const WSA_FLAG_OVERLAPPED: DWORD = 0x01;
 
 pub use winapi::ctypes::c_int;
 
@@ -151,7 +149,7 @@ pub(crate) fn socket(family: c_int, ty: c_int, protocol: c_int) -> io::Result<Sy
             protocol,
             ptr::null_mut(),
             0,
-            WSA_FLAG_OVERLAPPED,
+            sock::WSA_FLAG_OVERLAPPED,
         ),
         PartialEq::eq,
         sock::INVALID_SOCKET
@@ -209,6 +207,30 @@ pub(crate) fn getpeername(socket: SysSocket) -> io::Result<SockAddr> {
     .map(|_| unsafe { SockAddr::from_raw_parts(&storage as *const _ as *const _, len) })
 }
 
+pub(crate) fn try_clone(socket: SysSocket) -> io::Result<SysSocket> {
+    let mut info: MaybeUninit<sock::WSAPROTOCOL_INFOW> = MaybeUninit::uninit();
+    syscall!(
+        WSADuplicateSocketW(socket, GetCurrentProcessId(), info.as_mut_ptr()),
+        PartialEq::eq,
+        sock::SOCKET_ERROR
+    )?;
+    // Safety: `WSADuplicateSocketW` intialised `info` for us.
+    let mut info = unsafe { info.assume_init() };
+
+    syscall!(
+        WSASocketW(
+            info.iAddressFamily,
+            info.iSocketType,
+            info.iProtocol,
+            &mut info,
+            0,
+            sock::WSA_FLAG_OVERLAPPED | sock::WSA_FLAG_NO_HANDLE_INHERIT,
+        ),
+        PartialEq::eq,
+        sock::INVALID_SOCKET
+    )
+}
+
 /// Windows only API.
 impl crate::Socket {
     /// Sets `HANDLE_FLAG_INHERIT` to zero using `SetHandleInformation`.
@@ -233,30 +255,6 @@ pub struct Socket {
 }
 
 impl Socket {
-    pub fn try_clone(&self) -> io::Result<Socket> {
-        unsafe {
-            let mut info: sock::WSAPROTOCOL_INFOW = mem::zeroed();
-            let r = sock::WSADuplicateSocketW(self.socket, GetCurrentProcessId(), &mut info);
-            if r != 0 {
-                return Err(io::Error::last_os_error());
-            }
-            let socket = sock::WSASocketW(
-                info.iAddressFamily,
-                info.iSocketType,
-                info.iProtocol,
-                &mut info,
-                0,
-                WSA_FLAG_OVERLAPPED,
-            );
-            let socket = match socket {
-                sock::INVALID_SOCKET => return Err(last_error()),
-                n => Socket::from_raw_socket(n as RawSocket),
-            };
-            socket.set_no_inherit()?;
-            Ok(socket)
-        }
-    }
-
     pub fn take_error(&self) -> io::Result<Option<io::Error>> {
         unsafe {
             let raw: c_int = self.getsockopt(SOL_SOCKET, SO_ERROR)?;
@@ -844,17 +842,6 @@ impl Socket {
         }
     }
 
-    fn set_no_inherit(&self) -> io::Result<()> {
-        unsafe {
-            let r = SetHandleInformation(self.socket as HANDLE, HANDLE_FLAG_INHERIT, 0);
-            if r == 0 {
-                Err(io::Error::last_os_error())
-            } else {
-                Ok(())
-            }
-        }
-    }
-
     pub fn inner(self) -> SysSocket {
         self.socket
     }