From d8aeef1bfdc30cdcd6f045521aa03ddae56d5e99 Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Wed, 10 Mar 2021 20:23:25 +0100 Subject: [PATCH] ressurect `killed`, is_text -> type --- mitmproxy/addons/dumper.py | 7 +- mitmproxy/addons/save.py | 2 + mitmproxy/io/compat.py | 7 +- mitmproxy/proxy/layers/websocket.py | 15 +++-- mitmproxy/test/tflow.py | 7 +- mitmproxy/websocket.py | 66 ++++++++----------- test/mitmproxy/proxy/layers/test_websocket.py | 7 +- test/mitmproxy/test_websocket.py | 9 ++- 8 files changed, 58 insertions(+), 62 deletions(-) diff --git a/mitmproxy/addons/dumper.py b/mitmproxy/addons/dumper.py index 15e5adf33..80ae72730 100644 --- a/mitmproxy/addons/dumper.py +++ b/mitmproxy/addons/dumper.py @@ -284,22 +284,21 @@ class Dumper: ) def websocket_message(self, f: http.HTTPFlow): - assert f.websocket is not None + assert f.websocket is not None # satisfy type checker if self.match(f): message = f.websocket.messages[-1] direction = "->" if message.from_client else "<-" - typ = "text" if message.is_text else "binary" self.echo( f"{human.format_address(f.client_conn.peername)} " - f"{direction} WebSocket {typ} message " + f"{direction} WebSocket {message.type.name.lower()} message " f"{direction} {human.format_address(f.server_conn.address)}{f.request.path}" ) if ctx.options.flow_detail >= 3: self._echo_message(message, f) def websocket_end(self, f: http.HTTPFlow): - assert f.websocket is not None + assert f.websocket is not None # satisfy type checker if self.match(f): c = 'client' if f.websocket.close_by_client else 'server' self.echo(f"WebSocket connection closed by {c}: {f.websocket.close_code} {f.websocket.close_reason}") diff --git a/mitmproxy/addons/save.py b/mitmproxy/addons/save.py index c0d389643..49468f1ea 100644 --- a/mitmproxy/addons/save.py +++ b/mitmproxy/addons/save.py @@ -102,6 +102,8 @@ class Save: self.active_flows.add(flow) def response(self, flow: http.HTTPFlow): + # websocket flows will receive either websocket_end or websocket_error, + # we don't want to persist them here already if self.stream and flow.websocket is None: self.stream.add(flow) self.active_flows.discard(flow) diff --git a/mitmproxy/io/compat.py b/mitmproxy/io/compat.py index 15f04ef31..963048614 100644 --- a/mitmproxy/io/compat.py +++ b/mitmproxy/io/compat.py @@ -291,12 +291,7 @@ def convert_11_12(data): "and may appear duplicated." ) data["websocket"] = { - "messages": [ - # old: int(self.type), self.from_client, self.content, self.timestamp, self.killed - # new: self.from_client, self.is_text, self.content, self.timestamp - [from_client, typ == 0x1, strutils.always_bytes(content) if not killed else b"", timestamp] - for typ, from_client, content, timestamp, killed in ws_flow["messages"] - ], + "messages": ws_flow["messages"], "close_by_client": ws_flow["close_sender"] == "client", "close_code": ws_flow["close_code"], "close_reason": ws_flow["close_reason"], diff --git a/mitmproxy/proxy/layers/websocket.py b/mitmproxy/proxy/layers/websocket.py index 78ff0adc9..f39a93ddc 100644 --- a/mitmproxy/proxy/layers/websocket.py +++ b/mitmproxy/proxy/layers/websocket.py @@ -11,7 +11,7 @@ from mitmproxy.proxy.commands import StartHook from mitmproxy.proxy.context import Context from mitmproxy.proxy.utils import expect from wsproto import ConnectionState -from wsproto.frame_protocol import CloseReason +from wsproto.frame_protocol import CloseReason, Opcode @dataclass @@ -96,7 +96,7 @@ class WebsocketLayer(layer.Layer): server_extensions = [] # Parse extension headers. We only support deflate at the moment and ignore everything else. - assert self.flow.response + assert self.flow.response # satisfy type checker ext_header = self.flow.response.headers.get("Sec-WebSocket-Extensions", "") if ext_header: for ext in wsproto.utilities.split_comma_header(ext_header.encode("ascii", "replace")): @@ -122,7 +122,7 @@ class WebsocketLayer(layer.Layer): @expect(events.DataReceived, events.ConnectionClosed) def relay_messages(self, event: events.ConnectionEvent) -> layer.CommandGenerator[None]: - assert self.flow.websocket + assert self.flow.websocket # satisfy type checker from_client = event.connection == self.context.client from_str = 'client' if from_client else 'server' @@ -144,8 +144,10 @@ class WebsocketLayer(layer.Layer): if isinstance(ws_event, wsproto.events.Message): is_text = isinstance(ws_event.data, str) if is_text: + typ = Opcode.TEXT src_ws.frame_buf.append(ws_event.data.encode()) else: + typ = Opcode.BINARY src_ws.frame_buf.append(ws_event.data) if ws_event.message_finished: @@ -154,12 +156,13 @@ class WebsocketLayer(layer.Layer): fragmentizer = Fragmentizer(src_ws.frame_buf, is_text) src_ws.frame_buf.clear() - message = websocket.WebSocketMessage(from_client, is_text, content) + message = websocket.WebSocketMessage(typ, from_client, content) self.flow.websocket.messages.append(message) yield WebsocketMessageHook(self.flow) - for msg in fragmentizer(message.content): - yield dst_ws.send2(msg) + if not message.killed: + for msg in fragmentizer(message.content): + yield dst_ws.send2(msg) elif isinstance(ws_event, (wsproto.events.Ping, wsproto.events.Pong)): yield commands.Log( diff --git a/mitmproxy/test/tflow.py b/mitmproxy/test/tflow.py index 7f8d7850e..9f76d34eb 100644 --- a/mitmproxy/test/tflow.py +++ b/mitmproxy/test/tflow.py @@ -7,6 +7,7 @@ from mitmproxy import http from mitmproxy import tcp from mitmproxy import websocket from mitmproxy.test import tutils +from wsproto.frame_protocol import Opcode def ttcpflow(client_conn=True, server_conn=True, messages=True, err=None): @@ -69,9 +70,9 @@ def twebsocketflow(messages=True, err=None) -> http.HTTPFlow: if messages is True: flow.websocket.messages = [ - websocket.WebSocketMessage(True, False, b"hello binary", 946681203), - websocket.WebSocketMessage(True, True, b"hello text", 946681204), - websocket.WebSocketMessage(False, True, b"it's me", 946681205), + websocket.WebSocketMessage(Opcode.BINARY, True, b"hello binary", 946681203), + 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() diff --git a/mitmproxy/websocket.py b/mitmproxy/websocket.py index 5d84fb18e..f1e37542b 100644 --- a/mitmproxy/websocket.py +++ b/mitmproxy/websocket.py @@ -6,12 +6,14 @@ as HTTP flows as well. They can be distinguished from regular HTTP requests by h This module only defines the classes for individual `WebSocketMessage`s and the `WebSocketData` container. """ import time -import warnings -from typing import List +from typing import List, Tuple, Union from typing import Optional from mitmproxy import stateobject from mitmproxy.coretypes import serializable +from wsproto.frame_protocol import Opcode + +WebSocketMessageState = Tuple[int, bool, bytes, float, bool] class WebSocketMessage(serializable.Serializable): @@ -25,75 +27,63 @@ class WebSocketMessage(serializable.Serializable): text and binary messages. To avoid a whole class of nasty type confusion bugs, mitmproxy stores all message contents as binary. If you need text, you can decode the `content` property: - >>> if message.is_text: + >>> from wsproto.frame_protocol import Opcode + >>> if message.type == Opcode.TEXT: >>> text = message.content.decode() + + Per the WebSocket spec, text messages always use UTF-8 encoding. """ from_client: bool """True if this messages was sent by the client.""" - is_text: bool + type: Opcode """ - True if the message is a text message, False if the message is a binary message. + The message type, as per RFC 6455's [opcode](https://tools.ietf.org/html/rfc6455#section-5.2). - In either case, mitmproxy will store the message contents as *bytes*. + Note that mitmproxy will always store the message contents as *bytes*. + A dedicated `.text` property for text messages is planned, see https://github.com/mitmproxy/mitmproxy/pull/4486. """ content: bytes """A byte-string representing the content of this message.""" timestamp: float """Timestamp of when this message was received or created.""" + killed: bool + """True if the message has not been forwarded by mitmproxy, False otherwise.""" def __init__( self, + type: Union[int, Opcode], from_client: bool, - is_text: bool, content: bytes, timestamp: Optional[float] = None, + killed: bool = False, ) -> None: self.from_client = from_client - self.is_text = is_text + self.type = Opcode(type) self.content = content self.timestamp: float = timestamp or time.time() + self.killed = killed @classmethod - def from_state(cls, state): + def from_state(cls, state: WebSocketMessageState): return cls(*state) - def get_state(self): - return self.from_client, self.is_text, self.content, self.timestamp + def get_state(self) -> WebSocketMessageState: + return int(self.type), self.from_client, self.content, self.timestamp, self.killed - def set_state(self, state): - self.from_client, self.is_text, self.content, self.timestamp = state + def set_state(self, state: WebSocketMessageState) -> None: + typ, self.from_client, self.content, self.timestamp, self.killed = state + self.type = Opcode(typ) def __repr__(self): - if self.is_text: + if self.type == Opcode.TEXT: return repr(self.content.decode(errors="replace")) else: return repr(self.content) - def kill(self): # pragma: no cover - """ - Kill this message. - - It will not be sent to the other endpoint. - """ - warnings.warn( - "WebSocketMessage.kill is deprecated, set an empty content instead.", - DeprecationWarning, - stacklevel=2, - ) - self.content = b"" - - @property - def killed(self) -> bool: # pragma: no cover - """ - True if this messages was killed and should not be sent to the other endpoint. - """ - warnings.warn( - "WebSocketMessage.killed is deprecated, check for an empty content instead.", - DeprecationWarning, - stacklevel=2, - ) - return bool(self.content) + def kill(self): + # Likely to be replaced with .drop() in the future, see https://github.com/mitmproxy/mitmproxy/pull/4486 + self.killed = True class WebSocketData(stateobject.StateObject): diff --git a/test/mitmproxy/proxy/layers/test_websocket.py b/test/mitmproxy/proxy/layers/test_websocket.py index 9c6ebc9ed..5d372a2a3 100644 --- a/test/mitmproxy/proxy/layers/test_websocket.py +++ b/test/mitmproxy/proxy/layers/test_websocket.py @@ -13,6 +13,7 @@ from mitmproxy.proxy.events import DataReceived, ConnectionClosed from mitmproxy.proxy.layers import http, websocket from mitmproxy.websocket import WebSocketData from test.mitmproxy.proxy.tutils import Placeholder, Playbook, reply +from wsproto.frame_protocol import Opcode @dataclass @@ -97,10 +98,10 @@ def test_upgrade(tctx): assert len(flow().websocket.messages) == 2 assert flow().websocket.messages[0].content == b"hello world" assert flow().websocket.messages[0].from_client - assert flow().websocket.messages[0].is_text + assert flow().websocket.messages[0].type == Opcode.TEXT assert flow().websocket.messages[1].content == b"hello back" assert flow().websocket.messages[1].from_client is False - assert flow().websocket.messages[1].is_text is False + assert flow().websocket.messages[1].type == Opcode.BINARY @pytest.fixture() @@ -150,7 +151,7 @@ def test_drop_message(ws_testdata): >> DataReceived(tctx.server, b"\x81\x03foo") << websocket.WebsocketMessageHook(flow) ) - flow.websocket.messages[-1].content = "" + flow.websocket.messages[-1].kill() assert ( playbook >> reply() diff --git a/test/mitmproxy/test_websocket.py b/test/mitmproxy/test_websocket.py index 788219752..26d7b0f6a 100644 --- a/test/mitmproxy/test_websocket.py +++ b/test/mitmproxy/test_websocket.py @@ -1,6 +1,7 @@ from mitmproxy import http from mitmproxy import websocket from mitmproxy.test import tflow +from wsproto.frame_protocol import Opcode class TestWebSocketData: @@ -15,9 +16,13 @@ class TestWebSocketData: class TestWebSocketMessage: def test_basic(self): - m = websocket.WebSocketMessage(True, True, b"foo") + m = websocket.WebSocketMessage(Opcode.TEXT, True, b"foo") m.set_state(m.get_state()) assert m.content == b"foo" assert repr(m) == "'foo'" - m.is_text = False + m.type = Opcode.BINARY assert repr(m) == "b'foo'" + + assert not m.killed + m.kill() + assert m.killed