Repository: rust-lang/socket2 (rust)

Issue:
Support for connect_timeout in 0.4
As I saw in 0.4 changelogs `Socket::connect_timeout` was removed because it was not a low-level function but a combination of system calls. This is true, however a lot of people may need this functionality.

For someone to implement this and make it cross-platform, it will be troublesome. So can we have it back? Maybe we can have it in a new crate? For example in `socket2-util` crate and in `SocketExt` trait.

Pull request description:
Add back `Socket::connect_timeout`
Fixes #190

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index f87f1ba..abec0da 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -88,9 +88,6 @@
     * `Socket::into_unix_stream` => `UnixStream::from(socket)`.
     * `Socket::into_unix_listener` => `UnixListener::from(socket)`.
     * `Socket::into_unix_datagram` => `UnixDatagram::from(socket)`.
-* **BREAKING:** Removed `Socket::connect_timeout`: it didn't correspond to a
-  single system call and thus doesn't belong in a low-level library like
-  socket2.
 * Removed `cfg-if` dependency.
 * Remove `redox_syscall` depdency.
 
diff --git a/src/socket.rs b/src/socket.rs
index a999993..b33387d 100644
--- a/src/socket.rs
+++ b/src/socket.rs
@@ -164,6 +164,41 @@ impl Socket {
         sys::connect(self.inner, address)
     }
 
+    /// Initiate a connection on this socket to the specified address, only
+    /// only waiting for a certain period of time for the connection to be
+    /// established.
+    ///
+    /// Unlike many other methods on `Socket`, this does *not* correspond to a
+    /// single C function. It sets the socket to nonblocking mode, connects via
+    /// connect(2), and then waits for the connection to complete with poll(2)
+    /// on Unix and select on Windows. When the connection is complete, the
+    /// socket is set back to blocking mode. On Unix, this will loop over
+    /// `EINTR` errors.
+    ///
+    /// # Warnings
+    ///
+    /// The non-blocking state of the socket is overridden by this function -
+    /// it will be returned in blocking mode on success, and in an indeterminate
+    /// state on failure.
+    ///
+    /// If the connection request times out, it may still be processing in the
+    /// background - a second call to `connect` or `connect_timeout` may fail.
+    pub fn connect_timeout(&self, addr: &SockAddr, timeout: Duration) -> io::Result<()> {
+        self.set_nonblocking(true)?;
+        let res = self.connect(addr);
+        self.set_nonblocking(false)?;
+
+        match res {
+            Ok(()) => return Ok(()),
+            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {}
+            #[cfg(unix)]
+            Err(ref e) if e.raw_os_error() == Some(libc::EINPROGRESS) => {}
+            Err(e) => return Err(e),
+        }
+
+        sys::poll_connect(self, timeout)
+    }
+
     /// Mark a socket as ready to accept incoming connection requests using
     /// [`Socket::accept()`].
     ///
diff --git a/src/sys/unix.rs b/src/sys/unix.rs
index b62393d..1caceaf 100644
--- a/src/sys/unix.rs
+++ b/src/sys/unix.rs
@@ -22,7 +22,7 @@ use std::os::unix::net::{UnixDatagram, UnixListener, UnixStream};
 use std::path::Path;
 #[cfg(not(all(target_os = "redox", not(feature = "all"))))]
 use std::ptr;
-use std::time::Duration;
+use std::time::{Duration, Instant};
 use std::{io, slice};
 
 #[cfg(not(target_vendor = "apple"))]
@@ -426,6 +426,63 @@ pub(crate) fn connect(fd: Socket, addr: &SockAddr) -> io::Result<()> {
     syscall!(connect(fd, addr.as_ptr(), addr.len())).map(|_| ())
 }
 
+pub(crate) fn poll_connect(socket: &crate::Socket, timeout: Duration) -> io::Result<()> {
+    let start = Instant::now();
+
+    let mut pollfd = libc::pollfd {
+        fd: socket.inner,
+        events: libc::POLLIN | libc::POLLOUT,
+        revents: 0,
+    };
+
+    loop {
+        let elapsed = start.elapsed();
+        if elapsed >= timeout {
+            return Err(io::ErrorKind::TimedOut.into());
+        }
+
+        let timeout = (timeout - elapsed).as_millis();
+        let timeout = clamp(timeout, 1, c_int::max_value() as u128) as c_int;
+
+        match syscall!(poll(&mut pollfd, 1, timeout)) {
+            Ok(0) => return Err(io::ErrorKind::TimedOut.into()),
+            Ok(_) => {
+                // Error or hang up indicates an error (or failure to connect).
+                if (pollfd.revents & libc::POLLHUP) != 0 || (pollfd.revents & libc::POLLERR) != 0 {
+                    match socket.take_error() {
+                        Ok(Some(err)) => return Err(err),
+                        Ok(None) => {
+                            return Err(io::Error::new(
+                                io::ErrorKind::Other,
+                                "no error set after POLLHUP",
+                            ))
+                        }
+                        Err(err) => return Err(err),
+                    }
+                }
+                return Ok(());
+            }
+            // Got interrupted, try again.
+            Err(ref err) if err.kind() == io::ErrorKind::Interrupted => continue,
+            Err(err) => return Err(err),
+        }
+    }
+}
+
+// TODO: use clamp from std lib, stable since 1.50.
+fn clamp<T>(value: T, min: T, max: T) -> T
+where
+    T: Ord,
+{
+    if value <= min {
+        min
+    } else if value >= max {
+        max
+    } else {
+        value
+    }
+}
+
 pub(crate) fn listen(fd: Socket, backlog: c_int) -> io::Result<()> {
     syscall!(listen(fd, backlog)).map(|_| ())
 }
diff --git a/src/sys/windows.rs b/src/sys/windows.rs
index c8379ef..1330432 100644
--- a/src/sys/windows.rs
+++ b/src/sys/windows.rs
@@ -13,7 +13,7 @@ use std::mem::{self, size_of, MaybeUninit};
 use std::net::{self, Ipv4Addr, Ipv6Addr, Shutdown};
 use std::os::windows::prelude::*;
 use std::sync::Once;
-use std::time::Duration;
+use std::time::{Duration, Instant};
 use std::{ptr, slice};
 
 use winapi::ctypes::c_long;
@@ -28,7 +28,10 @@ use winapi::shared::ws2def::WSABUF;
 use winapi::um::handleapi::SetHandleInformation;
 use winapi::um::processthreadsapi::GetCurrentProcessId;
 use winapi::um::winbase::{self, INFINITE};
-use winapi::um::winsock2::{self as sock, u_long, SD_BOTH, SD_RECEIVE, SD_SEND};
+use winapi::um::winsock2::{
+    self as sock, u_long, POLLERR, POLLHUP, POLLRDNORM, POLLWRNORM, SD_BOTH, SD_RECEIVE, SD_SEND,
+    WSAPOLLFD,
+};
 
 use crate::{RecvFlags, SockAddr, TcpKeepalive, Type};
 
@@ -220,6 +223,67 @@ pub(crate) fn connect(socket: Socket, addr: &SockAddr) -> io::Result<()> {
     syscall!(connect(socket, addr.as_ptr(), addr.len()), PartialEq::ne, 0).map(|_| ())
 }
 
+pub(crate) fn poll_connect(socket: &crate::Socket, timeout: Duration) -> io::Result<()> {
+    let start = Instant::now();
+
+    let mut fd_array = WSAPOLLFD {
+        fd: socket.inner,
+        events: POLLRDNORM | POLLWRNORM,
+        revents: 0,
+    };
+
+    loop {
+        let elapsed = start.elapsed();
+        if elapsed >= timeout {
+            return Err(io::ErrorKind::TimedOut.into());
+        }
+
+        let timeout = (timeout - elapsed).as_millis();
+        let timeout = clamp(timeout, 1, c_int::max_value() as u128) as c_int;
+
+        match syscall!(
+            WSAPoll(&mut fd_array, 1, timeout),
+            PartialEq::eq,
+            sock::SOCKET_ERROR
+        ) {
+            Ok(0) => return Err(io::ErrorKind::TimedOut.into()),
+            Ok(_) => {
+                // Error or hang up indicates an error (or failure to connect).
+                if (fd_array.revents & POLLERR) != 0 || (fd_array.revents & POLLHUP) != 0 {
+                    match socket.take_error() {
+                        Ok(Some(err)) => return Err(err),
+                        Ok(None) => {
+                            return Err(io::Error::new(
+                                io::ErrorKind::Other,
+                                "no error set after POLLHUP",
+                            ))
+                        }
+                        Err(err) => return Err(err),
+                    }
+                }
+                return Ok(());
+            }
+            // Got interrupted, try again.
+            Err(ref err) if err.kind() == io::ErrorKind::Interrupted => continue,
+            Err(err) => return Err(err),
+        }
+    }
+}
+
+// TODO: use clamp from std lib, stable since 1.50.
+fn clamp<T>(value: T, min: T, max: T) -> T
+where
+    T: Ord,
+{
+    if value <= min {
+        min
+    } else if value >= max {
+        max
+    } else {
+        value
+    }
+}
+
 pub(crate) fn listen(socket: Socket, backlog: c_int) -> io::Result<()> {
     syscall!(listen(socket, backlog), PartialEq::ne, 0).map(|_| ())
 }


Test patch:
diff --git a/tests/socket.rs b/tests/socket.rs
index 0b5e027..af59fdc 100644
--- a/tests/socket.rs
+++ b/tests/socket.rs
@@ -1,4 +1,3 @@
-#[cfg(any(windows, target_vendor = "apple"))]
 use std::io;
 #[cfg(not(target_os = "redox"))]
 use std::io::IoSlice;
@@ -307,6 +306,54 @@ where
 
 const DATA: &[u8] = b"hello world";
 
+#[test]
+fn connect_timeout_unrouteable() {
+    // This IP is unroutable, so connections should always time out.
+    let addr = "10.255.255.1:80".parse::<SocketAddr>().unwrap().into();
+
+    let socket = Socket::new(Domain::IPV4, Type::STREAM, None).unwrap();
+    match socket.connect_timeout(&addr, Duration::from_millis(250)) {
+        Ok(_) => panic!("unexpected success"),
+        Err(ref e) if e.kind() == io::ErrorKind::TimedOut => {}
+        Err(e) => panic!("unexpected error {}", e),
+    }
+}
+
+#[test]
+fn connect_timeout_unbound() {
+    // Bind and drop a socket to track down a "probably unassigned" port.
+    let socket = Socket::new(Domain::IPV4, Type::STREAM, None).unwrap();
+    let addr = "127.0.0.1:0".parse::<SocketAddr>().unwrap().into();
+    socket.bind(&addr).unwrap();
+    let addr = socket.local_addr().unwrap();
+    drop(socket);
+
+    let socket = Socket::new(Domain::IPV4, Type::STREAM, None).unwrap();
+    match socket.connect_timeout(&addr, Duration::from_millis(250)) {
+        Ok(_) => panic!("unexpected success"),
+        Err(ref e)
+            if e.kind() == io::ErrorKind::ConnectionRefused
+                || e.kind() == io::ErrorKind::TimedOut => {}
+        Err(e) => panic!("unexpected error {}", e),
+    }
+}
+
+#[test]
+fn connect_timeout_valid() {
+    let socket = Socket::new(Domain::IPV4, Type::STREAM, None).unwrap();
+    socket
+        .bind(&"127.0.0.1:0".parse::<SocketAddr>().unwrap().into())
+        .unwrap();
+    socket.listen(128).unwrap();
+
+    let addr = socket.local_addr().unwrap();
+
+    let socket = Socket::new(Domain::IPV4, Type::STREAM, None).unwrap();
+    socket
+        .connect_timeout(&addr, Duration::from_millis(250))
+        .unwrap();
+}
+
 #[test]
 #[cfg(all(feature = "all", unix))]
 fn pair() {
@@ -768,7 +815,8 @@ test!(IPv4 ttl, set_ttl(40));
 test!(IPv4 broadcast, set_broadcast(true));
 
 test!(IPv6 unicast_hops_v6, set_unicast_hops_v6(20));
-#[cfg(not(windows))]
+#[cfg(not(any(windows, target_os = "freebsd")))]
 test!(IPv6 only_v6, set_only_v6(true));
-#[cfg(windows)] // IPv6 socket are already IPv6 only on Windows.
+// IPv6 socket are already IPv6 only on FreeBSD and Windows.
+#[cfg(any(windows, target_os = "freebsd"))]
 test!(IPv6 only_v6, set_only_v6(false));