From 59c277effd016fdf771a65aaaa87526ff5fe869e Mon Sep 17 00:00:00 2001 From: Thomas Kriechbaumer Date: Fri, 29 Dec 2017 21:55:41 +0100 Subject: websocket: add flow kill test --- test/mitmproxy/proxy/protocol/test_websocket.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) (limited to 'test') diff --git a/test/mitmproxy/proxy/protocol/test_websocket.py b/test/mitmproxy/proxy/protocol/test_websocket.py index a7acdc4d..490d31a0 100644 --- a/test/mitmproxy/proxy/protocol/test_websocket.py +++ b/test/mitmproxy/proxy/protocol/test_websocket.py @@ -221,6 +221,25 @@ class TestSimple(_WebSocketTest): assert frame.payload == b'foo' +class TestKillFlow(_WebSocketTest): + + @classmethod + def handle_websockets(cls, rfile, wfile): + wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.TEXT, payload=b'server-foobar'))) + wfile.flush() + + def test_kill(self): + class KillFlow: + def websocket_message(self, f): + f.kill() + + self.master.addons.add(KillFlow()) + self.setup_connection() + + frame = websockets.Frame.from_file(self.client.rfile) + assert frame.payload == b'foo' + + class TestSimpleTLS(_WebSocketTest): ssl = True -- cgit v1.2.3 From 62326227740d3808c0886ed381cb976a6b2f8b03 Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Fri, 29 Dec 2017 22:44:18 +0100 Subject: fix Flow.kill behaviour This now just sets a kill reply instead of committing directly. First, this seems like the more sane thing to do. Second, we have an iffy race condition where we call Reply.commit() before the addonmanager finishes its invocation, the proxy thread then progresses and sets a new flow.reply attribute, and the addonmanager then gets confused when finishing. This commit doesn't fix that, but mitigates it for Flow.kill which is now committed by the addonmanager. --- test/mitmproxy/proxy/protocol/test_websocket.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) (limited to 'test') diff --git a/test/mitmproxy/proxy/protocol/test_websocket.py b/test/mitmproxy/proxy/protocol/test_websocket.py index 490d31a0..d9389faf 100644 --- a/test/mitmproxy/proxy/protocol/test_websocket.py +++ b/test/mitmproxy/proxy/protocol/test_websocket.py @@ -236,8 +236,8 @@ class TestKillFlow(_WebSocketTest): self.master.addons.add(KillFlow()) self.setup_connection() - frame = websockets.Frame.from_file(self.client.rfile) - assert frame.payload == b'foo' + with pytest.raises(exceptions.TcpDisconnect): + websockets.Frame.from_file(self.client.rfile) class TestSimpleTLS(_WebSocketTest): -- cgit v1.2.3