Refactor Socket::set_nonblocking

Author: Thomasdezeeuw

src/socket.rs

@@ -250,10 +250,14 @@ impl Socket {
 
     /// Moves this TCP stream into or out of nonblocking mode.
     ///
-    /// On Unix this corresponds to calling fcntl, and on Windows this
-    /// corresponds to calling ioctlsocket.
+    /// # Notes
+    ///
+    /// On Unix this corresponds to calling `fcntl` (un)setting `O_NONBLOCK`.
+    ///
+    /// On Windows this corresponds to calling `ioctlsocket` (un)setting
+    /// `FIONBIO`.
     pub fn set_nonblocking(&self, nonblocking: bool) -> io::Result<()> {
-        self.inner().set_nonblocking(nonblocking)
+        sys::set_nonblocking(self.inner, nonblocking)
     }
 
     /// Shuts down the read, write, or both halves of this connection.

src/sys/unix.rs

@@ -357,6 +357,14 @@ pub(crate) fn take_error(fd: SysSocket) -> io::Result<Option<io::Error>> {
     }
 }
 
+pub(crate) fn set_nonblocking(fd: SysSocket, nonblocking: bool) -> io::Result<()> {
+    if nonblocking {
+        fcntl_add(fd, libc::O_NONBLOCK)
+    } else {
+        fcntl_remove(fd, libc::O_NONBLOCK)
+    }
+}
+
 /// Unix only API.
 impl crate::Socket {
     /// Accept a new incoming connection from this listener.
@@ -431,6 +439,18 @@ fn fcntl_add(fd: SysSocket, flag: c_int) -> io::Result<()> {
     }
 }
 
+/// Remove `flag` to the current set flags of `F_GETFD`.
+fn fcntl_remove(fd: SysSocket, flag: c_int) -> io::Result<()> {
+    let previous = syscall!(fcntl(fd, libc::F_GETFD))?;
+    let new = previous & !flag;
+    if new != previous {
+        syscall!(fcntl(fd, libc::F_SETFD, new)).map(|_| ())
+    } else {
+        // Flag was already set.
+        Ok(())
+    }
+}
+
 /// Caller must ensure `T` is the correct type for `opt` and `val`.
 unsafe fn getsockopt<T>(fd: SysSocket, opt: c_int, val: c_int) -> io::Result<T> {
     let mut payload: MaybeUninit<T> = MaybeUninit::uninit();
@@ -472,19 +492,6 @@ pub struct Socket {
 }
 
 impl Socket {
-    pub fn set_nonblocking(&self, nonblocking: bool) -> io::Result<()> {
-        let previous = syscall!(fcntl(self.fd, libc::F_GETFL))?;
-        let new = if nonblocking {
-            previous | libc::O_NONBLOCK
-        } else {
-            previous & !libc::O_NONBLOCK
-        };
-        if new != previous {
-            syscall!(fcntl(self.fd, libc::F_SETFL, new))?;
-        }
-        Ok(())
-    }
-
     pub fn shutdown(&self, how: Shutdown) -> io::Result<()> {
         let how = match how {
             Shutdown::Write => libc::SHUT_WR,

src/sys/windows.rs

@@ -18,7 +18,7 @@ use std::ptr;
 use std::sync::Once;
 use std::time::Duration;
 
-use winapi::ctypes::{c_char, c_ulong};
+use winapi::ctypes::{c_char, c_long, c_ulong};
 use winapi::shared::in6addr::*;
 use winapi::shared::inaddr::*;
 use winapi::shared::minwindef::DWORD;
@@ -32,7 +32,7 @@ use winapi::um::processthreadsapi::GetCurrentProcessId;
 #[cfg(feature = "all")]
 use winapi::um::winbase;
 use winapi::um::winbase::INFINITE;
-use winapi::um::winsock2 as sock;
+use winapi::um::winsock2::{self as sock, u_long};
 
 use crate::{RecvFlags, SockAddr, Type};
 
@@ -261,6 +261,12 @@ pub(crate) fn take_error(socket: SysSocket) -> io::Result<Option<io::Error>> {
     }
 }
 
+pub(crate) fn set_nonblocking(socket: SysSocket, nonblocking: bool) -> io::Result<()> {
+    let mut nonblocking = nonblocking as u_long;
+    ioctlsocket(socket, sock::FIONBIO, &mut nonblocking)
+}
+
+/// 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();
     let mut len = mem::size_of::<T>() as c_int;
@@ -276,6 +282,15 @@ unsafe fn getsockopt<T>(socket: SysSocket, opt: c_int, val: c_int) -> io::Result
     })
 }
 
+fn ioctlsocket(socket: SysSocket, cmd: c_long, payload: &mut u_long) -> io::Result<()> {
+    syscall!(
+        ioctlsocket(socket, cmd, payload),
+        PartialEq::eq,
+        sock::SOCKET_ERROR
+    )
+    .map(|_| ())
+}
+
 /// Windows only API.
 impl crate::Socket {
     /// Sets `HANDLE_FLAG_INHERIT` to zero using `SetHandleInformation`.
@@ -300,18 +315,6 @@ pub struct Socket {
 }
 
 impl Socket {
-    pub fn set_nonblocking(&self, nonblocking: bool) -> io::Result<()> {
-        unsafe {
-            let mut nonblocking = nonblocking as c_ulong;
-            let r = sock::ioctlsocket(self.socket, sock::FIONBIO as c_int, &mut nonblocking);
-            if r == 0 {
-                Ok(())
-            } else {
-                Err(io::Error::last_os_error())
-            }
-        }
-    }
-
     pub fn shutdown(&self, how: Shutdown) -> io::Result<()> {
         let how = match how {
             Shutdown::Write => SD_SEND,