From caabdf5bc9d61734062f7b1eb88817d8193c2b01 Mon Sep 17 00:00:00 2001
From: Steven Fackler <sfackler@gmail.com>
Date: Sat, 14 Oct 2017 13:25:42 -0700
Subject: [PATCH] Fix refused connections in connect_timeout

---
 src/socket.rs       | 18 ++++++++++++++++++
 src/sys/unix/mod.rs | 17 ++++++++++-------
 2 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/src/socket.rs b/src/socket.rs
index 1ba2fc5a..703f678b 100644
--- a/src/socket.rs
+++ b/src/socket.rs
@@ -827,6 +827,24 @@ mod test {
         }
     }
 
+     #[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();
diff --git a/src/sys/unix/mod.rs b/src/sys/unix/mod.rs
index 7757c4b9..49b21995 100644
--- a/src/sys/unix/mod.rs
+++ b/src/sys/unix/mod.rs
@@ -22,8 +22,7 @@ use std::time::{Duration, Instant};
 #[cfg(feature = "unix")]
 use std::os::unix::net::{UnixDatagram, UnixListener, UnixStream};
 
-use libc::{self, c_void, c_int};
-use libc::{sockaddr, socklen_t, ssize_t};
+use libc::{self, c_void, c_int, socklen_t, ssize_t};
 
 cfg_if! {
     if #[cfg(any(target_os = "dragonfly", target_os = "freebsd",
@@ -184,10 +183,14 @@ impl Socket {
                 }
                 0 => return Err(io::Error::new(io::ErrorKind::TimedOut, "connection timed out")),
                 _ => {
-                    if pollfd.revents & libc::POLLOUT == 0 {
-                        if let Some(e) = self.take_error()? {
-                            return Err(e);
-                        }
+                    // linux returns POLLOUT|POLLERR|POLLHUP for refused connections (!), so look
+                    // for POLLHUP rather than read readiness
+                    if pollfd.revents & libc::POLLHUP != 0 {
+                        let e = self.take_error()?
+                            .unwrap_or_else(|| {
+                                io::Error::new(io::ErrorKind::Other, "no error set after POLLHUP")
+                            });
+                        return Err(e);
                     }
                     return Ok(());
                 }
@@ -256,7 +259,7 @@ impl Socket {
         let mut socket = None;
         #[cfg(target_os = "linux")] {
             weak! {
-                fn accept4(c_int, *mut sockaddr, *mut socklen_t, c_int) -> c_int
+                fn accept4(c_int, *mut libc::sockaddr, *mut socklen_t, c_int) -> c_int
             }
             if let Some(f) = accept4.get() {
                 let res = cvt_r(|| unsafe {