From 4212a56f2567ab750137ac8315561d941741035f Mon Sep 17 00:00:00 2001
From: Alexander Prinzhorn <alexander@prinzhorn.it>
Date: Tue, 9 Feb 2021 19:37:46 +0100
Subject: [PATCH] Collect all flow filter matches before modifying headers,
 fixes #4245 (#4246)

---
 CHANGELOG.md                                |  9 +++++----
 mitmproxy/addons/modifyheaders.py           | 15 +++++++++++----
 test/mitmproxy/addons/test_modifyheaders.py | 13 +++++++++++++
 3 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 853f5f407..633dd2670 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -49,10 +49,11 @@ If you depend on these features, please raise your voice in
 * In reverse proxy mode, mitmproxy now does not assume TLS if no scheme
   is given but a custom port is provided (@mhils)
 * Remove the following options: `http2_priority`, `relax_http_form_validation`, `upstream_bind_address`,
-  `spoof_source_address`, and `stream_websockets`. If you depended on one of them please let us know. 
+  `spoof_source_address`, and `stream_websockets`. If you depended on one of them please let us know.
   mitmproxy never phones home, which means we don't know how prominently these options were used. (@mhils)
 * Fix IDNA host 'Bad HTTP request line' error (@grahamrobbins)
 * Pressing `?` now exits console help view (@abitrolly)
+* `--modify-headers` now works correctly when modifying a header that is also part of the filter expression (@Prinzhorn)
 * --- TODO: add new PRs above this line ---
 * ... and various other fixes, documentation improvements, dependency version bumps, etc.
 
@@ -88,11 +89,11 @@ If you depend on these features, please raise your voice in
 
 * Support for Python 3.9 (@mhils)
 * Add MsgPack content viewer (@tasn)
-* Use `@charset` to decode CSS files if available (@prinzhorn)
+* Use `@charset` to decode CSS files if available (@Prinzhorn)
 * Fix links to anticache docs in mitmweb and use HTTPS for links to documentation (@rugk)
-* Updated typing for WebsocketMessage.content (@prinzhorn)
+* Updated typing for WebsocketMessage.content (@Prinzhorn)
 * Add option `console_strip_trailing_newlines`, and no longer strip trailing newlines by default (@capt8bit)
-* Prevent transparent mode from connecting to itself in the basic cases (@prinzhorn)
+* Prevent transparent mode from connecting to itself in the basic cases (@Prinzhorn)
 * Display HTTP trailers in mitmweb (@sanlengjingvv)
 * Revamp onboarding app (@mhils)
 * Add ASGI support for embedded apps (@mhils)
diff --git a/mitmproxy/addons/modifyheaders.py b/mitmproxy/addons/modifyheaders.py
index 9801f4517..38545d30d 100644
--- a/mitmproxy/addons/modifyheaders.py
+++ b/mitmproxy/addons/modifyheaders.py
@@ -83,14 +83,21 @@ class ModifyHeaders:
         self.run(flow, flow.response.headers)
 
     def run(self, flow: http.HTTPFlow, hdrs: Headers) -> None:
-        # unset all specified headers
+        matches = []
+
+        # first check all the filters against the original, unmodified flow
         for spec in self.replacements:
-            if spec.matches(flow):
+            matches.append(spec.matches(flow))
+
+        # unset all specified headers
+        for i, spec in enumerate(self.replacements):
+            if matches[i]:
                 hdrs.pop(spec.subject, None)
 
         # set all specified headers if the replacement string is not empty
-        for spec in self.replacements:
-            if spec.matches(flow):
+
+        for i, spec in enumerate(self.replacements):
+            if matches[i]:
                 try:
                     replacement = spec.read_replacement()
                 except OSError as e:
diff --git a/test/mitmproxy/addons/test_modifyheaders.py b/test/mitmproxy/addons/test_modifyheaders.py
index 80ddad7a7..8718510ae 100644
--- a/test/mitmproxy/addons/test_modifyheaders.py
+++ b/test/mitmproxy/addons/test_modifyheaders.py
@@ -113,6 +113,19 @@ class TestModifyHeaders:
             mh.response(f)
             assert "one" not in f.response.headers
 
+            # test modifying a header that is also part of the filter expression
+            # https://github.com/mitmproxy/mitmproxy/issues/4245
+            tctx.configure(
+                mh,
+                modify_headers=[
+                    "/~hq ^user-agent:.+Mozilla.+$/user-agent/Definitely not Mozilla ;)"
+                ]
+            )
+            f = tflow.tflow()
+            f.request.headers["user-agent"] = "Hello, it's me, Mozilla"
+            mh.request(f)
+            assert "Definitely not Mozilla ;)" == f.request.headers["user-agent"]
+
     @pytest.mark.parametrize("take", [True, False])
     def test_taken(self, take):
         mh = ModifyHeaders()