diff options
| -rw-r--r-- | mitmproxy/dump.py | 2 | ||||
| -rw-r--r-- | mitmproxy/flow.py | 22 | ||||
| -rw-r--r-- | test/mitmproxy/tservers.py | 5 |
3 files changed, 13 insertions, 16 deletions
diff --git a/mitmproxy/dump.py b/mitmproxy/dump.py index aae397cd..f1eabdb8 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 77ed8747..51ef7193 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 57225ee1..c9d68cfd 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: |