From f96697646ca495ddc045ab18134b8f70052c86ff Mon Sep 17 00:00:00 2001
From: Maximilian Hils <git@maximilianhils.com>
Date: Thu, 19 May 2016 22:09:00 -0700
Subject: [PATCH] add FlowMaster.active_flows

---
 mitmproxy/dump.py          |  2 +-
 mitmproxy/flow.py          | 22 +++++++++++-----------
 test/mitmproxy/tservers.py |  5 +----
 3 files changed, 13 insertions(+), 16 deletions(-)

diff --git a/mitmproxy/dump.py b/mitmproxy/dump.py
index aae397cd4..f1eabdb8c 100644
--- a/mitmproxy/dump.py
+++ b/mitmproxy/dump.py
@@ -320,7 +320,6 @@ class DumpMaster(flow.FlowMaster):
             self.outfile.flush()
 
     def _process_flow(self, f):
-        self.state.delete_flow(f)
         if self.filt and not f.match(self.filt):
             return
 
@@ -328,6 +327,7 @@ class DumpMaster(flow.FlowMaster):
 
     def handle_request(self, f):
         flow.FlowMaster.handle_request(self, f)
+        self.state.delete_flow(f)
         if f:
             f.reply()
         return f
diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py
index 77ed87470..51ef71933 100644
--- a/mitmproxy/flow.py
+++ b/mitmproxy/flow.py
@@ -653,8 +653,9 @@ class FlowMaster(controller.ServerMaster):
         if server:
             self.add_server(server)
         self.state = state
-        self.server_playback = None
-        self.client_playback = None
+        self.active_flows = set()  # type: Set[Flow]
+        self.server_playback = None  # type: Optional[ServerPlaybackState]
+        self.client_playback = None  # type: Optional[ClientPlaybackState]
         self.kill_nonreplay = False
         self.scripts = []  # type: List[script.Script]
         self.pause_scripts = False
@@ -1033,6 +1034,7 @@ class FlowMaster(controller.ServerMaster):
                 return
         if f not in self.state.flows:  # don't add again on replay
             self.state.add_flow(f)
+        self.active_flows.add(f)
         self.replacehooks.run(f)
         self.setheaders.run(f)
         self.process_new_request(f)
@@ -1053,6 +1055,7 @@ class FlowMaster(controller.ServerMaster):
         return f
 
     def handle_response(self, f):
+        self.active_flows.discard(f)
         self.state.update_flow(f)
         self.replacehooks.run(f)
         self.setheaders.run(f)
@@ -1099,7 +1102,9 @@ class FlowMaster(controller.ServerMaster):
         return ok
 
     def handle_tcp_open(self, flow):
-        self.state.add_flow(flow)
+        # TODO: This would break mitmproxy currently.
+        # self.state.add_flow(flow)
+        self.active_flows.add(flow)
         self.run_script_hook("tcp_open", flow)
         flow.reply()
 
@@ -1124,7 +1129,7 @@ class FlowMaster(controller.ServerMaster):
         flow.reply()
 
     def handle_tcp_close(self, flow):
-        self.state.delete_flow(flow)
+        self.active_flows.discard(flow)
         if self.stream:
             self.stream.add(flow)
         self.run_script_hook("tcp_close", flow)
@@ -1135,13 +1140,8 @@ class FlowMaster(controller.ServerMaster):
 
         # Add all flows that are still active
         if self.stream:
-            for flow in self.state.flows:
-                # FIXME: We actually need to keep track of which flows are still active.
-                if isinstance(flow, HTTPFlow) and not flow.response:
-                    self.stream.add(flow)
-                if isinstance(flow, TCPFlow):
-                    # (assuming mitmdump only, this must be still active)
-                    self.stream.add(flow)
+            for flow in self.active_flows:
+                self.stream.add(flow)
             self.stop_stream()
 
         self.unload_scripts()
diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py
index 57225ee1b..c9d68cfd0 100644
--- a/test/mitmproxy/tservers.py
+++ b/test/mitmproxy/tservers.py
@@ -161,10 +161,7 @@ class HTTPProxyTest(ProxyTestBase):
             q = "get:'/p/%s'" % spec
         else:
             q = "get:'%s/p/%s'" % (self.server.urlbase, spec)
-        resp = p.request(q)
-        p.finish()
-        p.close()
-        return resp
+        return p.request(q)
 
     def app(self, page):
         if self.ssl: