Complete upstream authentication module

- Handles upstream CONNECT and regular requests, plus HTTP Basic for reverse
proxy
- Add some tests to make sure we can rely on the .via attribute on server
connections.

5 files changed, 62 insertions, 11 deletions

diff --git a/mitmproxy/addons/upstream_proxy_auth.py b/mitmproxy/addons/upstream_proxy_auth.py
index 2ee51fcb..8b31c10a 100644
--- a/mitmproxy/addons/upstream_proxy_auth.py
+++ b/mitmproxy/addons/upstream_proxy_auth.py
@@ -15,16 +15,39 @@ def parse_upstream_auth(auth):
 
 
 class UpstreamProxyAuth():
+    """
+        This addon handles authentication to systems upstream from us for the
+        upstream proxy and reverse proxy mode. There are 3 cases:
+
+        - Upstream proxy CONNECT requests should have authentication added, and
+          subsequent already connected requests should not.
+        - Upstream proxy regular requests
+        - Reverse proxy regular requests (CONNECT is invalid in this mode)
+    """
     def __init__(self):
         self.auth = None
+        self.root_mode = None
 
     def configure(self, options, updated):
+        # FIXME: We're doing this because our proxy core is terminally confused
+        # at the moment. Ideally, we should be able to check if we're in
+        # reverse proxy mode at the HTTP layer, so that scripts can put the
+        # proxy in reverse proxy mode for specific reuests.
+        if "mode" in updated:
+            self.root_mode = options.mode
         if "upstream_auth" in updated:
             if options.upstream_auth is None:
                 self.auth = None
             else:
                 self.auth = parse_upstream_auth(options.upstream_auth)
 
-    def requestheaders(self, f):
+    def http_connect(self, f):
         if self.auth and f.mode == "upstream":
             f.request.headers["Proxy-Authorization"] = self.auth
+
+    def requestheaders(self, f):
+        if self.auth:
+            if f.mode == "upstream" and not f.server_conn.via:
+                f.request.headers["Proxy-Authorization"] = self.auth
+            elif self.root_mode == "reverse":
+                f.request.headers["Proxy-Authorization"] = self.auth
diff --git a/mitmproxy/proxy/protocol/http.py b/mitmproxy/proxy/protocol/http.py
index 15f3f7cf..5165018d 100644
--- a/mitmproxy/proxy/protocol/http.py
+++ b/mitmproxy/proxy/protocol/http.py
@@ -135,7 +135,9 @@ MODE_REQUEST_FORMS = {
 
 def validate_request_form(mode, request):
     if request.first_line_format == "absolute" and request.scheme != "http":
-        raise exceptions.HttpException("Invalid request scheme: %s" % request.scheme)
+        raise exceptions.HttpException(
+            "Invalid request scheme: %s" % request.scheme
+        )
     allowed_request_forms = MODE_REQUEST_FORMS[mode]
     if request.first_line_format not in allowed_request_forms:
         err_message = "Invalid HTTP request form (expected: %s, got: %s)" % (
@@ -275,8 +277,6 @@ class HttpLayer(base.Layer):
         if not self.connect_request and not self.authenticate(request):
             return False
 
-        f.request = request
-
         # update host header in reverse proxy mode
         if self.config.options.mode == "reverse":
             f.request.headers["Host"] = self.config.upstream_server.address.host
@@ -389,10 +389,8 @@ class HttpLayer(base.Layer):
 
             # Handle 101 Switching Protocols
             if f.response.status_code == 101:
-                """
-                    Handle a successful HTTP 101 Switching Protocols Response, received after
-                    e.g. a WebSocket upgrade request.
-                """
+                # Handle a successful HTTP 101 Switching Protocols Response,
+                # received after e.g. a WebSocket upgrade request.
                 # Check for WebSockets handshake
                 is_websockets = (
                     websockets.check_handshake(f.request.headers) and
@@ -467,13 +465,17 @@ class HttpLayer(base.Layer):
                     self.send_response(http.make_error_response(
                         401,
                         "Authentication Required",
-                        mitmproxy.net.http.Headers(**self.config.authenticator.auth_challenge_headers())
+                        mitmproxy.net.http.Headers(
+                            **self.config.authenticator.auth_challenge_headers()
+                        )
                     ))
                 else:
                     self.send_response(http.make_error_response(
                         407,
                         "Proxy Authentication Required",
-                        mitmproxy.net.http.Headers(**self.config.authenticator.auth_challenge_headers())
+                        mitmproxy.net.http.Headers(
+                            **self.config.authenticator.auth_challenge_headers()
+                        )
                     ))
                 return False
         return True
diff --git a/mitmproxy/proxy/root_context.py b/mitmproxy/proxy/root_context.py
index 50dbe79e..f38f2a8c 100644
--- a/mitmproxy/proxy/root_context.py
+++ b/mitmproxy/proxy/root_context.py
@@ -64,7 +64,12 @@ class RootContext:
         # An inline script may upgrade from http to https,
         # in which case we need some form of TLS layer.
         if isinstance(top_layer, modes.ReverseProxy):
-            return protocol.TlsLayer(top_layer, client_tls, top_layer.server_tls, top_layer.server_conn.address.host)
+            return protocol.TlsLayer(
+                top_layer,
+                client_tls,
+                top_layer.server_tls,
+                top_layer.server_conn.address.host
+            )
         if isinstance(top_layer, protocol.ServerConnectionMixin) or isinstance(top_layer, protocol.UpstreamConnectLayer):
             return protocol.TlsLayer(top_layer, client_tls, client_tls)
 
diff --git a/test/mitmproxy/addons/test_upstream_proxy_auth.py b/test/mitmproxy/addons/test_upstream_proxy_auth.py
index e9a7f4ef..d5d6a3e3 100644
--- a/test/mitmproxy/addons/test_upstream_proxy_auth.py
+++ b/test/mitmproxy/addons/test_upstream_proxy_auth.py
@@ -52,3 +52,14 @@ def test_simple():
         f = tflow.tflow()
         up.requestheaders(f)
         assert "proxy-authorization" not in f.request.headers
+
+        tctx.configure(up, mode="reverse")
+        f = tflow.tflow()
+        f.mode = "transparent"
+        up.requestheaders(f)
+        assert "proxy-authorization" in f.request.headers
+
+        f = tflow.tflow()
+        f.mode = "upstream"
+        up.http_connect(f)
+        assert "proxy-authorization" in f.request.headers
diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py
index dab47c9c..5a5b6817 100644
--- a/test/mitmproxy/test_server.py
+++ b/test/mitmproxy/test_server.py
@@ -669,6 +669,13 @@ class TestProxySSL(tservers.HTTPProxyTest):
         first_flow = self.master.state.flows[0]
         assert first_flow.server_conn.timestamp_ssl_setup
 
+    def test_via(self):
+        # tests that the ssl timestamp is present when ssl is used
+        f = self.pathod("200:b@10")
+        assert f.status_code == 200
+        first_flow = self.master.state.flows[0]
+        assert not first_flow.server_conn.via
+
 
 class MasterRedirectRequest(tservers.TestMaster):
     redirect_port = None  # Set by TestRedirectRequest
@@ -950,11 +957,14 @@ class TestUpstreamProxySSL(
 
         # CONNECT from pathoc to chain[0],
         assert self.proxy.tmaster.state.flow_count() == 1
+        assert self.proxy.tmaster.state.flows[0].server_conn.via
         # request from pathoc to chain[0]
         # CONNECT from proxy to chain[1],
         assert self.chain[0].tmaster.state.flow_count() == 1
+        assert self.chain[0].tmaster.state.flows[0].server_conn.via
         # request from proxy to chain[1]
         # request from chain[0] (regular proxy doesn't store CONNECTs)
+        assert not self.chain[1].tmaster.state.flows[0].server_conn.via
         assert self.chain[1].tmaster.state.flow_count() == 1