fix change_upstream_proxy.py example, fix #4981 (#5007)

examples/contrib/change_upstream_proxy.py

@@ -1,10 +1,18 @@
-from mitmproxy import http
 import typing
 
+from mitmproxy import http
+from mitmproxy.connection import Server
+from mitmproxy.net.server_spec import ServerSpec
+
+
 # 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
+# Usage: mitmdump
+#   -s change_upstream_proxy.py
+#   --mode upstream:http://default-upstream-proxy:8080/
+#   --set connection_strategy=lazy
+#   --set upstream_cert=false
 #
 # If you want to change the target server, you should modify flow.request.host and flow.request.port
 
@@ -18,10 +26,12 @@ def proxy_address(flow: http.HTTPFlow) -> typing.Tuple[str, int]:
 
 
 def request(flow: http.HTTPFlow) -> None:
-    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:
-        flow.live.change_upstream_proxy_server(address)  # type: ignore
+
+    is_proxy_change = address != flow.server_conn.via.address
+    server_connection_already_open = flow.server_conn.timestamp_start is not None
+    if is_proxy_change and server_connection_already_open:
+        # server_conn already refers to an existing connection (which cannot be modified),
+        # so we need to replace it with a new server connection object.
+        flow.server_conn = Server(flow.server_conn.address)
+    flow.server_conn.via = ServerSpec("http", address)

mitmproxy/addons/clientplayback.py

@@ -74,7 +74,7 @@ class ReplayHandler(server.ConnectionHandler):
         )
         context.server.tls = flow.request.scheme == "https"
         if options.mode.startswith("upstream:"):
-            context.server.via = server_spec.parse_with_mode(options.mode)[1]
+            context.server.via = flow.server_conn.via = server_spec.parse_with_mode(options.mode)[1]
 
         super().__init__(context)
 

mitmproxy/proxy/layers/http/__init__.py

@@ -541,7 +541,7 @@ class HttpStream(layer.Layer):
         connection, err = yield GetHttpConnection(
             (self.flow.request.host, self.flow.request.port),
             self.flow.request.scheme == "https",
-            self.context.server.via,
+            self.flow.server_conn.via,
         )
         if err:
             yield from self.handle_protocol_error(ResponseProtocolError(self.stream_id, err))

mitmproxy/proxy/tunnel.py

@@ -1,3 +1,4 @@
+import time
 from enum import Enum, auto
 from typing import List, Optional, Tuple, Union
 
@@ -62,16 +63,20 @@ class TunnelLayer(layer.Layer):
                     if done:
                         if self.conn != self.tunnel_connection:
                             self.conn.state = connection.ConnectionState.OPEN
+                            self.conn.timestamp_start = time.time()
                     if err:
                         if self.conn != self.tunnel_connection:
                             self.conn.state = connection.ConnectionState.CLOSED
+                            self.conn.timestamp_start = time.time()
                         yield from self.on_handshake_error(err)
                     if done or err:
                         yield from self._handshake_finished(err)
                 else:
                     yield from self.receive_data(event.data)
             elif isinstance(event, events.ConnectionClosed):
-                self.conn.state &= ~connection.ConnectionState.CAN_READ
+                if self.conn != self.tunnel_connection:
+                    self.conn.state &= ~connection.ConnectionState.CAN_READ
+                    self.conn.timestamp_end = time.time()
                 if self.tunnel_state is TunnelState.OPEN:
                     yield from self.receive_close()
                 elif self.tunnel_state is TunnelState.ESTABLISHING: