remove depth attribute from set_server

4 files changed, 36 insertions, 39 deletions

diff --git a/examples/change_upstream_proxy.py b/examples/change_upstream_proxy.py
index 8f58e1f2..9c454897 100644
--- a/examples/change_upstream_proxy.py
+++ b/examples/change_upstream_proxy.py
@@ -4,7 +4,6 @@
 # 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 proxy_address(flow):
@@ -22,13 +21,4 @@ def request(context, flow):
         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)
+        flow.live.change_upstream_proxy_server(address)
\ No newline at end of file
diff --git a/libmproxy/protocol/base.py b/libmproxy/protocol/base.py
index f1718065..f27cb04b 100644
--- a/libmproxy/protocol/base.py
+++ b/libmproxy/protocol/base.py
@@ -116,19 +116,16 @@ class ServerConnectionMixin(object):
                     "The proxy shall not connect to itself.".format(repr(address))
                 )
 
-    def set_server(self, address, server_tls=None, sni=None, depth=1):
-        if depth == 1:
-            if self.server_conn:
-                self.disconnect()
-            self.log("Set new server address: " + repr(address), "debug")
-            self.server_conn.address = address
-            self.__check_self_connect()
-            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)
+    def set_server(self, address, server_tls=None, sni=None):
+        if self.server_conn:
+            self.disconnect()
+        self.log("Set new server address: " + repr(address), "debug")
+        self.server_conn.address = address
+        self.__check_self_connect()
+        if server_tls:
+            raise ProtocolException(
+                "Cannot upgrade to TLS, no TLS layer on the protocol stack."
+            )
 
     def disconnect(self):
         """
diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py
index 3c934393..f2265c34 100644
--- a/libmproxy/protocol/http.py
+++ b/libmproxy/protocol/http.py
@@ -304,16 +304,22 @@ class UpstreamConnectLayer(Layer):
         else:
             pass  # swallow the message
 
-    def set_server(self, address, server_tls=None, sni=None, depth=1):
-        if depth == 1:
-            if self.ctx.server_conn:
-                self.ctx.disconnect()
-            address = Address.wrap(address)
-            self.connect_request.host = address.host
-            self.connect_request.port = address.port
-            self.server_conn.address = address
-        else:
-            self.ctx.set_server(address, server_tls, sni, depth - 1)
+    def change_upstream_proxy_server(self, address):
+        if address != self.server_conn.via.address:
+            self.ctx.set_server(address)
+
+    def set_server(self, address, server_tls=None, sni=None):
+        if self.ctx.server_conn:
+            self.ctx.disconnect()
+        address = Address.wrap(address)
+        self.connect_request.host = address.host
+        self.connect_request.port = address.port
+        self.server_conn.address = address
+
+        if server_tls:
+            raise ProtocolException(
+                "Cannot upgrade to TLS, no TLS layer on the protocol stack."
+            )
 
 
 class HttpLayer(Layer):
@@ -388,6 +394,12 @@ class HttpLayer(Layer):
             finally:
                 flow.live = False
 
+    def change_upstream_proxy_server(self, address):
+        # Make set_upstream_proxy_server always available,
+        # even if there's no UpstreamConnectLayer
+        if address != self.server_conn.address:
+            return self.set_server(address)
+
     def handle_regular_mode_connect(self, request):
         self.set_server((request.host, request.port))
         self.send_response(make_connect_response(request.httpversion))
diff --git a/libmproxy/protocol/tls.py b/libmproxy/protocol/tls.py
index 00e016ea..a62b1a22 100644
--- a/libmproxy/protocol/tls.py
+++ b/libmproxy/protocol/tls.py
@@ -338,13 +338,11 @@ class TlsLayer(Layer):
         if self._server_tls and not self.server_conn.tls_established:
             self._establish_tls_with_server()
 
-    def set_server(self, address, server_tls=None, sni=None, depth=1):
-        if depth == 1 and server_tls is not None:
-            self.ctx.set_server(address, None, None, 1)
+    def set_server(self, address, server_tls=None, sni=None):
+        if server_tls is not None:
             self._sni_from_server_change = sni
             self._server_tls = server_tls
-        else:
-            self.ctx.set_server(address, server_tls, sni, depth)
+        self.ctx.set_server(address, None, None)
 
     @property
     def sni_for_server_connection(self):