remove depth attribute from set_server

3 files changed, 35 insertions, 28 deletions

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):