Fix non-blocking connect() for Windows. Refactored the code

that retries the connect() call in timeout mode so it can be shared
between connect() and connect_ex(), and needs only a single #ifdef.

The test for this was doing funky stuff I don't approve of,
so I removed it in favor of a simpler test.  This allowed me
to implement a simpler, "purer" form of the timeout retry code.
Hopefully that's enough (if you want to be fancy, use non-blocking
mode and decode the errors yourself, like before).

2 files changed, 42 insertions, 39 deletions

diff --git a/Lib/test/test_socket.py b/Lib/test/test_socket.py
index 485e038..a2db64a 100644
--- a/Lib/test/test_socket.py
+++ b/Lib/test/test_socket.py
@@ -416,18 +416,8 @@ class NonBlockingTCPTests(ThreadedTCPSocketTest):
         conn, addr = self.serv.accept()
 
     def _testConnect(self):
-        self.cli.setblocking(0)
-        try:
-            self.cli.connect((HOST, PORT))
-        except socket.error:
-            pass
-        else:
-            self.fail("Error trying to do non-blocking connect.")
-        read, write, err = select.select([self.cli], [], [])
-        if self.cli in read:
-            self.cli.connect((HOST, PORT))
-        else:
-            self.fail("Error trying to do connect after select.")
+        self.cli.settimeout(10)
+        self.cli.connect((HOST, PORT))
 
     def testRecv(self):
         """Testing non-blocking recv."""
diff --git a/Modules/socketmodule.c b/Modules/socketmodule.c
index 3896831..1da9afb 100644
--- a/Modules/socketmodule.c
+++ b/Modules/socketmodule.c
@@ -1270,6 +1270,43 @@ static char close_doc[] =
 \n\
 Close the socket.  It cannot be used after this call.";
 
+static int
+internal_connect(PySocketSockObject *s, struct sockaddr *addr, int addrlen)
+{
+	int res;
+
+	res = connect(s->sock_fd, addr, addrlen);
+
+#ifdef MS_WINDOWS
+
+	if (s->sock_timeout > 0.0) {
+		if (res < 0 && WSAGetLastError() == WSAEWOULDBLOCK) {
+			internal_select(s, 1);
+			res = connect(s->sock_fd, addr, addrlen);
+			if (res < 0 && WSAGetLastError() == WSAEISCONN)
+				res = 0;
+		}
+	}
+
+	if (res < 0)
+		res = WSAGetLastError();
+
+#else
+
+	if (s->sock_timeout > 0.0) {
+		if (res < 0 && errno == EINPROGRESS) {
+			internal_select(s, 1);
+			res = connect(s->sock_fd, addr, addrlen);
+		}
+	}
+
+	if (res < 0)
+		res = errno;
+
+#endif
+
+	return res;
+}
 
 /* s.connect(sockaddr) method */
 
@@ -1284,18 +1321,10 @@ sock_connect(PySocketSockObject *s, PyObject *addro)
 		return NULL;
 
 	Py_BEGIN_ALLOW_THREADS
-	if (s->sock_timeout > 0.0) {
-		res = connect(s->sock_fd, addr, addrlen);
-		if (res == EINPROGRESS) {
-			internal_select(s, 1);
-			res = connect(s->sock_fd, addr, addrlen);
-		}
-	}
-	else
-		res = connect(s->sock_fd, addr, addrlen);
+	res = internal_connect(s, addr, addrlen);
 	Py_END_ALLOW_THREADS
 
-	if (res < 0)
+	if (res != 0)
 		return s->errorhandler();
 	Py_INCREF(Py_None);
 	return Py_None;
@@ -1321,25 +1350,9 @@ sock_connect_ex(PySocketSockObject *s, PyObject *addro)
 		return NULL;
 
 	Py_BEGIN_ALLOW_THREADS
-	if (s->sock_timeout > 0.0) {
-		res = connect(s->sock_fd, addr, addrlen);
-		if (res == EINPROGRESS) {
-			internal_select(s, 1);
-			res = connect(s->sock_fd, addr, addrlen);
-		}
-	}
-	else
-		res = connect(s->sock_fd, addr, addrlen);
+	res = internal_connect(s, addr, addrlen);
 	Py_END_ALLOW_THREADS
 
-	if (res != 0) {
-#ifdef MS_WINDOWS
-		res = WSAGetLastError();
-#else
-		res = errno;
-#endif
-	}
-
 	return PyInt_FromLong((long) res);
 }