various SSL fixes; issues 1251, 3162, 3212

1 files changed, 105 insertions, 256 deletions

diff --git a/Lib/ssl.py b/Lib/ssl.py
index 24502e4..e45e16b 100644
--- a/Lib/ssl.py
+++ b/Lib/ssl.py
@@ -74,7 +74,7 @@ from _ssl import \
      SSL_ERROR_EOF, \
      SSL_ERROR_INVALID_ERROR_CODE
 
-from socket import socket
+from socket import socket, _fileobject
 from socket import getnameinfo as _getnameinfo
 import base64        # for DER-to-PEM translation
 
@@ -86,8 +86,16 @@ class SSLSocket (socket):
 
     def __init__(self, sock, keyfile=None, certfile=None,
                  server_side=False, cert_reqs=CERT_NONE,
-                 ssl_version=PROTOCOL_SSLv23, ca_certs=None):
+                 ssl_version=PROTOCOL_SSLv23, ca_certs=None,
+                 do_handshake_on_connect=True,
+                 suppress_ragged_eofs=True):
         socket.__init__(self, _sock=sock._sock)
+        # the initializer for socket trashes the methods (tsk, tsk), so...
+        self.send = lambda x, flags=0: SSLSocket.send(self, x, flags)
+        self.recv = lambda x, flags=0: SSLSocket.recv(self, x, flags)
+        self.sendto = lambda data, addr, flags=0: SSLSocket.sendto(self, data, addr, flags)
+        self.recvfrom = lambda addr, buflen, flags: SSLSocket.recvfrom(self, addr, buflen, flags)
+
         if certfile and not keyfile:
             keyfile = certfile
         # see if it's connected
@@ -101,18 +109,34 @@ class SSLSocket (socket):
             self._sslobj = _ssl.sslwrap(self._sock, server_side,
                                         keyfile, certfile,
                                         cert_reqs, ssl_version, ca_certs)
+            if do_handshake_on_connect:
+                timeout = self.gettimeout()
+                try:
+                    self.settimeout(None)
+                    self.do_handshake()
+                finally:
+                    self.settimeout(timeout)
         self.keyfile = keyfile
         self.certfile = certfile
         self.cert_reqs = cert_reqs
         self.ssl_version = ssl_version
         self.ca_certs = ca_certs
+        self.do_handshake_on_connect = do_handshake_on_connect
+        self.suppress_ragged_eofs = suppress_ragged_eofs
+        self._makefile_refs = 0
 
     def read(self, len=1024):
 
         """Read up to LEN bytes and return them.
         Return zero-length string on EOF."""
 
-        return self._sslobj.read(len)
+        try:
+            return self._sslobj.read(len)
+        except SSLError, x:
+            if x.args[0] == SSL_ERROR_EOF and self.suppress_ragged_eofs:
+                return ''
+            else:
+                raise
 
     def write(self, data):
 
@@ -143,16 +167,27 @@ class SSLSocket (socket):
                 raise ValueError(
                     "non-zero flags not allowed in calls to send() on %s" %
                     self.__class__)
-            return self._sslobj.write(data)
+            while True:
+                try:
+                    v = self._sslobj.write(data)
+                except SSLError, x:
+                    if x.args[0] == SSL_ERROR_WANT_READ:
+                        return 0
+                    elif x.args[0] == SSL_ERROR_WANT_WRITE:
+                        return 0
+                    else:
+                        raise
+                else:
+                    return v
         else:
             return socket.send(self, data, flags)
 
-    def send_to (self, data, addr, flags=0):
+    def sendto (self, data, addr, flags=0):
         if self._sslobj:
-            raise ValueError("send_to not allowed on instances of %s" %
+            raise ValueError("sendto not allowed on instances of %s" %
                              self.__class__)
         else:
-            return socket.send_to(self, data, addr, flags)
+            return socket.sendto(self, data, addr, flags)
 
     def sendall (self, data, flags=0):
         if self._sslobj:
@@ -160,7 +195,12 @@ class SSLSocket (socket):
                 raise ValueError(
                     "non-zero flags not allowed in calls to sendall() on %s" %
                     self.__class__)
-            return self._sslobj.write(data)
+            amount = len(data)
+            count = 0
+            while (count < amount):
+                v = self.send(data[count:])
+                count += v
+            return amount
         else:
             return socket.sendall(self, data, flags)
 
@@ -170,25 +210,51 @@ class SSLSocket (socket):
                 raise ValueError(
                     "non-zero flags not allowed in calls to sendall() on %s" %
                     self.__class__)
-            return self._sslobj.read(data, buflen)
+            while True:
+                try:
+                    return self.read(buflen)
+                except SSLError, x:
+                    if x.args[0] == SSL_ERROR_WANT_READ:
+                        continue
+                    else:
+                        raise x
         else:
             return socket.recv(self, buflen, flags)
 
-    def recv_from (self, addr, buflen=1024, flags=0):
+    def recvfrom (self, addr, buflen=1024, flags=0):
         if self._sslobj:
-            raise ValueError("recv_from not allowed on instances of %s" %
+            raise ValueError("recvfrom not allowed on instances of %s" %
                              self.__class__)
         else:
-            return socket.recv_from(self, addr, buflen, flags)
+            return socket.recvfrom(self, addr, buflen, flags)
 
-    def shutdown(self, how):
+    def pending (self):
+        if self._sslobj:
+            return self._sslobj.pending()
+        else:
+            return 0
+
+    def shutdown (self, how):
         self._sslobj = None
         socket.shutdown(self, how)
 
-    def close(self):
+    def close (self):
         self._sslobj = None
         socket.close(self)
 
+    def close (self):
+        if self._makefile_refs < 1:
+            self._sslobj = None
+            socket.close(self)
+        else:
+            self._makefile_refs -= 1
+
+    def do_handshake (self):
+
+        """Perform a TLS/SSL handshake."""
+
+        self._sslobj.do_handshake()
+
     def connect(self, addr):
 
         """Connects to remote ADDR, and then wraps the connection in
@@ -202,6 +268,8 @@ class SSLSocket (socket):
         self._sslobj = _ssl.sslwrap(self._sock, False, self.keyfile, self.certfile,
                                     self.cert_reqs, self.ssl_version,
                                     self.ca_certs)
+        if self.do_handshake_on_connect:
+            self.do_handshake()
 
     def accept(self):
 
@@ -210,260 +278,39 @@ class SSLSocket (socket):
         SSL channel, and the address of the remote client."""
 
         newsock, addr = socket.accept(self)
-        return (SSLSocket(newsock, True, self.keyfile, self.certfile,
-                          self.cert_reqs, self.ssl_version,
-                          self.ca_certs), addr)
-
+        return (SSLSocket(newsock,
+                          keyfile=self.keyfile,
+                          certfile=self.certfile,
+                          server_side=True,
+                          cert_reqs=self.cert_reqs,
+                          ssl_version=self.ssl_version,
+                          ca_certs=self.ca_certs,
+                          do_handshake_on_connect=self.do_handshake_on_connect,
+                          suppress_ragged_eofs=self.suppress_ragged_eofs),
+                addr)
 
     def makefile(self, mode='r', bufsize=-1):
 
         """Ouch.  Need to make and return a file-like object that
         works with the SSL connection."""
 
-        if self._sslobj:
-            return SSLFileStream(self._sslobj, mode, bufsize)
-        else:
-            return socket.makefile(self, mode, bufsize)
-
-
-class SSLFileStream:
-
-    """A class to simulate a file stream on top of a socket.
-    Most of this is just lifted from the socket module, and
-    adjusted to work with an SSL stream instead of a socket."""
-
-
-    default_bufsize = 8192
-    name = "<SSL stream>"
-
-    __slots__ = ["mode", "bufsize", "softspace",
-                 # "closed" is a property, see below
-                 "_sslobj", "_rbufsize", "_wbufsize", "_rbuf", "_wbuf",
-                 "_close", "_fileno"]
-
-    def __init__(self, sslobj, mode='rb', bufsize=-1, close=False):
-        self._sslobj = sslobj
-        self.mode = mode # Not actually used in this version
-        if bufsize < 0:
-            bufsize = self.default_bufsize
-        self.bufsize = bufsize
-        self.softspace = False
-        if bufsize == 0:
-            self._rbufsize = 1
-        elif bufsize == 1:
-            self._rbufsize = self.default_bufsize
-        else:
-            self._rbufsize = bufsize
-        self._wbufsize = bufsize
-        self._rbuf = "" # A string
-        self._wbuf = [] # A list of strings
-        self._close = close
-        self._fileno = -1
-
-    def _getclosed(self):
-        return self._sslobj is None
-    closed = property(_getclosed, doc="True if the file is closed")
-
-    def fileno(self):
-        return self._fileno
-
-    def close(self):
-        try:
-            if self._sslobj:
-                self.flush()
-        finally:
-            if self._close and self._sslobj:
-                self._sslobj.close()
-            self._sslobj = None
-
-    def __del__(self):
-        try:
-            self.close()
-        except:
-            # close() may fail if __init__ didn't complete
-            pass
-
-    def flush(self):
-        if self._wbuf:
-            buffer = "".join(self._wbuf)
-            self._wbuf = []
-            count = 0
-            while (count < len(buffer)):
-                written = self._sslobj.write(buffer)
-                count += written
-                buffer = buffer[written:]
-
-    def write(self, data):
-        data = str(data) # XXX Should really reject non-string non-buffers
-        if not data:
-            return
-        self._wbuf.append(data)
-        if (self._wbufsize == 0 or
-            self._wbufsize == 1 and '\n' in data or
-            self._get_wbuf_len() >= self._wbufsize):
-            self.flush()
-
-    def writelines(self, list):
-        # XXX We could do better here for very long lists
-        # XXX Should really reject non-string non-buffers
-        self._wbuf.extend(filter(None, map(str, list)))
-        if (self._wbufsize <= 1 or
-            self._get_wbuf_len() >= self._wbufsize):
-            self.flush()
-
-    def _get_wbuf_len(self):
-        buf_len = 0
-        for x in self._wbuf:
-            buf_len += len(x)
-        return buf_len
-
-    def read(self, size=-1):
-        data = self._rbuf
-        if size < 0:
-            # Read until EOF
-            buffers = []
-            if data:
-                buffers.append(data)
-            self._rbuf = ""
-            if self._rbufsize <= 1:
-                recv_size = self.default_bufsize
-            else:
-                recv_size = self._rbufsize
-            while True:
-                data = self._sslobj.read(recv_size)
-                if not data:
-                    break
-                buffers.append(data)
-            return "".join(buffers)
-        else:
-            # Read until size bytes or EOF seen, whichever comes first
-            buf_len = len(data)
-            if buf_len >= size:
-                self._rbuf = data[size:]
-                return data[:size]
-            buffers = []
-            if data:
-                buffers.append(data)
-            self._rbuf = ""
-            while True:
-                left = size - buf_len
-                recv_size = max(self._rbufsize, left)
-                data = self._sslobj.read(recv_size)
-                if not data:
-                    break
-                buffers.append(data)
-                n = len(data)
-                if n >= left:
-                    self._rbuf = data[left:]
-                    buffers[-1] = data[:left]
-                    break
-                buf_len += n
-            return "".join(buffers)
-
-    def readline(self, size=-1):
-        data = self._rbuf
-        if size < 0:
-            # Read until \n or EOF, whichever comes first
-            if self._rbufsize <= 1:
-                # Speed up unbuffered case
-                assert data == ""
-                buffers = []
-                while data != "\n":
-                    data = self._sslobj.read(1)
-                    if not data:
-                        break
-                    buffers.append(data)
-                return "".join(buffers)
-            nl = data.find('\n')
-            if nl >= 0:
-                nl += 1
-                self._rbuf = data[nl:]
-                return data[:nl]
-            buffers = []
-            if data:
-                buffers.append(data)
-            self._rbuf = ""
-            while True:
-                data = self._sslobj.read(self._rbufsize)
-                if not data:
-                    break
-                buffers.append(data)
-                nl = data.find('\n')
-                if nl >= 0:
-                    nl += 1
-                    self._rbuf = data[nl:]
-                    buffers[-1] = data[:nl]
-                    break
-            return "".join(buffers)
-        else:
-            # Read until size bytes or \n or EOF seen, whichever comes first
-            nl = data.find('\n', 0, size)
-            if nl >= 0:
-                nl += 1
-                self._rbuf = data[nl:]
-                return data[:nl]
-            buf_len = len(data)
-            if buf_len >= size:
-                self._rbuf = data[size:]
-                return data[:size]
-            buffers = []
-            if data:
-                buffers.append(data)
-            self._rbuf = ""
-            while True:
-                data = self._sslobj.read(self._rbufsize)
-                if not data:
-                    break
-                buffers.append(data)
-                left = size - buf_len
-                nl = data.find('\n', 0, left)
-                if nl >= 0:
-                    nl += 1
-                    self._rbuf = data[nl:]
-                    buffers[-1] = data[:nl]
-                    break
-                n = len(data)
-                if n >= left:
-                    self._rbuf = data[left:]
-                    buffers[-1] = data[:left]
-                    break
-                buf_len += n
-            return "".join(buffers)
-
-    def readlines(self, sizehint=0):
-        total = 0
-        list = []
-        while True:
-            line = self.readline()
-            if not line:
-                break
-            list.append(line)
-            total += len(line)
-            if sizehint and total >= sizehint:
-                break
-        return list
-
-    # Iterator protocols
-
-    def __iter__(self):
-        return self
-
-    def next(self):
-        line = self.readline()
-        if not line:
-            raise StopIteration
-        return line
-
+        self._makefile_refs += 1
+        return _fileobject(self, mode, bufsize)
 
 
 
 def wrap_socket(sock, keyfile=None, certfile=None,
                 server_side=False, cert_reqs=CERT_NONE,
-                ssl_version=PROTOCOL_SSLv23, ca_certs=None):
+                ssl_version=PROTOCOL_SSLv23, ca_certs=None,
+                do_handshake_on_connect=True,
+                suppress_ragged_eofs=True):
 
     return SSLSocket(sock, keyfile=keyfile, certfile=certfile,
                      server_side=server_side, cert_reqs=cert_reqs,
-                     ssl_version=ssl_version, ca_certs=ca_certs)
+                     ssl_version=ssl_version, ca_certs=ca_certs,
+                     do_handshake_on_connect=do_handshake_on_connect,
+                     suppress_ragged_eofs=suppress_ragged_eofs)
+
 
 # some utility functions
 
@@ -549,5 +396,7 @@ def sslwrap_simple (sock, keyfile=None, certfile=None):
     for compability with Python 2.5 and earlier.  Will disappear in
     Python 3.0."""
 
-    return _ssl.sslwrap(sock._sock, 0, keyfile, certfile, CERT_NONE,
-                        PROTOCOL_SSLv23, None)
+    ssl_sock = _ssl.sslwrap(sock._sock, 0, keyfile, certfile, CERT_NONE,
+                            PROTOCOL_SSLv23, None)
+    ssl_sock.do_handshake()
+    return ssl_sock