Fix header replacement

bubble-server/src/main/resources/packer/roles/mitmproxy/files/bubble_api.py

@@ -450,7 +450,7 @@ def original_flex_ip(client_addr, fqdns):
     return None
 
 
-def _replace_in_headers(headers: nheaders.Headers, pattern: str, replacement: str):
+def _replace_in_headers(headers: nheaders.Headers, modifiers_dict: dict) -> int:
     """
     Taken from original mitmproxy's Header class implementation with sligh change to allow replacement with empty string
     (resulting with actual removal/skip of the header line).
@@ -461,26 +461,34 @@ def _replace_in_headers(headers: nheaders.Headers, pattern: str, replacement: st
     Returns:
         The number of replacements made.
     """
-    pattern = re.compile(strutils.escaped_str_to_bytes(pattern))
-    replacement = strutils.escaped_str_to_bytes(replacement)
     repl_count = 0
     fields = []
 
     for name, value in headers.fields:
-        line, n = pattern.subn(replacement, name + b": " + value)
+
+        line = name + b": " + value
+        inner_repl_count = 0
+        for pattern, replacement in modifiers_dict.items():
+            line, n = pattern.subn(replacement, line)
+            inner_repl_count += n
+            if len(line) == 0:
+                # No need to go though other patterns for this line
+                break
 
         if len(line) == 0:
-            # Skip/remove this header line and go with the next one:
-            continue
+            # Skip (remove) this header line in this case
+            break
 
-        try:
-            name, value = line.split(b": ", 1)
-        except ValueError:
-            # We get a ValueError if the replacement removed the ": "
-            # There's not much we can do about this, so we just keep the header as-is.
-            pass
-        else:
-            repl_count += n
+        if inner_repl_count > 0:
+            # only in case when there were some replacements:
+            try:
+                name, value = line.split(b": ", 1)
+            except ValueError:
+                # We get a ValueError if the replacement removed the ": "
+                # There's not much we can do about this, so we just keep the header as-is.
+                pass
+            else:
+                repl_count += inner_repl_count
 
         fields.append((name, value))
 
@@ -488,21 +496,25 @@ def _replace_in_headers(headers: nheaders.Headers, pattern: str, replacement: st
     return repl_count
 
 
-def response_header_modify(flow):
+def response_header_modify(flow) -> int:
     if flow.response is None:
         return None
-    return _header_modify(flow.client_conn.address[0], flow.server_conn.address[0], flow.response.headers)
+
+    ctx = {'fqdn': flow.server_conn.address[0]}
+    return _header_modify(flow.client_conn.address[0], ctx, flow.response.headers)
 
 
-def _header_modify(client_addr: str, server_addr: str, headers: nheaders.Headers):
+def _header_modify(client_addr: str, ctx: dict, headers: nheaders.Headers) -> int:
     modifiers_set = 'responseHeaderModifierLists~' + client_addr + '~UNION'
     modifiers = REDIS.smembers(modifiers_set)
 
     repl_count = 0
     if modifiers:
+        modifiers_dict = {}
         for modifier in modifiers:
-            modifier_config = _extract_modifier_config(modifier, server_addr)
-            repl_count += _replace_in_headers(headers, modifier_config['regex'], modifier_config['replacement'])
+            regex, replacement = _extract_modifier_config(modifier, ctx)
+            modifiers_dict[regex] = replacement
+        repl_count += _replace_in_headers(headers, modifiers_dict)
 
     if bubble_log.isEnabledFor(DEBUG):
         bubble_log.debug('_header_modify: replacing headers - replacements count: ' + repl_count)
@@ -510,9 +522,22 @@ def _header_modify(client_addr: str, server_addr: str, headers: nheaders.Headers
     return repl_count
 
 
-def _extract_modifier_config(modifier: str, server_addr: str):
-    modifier.replace('{{fqdn}}', re.escape(server_addr))
-    return json.loads(modifier)
+def _extract_modifier_config(modifier: bytes, ctx: dict) -> tuple:
+    modifier_obj = json.loads(modifier)
+
+    regex = _replace_modifier_values(modifier_obj['regex'], ctx)
+    replacement = _replace_modifier_values(modifier_obj['replacement'], ctx)
+
+    regex = re.compile(strutils.escaped_str_to_bytes(regex))
+    replacement = strutils.escaped_str_to_bytes(replacement)
+
+    return regex, replacement
+
+
+def _replace_modifier_values(s: str, ctx: dict) -> str:
+    # no loop over ctx currently to speed up as there's just 1 variable inside
+    s = s.replace('{{fqdn}}', re.escape(ctx['fqdn']))
+    return s
 
 
 def health_check_response(flow):