fix temporary server change, coverage++

libmproxy/protocol/primitives.py

@@ -1,8 +1,9 @@
 from __future__ import absolute_import
+import copy
+import netlib.tcp
 from .. import stateobject, utils, version
 from ..proxy.primitives import AddressPriority
 from ..proxy.connection import ClientConnection, ServerConnection
-import copy
 
 
 KILL = 0  # const for killed requests
@@ -161,16 +162,12 @@ class TemporaryServerChangeMixin(object):
     without any need to expose the ConnectionHandler to the Flow.
     """
     def change_server(self, address, ssl):
+        address = netlib.tcp.Address.wrap(address)
         if address == self.c.server_conn.address():
             return
         priority = AddressPriority.MANUALLY_CHANGED
 
-        if self.c.server_conn.priority > priority:
-            self.log("Attempt to change server address, "
-                     "but priority is too low (is: %s, got: %s)" % (self.server_conn.priority, priority))
-            return
-
-        self.log("Temporarily change server connection: %s:%s -> %s:%s" % (
+        self.c.log("Temporarily change server connection: %s:%s -> %s:%s" % (
             self.c.server_conn.address.host,
             self.c.server_conn.address.port,
             address.host,
@@ -191,11 +188,11 @@ class TemporaryServerChangeMixin(object):
         if not hasattr(self, "_backup_server_conn"):
             return
 
-        self.log("Restore original server connection: %s:%s -> %s:%s" % (
+        self.c.log("Restore original server connection: %s:%s -> %s:%s" % (
             self.c.server_conn.address.host,
             self.c.server_conn.address.port,
-            self._backup_server_conn.host,
-            self._backup_server_conn.port
+            self._backup_server_conn.address.host,
+            self._backup_server_conn.address.port
         ))
 
         self.c.del_server_connection()

libmproxy/proxy/primitives.py

@@ -25,7 +25,7 @@ class UpstreamServerResolver(object):
         """
         Returns the address of the server to connect to.
         """
-        raise NotImplementedError
+        raise NotImplementedError  # pragma: nocover
 
 
 class ConstUpstreamServerResolver(UpstreamServerResolver):

libmproxy/proxy/server.py

@@ -90,13 +90,15 @@ class ConnectionHandler:
                         continue
 
             # FIXME: Do we want to persist errors?
-            except (ProxyError, tcp.NetLibError), e:
+            except (ProxyError, tcp.NetLibError, IOError), e:
                 handle_error(self.conntype, self, e)
         except Exception, e:
-            self.log(e.__class__)
-            import traceback
+            import traceback, sys
             self.log(traceback.format_exc())
-            self.log(str(e))
+            print >> sys.stderr, traceback.format_exc()
+            print >> sys.stderr, "mitmproxy has crashed!"
+            print >> sys.stderr, "Please lodge a bug report at: https://github.com/mitmproxy/mitmproxy"
+            raise e
 
         self.del_server_connection()
         self.log("clientdisconnect")

test/test_server.py

@@ -3,6 +3,7 @@ import mock
 from netlib import tcp, http_auth, http
 from libpathod import pathoc, pathod
 import tutils, tservers
+from libmproxy import flow
 from libmproxy.protocol import KILL
 from libmproxy.protocol.http import CONTENT_MISSING
 
@@ -328,6 +329,61 @@ class TestProxySSL(tservers.HTTPProxTest):
         first_request = self.master.state.view[0].request
         assert first_request.flow.server_conn.timestamp_ssl_setup
 
+
+class MasterRedirectRequest(tservers.TestMaster):
+    def handle_request(self, request):
+        if request.path == "/p/201":
+            url = request.get_url()
+            new = "http://127.0.0.1:%s/p/201" % self.redirect_port
+
+            request.set_url(new)
+            request.set_url(new)
+            request.flow.change_server(("127.0.0.1", self.redirect_port), False)
+            request.set_url(url)
+            tutils.raises("SSL handshake error", request.flow.change_server, ("127.0.0.1", self.redirect_port), True)
+            request.set_url(new)
+            request.set_url(url)
+            request.set_url(new)
+        tservers.TestMaster.handle_request(self, request)
+
+    def handle_response(self, response):
+        response.content = str(response.flow.client_conn.address.port)
+        tservers.TestMaster.handle_response(self, response)
+
+
+class TestRedirectRequest(tservers.HTTPProxTest):
+    masterclass = MasterRedirectRequest
+
+    def test_redirect(self):
+        self.master.redirect_port = self.server2.port
+
+        p = self.pathoc()
+
+        self.server.clear_log()
+        self.server2.clear_log()
+        r1 = p.request("get:'%s/p/200'"%self.server.urlbase)
+        assert r1.status_code == 200
+        assert self.server.last_log()
+        assert not self.server2.last_log()
+
+        self.server.clear_log()
+        self.server2.clear_log()
+        r2 = p.request("get:'%s/p/201'"%self.server.urlbase)
+        assert r2.status_code == 201
+        assert not self.server.last_log()
+        assert self.server2.last_log()
+
+        self.server.clear_log()
+        self.server2.clear_log()
+        r3 = p.request("get:'%s/p/202'"%self.server.urlbase)
+        assert r3.status_code == 202
+        assert self.server.last_log()
+        assert not self.server2.last_log()
+
+        assert r3.content == r2.content == r1.content
+        # Make sure that we actually use the same connection in this test case
+
+
 class MasterFakeResponse(tservers.TestMaster):
     def handle_request(self, m):
         resp = tutils.tresp()