diff --git a/docs/scripts/api-events.py b/docs/scripts/api-events.py index 7802f603f..2429c751f 100644 --- a/docs/scripts/api-events.py +++ b/docs/scripts/api-events.py @@ -134,7 +134,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..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,8 +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: - flow.error = terr() + + flow.websocket.close_reason = close_reason + + if close_code is not None: + flow.websocket.close_code = close_code + else: + 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 c4fa19ae4..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_error(f) + 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(): 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..84c1de058 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,14 +245,16 @@ 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 + # 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): +def test_close_code(ws_testdata): tctx, playbook, flow = ws_testdata assert ( playbook @@ -263,10 +265,10 @@ 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 + assert flow.websocket.close_code == 4000 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)