From d3f6ea1d1e53f5ff35ebd149ddc57a1a0cb5e543 Mon Sep 17 00:00:00 2001
From: Antoine Pitrou <solipsis@pitrou.net>
Date: Sat, 26 Feb 2011 23:35:27 +0000
Subject: [PATCH] Merged revisions 88664 via svnmerge from
 svn+ssh://pythondev@svn.python.org/python/branches/py3k
MIME-Version: 1.0
Content-Type: text/plain; charset=utf8
Content-Transfer-Encoding: 8bit

........
  r88664 | antoine.pitrou | 2011-02-27 00:24:06 +0100 (dim., 27 févr. 2011) | 4 lines

  Issue #11326: Add the missing connect_ex() implementation for SSL sockets,
  and make it work for non-blocking connects.
........
---
 Lib/ssl.py           | 35 ++++++++++++++++++++++++++---------
 Lib/test/test_ssl.py | 43 +++++++++++++++++++++++++++++++++++++++++++
 Misc/NEWS            |  3 +++
 3 files changed, 72 insertions(+), 9 deletions(-)

diff --git a/Lib/ssl.py b/Lib/ssl.py
index 8c2cdf1bb4..72fbae57ff 100644
--- a/Lib/ssl.py
+++ b/Lib/ssl.py
@@ -110,9 +110,11 @@ class SSLSocket(socket):
             if e.errno != errno.ENOTCONN:
                 raise
             # no, no connection yet
+            self._connected = False
             self._sslobj = None
         else:
             # yes, create the SSL object
+            self._connected = True
             self._sslobj = _ssl.sslwrap(self._sock, server_side,
                                         keyfile, certfile,
                                         cert_reqs, ssl_version, ca_certs,
@@ -282,21 +284,36 @@ class SSLSocket(socket):
 
         self._sslobj.do_handshake()
 
-    def connect(self, addr):
-
-        """Connects to remote ADDR, and then wraps the connection in
-        an SSL channel."""
-
+    def _real_connect(self, addr, return_errno):
         # Here we assume that the socket is client-side, and not
         # connected at the time of the call.  We connect it, then wrap it.
-        if self._sslobj:
+        if self._connected:
             raise ValueError("attempt to connect already-connected SSLSocket!")
-        socket.connect(self, addr)
         self._sslobj = _ssl.sslwrap(self._sock, False, self.keyfile, self.certfile,
                                     self.cert_reqs, self.ssl_version,
                                     self.ca_certs, self.ciphers)
-        if self.do_handshake_on_connect:
-            self.do_handshake()
+        try:
+            socket.connect(self, addr)
+            if self.do_handshake_on_connect:
+                self.do_handshake()
+        except socket_error as e:
+            if return_errno:
+                return e.errno
+            else:
+                self._sslobj = None
+                raise e
+        self._connected = True
+        return 0
+
+    def connect(self, addr):
+        """Connects to remote ADDR, and then wraps the connection in
+        an SSL channel."""
+        self._real_connect(addr, False)
+
+    def connect_ex(self, addr):
+        """Connects to remote ADDR, and then wraps the connection in
+        an SSL channel."""
+        return self._real_connect(addr, True)
 
     def accept(self):
 
diff --git a/Lib/test/test_ssl.py b/Lib/test/test_ssl.py
index bc2bfccf41..20cfe38e36 100644
--- a/Lib/test/test_ssl.py
+++ b/Lib/test/test_ssl.py
@@ -225,6 +225,49 @@ class NetworkedTests(unittest.TestCase):
             finally:
                 s.close()
 
+    def test_connect_ex(self):
+        # Issue #11326: check connect_ex() implementation
+        with test_support.transient_internet("svn.python.org"):
+            s = ssl.wrap_socket(socket.socket(socket.AF_INET),
+                                cert_reqs=ssl.CERT_REQUIRED,
+                                ca_certs=SVN_PYTHON_ORG_ROOT_CERT)
+            try:
+                self.assertEqual(0, s.connect_ex(("svn.python.org", 443)))
+                self.assertTrue(s.getpeercert())
+            finally:
+                s.close()
+
+    def test_non_blocking_connect_ex(self):
+        # Issue #11326: non-blocking connect_ex() should allow handshake
+        # to proceed after the socket gets ready.
+        with test_support.transient_internet("svn.python.org"):
+            s = ssl.wrap_socket(socket.socket(socket.AF_INET),
+                                cert_reqs=ssl.CERT_REQUIRED,
+                                ca_certs=SVN_PYTHON_ORG_ROOT_CERT,
+                                do_handshake_on_connect=False)
+            try:
+                s.setblocking(False)
+                rc = s.connect_ex(('svn.python.org', 443))
+                self.assertIn(rc, (0, errno.EINPROGRESS))
+                # Wait for connect to finish
+                select.select([], [s], [], 5.0)
+                # Non-blocking handshake
+                while True:
+                    try:
+                        s.do_handshake()
+                        break
+                    except ssl.SSLError as err:
+                        if err.args[0] == ssl.SSL_ERROR_WANT_READ:
+                            select.select([s], [], [], 5.0)
+                        elif err.args[0] == ssl.SSL_ERROR_WANT_WRITE:
+                            select.select([], [s], [], 5.0)
+                        else:
+                            raise
+                # SSL established
+                self.assertTrue(s.getpeercert())
+            finally:
+                s.close()
+
     @unittest.skipIf(os.name == "nt", "Can't use a socket as a file under Windows")
     def test_makefile_close(self):
         # Issue #5238: creating a file-like object with makefile() shouldn't
diff --git a/Misc/NEWS b/Misc/NEWS
index d4f8b31fdc..6ec6bf32fe 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -37,6 +37,9 @@ Core and Builtins
 Library
 -------
 
+- Issue #11326: Add the missing connect_ex() implementation for SSL sockets,
+  and make it work for non-blocking connects.
+
 - Issue #10956: Buffered I/O classes retry reading or writing after a signal
   has arrived and the handler returned successfully.
 
-- 
2.50.1