Add a way for users to specify header significance in server replay.

Also add the --rheader command-line option to mitmdump to let the user specify
an arbitrary number of significant headers. The default is to treat no headers
as significant.

1 files changed, 17 insertions, 3 deletions

diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 80a88708..42870f17 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -32,7 +32,12 @@ class RequestReplayThread(threading.Thread):
 
 
 class ServerPlaybackState:
-    def __init__(self):
+    def __init__(self, headers):
+        """
+            headers: A case-insensitive list of request headers that should be
+            included in request-response matching.
+        """
+        self.headers = headers
         self.fmap = {}
 
     def count(self):
@@ -62,6 +67,15 @@ class ServerPlaybackState:
             str(r.path),
             str(r.content),
         ]
+        if self.headers:
+            hdrs = []
+            for i in self.headers:
+                v = r.headers.get(i, [])
+                # Slightly subtle: we need to convert everything to strings
+                # to prevent a mismatch between unicode/non-unicode.
+                v = [str(x) for x in v]
+                hdrs.append((i, v))
+            key.append(repr(hdrs))
         return hashlib.sha256(repr(key)).digest()
 
     def next_flow(self, request):
@@ -342,12 +356,12 @@ class FlowMaster(controller.Master):
     def set_request_script(self, s):
         self.scripts["request"] = s
 
-    def start_playback(self, flows, kill):
+    def start_playback(self, flows, kill, headers):
         """
             flows: A list of flows.
             kill: Boolean, should we kill requests not part of the replay?
         """
-        self.playback = ServerPlaybackState()
+        self.playback = ServerPlaybackState(headers)
         self.playback.load(flows)
         self.kill_nonreplay = kill