fix upstream proxy server change, update example

examples/change_upstream_proxy.py

@@ -1,29 +1,34 @@
 # This scripts demonstrates how mitmproxy can switch to a second/different upstream proxy
 # in upstream proxy mode.
 #
-# Usage: mitmdump -U http://default-upstream-proxy.local:8080/ -s
-# "change_upstream_proxy.py host"
-from libmproxy.protocol.http import send_connect_request
-
-alternative_upstream_proxy = ("localhost", 8082)
+# Usage: mitmdump -U http://default-upstream-proxy.local:8080/ -s change_upstream_proxy.py
+#
+# If you want to change the target server, you should modify flow.request.host and flow.request.port
+# flow.live.set_server should only be used by inline scripts to change the upstream proxy.
 
 
-def should_redirect(flow):
-    return flow.request.host == "example.com"
+def proxy_address(flow):
+    # Poor man's loadbalancing: route every second domain through the alternative proxy.
+    if hash(flow.request.host) % 2 == 1:
+        return ("localhost", 8082)
+    else:
+        return ("localhost", 8081)
 
 
 def request(context, flow):
-    if flow.live and should_redirect(flow):
-
-        # If you want to change the target server, you should modify flow.request.host and flow.request.port
-        # flow.live.change_server should only be used by inline scripts to change the upstream proxy,
-        # unless you are sure that you know what you are doing.
-        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)
+    if flow.request.method == "CONNECT":
+        # If the decision is done by domain, one could also modify the server address here.
+        # We do it after CONNECT here to have the request data available as well.
+        return
+    address = proxy_address(flow)
+    if flow.live:
+        if flow.request.scheme == "http":
+            # For a normal HTTP request, we just change the proxy server and we're done!
+            if address != flow.live.server_conn.address:
+                flow.live.set_server(address, depth=1)
+        else:
+            # If we have CONNECTed (and thereby established "destination state"), the story is
+            # a bit more complex. Now we don't want to change the top level address (which is
+            # the connect destination) but the address below that. (Notice the `.via` and depth=2).
+            if address != flow.live.server_conn.via.address:
+                flow.live.set_server(address, depth=2)

libmproxy/protocol/base.py

@@ -116,6 +116,10 @@ def set_server(self, address, server_tls=None, sni=None, depth=1):
                 self._disconnect()
             self.log("Set new server address: " + repr(address), "debug")
             self.server_conn.address = address
+            if server_tls:
+                raise ProtocolException(
+                    "Cannot upgrade to TLS, no TLS layer on the protocol stack."
+                )
         else:
             self.ctx.set_server(address, server_tls, sni, depth - 1)
 

libmproxy/protocol/http.py

@@ -207,6 +207,9 @@ def via(self):
     def __getattr__(self, item):
         return getattr(self.via, item)
 
+    def __nonzero__(self):
+        return bool(self.via)
+
 
 class UpstreamConnectLayer(Layer):
     def __init__(self, ctx, connect_request):
@@ -221,19 +224,22 @@ def __call__(self):
         layer = self.ctx.next_layer(self)
         layer()
 
+    def _send_connect_request(self):
+        self.send_request(self.connect_request)
+        resp = self.read_response("CONNECT")
+        if resp.code != 200:
+            raise ProtocolException("Reconnect: Upstream server refuses CONNECT request")
+
     def connect(self):
         if not self.server_conn:
             self.ctx.connect()
-            self.send_request(self.connect_request)
+            self._send_connect_request()
         else:
             pass  # swallow the message
 
     def reconnect(self):
         self.ctx.reconnect()
-        self.send_request(self.connect_request)
-        resp = self.read_response("CONNECT")
-        if resp.code != 200:
-            raise ProtocolException("Reconnect: Upstream server refuses CONNECT request")
+        self._send_connect_request()
 
     def set_server(self, address, server_tls=None, sni=None, depth=1):
         if depth == 1:
@@ -386,7 +392,7 @@ def get_response():
             if self.supports_streaming:
                 flow.response = self.read_response_headers()
             else:
-                flow.response = self.read_response()
+                flow.response = self.read_response(flow.request.method)
 
         try:
             get_response()
@@ -473,13 +479,6 @@ def establish_server_connection(self, flow):
             # Establish connection is neccessary.
             if not self.server_conn:
                 self.connect()
-
-            # SetServer is not guaranteed to work with TLS:
-            # If there's not TlsLayer below which could catch the exception,
-            # TLS will not be established.
-            if tls and not self.server_conn.tls_established:
-                raise ProtocolException(
-                    "Cannot upgrade to SSL, no TLS layer on the protocol stack.")
         else:
             if not self.server_conn:
                 self.connect()

libmproxy/protocol/tls.py

@@ -152,10 +152,12 @@ def reconnect(self):
             self._establish_tls_with_server()
 
     def set_server(self, address, server_tls=None, sni=None, depth=1):
-        self.ctx.set_server(address, server_tls, sni, depth)
         if depth == 1 and server_tls is not None:
+            self.ctx.set_server(address, None, None, 1)
             self._sni_from_server_change = sni
             self._server_tls = server_tls
+        else:
+            self.ctx.set_server(address, server_tls, sni, depth)
 
     @property
     def sni_for_server_connection(self):