From: Aron Xu <aron@debian.org>
Date: Mon, 13 Feb 2012 14:43:56 +0800
Subject: Fix connect() timeout

---
 netcat.c | 82 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 77 insertions(+), 5 deletions(-)

diff --git a/netcat.c b/netcat.c
index bec6e46..461e2bd 100644
--- a/netcat.c
+++ b/netcat.c
@@ -90,6 +90,7 @@
 #include <ctype.h>
 #include <err.h>
 #include <errno.h>
+#include <fcntl.h>
 #include <limits.h>
 #include <netdb.h>
 #include <poll.h>
@@ -124,6 +125,10 @@
 # define TLS_MUSTSTAPLE	(1 << 4)
 #endif
 
+#define CONNECTION_SUCCESS 0
+#define CONNECTION_FAILED 1
+#define CONNECTION_TIMEOUT 2
+
 /* Command Line Options */
 int	dflag;					/* detached, no stdin */
 int	Fflag;					/* fdpass sock to stdout */
@@ -215,6 +220,8 @@ ssize_t drainbuf(int, unsigned char *, size_t *);
 ssize_t fillbuf(int, unsigned char *, size_t *);
 #endif
 
+static int connect_with_timeout(int, const struct sockaddr *, socklen_t, int);
+
 int
 main(int argc, char *argv[])
 {
@@ -1088,18 +1095,21 @@ remote_connect(const char *host, const char *port, struct addrinfo hints,
 			}
 		}
 
-		if (timeout_connect(s, res->ai_addr, res->ai_addrlen) == 0)
+		if ((error = connect_with_timeout(s, res->ai_addr, res->ai_addrlen,
+		    timeout)) == CONNECTION_SUCCESS)
 			break;
 
 		if (vflag) {
 			/* only print IP if there is something to report */
 			if (nflag || ipaddr == NULL ||
 			    (strncmp(host, ipaddr, NI_MAXHOST) == 0))
-				warn("connect to %s port %s (%s) failed", host,
-				    port, uflag ? "udp" : "tcp");
+				warn("connect to %s port %s (%s) %s", host,
+				    port, uflag ? "udp" : "tcp",
+				    error == CONNECTION_TIMEOUT ? "timed out" : "failed");
 			else
-				warn("connect to %s (%s) port %s (%s) failed",
-				    host, ipaddr, port, uflag ? "udp" : "tcp");
+				warn("connect to %s (%s) port %s (%s) %s",
+				    host, ipaddr, port, uflag ? "udp" : "tcp",
+				    error == CONNECTION_TIMEOUT ? "timed out" : "failed");
 		}
 
 		save_errno = errno;
@@ -1141,6 +1151,68 @@ timeout_connect(int s, const struct sockaddr *name, socklen_t namelen)
 	return ret;
 }
 
+static int connect_with_timeout(int fd, const struct sockaddr *sa,
+    socklen_t salen, int ctimeout)
+{
+	int err;
+	struct timeval tv, *tvp = NULL;
+	fd_set connect_fdset;
+	socklen_t len;
+	int orig_flags;
+
+	orig_flags = fcntl(fd, F_GETFL, 0);
+	if (fcntl(fd, F_SETFL, orig_flags | O_NONBLOCK) < 0 ) {
+		warn("can't set O_NONBLOCK - timeout not available");
+		if (connect(fd, sa, salen) == 0)
+			return CONNECTION_SUCCESS;
+		else
+			return CONNECTION_FAILED;
+	}
+
+	/* set connect timeout */
+	if (ctimeout > 0) {
+		tv.tv_sec = (time_t)ctimeout/1000;
+		tv.tv_usec = 0;
+		tvp = &tv;
+	}
+
+	/* attempt the connection */
+	err = connect(fd, sa, salen);
+	if (err != 0 && errno == EINPROGRESS) {
+		/* connection is proceeding
+		 * it is complete (or failed) when select returns */
+
+		/* initialize connect_fdset */
+		FD_ZERO(&connect_fdset);
+		FD_SET(fd, &connect_fdset);
+
+		/* call select */
+		do {
+			err = select(fd + 1, NULL, &connect_fdset, NULL, tvp);
+		} while (err < 0 && errno == EINTR);
+
+		/* select error */
+		if (err < 0)
+			errx(1,"select error: %s", strerror(errno));
+		/* we have reached a timeout */
+		if (err == 0)
+			return CONNECTION_TIMEOUT;
+		/* select returned successfully, but we must test socket
+		 * error for result */
+		len = sizeof(err);
+		if (getsockopt(fd, SOL_SOCKET, SO_ERROR, &err, &len) < 0)
+			errx(1, "getsockopt error: %s", strerror(errno));
+		/* setup errno according to the result returned by
+		 * getsockopt */
+		if (err != 0)
+			errno = err;
+	}
+
+	/* return aborted if an error occured, and valid otherwise */
+	fcntl(fd, F_SETFL, orig_flags);
+	return (err != 0)? CONNECTION_FAILED : CONNECTION_SUCCESS;
+}
+
 /*
  * local_listen()
  * Returns a socket listening on a local port, binds to specified source