organize exceptions, improve content-length handling

5 files changed, 107 insertions, 72 deletions

diff --git a/netlib/exceptions.py b/netlib/exceptions.py
index e13af473..e30235af 100644
--- a/netlib/exceptions.py
+++ b/netlib/exceptions.py
@@ -16,7 +16,7 @@ class NetlibException(Exception):
         super(NetlibException, self).__init__(message)
 
 
-class ReadDisconnect(object):
+class Disconnect(object):
     """Immediate EOF"""
 
 
@@ -24,9 +24,35 @@ class HttpException(NetlibException):
     pass
 
 
-class HttpReadDisconnect(HttpException, ReadDisconnect):
+class HttpReadDisconnect(HttpException, Disconnect):
     pass
 
 
 class HttpSyntaxException(HttpException):
     pass
+
+
+class TcpException(NetlibException):
+    pass
+
+
+class TcpDisconnect(TcpException, Disconnect):
+    pass
+
+
+
+
+class TcpReadIncomplete(TcpException):
+    pass
+
+
+class TcpTimeout(TcpException):
+    pass
+
+
+class TlsException(NetlibException):
+    pass
+
+
+class InvalidCertificateException(TlsException):
+    pass
diff --git a/netlib/http/http1/assemble.py b/netlib/http/http1/assemble.py
index ace25d79..33b9ef25 100644
--- a/netlib/http/http1/assemble.py
+++ b/netlib/http/http1/assemble.py
@@ -95,9 +95,9 @@ def _assemble_response_headers(response, preserve_transfer_encoding=False):
     if not preserve_transfer_encoding:
         headers.pop(b"Transfer-Encoding", None)
 
-    # If body is defined (i.e. not None or CONTENT_MISSING), we always
-    # add a content-length header.
-    if response.body or response.body == b"":
-        headers[b"Content-Length"] = str(len(response.body)).encode("ascii")
+        # If body is defined (i.e. not None or CONTENT_MISSING),
+        # we now need to set a content-length header.
+        if response.body or response.body == b"":
+            headers[b"Content-Length"] = str(len(response.body)).encode("ascii")
 
     return bytes(headers)
diff --git a/netlib/http/http1/read.py b/netlib/http/http1/read.py
index 62025d15..7f2b7bab 100644
--- a/netlib/http/http1/read.py
+++ b/netlib/http/http1/read.py
@@ -4,15 +4,14 @@ import sys
 import re
 
 from ... import utils
-from ...exceptions import HttpReadDisconnect, HttpSyntaxException, HttpException
+from ...exceptions import HttpReadDisconnect, HttpSyntaxException, HttpException, TcpDisconnect
 from .. import Request, Response, Headers
-from netlib.tcp import NetLibDisconnect
 
 
 def read_request(rfile, body_size_limit=None):
     request = read_request_head(rfile)
     expected_body_size = expected_http_body_size(request)
-    request.body = b"".join(read_body(rfile, expected_body_size, limit=body_size_limit))
+    request._body = b"".join(read_body(rfile, expected_body_size, limit=body_size_limit))
     request.timestamp_end = time.time()
     return request
 
@@ -51,7 +50,7 @@ def read_request_head(rfile):
 def read_response(rfile, request, body_size_limit=None):
     response = read_response_head(rfile)
     expected_body_size = expected_http_body_size(request, response)
-    response.body = b"".join(read_body(rfile, expected_body_size, body_size_limit))
+    response._body = b"".join(read_body(rfile, expected_body_size, body_size_limit))
     response.timestamp_end = time.time()
     return response
 
@@ -215,7 +214,7 @@ def _get_first_line(rfile):
         if line == b"\r\n" or line == b"\n":
             # Possible leftover from previous message
             line = rfile.readline()
-    except NetLibDisconnect:
+    except TcpDisconnect:
         raise HttpReadDisconnect()
     if not line:
         raise HttpReadDisconnect()
diff --git a/netlib/http/models.py b/netlib/http/models.py
index 2d09535c..b4446ecb 100644
--- a/netlib/http/models.py
+++ b/netlib/http/models.py
@@ -231,7 +231,7 @@ class Request(object):
         self.path = path
         self.httpversion = httpversion
         self.headers = headers
-        self.body = body
+        self._body = body
         self.timestamp_start = timestamp_start
         self.timestamp_end = timestamp_end
         self.form_out = form_out or form_in
@@ -453,6 +453,16 @@ class Request(object):
         self.scheme, self.host, self.port, self.path = parts
 
     @property
+    def body(self):
+        return self._body
+
+    @body.setter
+    def body(self, body):
+        self._body = body
+        if isinstance(body, bytes):
+            self.headers["Content-Length"] = str(len(body)).encode()
+
+    @property
     def content(self):  # pragma: no cover
         # TODO: remove deprecated getter
         return self.body
@@ -488,7 +498,7 @@ class Response(object):
         self.status_code = status_code
         self.msg = msg
         self.headers = headers
-        self.body = body
+        self._body = body
         self.timestamp_start = timestamp_start
         self.timestamp_end = timestamp_end
 
@@ -552,6 +562,16 @@ class Response(object):
         self.headers.set_all("Set-Cookie", values)
 
     @property
+    def body(self):
+        return self._body
+
+    @body.setter
+    def body(self, body):
+        self._body = body
+        if isinstance(body, bytes):
+            self.headers["Content-Length"] = str(len(body)).encode()
+
+    @property
     def content(self):  # pragma: no cover
         # TODO: remove deprecated getter
         return self.body
diff --git a/netlib/tcp.py b/netlib/tcp.py
index 1eb417b4..707e11e0 100644
--- a/netlib/tcp.py
+++ b/netlib/tcp.py
@@ -16,6 +16,9 @@ from . import certutils, version_check
 
 # This is a rather hackish way to make sure that
 # the latest version of pyOpenSSL is actually installed.
+from netlib.exceptions import InvalidCertificateException, TcpReadIncomplete, TlsException, \
+    TcpTimeout, TcpDisconnect, TcpException
+
 version_check.check_pyopenssl_version()
 
 
@@ -24,11 +27,17 @@ EINTR = 4
 # To enable all SSL methods use: SSLv23
 # then add options to disable certain methods
 # https://bugs.launchpad.net/pyopenssl/+bug/1020632/comments/3
+SSL_BASIC_OPTIONS = (
+    SSL.OP_CIPHER_SERVER_PREFERENCE
+)
+if hasattr(SSL, "OP_NO_COMPRESSION"):
+    SSL_BASIC_OPTIONS |= SSL.OP_NO_COMPRESSION
+
 SSL_DEFAULT_METHOD = SSL.SSLv23_METHOD
 SSL_DEFAULT_OPTIONS = (
     SSL.OP_NO_SSLv2 |
     SSL.OP_NO_SSLv3 |
-    SSL.OP_CIPHER_SERVER_PREFERENCE
+    SSL_BASIC_OPTIONS
 )
 if hasattr(SSL, "OP_NO_COMPRESSION"):
     SSL_DEFAULT_OPTIONS |= SSL.OP_NO_COMPRESSION
@@ -39,42 +48,17 @@ Don't ask...
 https://bugs.launchpad.net/pyopenssl/+bug/1020632/comments/3
 """
 sslversion_choices = {
-    "all": (SSL.SSLv23_METHOD, 0),
+    "all": (SSL.SSLv23_METHOD, SSL_BASIC_OPTIONS),
     # SSLv23_METHOD + NO_SSLv2 + NO_SSLv3 == TLS 1.0+
     # TLSv1_METHOD would be TLS 1.0 only
-    "secure": (SSL.SSLv23_METHOD, (SSL.OP_NO_SSLv2 | SSL.OP_NO_SSLv3)),
-    "SSLv2": (SSL.SSLv2_METHOD, 0),
-    "SSLv3": (SSL.SSLv3_METHOD, 0),
-    "TLSv1": (SSL.TLSv1_METHOD, 0),
-    "TLSv1_1": (SSL.TLSv1_1_METHOD, 0),
-    "TLSv1_2": (SSL.TLSv1_2_METHOD, 0),
+    "secure": (SSL.SSLv23_METHOD, (SSL.OP_NO_SSLv2 | SSL.OP_NO_SSLv3 | SSL_BASIC_OPTIONS)),
+    "SSLv2": (SSL.SSLv2_METHOD, SSL_BASIC_OPTIONS),
+    "SSLv3": (SSL.SSLv3_METHOD, SSL_BASIC_OPTIONS),
+    "TLSv1": (SSL.TLSv1_METHOD, SSL_BASIC_OPTIONS),
+    "TLSv1_1": (SSL.TLSv1_1_METHOD, SSL_BASIC_OPTIONS),
+    "TLSv1_2": (SSL.TLSv1_2_METHOD, SSL_BASIC_OPTIONS),
 }
 
-
-class NetLibError(Exception):
-    pass
-
-
-class NetLibDisconnect(NetLibError):
-    pass
-
-
-class NetLibIncomplete(NetLibError):
-    pass
-
-
-class NetLibTimeout(NetLibError):
-    pass
-
-
-class NetLibSSLError(NetLibError):
-    pass
-
-
-class NetLibInvalidCertificateError(NetLibSSLError):
-    pass
-
-
 class SSLKeyLogger(object):
 
     def __init__(self, filename):
@@ -168,17 +152,17 @@ class Writer(_FileLike):
 
     def flush(self):
         """
-            May raise NetLibDisconnect
+            May raise TcpDisconnect
         """
         if hasattr(self.o, "flush"):
             try:
                 self.o.flush()
             except (socket.error, IOError) as v:
-                raise NetLibDisconnect(str(v))
+                raise TcpDisconnect(str(v))
 
     def write(self, v):
         """
-            May raise NetLibDisconnect
+            May raise TcpDisconnect
         """
         if v:
             self.first_byte_timestamp = self.first_byte_timestamp or time.time()
@@ -191,7 +175,7 @@ class Writer(_FileLike):
                     self.add_log(v[:r])
                     return r
             except (SSL.Error, socket.error) as e:
-                raise NetLibDisconnect(str(e))
+                raise TcpDisconnect(str(e))
 
 
 class Reader(_FileLike):
@@ -210,23 +194,29 @@ class Reader(_FileLike):
             try:
                 data = self.o.read(rlen)
             except SSL.ZeroReturnError:
+                # TLS connection was shut down cleanly
                 break
-            except SSL.WantReadError:
+            except (SSL.WantWriteError, SSL.WantReadError):
+                # From the OpenSSL docs:
+                # If the underlying BIO is non-blocking, SSL_read() will also return when the
+                # underlying BIO could not satisfy the needs of SSL_read() to continue the
+                # operation. In this case a call to SSL_get_error with the return value of
+                # SSL_read() will yield SSL_ERROR_WANT_READ or SSL_ERROR_WANT_WRITE.
                 if (time.time() - start) < self.o.gettimeout():
                     time.sleep(0.1)
                     continue
                 else:
-                    raise NetLibTimeout
+                    raise TcpTimeout()
             except socket.timeout:
-                raise NetLibTimeout
-            except socket.error:
-                raise NetLibDisconnect
+                raise TcpTimeout()
+            except socket.error as e:
+                raise TcpDisconnect(str(e))
             except SSL.SysCallError as e:
                 if e.args == (-1, 'Unexpected EOF'):
                     break
-                raise NetLibSSLError(e.message)
+                raise TlsException(e.message)
             except SSL.Error as e:
-                raise NetLibSSLError(e.message)
+                raise TlsException(e.message)
             self.first_byte_timestamp = self.first_byte_timestamp or time.time()
             if not data:
                 break
@@ -260,9 +250,9 @@ class Reader(_FileLike):
         result = self.read(length)
         if length != -1 and len(result) != length:
             if not result:
-                raise NetLibDisconnect()
+                raise TcpDisconnect()
             else:
-                raise NetLibIncomplete(
+                raise TcpReadIncomplete(
                     "Expected %s bytes, got %s" % (length, len(result))
                 )
         return result
@@ -275,15 +265,15 @@ class Reader(_FileLike):
             Up to the next N bytes if peeking is successful.
 
         Raises:
-            NetLibError if there was an error with the socket
-            NetLibSSLError if there was an error with pyOpenSSL.
+            TcpException if there was an error with the socket
+            TlsException if there was an error with pyOpenSSL.
             NotImplementedError if the underlying file object is not a (pyOpenSSL) socket
         """
         if isinstance(self.o, socket._fileobject):
             try:
                 return self.o._sock.recv(length, socket.MSG_PEEK)
             except socket.error as e:
-                raise NetLibError(repr(e))
+                raise TcpException(repr(e))
         elif isinstance(self.o, SSL.Connection):
             try:
                 if tuple(int(x) for x in OpenSSL.__version__.split(".")[:2]) > (0, 15):
@@ -296,7 +286,7 @@ class Reader(_FileLike):
                     self.o._raise_ssl_error(self.o._ssl, result)
                     return SSL._ffi.buffer(buf, result)[:]
             except SSL.Error as e:
-                six.reraise(NetLibSSLError, NetLibSSLError(str(e)), sys.exc_info()[2])
+                six.reraise(TlsException, TlsException(str(e)), sys.exc_info()[2])
         else:
             raise NotImplementedError("Can only peek into (pyOpenSSL) sockets")
 
@@ -461,7 +451,7 @@ class _Connection(object):
                 try:
                     self.wfile.flush()
                     self.wfile.close()
-                except NetLibDisconnect:
+                except TcpDisconnect:
                     pass
 
             self.rfile.close()
@@ -525,7 +515,7 @@ class _Connection(object):
                 # TODO: maybe change this to with newer pyOpenSSL APIs
                 context.set_tmp_ecdh(OpenSSL.crypto.get_elliptic_curve('prime256v1'))
             except SSL.Error as v:
-                raise NetLibError("SSL cipher specification error: %s" % str(v))
+                raise TlsException("SSL cipher specification error: %s" % str(v))
 
         # SSLKEYLOGFILE
         if log_ssl_key:
@@ -546,7 +536,7 @@ class _Connection(object):
             elif alpn_select_callback is not None and alpn_select is None:
                 context.set_alpn_select_callback(alpn_select_callback)
             elif alpn_select_callback is not None and alpn_select is not None:
-                raise NetLibError("ALPN error: only define alpn_select (string) OR alpn_select_callback (method).")
+                raise TlsException("ALPN error: only define alpn_select (string) OR alpn_select_callback (method).")
 
         return context
 
@@ -594,7 +584,7 @@ class TCPClient(_Connection):
                 context.use_privatekey_file(cert)
                 context.use_certificate_file(cert)
             except SSL.Error as v:
-                raise NetLibError("SSL client certificate error: %s" % str(v))
+                raise TlsException("SSL client certificate error: %s" % str(v))
         return context
 
     def convert_to_ssl(self, sni=None, alpn_protos=None, **sslctx_kwargs):
@@ -618,15 +608,15 @@ class TCPClient(_Connection):
             self.connection.do_handshake()
         except SSL.Error as v:
             if self.ssl_verification_error:
-                raise NetLibInvalidCertificateError("SSL handshake error: %s" % repr(v))
+                raise InvalidCertificateException("SSL handshake error: %s" % repr(v))
             else:
-                raise NetLibError("SSL handshake error: %s" % repr(v))
+                raise TlsException("SSL handshake error: %s" % repr(v))
 
         # Fix for pre v1.0 OpenSSL, which doesn't throw an exception on
         # certificate validation failure
         verification_mode = sslctx_kwargs.get('verify_options', None)
         if self.ssl_verification_error is not None and verification_mode == SSL.VERIFY_PEER:
-            raise NetLibInvalidCertificateError("SSL handshake error: certificate verify failed")
+            raise InvalidCertificateException("SSL handshake error: certificate verify failed")
 
         self.ssl_established = True
         self.cert = certutils.SSLCert(self.connection.get_peer_certificate())
@@ -644,7 +634,7 @@ class TCPClient(_Connection):
             self.rfile = Reader(connection.makefile('rb', self.rbufsize))
             self.wfile = Writer(connection.makefile('wb', self.wbufsize))
         except (socket.error, IOError) as err:
-            raise NetLibError(
+            raise TcpException(
                 'Error connecting to "%s": %s' %
                 (self.address.host, err))
         self.connection = connection
@@ -750,7 +740,7 @@ class BaseHandler(_Connection):
         try:
             self.connection.do_handshake()
         except SSL.Error as v:
-            raise NetLibError("SSL handshake error: %s" % repr(v))
+            raise TlsException("SSL handshake error: %s" % repr(v))
         self.ssl_established = True
         self.rfile.set_descriptor(self.connection)
         self.wfile.set_descriptor(self.connection)