From eef6522a5e4bc7702eb3a8b021a4d3364f3ea450 Mon Sep 17 00:00:00 2001 From: Alexander Prinzhorn Date: Mon, 12 Jul 2021 16:35:25 +0200 Subject: [PATCH 1/4] remove websocket_error event, fixes #4674 --- docs/scripts/api-events.py | 1 - examples/contrib/jsondump.py | 6 ---- mitmproxy/addons/dumper.py | 30 ++++++++++++------ mitmproxy/addons/save.py | 5 +-- mitmproxy/eventsequence.py | 5 +-- mitmproxy/proxy/layers/websocket.py | 31 +++---------------- mitmproxy/test/tflow.py | 6 +++- test/mitmproxy/addons/test_dumper.py | 2 +- test/mitmproxy/addons/test_save.py | 2 +- test/mitmproxy/proxy/layers/test_websocket.py | 8 ++--- test/mitmproxy/test_eventsequence.py | 10 ++---- test/mitmproxy/test_flowfilter.py | 7 ----- 12 files changed, 40 insertions(+), 73 deletions(-) diff --git a/docs/scripts/api-events.py b/docs/scripts/api-events.py index f365d8392..05f5126fe 100644 --- a/docs/scripts/api-events.py +++ b/docs/scripts/api-events.py @@ -133,7 +133,6 @@ with outfile.open("w") as f, contextlib.redirect_stdout(f): websocket.WebsocketStartHook, websocket.WebsocketMessageHook, websocket.WebsocketEndHook, - websocket.WebsocketErrorHook, ] ) diff --git a/examples/contrib/jsondump.py b/examples/contrib/jsondump.py index 615a42d82..ab808a2d7 100644 --- a/examples/contrib/jsondump.py +++ b/examples/contrib/jsondump.py @@ -247,11 +247,5 @@ class JSONDumper: """ self.queue.put(flow.get_state()) - def websocket_error(self, flow): - """ - Dump websocket errors. - """ - self.queue.put(flow.get_state()) - addons = [JSONDumper()] # pylint: disable=invalid-name diff --git a/mitmproxy/addons/dumper.py b/mitmproxy/addons/dumper.py index 7ebb128b5..335d80344 100644 --- a/mitmproxy/addons/dumper.py +++ b/mitmproxy/addons/dumper.py @@ -10,10 +10,12 @@ from mitmproxy import ctx from mitmproxy import exceptions from mitmproxy import flowfilter from mitmproxy import http +from mitmproxy import flow from mitmproxy.tcp import TCPFlow, TCPMessage from mitmproxy.utils import human from mitmproxy.utils import strutils -from mitmproxy.websocket import WebSocketMessage +from mitmproxy.websocket import WebSocketMessage, WebSocketData +from wsproto.frame_protocol import CloseReason def indent(n: int, text: str) -> str: @@ -277,12 +279,6 @@ class Dumper: if self.match(f): self.echo_flow(f) - def websocket_error(self, f: http.HTTPFlow): - self.echo_error( - f"Error in WebSocket connection to {human.format_address(f.server_conn.address)}: {f.error}", - fg="red" - ) - def websocket_message(self, f: http.HTTPFlow): assert f.websocket is not None # satisfy type checker if self.match(f): @@ -300,8 +296,24 @@ class Dumper: def websocket_end(self, f: http.HTTPFlow): assert f.websocket is not None # satisfy type checker if self.match(f): - c = 'client' if f.websocket.closed_by_client else 'server' - self.echo(f"WebSocket connection closed by {c}: {f.websocket.close_code} {f.websocket.close_reason}") + if f.websocket.close_code in {1000, 1001, 1005}: + c = 'client' if f.websocket.closed_by_client else 'server' + self.echo(f"WebSocket connection closed by {c}: {f.websocket.close_code} {f.websocket.close_reason}") + else: + error = flow.Error(f"WebSocket Error: {self.format_websocket_error(f.websocket)}") + self.echo_error( + f"Error in WebSocket connection to {human.format_address(f.server_conn.address)}: {error}", + fg="red" + ) + + def format_websocket_error(self, websocket: WebSocketData) -> str: + try: + ret = CloseReason(websocket.close_code).name + except ValueError: + ret = f"UNKNOWN_ERROR={websocket.close_code}" + if websocket.close_reason: + ret += f" (reason: {websocket.close_reason})" + return ret def tcp_error(self, f): if self.match(f): diff --git a/mitmproxy/addons/save.py b/mitmproxy/addons/save.py index 49468f1ea..9297e7399 100644 --- a/mitmproxy/addons/save.py +++ b/mitmproxy/addons/save.py @@ -94,15 +94,12 @@ class Save: self.stream.add(flow) self.active_flows.discard(flow) - def websocket_error(self, flow: http.HTTPFlow): - self.websocket_end(flow) - def request(self, flow: http.HTTPFlow): if self.stream: self.active_flows.add(flow) def response(self, flow: http.HTTPFlow): - # websocket flows will receive either websocket_end or websocket_error, + # websocket flows will receive a websocket_end, # we don't want to persist them here already if self.stream and flow.websocket is None: self.stream.add(flow) diff --git a/mitmproxy/eventsequence.py b/mitmproxy/eventsequence.py index 0efb916b9..8cd7ce69a 100644 --- a/mitmproxy/eventsequence.py +++ b/mitmproxy/eventsequence.py @@ -24,10 +24,7 @@ def _iterate_http(f: http.HTTPFlow) -> TEventGenerator: for m in message_queue: f.websocket.messages.append(m) yield layers.websocket.WebsocketMessageHook(f) - if f.error: - yield layers.websocket.WebsocketErrorHook(f) - else: - yield layers.websocket.WebsocketEndHook(f) + yield layers.websocket.WebsocketEndHook(f) elif f.error: yield layers.http.HttpErrorHook(f) diff --git a/mitmproxy/proxy/layers/websocket.py b/mitmproxy/proxy/layers/websocket.py index cdf10ab1d..5e78c2267 100644 --- a/mitmproxy/proxy/layers/websocket.py +++ b/mitmproxy/proxy/layers/websocket.py @@ -6,14 +6,14 @@ import wsproto import wsproto.extensions import wsproto.frame_protocol import wsproto.utilities -from mitmproxy import connection, flow, http, websocket +from mitmproxy import connection, http, websocket from mitmproxy.proxy import commands, events, layer from mitmproxy.proxy.commands import StartHook from mitmproxy.proxy.context import Context from mitmproxy.proxy.events import MessageInjected from mitmproxy.proxy.utils import expect from wsproto import ConnectionState -from wsproto.frame_protocol import CloseReason, Opcode +from wsproto.frame_protocol import Opcode @dataclass @@ -39,21 +39,12 @@ class WebsocketMessageHook(StartHook): class WebsocketEndHook(StartHook): """ A WebSocket connection has ended. + You can check `flow.websocket.close_code` to determine why it ended. """ flow: http.HTTPFlow -@dataclass -class WebsocketErrorHook(StartHook): - """ - A WebSocket connection has had an error. - - Every WebSocket flow will receive either a websocket_error or a websocket_end event, but not both. - """ - flow: http.HTTPFlow - - class WebSocketMessageInjected(MessageInjected[websocket.WebSocketMessage]): """ The user has injected a custom WebSocket message. @@ -200,11 +191,7 @@ class WebsocketLayer(layer.Layer): # response == original event, so no need to differentiate here. yield ws.send2(ws_event) yield commands.CloseConnection(ws.conn) - if ws_event.code in {1000, 1001, 1005}: - yield WebsocketEndHook(self.flow) - else: - self.flow.error = flow.Error(f"WebSocket Error: {format_close_event(ws_event)}") - yield WebsocketErrorHook(self.flow) + yield WebsocketEndHook(self.flow) self._handle_event = self.done else: # pragma: no cover raise AssertionError(f"Unexpected WebSocket event: {ws_event}") @@ -214,16 +201,6 @@ class WebsocketLayer(layer.Layer): yield from () -def format_close_event(event: wsproto.events.CloseConnection) -> str: - try: - ret = CloseReason(event.code).name - except ValueError: - ret = f"UNKNOWN_ERROR={event.code}" - if event.reason: - ret += f" (reason: {event.reason})" - return ret - - class Fragmentizer: """ Theory (RFC 6455): diff --git a/mitmproxy/test/tflow.py b/mitmproxy/test/tflow.py index 7189adcef..10d6e03cd 100644 --- a/mitmproxy/test/tflow.py +++ b/mitmproxy/test/tflow.py @@ -75,7 +75,11 @@ def twebsocketflow(messages=True, err=None) -> http.HTTPFlow: websocket.WebSocketMessage(Opcode.TEXT, False, b"it's me", 946681205), ] if err is True: - flow.error = terr() + # ABNORMAL_CLOSURE + flow.websocket.close_code = 1006 + else: + # NORMAL_CLOSURE + flow.websocket.close_code = 1000 flow.reply = controller.DummyReply() return flow diff --git a/test/mitmproxy/addons/test_dumper.py b/test/mitmproxy/addons/test_dumper.py index c4fa19ae4..bffe721f2 100644 --- a/test/mitmproxy/addons/test_dumper.py +++ b/test/mitmproxy/addons/test_dumper.py @@ -234,7 +234,7 @@ def test_websocket(): assert "WebSocket connection closed by" in sio.getvalue() f = tflow.twebsocketflow(err=True) - d.websocket_error(f) + d.websocket_end(f) assert "Error in WebSocket" in sio_err.getvalue() diff --git a/test/mitmproxy/addons/test_save.py b/test/mitmproxy/addons/test_save.py index 754c755d1..5067562cd 100644 --- a/test/mitmproxy/addons/test_save.py +++ b/test/mitmproxy/addons/test_save.py @@ -60,7 +60,7 @@ def test_websocket(tmpdir): f = tflow.twebsocketflow() sa.request(f) - sa.websocket_error(f) + sa.websocket_end(f) tctx.configure(sa, save_stream_file=None) assert len(rd(p)) == 2 diff --git a/test/mitmproxy/proxy/layers/test_websocket.py b/test/mitmproxy/proxy/layers/test_websocket.py index 7e3ef5cd8..b58c66e50 100644 --- a/test/mitmproxy/proxy/layers/test_websocket.py +++ b/test/mitmproxy/proxy/layers/test_websocket.py @@ -188,7 +188,7 @@ def test_protocol_error(ws_testdata): << CloseConnection(tctx.server) << SendData(tctx.client, b"\x88/\x03\xeaexpected CONTINUATION, got ") << CloseConnection(tctx.client) - << websocket.WebsocketErrorHook(flow) + << websocket.WebsocketEndHook(flow) >> reply() ) @@ -245,11 +245,10 @@ def test_close_disconnect(ws_testdata): << CloseConnection(tctx.server) << SendData(tctx.client, b"\x88\x02\x03\xe8") << CloseConnection(tctx.client) - << websocket.WebsocketErrorHook(flow) + << websocket.WebsocketEndHook(flow) >> reply() >> ConnectionClosed(tctx.client) ) - assert "ABNORMAL_CLOSURE" in flow.error.msg def test_close_error(ws_testdata): @@ -263,10 +262,9 @@ def test_close_error(ws_testdata): << CloseConnection(tctx.server) << SendData(tctx.client, b"\x88\x02\x0f\xa0") << CloseConnection(tctx.client) - << websocket.WebsocketErrorHook(flow) + << websocket.WebsocketEndHook(flow) >> reply() ) - assert "UNKNOWN_ERROR=4000" in flow.error.msg def test_deflate(ws_testdata): diff --git a/test/mitmproxy/test_eventsequence.py b/test/mitmproxy/test_eventsequence.py index cbfbb7087..03ad5da57 100644 --- a/test/mitmproxy/test_eventsequence.py +++ b/test/mitmproxy/test_eventsequence.py @@ -23,9 +23,8 @@ def test_http_flow(resp, err): assert isinstance(next(i), layers.http.HttpErrorHook) -@pytest.mark.parametrize("err", [False, True]) -def test_websocket_flow(err): - f = tflow.twebsocketflow(err=err) +def test_websocket_flow(): + f = tflow.twebsocketflow() i = eventsequence.iterate(f) assert isinstance(next(i), layers.http.HttpRequestHeadersHook) @@ -41,10 +40,7 @@ def test_websocket_flow(err): assert len(f.websocket.messages) == 2 assert isinstance(next(i), layers.websocket.WebsocketMessageHook) assert len(f.websocket.messages) == 3 - if err: - assert isinstance(next(i), layers.websocket.WebsocketErrorHook) - else: - assert isinstance(next(i), layers.websocket.WebsocketEndHook) + assert isinstance(next(i), layers.websocket.WebsocketEndHook) @pytest.mark.parametrize("err", [False, True]) diff --git a/test/mitmproxy/test_flowfilter.py b/test/mitmproxy/test_flowfilter.py index 5ab2e16c4..9105b4243 100644 --- a/test/mitmproxy/test_flowfilter.py +++ b/test/mitmproxy/test_flowfilter.py @@ -461,9 +461,6 @@ class TestMatchingWebSocketFlow: def flow(self) -> http.HTTPFlow: return tflow.twebsocketflow() - def err(self) -> http.HTTPFlow: - return tflow.twebsocketflow(err=True) - def q(self, q, o): return flowfilter.parse(q)(o) @@ -484,10 +481,6 @@ class TestMatchingWebSocketFlow: f = tflow.tflow(resp=True) assert not self.q("~websocket", f) - def test_ferr(self): - e = self.err() - assert self.q("~e", e) - def test_domain(self): q = self.flow() assert self.q("~d example.com", q) From 793f5c5681ade96d6650130361229dd05b6a05b0 Mon Sep 17 00:00:00 2001 From: Alexander Prinzhorn Date: Tue, 13 Jul 2021 12:51:24 +0200 Subject: [PATCH 2/4] fix dumper test cov --- mitmproxy/test/tflow.py | 18 ++++++++++++------ test/mitmproxy/addons/test_dumper.py | 20 ++++++++++++++++++++ 2 files changed, 32 insertions(+), 6 deletions(-) diff --git a/mitmproxy/test/tflow.py b/mitmproxy/test/tflow.py index 10d6e03cd..615e92434 100644 --- a/mitmproxy/test/tflow.py +++ b/mitmproxy/test/tflow.py @@ -30,7 +30,7 @@ def ttcpflow(client_conn=True, server_conn=True, messages=True, err=None): return f -def twebsocketflow(messages=True, err=None) -> http.HTTPFlow: +def twebsocketflow(messages=True, err=None, close_code=None, close_reason='') -> http.HTTPFlow: flow = http.HTTPFlow(tclient_conn(), tserver_conn()) flow.request = http.Request( "example.com", @@ -74,12 +74,18 @@ def twebsocketflow(messages=True, err=None) -> http.HTTPFlow: websocket.WebSocketMessage(Opcode.TEXT, True, b"hello text", 946681204), websocket.WebSocketMessage(Opcode.TEXT, False, b"it's me", 946681205), ] - if err is True: - # ABNORMAL_CLOSURE - flow.websocket.close_code = 1006 + + flow.websocket.close_reason = close_reason + + if close_code is not None: + flow.websocket.close_code = close_code else: - # NORMAL_CLOSURE - flow.websocket.close_code = 1000 + if err is True: + # ABNORMAL_CLOSURE + flow.websocket.close_code = 1006 + else: + # NORMAL_CLOSURE + flow.websocket.close_code = 1000 flow.reply = controller.DummyReply() return flow diff --git a/test/mitmproxy/addons/test_dumper.py b/test/mitmproxy/addons/test_dumper.py index bffe721f2..2f6a86746 100644 --- a/test/mitmproxy/addons/test_dumper.py +++ b/test/mitmproxy/addons/test_dumper.py @@ -232,10 +232,30 @@ def test_websocket(): d.websocket_end(f) assert "WebSocket connection closed by" in sio.getvalue() + sio_err.truncate(0) f = tflow.twebsocketflow(err=True) d.websocket_end(f) assert "Error in WebSocket" in sio_err.getvalue() + assert "(reason:" not in sio_err.getvalue() + sio_err.truncate(0) + + f = tflow.twebsocketflow(err=True, close_reason='Some lame excuse') + d.websocket_end(f) + assert "Error in WebSocket" in sio_err.getvalue() + assert "(reason: Some lame excuse)" in sio_err.getvalue() + sio_err.truncate(0) + + f = tflow.twebsocketflow(close_code=4000) + d.websocket_end(f) + assert "UNKNOWN_ERROR=4000" in sio_err.getvalue() + assert "(reason:" not in sio_err.getvalue() + sio_err.truncate(0) + + f = tflow.twebsocketflow(close_code=4000, close_reason='I swear I had a reason') + d.websocket_end(f) + assert "UNKNOWN_ERROR=4000" in sio_err.getvalue() + assert "(reason: I swear I had a reason)" in sio_err.getvalue() def test_http2(): From 13f115451b85e1dab6d2b35d4f5c8990cde96b38 Mon Sep 17 00:00:00 2001 From: Alexander Prinzhorn Date: Tue, 13 Jul 2021 16:22:10 +0200 Subject: [PATCH 3/4] update tests --- test/mitmproxy/proxy/layers/test_websocket.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/test/mitmproxy/proxy/layers/test_websocket.py b/test/mitmproxy/proxy/layers/test_websocket.py index b58c66e50..7322aa1b6 100644 --- a/test/mitmproxy/proxy/layers/test_websocket.py +++ b/test/mitmproxy/proxy/layers/test_websocket.py @@ -249,6 +249,9 @@ def test_close_disconnect(ws_testdata): >> reply() >> ConnectionClosed(tctx.client) ) + # The \x03\xe8 above is code 1000 (normal closure). + # But 1006 (ABNORMAL_CLOSURE) is expected, because the connection was already closed. + assert flow.websocket.close_code == 1006 def test_close_error(ws_testdata): @@ -265,6 +268,7 @@ def test_close_error(ws_testdata): << websocket.WebsocketEndHook(flow) >> reply() ) + assert flow.websocket.close_code == 4000 def test_deflate(ws_testdata): From 1f7745dbb8cf3e25c1fb7e617e72f1050ac8e3f0 Mon Sep 17 00:00:00 2001 From: Alexander Prinzhorn Date: Wed, 14 Jul 2021 09:04:57 +0200 Subject: [PATCH 4/4] rename test --- test/mitmproxy/proxy/layers/test_websocket.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/mitmproxy/proxy/layers/test_websocket.py b/test/mitmproxy/proxy/layers/test_websocket.py index 7322aa1b6..84c1de058 100644 --- a/test/mitmproxy/proxy/layers/test_websocket.py +++ b/test/mitmproxy/proxy/layers/test_websocket.py @@ -254,7 +254,7 @@ def test_close_disconnect(ws_testdata): assert flow.websocket.close_code == 1006 -def test_close_error(ws_testdata): +def test_close_code(ws_testdata): tctx, playbook, flow = ws_testdata assert ( playbook