add request.host_header, fix #2036

2 files changed, 83 insertions, 1 deletions

diff --git a/mitmproxy/net/http/request.py b/mitmproxy/net/http/request.py
index 68a11ce7..9527bb2d 100644
--- a/mitmproxy/net/http/request.py
+++ b/mitmproxy/net/http/request.py
@@ -1,5 +1,6 @@
 import re
 import urllib
+from typing import Optional
 
 from mitmproxy.types import multidict
 from mitmproxy.utils import strutils
@@ -171,6 +172,42 @@ class Request(message.Message):
                 self.headers.pop("host")
 
     @property
+    def host_header(self) -> Optional[str]:
+        """
+        The request's host/authority header.
+
+        This property maps to either ``request.headers["Host"]`` or
+        ``request.headers[":authority"]``, depending on whether it's HTTP/1.x or HTTP/2.0.
+        """
+        if ":authority" in self.headers:
+            return self.headers[":authority"]
+        if "Host" in self.headers:
+            return self.headers["Host"]
+        return None
+
+    @host_header.setter
+    def host_header(self, val: Optional[str]) -> None:
+        if val is None:
+            self.headers.pop("Host", None)
+            self.headers.pop(":authority", None)
+        elif self.host_header is not None:
+            # Update any existing headers.
+            if ":authority" in self.headers:
+                self.headers[":authority"] = val
+            if "Host" in self.headers:
+                self.headers["Host"] = val
+        else:
+            # Only add the correct new header.
+            if self.http_version.upper().startswith("HTTP/2"):
+                self.headers[":authority"] = val
+            else:
+                self.headers["Host"] = val
+
+    @host_header.deleter
+    def host_header(self):
+        self.host_header = None
+
+    @property
     def port(self):
         """
         Target port
diff --git a/test/mitmproxy/net/http/test_request.py b/test/mitmproxy/net/http/test_request.py
index 6fe57010..90ec31fe 100644
--- a/test/mitmproxy/net/http/test_request.py
+++ b/test/mitmproxy/net/http/test_request.py
@@ -97,7 +97,7 @@ class TestRequestCore:
         request.host = d
         assert request.data.host == b"foo\xFF\x00bar"
 
-    def test_host_header_update(self):
+    def test_host_update_also_updates_header(self):
         request = treq()
         assert "host" not in request.headers
         request.host = "example.com"
@@ -107,6 +107,51 @@ class TestRequestCore:
         request.host = "example.org"
         assert request.headers["Host"] == "example.org"
 
+    def test_get_host_header(self):
+        no_hdr = treq()
+        assert no_hdr.host_header is None
+
+        h1 = treq(headers=(
+            (b"host", b"example.com"),
+        ))
+        assert h1.host_header == "example.com"
+
+        h2 = treq(headers=(
+            (b":authority", b"example.org"),
+        ))
+        assert h2.host_header == "example.org"
+
+        both_hdrs = treq(headers=(
+            (b"host", b"example.org"),
+            (b":authority", b"example.com"),
+        ))
+        assert both_hdrs.host_header == "example.com"
+
+    def test_modify_host_header(self):
+        h1 = treq()
+        assert "host" not in h1.headers
+        assert ":authority" not in h1.headers
+        h1.host_header = "example.com"
+        assert "host" in h1.headers
+        assert ":authority" not in h1.headers
+        h1.host_header = None
+        assert "host" not in h1.headers
+
+        h2 = treq(http_version=b"HTTP/2.0")
+        h2.host_header = "example.org"
+        assert "host" not in h2.headers
+        assert ":authority" in h2.headers
+        del h2.host_header
+        assert ":authority" not in h2.headers
+
+        both_hdrs = treq(headers=(
+            (b":authority", b"example.com"),
+            (b"host", b"example.org"),
+        ))
+        both_hdrs.host_header = "foo.example.com"
+        assert both_hdrs.headers["Host"] == "foo.example.com"
+        assert both_hdrs.headers[":authority"] == "foo.example.com"
+
 
 class TestRequestUtils:
     """