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.

2 files changed, 19 insertions, 4 deletions

diff --git a/libmproxy/dump.py b/libmproxy/dump.py
index 66bb5206..396dc31d 100644
--- a/libmproxy/dump.py
+++ b/libmproxy/dump.py
@@ -12,6 +12,7 @@ class Options(object):
         "replay",
         "verbosity",
         "wfile",
+        "rheaders",
     ]
     def __init__(self, **kwargs):
         for k, v in kwargs.items():
@@ -52,7 +53,7 @@ class DumpMaster(flow.FlowMaster):
                 flows = list(flow.FlowReader(f).stream())
             except IOError, v:
                 raise DumpError(v.strerror)
-            self.start_playback(flows, options.kill)
+            self.start_playback(flows, options.kill, options.rheaders)
 
     def _runscript(self, f, script):
         try:
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