Merge pull request #2621 from mhils/clientreplay

Improve Clientreplay

6 files changed, 35 insertions, 14 deletions

diff --git a/mitmproxy/addons/clientplayback.py b/mitmproxy/addons/clientplayback.py
index 0db6d336..9e012b67 100644
--- a/mitmproxy/addons/clientplayback.py
+++ b/mitmproxy/addons/clientplayback.py
@@ -9,15 +9,16 @@ import typing
 
 class ClientPlayback:
     def __init__(self):
-        self.flows = None
+        self.flows = []  # type: typing.List[flow.Flow]
         self.current_thread = None
-        self.has_replayed = False
         self.configured = False
 
     def count(self) -> int:
-        if self.flows:
-            return len(self.flows)
-        return 0
+        if self.current_thread:
+            current = 1
+        else:
+            current = 0
+        return current + len(self.flows)
 
     @command.command("replay.client.stop")
     def stop_replay(self) -> None:
@@ -32,7 +33,7 @@ class ClientPlayback:
         """
             Replay requests from flows.
         """
-        self.flows = flows
+        self.flows = list(flows)
         ctx.master.addons.trigger("update", [])
 
     @command.command("replay.client.file")
@@ -54,13 +55,16 @@ class ClientPlayback:
             self.start_replay(flows)
 
     def tick(self):
-        if self.current_thread and not self.current_thread.is_alive():
+        current_is_done = self.current_thread and not self.current_thread.is_alive()
+        can_start_new = not self.current_thread or current_is_done
+        will_start_new = can_start_new and self.flows
+
+        if current_is_done:
             self.current_thread = None
-        if self.flows and not self.current_thread:
+            ctx.master.addons.trigger("update", [])
+        if will_start_new:
             f = self.flows.pop(0)
             self.current_thread = ctx.master.replay_request(f)
             ctx.master.addons.trigger("update", [f])
-            self.has_replayed = True
-        if self.has_replayed:
-            if not self.flows and not self.current_thread:
-                ctx.master.addons.trigger("processing_complete")
+        if current_is_done and not will_start_new:
+            ctx.master.addons.trigger("processing_complete")
diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py
index 294aba26..dc778404 100644
--- a/mitmproxy/flow.py
+++ b/mitmproxy/flow.py
@@ -99,6 +99,7 @@ class Flow(stateobject.StateObject):
         return d
 
     def set_state(self, state):
+        state = state.copy()
         state.pop("version")
         if "backup" in state:
             self._backup = state.pop("backup")
diff --git a/mitmproxy/http.py b/mitmproxy/http.py
index c09778fe..7762647b 100644
--- a/mitmproxy/http.py
+++ b/mitmproxy/http.py
@@ -56,6 +56,7 @@ class HTTPRequest(http.Request):
         return state
 
     def set_state(self, state):
+        state = state.copy()
         self.is_replay = state.pop("is_replay")
         super().set_state(state)
 
diff --git a/mitmproxy/proxy/protocol/http_replay.py b/mitmproxy/proxy/protocol/http_replay.py
index fd673a6f..00bb31c9 100644
--- a/mitmproxy/proxy/protocol/http_replay.py
+++ b/mitmproxy/proxy/protocol/http_replay.py
@@ -42,6 +42,7 @@ class RequestReplayThread(basethread.BaseThread):
         super().__init__(
             "RequestReplay (%s)" % f.request.url
         )
+        self.daemon = True
 
     def run(self):
         r = self.f.request
diff --git a/test/mitmproxy/addons/test_clientplayback.py b/test/mitmproxy/addons/test_clientplayback.py
index 6089b2d5..2dc7eb92 100644
--- a/test/mitmproxy/addons/test_clientplayback.py
+++ b/test/mitmproxy/addons/test_clientplayback.py
@@ -36,9 +36,12 @@ class TestClientPlayback:
                 assert rp.called
                 assert cp.current_thread
 
-            cp.flows = None
-            cp.current_thread = None
+            cp.flows = []
+            cp.current_thread.is_alive.return_value = False
+            assert cp.count() == 1
             cp.tick()
+            assert cp.count() == 0
+            assert tctx.master.has_event("update")
             assert tctx.master.has_event("processing_complete")
 
             cp.current_thread = MockThread()
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index 7f9d577b..fcc766b5 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -84,6 +84,17 @@ class TestSerialize:
         with pytest.raises(Exception, match="version"):
             list(r.stream())
 
+    def test_copy(self):
+        """
+        _backup may be shared across instances. That should not raise errors.
+        """
+        f = tflow.tflow()
+        f.backup()
+        f.request.path = "/foo"
+        f2 = f.copy()
+        f2.revert()
+        f.revert()
+
 
 class TestFlowMaster:
     def test_load_flow_reverse(self):