From dcdee7dc461677d8beef0ffba1eb2b30ae064609 Mon Sep 17 00:00:00 2001
From: Maximilian Hils <git@maximilianhils.com>
Date: Sat, 9 Apr 2016 16:47:05 -0700
Subject: [PATCH] server replay: keep --kill active

---
 mitmproxy/flow.py           |  7 +++----
 test/mitmproxy/test_dump.py |  5 +++--
 test/mitmproxy/test_flow.py | 10 ----------
 3 files changed, 6 insertions(+), 16 deletions(-)

diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py
index f5e238a26..246e9f82b 100644
--- a/mitmproxy/flow.py
+++ b/mitmproxy/flow.py
@@ -789,8 +789,6 @@ class FlowMaster(controller.Master):
         self.kill_nonreplay = kill
 
     def stop_server_playback(self):
-        if self.server_playback.exit:
-            self.shutdown()
         self.server_playback = None
 
     def do_server_playback(self, flow):
@@ -827,7 +825,8 @@ class FlowMaster(controller.Master):
         if self.server_playback:
             stop = (
                 self.server_playback.count() == 0 and
-                self.state.active_flow_count() == 0
+                self.state.active_flow_count() == 0 and
+                not self.kill_nonreplay
             )
             exit = self.server_playback.exit
             if stop:
@@ -935,7 +934,7 @@ class FlowMaster(controller.Master):
         if self.server_playback:
             pb = self.do_server_playback(f)
             if not pb and self.kill_nonreplay:
-                    f.kill(self)
+                f.kill(self)
 
     def process_new_response(self, f):
         if self.stickycookie_state:
diff --git a/test/mitmproxy/test_dump.py b/test/mitmproxy/test_dump.py
index 7e7728814..ad4cee539 100644
--- a/test/mitmproxy/test_dump.py
+++ b/test/mitmproxy/test_dump.py
@@ -68,8 +68,9 @@ class TestDumpMaster:
         m.handle_clientconnect(f.client_conn)
         m.handle_serverconnect(f.server_conn)
         m.handle_request(f)
-        f.response = HTTPResponse.wrap(netlib.tutils.tresp(content=content))
-        f = m.handle_response(f)
+        if not f.error:
+            f.response = HTTPResponse.wrap(netlib.tutils.tresp(content=content))
+            f = m.handle_response(f)
         m.handle_clientdisconnect(f.client_conn)
         return f
 
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index 1fadb23dc..7ede7a81c 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -904,16 +904,6 @@ class TestFlowMaster:
         assert not fm.do_server_playback(r)
         assert fm.do_server_playback(tutils.tflow())
 
-        fm.start_server_playback(
-            pb,
-            False,
-            [],
-            True,
-            False,
-            None,
-            False,
-            None,
-            False)
         q = queue.Queue()
         fm.tick(q, 0)
         assert fm.should_exit.is_set()