improve error handling

5 files changed, 76 insertions, 70 deletions

diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py
index 82ee2b3c..038fde97 100644
--- a/libmproxy/protocol/http.py
+++ b/libmproxy/protocol/http.py
@@ -22,7 +22,7 @@ def get_line(fp):
     if line == "\r\n" or line == "\n":  # Possible leftover from previous message
         line = fp.readline()
     if line == "":
-        raise tcp.NetLibDisconnect
+        raise tcp.NetLibDisconnect()
     return line
 
 
@@ -927,7 +927,7 @@ class HTTPHandler(ProtocolHandler):
                                                body_size_limit=self.c.config.body_size_limit, include_body=include_body)
                 return res
             except (tcp.NetLibDisconnect, http.HttpErrorConnClosed), v:
-                self.c.log("error in server communication: %s" % str(v), level="debug")
+                self.c.log("error in server communication: %s" % repr(v), level="debug")
                 if i < 1:
                     # In any case, we try to reconnect at least once.
                     # This is necessary because it might be possible that we already initiated an upstream connection
@@ -1064,23 +1064,23 @@ class HTTPHandler(ProtocolHandler):
     def handle_error(self, error, flow=None):
 
         message = repr(error)
-        code = getattr(error, "code", 502)
-        headers = getattr(error, "headers", None)
-
         if "tlsv1 alert unknown ca" in message:
-            message = message + " \nThe client does not trust the proxy's certificate."
+            message += " The client does not trust the proxy's certificate."
 
-        self.c.log("error: %s" % message, level="info")
+        if isinstance(error, tcp.NetLibDisconnect):
+            self.c.log("TCP connection closed unexpectedly.", "debug", [repr(error)])
+        else:
+            self.c.log("error: %s" % message, level="info")
 
         if flow:
-            flow.error = Error(message)
-            # FIXME: no flows without request or with both request and response at the moment.
+            # TODO: no flows without request or with both request and response at the moment.
             if flow.request and not flow.response:
+                flow.error = Error(message)
                 self.c.channel.ask("error", flow.error)
-        else:
-            pass
 
         try:
+            code = getattr(error, "code", 502)
+            headers = getattr(error, "headers", None)
             self.send_error(code, message, headers)
         except:
             pass
diff --git a/libmproxy/protocol/tcp.py b/libmproxy/protocol/tcp.py
index 00dbf4b3..8411ba6c 100644
--- a/libmproxy/protocol/tcp.py
+++ b/libmproxy/protocol/tcp.py
@@ -16,51 +16,55 @@ class TCPHandler(ProtocolHandler):
 
         server = "%s:%s" % self.c.server_conn.address()[:2]
         buf = memoryview(bytearray(self.chunk_size))
-
         conns = [self.c.client_conn.rfile, self.c.server_conn.rfile]
-        while True:
-            r, _, _ = select.select(conns, [], [], 10)
-            for rfile in r:
-                if self.c.client_conn.rfile == rfile:
-                    src, dst = self.c.client_conn, self.c.server_conn
-                    direction = "-> tcp ->"
-                    src_str, dst_str = "client", server
-                else:
-                    dst, src = self.c.client_conn, self.c.server_conn
-                    direction = "<- tcp <-"
-                    dst_str, src_str = "client", server
 
-                closed = False
-                if src.ssl_established:
-                    # Unfortunately, pyOpenSSL lacks a recv_into function.
-                    contents = src.rfile.read(1)  # We need to read a single byte before .pending() becomes usable
-                    contents += src.rfile.read(src.connection.pending())
-                    if not contents:
-                        closed = True
-                else:
-                    size = src.connection.recv_into(buf)
-                    if not size:
-                        closed = True
+        try:
+            while True:
+                r, _, _ = select.select(conns, [], [], 10)
+                for rfile in r:
+                    if self.c.client_conn.rfile == rfile:
+                        src, dst = self.c.client_conn, self.c.server_conn
+                        direction = "-> tcp ->"
+                        src_str, dst_str = "client", server
+                    else:
+                        dst, src = self.c.client_conn, self.c.server_conn
+                        direction = "<- tcp <-"
+                        dst_str, src_str = "client", server
 
-                if closed:
-                    conns.remove(src.rfile)
-                    # Shutdown connection to the other peer
-                    if dst.ssl_established:
-                        dst.connection.shutdown()
+                    closed = False
+                    if src.ssl_established:
+                        # Unfortunately, pyOpenSSL lacks a recv_into function.
+                        contents = src.rfile.read(1)  # We need to read a single byte before .pending() becomes usable
+                        contents += src.rfile.read(src.connection.pending())
+                        if not contents:
+                            closed = True
                     else:
-                        dst.connection.shutdown(socket.SHUT_WR)
+                        size = src.connection.recv_into(buf)
+                        if not size:
+                            closed = True
 
-                    if len(conns) == 0:
-                        return
-                    continue
+                    if closed:
+                        conns.remove(src.rfile)
+                        # Shutdown connection to the other peer
+                        if dst.ssl_established:
+                            dst.connection.shutdown()
+                        else:
+                            dst.connection.shutdown(socket.SHUT_WR)
 
-                if src.ssl_established or dst.ssl_established:
-                    # if one of the peers is over SSL, we need to send bytes/strings
-                    if not src.ssl_established:  # only ssl to dst, i.e. we revc'd into buf but need bytes/string now.
-                        contents = buf[:size].tobytes()
-                    self.c.log("%s %s\r\n%s" % (direction, dst_str, cleanBin(contents)), "debug")
-                    dst.connection.send(contents)
-                else:
-                    # socket.socket.send supports raw bytearrays/memoryviews
-                    self.c.log("%s %s\r\n%s" % (direction, dst_str, cleanBin(buf.tobytes())), "debug")
-                    dst.connection.send(buf[:size])
\ No newline at end of file
+                        if len(conns) == 0:
+                            return
+                        continue
+
+                    if src.ssl_established or dst.ssl_established:
+                        # if one of the peers is over SSL, we need to send bytes/strings
+                        if not src.ssl_established:  # only ssl to dst, i.e. we revc'd into buf but need bytes/string now.
+                            contents = buf[:size].tobytes()
+                        # self.c.log("%s %s\r\n%s" % (direction, dst_str, cleanBin(contents)), "debug")
+                        dst.connection.send(contents)
+                    else:
+                        # socket.socket.send supports raw bytearrays/memoryviews
+                        # self.c.log("%s %s\r\n%s" % (direction, dst_str, cleanBin(buf.tobytes())), "debug")
+                        dst.connection.send(buf[:size])
+        except socket.error as e:
+            self.c.log("TCP connection closed unexpectedly.", "debug", [repr(e)])
+            return
\ No newline at end of file
diff --git a/libmproxy/proxy/connection.py b/libmproxy/proxy/connection.py
index 07129d6a..d99ffa9b 100644
--- a/libmproxy/proxy/connection.py
+++ b/libmproxy/proxy/connection.py
@@ -3,7 +3,6 @@ import copy
 import os
 from netlib import tcp, certutils
 from .. import stateobject, utils
-from .primitives import ProxyError
 
 
 class ClientConnection(tcp.BaseHandler, stateobject.SimpleStateObject):
@@ -156,11 +155,8 @@ class ServerConnection(tcp.TCPClient, stateobject.SimpleStateObject):
             path = os.path.join(clientcerts, self.address.host.encode("idna")) + ".pem"
             if os.path.exists(path):
                 clientcert = path
-        try:
-            self.convert_to_ssl(cert=clientcert, sni=sni)
-            self.timestamp_ssl_setup = utils.timestamp()
-        except tcp.NetLibError, v:
-            raise ProxyError(400, repr(v))
+        self.convert_to_ssl(cert=clientcert, sni=sni)
+        self.timestamp_ssl_setup = utils.timestamp()
 
     def finish(self):
         tcp.TCPClient.finish(self)
diff --git a/libmproxy/proxy/server.py b/libmproxy/proxy/server.py
index 946a301a..092eae54 100644
--- a/libmproxy/proxy/server.py
+++ b/libmproxy/proxy/server.py
@@ -95,7 +95,7 @@ class ConnectionHandler:
             # Delegate handling to the protocol handler
             protocol_handler(self.conntype)(self).handle_messages()
 
-        except (ProxyError, tcp.NetLibError), e:
+        except ProxyError as e:
             protocol_handler(self.conntype)(self).handle_error(e)
         except Exception:
             import traceback, sys
@@ -190,18 +190,24 @@ class ConnectionHandler:
                 raise ProxyError(502, "No server connection.")
             if self.server_conn.ssl_established:
                 raise ProxyError(502, "SSL to Server already established.")
-            self.server_conn.establish_ssl(self.config.clientcerts, self.sni)
+            try:
+                self.server_conn.establish_ssl(self.config.clientcerts, self.sni)
+            except tcp.NetLibError as v:
+                raise ProxyError(502, repr(v))
         if client:
             if self.client_conn.ssl_established:
                 raise ProxyError(502, "SSL to Client already established.")
             cert, key = self.find_cert()
-            self.client_conn.convert_to_ssl(
-                cert, key,
-                handle_sni=self.handle_sni,
-                cipher_list=self.config.ciphers,
-                dhparams=self.config.certstore.dhparams,
-                ca_file=self.config.ca_file
-            )
+            try:
+                self.client_conn.convert_to_ssl(
+                    cert, key,
+                    handle_sni=self.handle_sni,
+                    cipher_list=self.config.ciphers,
+                    dhparams=self.config.certstore.dhparams,
+                    ca_file=self.config.ca_file
+                )
+            except tcp.NetLibError as v:
+                raise ProxyError(400, repr(v))
 
     def server_reconnect(self):
         address = self.server_conn.address
diff --git a/test/test_server.py b/test/test_server.py
index 52efa5f2..a570f10f 100644
--- a/test/test_server.py
+++ b/test/test_server.py
@@ -179,7 +179,7 @@ class TestHTTPConnectSSLError(tservers.HTTPProxTest):
         p = self.pathoc_raw()
         dst = ("localhost", self.proxy.port)
         p.connect(connect_to=dst)
-        tutils.raises("400 - Bad Request", p.http_connect, dst)
+        tutils.raises("502 - Bad Gateway", p.http_connect, dst)
 
 
 class TestHTTPS(tservers.HTTPProxTest, CommonMixin):
@@ -244,7 +244,7 @@ class TestTransparentSSL(tservers.TransparentProxTest, CommonMixin):
         p = pathoc.Pathoc(("localhost", self.proxy.port))
         p.connect()
         r = p.request("get:/")
-        assert r.status_code == 502
+        assert r.status_code == 400
 
 
 class TestProxy(tservers.HTTPProxTest):