diff --git a/mitmproxy/addons/save.py b/mitmproxy/addons/save.py index b7f329c4f..60118aff6 100644 --- a/mitmproxy/addons/save.py +++ b/mitmproxy/addons/save.py @@ -85,6 +85,9 @@ class Save: self.stream.add(flow) self.active_flows.discard(flow) + def tcp_error(self, flow): + self.tcp_end(flow) + def websocket_start(self, flow): if self.stream: self.active_flows.add(flow) @@ -94,14 +97,20 @@ class Save: self.stream.add(flow) self.active_flows.discard(flow) + def websocket_error(self, flow): + self.websocket_end(flow) + + def request(self, flow): + if self.stream: + self.active_flows.add(flow) + def response(self, flow): if self.stream: self.stream.add(flow) self.active_flows.discard(flow) - def request(self, flow): - if self.stream: - self.active_flows.add(flow) + def error(self, flow): + self.response(flow) def done(self): if self.stream: diff --git a/test/mitmproxy/addons/test_save.py b/test/mitmproxy/addons/test_save.py index 6727a96f7..65ec0d76d 100644 --- a/test/mitmproxy/addons/test_save.py +++ b/test/mitmproxy/addons/test_save.py @@ -39,8 +39,13 @@ def test_tcp(tmpdir): tt = tflow.ttcpflow() sa.tcp_start(tt) sa.tcp_end(tt) + + tt = tflow.ttcpflow() + sa.tcp_start(tt) + sa.tcp_error(tt) + tctx.configure(sa, save_stream_file=None) - assert rd(p) + assert len(rd(p)) == 2 def test_websocket(tmpdir): @@ -52,8 +57,13 @@ def test_websocket(tmpdir): f = tflow.twebsocketflow() sa.websocket_start(f) sa.websocket_end(f) + + f = tflow.twebsocketflow() + sa.websocket_start(f) + sa.websocket_error(f) + tctx.configure(sa, save_stream_file=None) - assert rd(p) + assert len(rd(p)) == 2 def test_save_command(tmpdir): @@ -89,8 +99,15 @@ def test_simple(tmpdir): tctx.configure(sa, save_stream_file=None) assert rd(p)[0].response + tctx.configure(sa, save_stream_file="+" + p) + f = tflow.tflow(err=True) + sa.request(f) + sa.error(f) + tctx.configure(sa, save_stream_file=None) + assert rd(p)[1].error + tctx.configure(sa, save_stream_file="+" + p) f = tflow.tflow() sa.request(f) tctx.configure(sa, save_stream_file=None) - assert not rd(p)[1].response + assert not rd(p)[2].response