From 9b97b6389127e6166fa0f1ad08951ccdb4a36290 Mon Sep 17 00:00:00 2001
From: Maximilian Hils <git@maximilianhils.com>
Date: Wed, 8 Feb 2017 11:49:22 +0100
Subject: [PATCH] fix change_upstream_proxy_server

---
 mitmproxy/proxy/protocol/http.py | 12 ++++++++++--
 test/mitmproxy/test_server.py    | 23 +++++++++++++++++++++++
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/mitmproxy/proxy/protocol/http.py b/mitmproxy/proxy/protocol/http.py
index 59de070f6..b70afe333 100644
--- a/mitmproxy/proxy/protocol/http.py
+++ b/mitmproxy/proxy/protocol/http.py
@@ -88,6 +88,10 @@ class UpstreamConnectLayer(base.Layer):
         layer()
 
     def _send_connect_request(self):
+        self.log("Sending CONNECT request", "debug", [
+            "Proxy Server: {}".format(self.ctx.server_conn.address),
+            "Connect to: {}:{}".format(self.connect_request.host, self.connect_request.port)
+        ])
         self.send_request(self.connect_request)
         resp = self.read_response(self.connect_request)
         if resp.status_code != 200:
@@ -101,6 +105,7 @@ class UpstreamConnectLayer(base.Layer):
             pass  # swallow the message
 
     def change_upstream_proxy_server(self, address):
+        self.log("Changing upstream proxy to {} (CONNECTed)".format(repr(address)), "debug")
         if address != self.server_conn.via.address:
             self.ctx.set_server(address)
 
@@ -432,10 +437,13 @@ class HttpLayer(base.Layer):
         except (exceptions.NetlibException, h2.exceptions.H2Error, exceptions.Http2ProtocolException):
             self.log("Failed to send error response to client: {}".format(message), "debug")
 
-    def change_upstream_proxy_server(self, address) -> None:
+    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:
+        if hasattr(self.ctx, "change_upstream_proxy_server"):
+            self.ctx.change_upstream_proxy_server(address)
+        elif address != self.server_conn.address:
+            self.log("Changing upstream proxy to {} (not CONNECTed)".format(repr(address)), "debug")
             self.set_server(address)
 
     def establish_server_connection(self, host: str, port: int, scheme: str):
diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py
index 6e7ca2754..af121534a 100644
--- a/test/mitmproxy/test_server.py
+++ b/test/mitmproxy/test_server.py
@@ -995,6 +995,29 @@ class TestUpstreamProxySSL(
         assert not self.chain[1].tmaster.state.flows[0].server_conn.via
         assert self.chain[1].tmaster.state.flow_count() == 1
 
+    def test_change_upstream_proxy_connect(self):
+        # skip chain[0].
+        self.proxy.tmaster.addons.add(
+            UpstreamProxyChanger(
+                ("127.0.0.1", self.chain[1].port)
+            )
+        )
+        p = self.pathoc()
+        with p.connect():
+            req = p.request("get:'/p/418'")
+
+        assert req.status_code == 418
+        assert self.chain[0].tmaster.state.flow_count() == 0
+        assert self.chain[1].tmaster.state.flow_count() == 1
+
+
+class UpstreamProxyChanger:
+    def __init__(self, addr):
+        self.address = addr
+
+    def request(self, f):
+        f.live.change_upstream_proxy_server(self.address)
+
 
 class RequestKiller:
     def __init__(self, exclude):