coverage++

3 files changed, 30 insertions, 7 deletions

diff --git a/libmproxy/proxy/server.py b/libmproxy/proxy/server.py
index 092eae54..d647ea9d 100644
--- a/libmproxy/proxy/server.py
+++ b/libmproxy/proxy/server.py
@@ -69,9 +69,9 @@ class ConnectionHandler:
         self.sni = None
 
     def handle(self):
-        self.log("clientconnect", "info")
-
         try:
+            self.log("clientconnect", "info")
+
             # Can we already identify the target server and connect to it?
             client_ssl, server_ssl = False, False
             if self.config.get_upstream_server:
@@ -95,6 +95,10 @@ class ConnectionHandler:
             # Delegate handling to the protocol handler
             protocol_handler(self.conntype)(self).handle_messages()
 
+            self.del_server_connection()
+            self.log("clientdisconnect", "info")
+            self.channel.tell("clientdisconnect", self)
+
         except ProxyError as e:
             protocol_handler(self.conntype)(self).handle_error(e)
         except Exception:
@@ -105,10 +109,6 @@ class ConnectionHandler:
             print >> sys.stderr, "mitmproxy has crashed!"
             print >> sys.stderr, "Please lodge a bug report at: https://github.com/mitmproxy/mitmproxy"
 
-        self.del_server_connection()
-        self.log("clientdisconnect", "info")
-        self.channel.tell("clientdisconnect", self)
-
     def del_server_connection(self):
         """
         Deletes (and closes) an existing server connection.
diff --git a/test/test_proxy.py b/test/test_proxy.py
index 614bb8b9..b33cdcfd 100644
--- a/test/test_proxy.py
+++ b/test/test_proxy.py
@@ -3,7 +3,7 @@ from libmproxy import cmdline
 from libmproxy.proxy.config import process_proxy_options
 from libmproxy.proxy.connection import ServerConnection
 from libmproxy.proxy.primitives import ProxyError
-from libmproxy.proxy.server import DummyServer, ProxyServer
+from libmproxy.proxy.server import DummyServer, ProxyServer, ConnectionHandler
 import tutils
 from libpathod import test
 from netlib import http, tcp
@@ -119,6 +119,12 @@ class TestProxyServer:
         opts = parser.parse_args(args=[])
         tutils.raises("error starting proxy server", ProxyServer, opts, 1)
 
+    def test_err_2(self):
+        parser = argparse.ArgumentParser()
+        cmdline.common_options(parser)
+        opts = parser.parse_args(args=[])
+        tutils.raises("error starting proxy server", ProxyServer, opts, 8080, "invalidhost")
+
 
 class TestDummyServer:
     def test_simple(self):
@@ -126,3 +132,10 @@ class TestDummyServer:
         d.start_slave()
         d.shutdown()
 
+
+class TestConnectionHandler:
+    def test_fatal_error(self):
+        config = dict(get_upstream_server=mock.Mock(side_effect=RuntimeError))
+        c = ConnectionHandler(config, mock.MagicMock(), ("127.0.0.1", 8080), None, mock.MagicMock(), None)
+        with tutils.capture_stderr(c.handle) as output:
+            assert "mitmproxy has crashed" in output
\ No newline at end of file
diff --git a/test/tutils.py b/test/tutils.py
index dc049adb..05f65a21 100644
--- a/test/tutils.py
+++ b/test/tutils.py
@@ -1,5 +1,7 @@
+from cStringIO import StringIO
 import os, shutil, tempfile, argparse
 from contextlib import contextmanager
+import sys
 from libmproxy import flow, utils, controller
 from libmproxy.protocol import http
 from libmproxy.proxy.connection import ClientConnection, ServerConnection
@@ -185,4 +187,12 @@ def raises(exc, obj, *args, **kwargs):
                 )
     raise AssertionError("No exception raised.")
 
+
+@contextmanager
+def capture_stderr(command, *args, **kwargs):
+    out, sys.stderr = sys.stderr, StringIO()
+    command(*args, **kwargs)
+    yield sys.stderr.getvalue()
+    sys.stderr = out
+
 test_data = utils.Data(__name__)