improve change_server api, add example how to change the upstream server

3 files changed, 41 insertions, 13 deletions

diff --git a/examples/change_upstream_proxy.py b/examples/change_upstream_proxy.py
new file mode 100644
index 00000000..2e4c6080
--- /dev/null
+++ b/examples/change_upstream_proxy.py
@@ -0,0 +1,17 @@
+# This scripts demonstrates how mitmproxy can switch to a different upstream proxy
+# in upstream proxy mode.
+#
+# Usage: mitmdump -s "change_upstream_proxy.py host"
+from libmproxy.protocol.http import send_connect_request
+
+def should_redirect(flow):
+	return (flow.request.host == "example.com")
+alternative_upstream_proxy = ("localhost",8082)
+
+def request(ctx, flow):
+	if flow.live and should_redirect(flow):
+
+		server_changed = flow.live.change_server(alternative_upstream_proxy, persistent_change=True)
+		if flow.request.scheme == "https" and server_changed:
+			send_connect_request(flow.live.c.server_conn, flow.request.host, flow.request.port)
+			flow.live.c.establish_ssl(server=True)
diff --git a/libmproxy/protocol/primitives.py b/libmproxy/protocol/primitives.py
index ecad9d9e..160c50c7 100644
--- a/libmproxy/protocol/primitives.py
+++ b/libmproxy/protocol/primitives.py
@@ -165,9 +165,22 @@ class LiveConnection(object):
         self._backup_server_conn = None
         """@type: libmproxy.proxy.connection.ServerConnection"""
 
-    def change_server(self, address, ssl=False, force=False, persistent_change=False):
+    def change_server(self, address, ssl=None, force=False, persistent_change=False):
+        """
+        Change the server connection to the specified address.
+        @returns:
+        True, if a new connection has been established,
+        False, if an existing connection has been used
+        """
         address = netlib.tcp.Address.wrap(address)
-        if force or address != self.c.server_conn.address or ssl != self.c.server_conn.ssl_established:
+
+        ssl_mismatch = (ssl is not None and ssl != self.c.server_conn.ssl_established)
+        address_mismatch = (address != self.c.server_conn.address)
+
+        if persistent_change:
+            self._backup_server_conn = None
+
+        if ssl_mismatch or address_mismatch or force:
 
             self.c.log("Change server connection: %s:%s -> %s:%s [persistent: %s]" % (
                 self.c.server_conn.address.host,
@@ -177,7 +190,7 @@ class LiveConnection(object):
                 persistent_change
             ), "debug")
 
-            if not self._backup_server_conn:
+            if not self._backup_server_conn and not persistent_change:
                 self._backup_server_conn = self.c.server_conn
                 self.c.server_conn = None
             else:  # This is at least the second temporary change. We can kill the current connection.
@@ -187,8 +200,8 @@ class LiveConnection(object):
             self.c.establish_server_connection(ask=False)
             if ssl:
                 self.c.establish_ssl(server=True)
-        if persistent_change:
-            self._backup_server_conn = None
+            return True
+        return False
 
     def restore_server(self):
         # TODO: Similar to _backup_server_conn, introduce _cache_server_conn, which keeps the changed connection open
diff --git a/test/test_server.py b/test/test_server.py
index d33bcc89..b128d0a2 100644
--- a/test/test_server.py
+++ b/test/test_server.py
@@ -387,14 +387,12 @@ class MasterRedirectRequest(tservers.TestMaster):
     def handle_request(self, f):
         request = f.request
         if request.path == "/p/201":
-            url = request.url
-            new = "http://127.0.0.1:%s/p/201" % self.redirect_port
-
-            request.url = new
-            f.live.change_server(("127.0.0.1", self.redirect_port), False)
-            request.url = url
-            tutils.raises("SSL handshake error", f.live.change_server, ("127.0.0.1", self.redirect_port), True)
-            request.url = new
+            addr = f.live.c.server_conn.address
+            assert f.live.change_server(("127.0.0.1", self.redirect_port), ssl=False)
+            assert not f.live.change_server(("127.0.0.1", self.redirect_port), ssl=False)
+            tutils.raises("SSL handshake error", f.live.change_server, ("127.0.0.1", self.redirect_port), ssl=True)
+            assert f.live.change_server(addr, ssl=False)
+            request.url = "http://127.0.0.1:%s/p/201" % self.redirect_port
         tservers.TestMaster.handle_request(self, f)
 
     def handle_response(self, f):