most recent changes to SSL module to support non-blocking sockets properly

2 files changed, 65 insertions, 12 deletions

diff --git a/Doc/library/ssl.rst b/Doc/library/ssl.rst
index d700229..5675baa 100644
--- a/Doc/library/ssl.rst
+++ b/Doc/library/ssl.rst
@@ -51,7 +51,7 @@ Functions, Constants, and Exceptions
    network connection.  This error is a subtype of :exc:`socket.error`, which
    in turn is a subtype of :exc:`IOError`.
 
-.. function:: wrap_socket (sock, keyfile=None, certfile=None, server_side=False, cert_reqs=CERT_NONE, ssl_version={see docs}, ca_certs=None)
+.. function:: wrap_socket (sock, keyfile=None, certfile=None, server_side=False, cert_reqs=CERT_NONE, ssl_version={see docs}, ca_certs=None, do_handshake_on_connect=True, suppress_ragged_eofs=True)
 
    Takes an instance ``sock`` of :class:`socket.socket`, and returns an instance of :class:`ssl.SSLSocket`, a subtype
    of :class:`socket.socket`, which wraps the underlying socket in an SSL context.
@@ -118,6 +118,18 @@ Functions, Constants, and Exceptions
    `*` In some older versions of OpenSSL (for instance, 0.9.7l on OS X 10.4),
    an SSLv2 client could not connect to an SSLv23 server.
 
+   The parameter ``do_handshake_on_connect`` specifies whether to do the SSL
+   handshake automatically after doing a :meth:`socket.connect`, or whether the
+   application program will call it explicitly, by invoking the :meth:`SSLSocket.do_handshake`
+   method.  Calling :meth:`SSLSocket.do_handshake` explicitly gives the program control over
+   the blocking behavior of the socket I/O involved in the handshake.
+
+   The parameter ``suppress_ragged_eofs`` specifies how the :meth:`SSLSocket.read`
+   method should signal unexpected EOF from the other end of the connection.  If specified
+   as :const:`True` (the default), it returns a normal EOF in response to unexpected
+   EOF errors raised from the underlying socket; if :const:`False`, it will raise
+   the exceptions back the caller.
+
 .. function:: RAND_status()
 
    Returns True if the SSL pseudo-random number generator has been
@@ -231,15 +243,41 @@ Functions, Constants, and Exceptions
 SSLSocket Objects
 -----------------
 
-.. method:: SSLSocket.read([nbytes=1024])
+.. method:: SSLSocket.read(nbytes=1024, buffer=None)
 
    Reads up to ``nbytes`` bytes from the SSL-encrypted channel and returns them.
+   If the ``buffer`` is specified, it will attempt to read into the buffer
+   the minimum of the size of the buffer and ``nbytes``, if that is specified.
+   If no buffer is specified, an immutable buffer is allocated and returned
+   with the data read from the socket.
 
 .. method:: SSLSocket.write(data)
 
    Writes the ``data`` to the other side of the connection, using the
    SSL channel to encrypt.  Returns the number of bytes written.
 
+.. method:: SSLSocket.do_handshake()
+
+   Performs the SSL setup handshake.  If the socket is non-blocking,
+   this method may raise :exc:`SSLError` with the value of the exception
+   instance's ``args[0]``
+   being either :const:`SSL_ERROR_WANT_READ` or
+   :const:`SSL_ERROR_WANT_WRITE`, and should be called again until
+   it stops raising those exceptions.  Here's an example of how to do
+   that::
+
+        while True:
+            try:
+                sock.do_handshake()
+                break
+            except ssl.SSLError as err:
+                if err.args[0] == ssl.SSL_ERROR_WANT_READ:
+                    select.select([sock], [], [])
+                elif err.args[0] == ssl.SSL_ERROR_WANT_WRITE:
+                    select.select([], [sock], [])
+                else:
+                    raise
+
 .. method:: SSLSocket.getpeercert(binary_form=False)
 
    If there is no certificate for the peer on the other end of the
diff --git a/Lib/ssl.py b/Lib/ssl.py
index 0161efa..be13866 100644
--- a/Lib/ssl.py
+++ b/Lib/ssl.py
@@ -126,12 +126,20 @@ class SSLSocket(socket):
                                             keyfile, certfile,
                                             cert_reqs, ssl_version, ca_certs)
                 if do_handshake_on_connect:
+                    timeout = self.gettimeout()
+                    if timeout == 0.0:
+                        # non-blocking
+                        raise ValueError("do_handshake_on_connect should not be specified for non-blocking sockets")
                     self.do_handshake()
+
             except socket_error as x:
                 self.close()
                 raise x
 
-        self._base = sock
+        if sock and (self.fileno() != sock.fileno()):
+            self._base = sock
+        else:
+            self._base = None
         self.keyfile = keyfile
         self.certfile = certfile
         self.cert_reqs = cert_reqs
@@ -148,7 +156,7 @@ class SSLSocket(socket):
         # raise an exception here if you wish to check for spurious closes
         pass
 
-    def read(self, len=1024, buffer=None):
+    def read(self, len=None, buffer=None):
         """Read up to LEN bytes and return them.
         Return zero-length string on EOF."""
 
@@ -157,7 +165,7 @@ class SSLSocket(socket):
             if buffer:
                 return self._sslobj.read(buffer, len)
             else:
-                return self._sslobj.read(len)
+                return self._sslobj.read(len or 1024)
         except SSLError as x:
             if x.args[0] == SSL_ERROR_EOF and self.suppress_ragged_eofs:
                 return b''
@@ -296,16 +304,18 @@ class SSLSocket(socket):
         # self._closed = True
         if self._base:
             self._base.close()
-        socket._real_close(self)
+        socket.close(self)
 
-    def do_handshake(self):
+    def do_handshake(self, block=False):
         """Perform a TLS/SSL handshake."""
 
+        timeout = self.gettimeout()
         try:
+            if timeout == 0.0 and block:
+                self.settimeout(None)
             self._sslobj.do_handshake()
-        except:
-            self._sslobj = None
-            raise
+        finally:
+            self.settimeout(timeout)
 
     def connect(self, addr):
         """Connects to remote ADDR, and then wraps the connection in
@@ -339,15 +349,20 @@ class SSLSocket(socket):
                 addr)
 
 
+    def __del__(self):
+        self._real_close()
+
 def wrap_socket(sock, keyfile=None, certfile=None,
                 server_side=False, cert_reqs=CERT_NONE,
                 ssl_version=PROTOCOL_SSLv23, ca_certs=None,
-                do_handshake_on_connect=True):
+                do_handshake_on_connect=True,
+                suppress_ragged_eofs=True):
 
     return SSLSocket(sock=sock, keyfile=keyfile, certfile=certfile,
                      server_side=server_side, cert_reqs=cert_reqs,
                      ssl_version=ssl_version, ca_certs=ca_certs,
-                     do_handshake_on_connect=do_handshake_on_connect)
+                     do_handshake_on_connect=do_handshake_on_connect,
+                     suppress_ragged_eofs=suppress_ragged_eofs)
 
 # some utility functions