add WebSocketMessage.text

docs/scripts/pdoc-template/module.html.jinja2

@@ -55,6 +55,10 @@ To document all event hooks, we do a bit of hackery:
             {% if doc.qualname.startswith("ServerConnectionHookData") and doc.name != "__init__" %}
                 {{ default_is_public(doc) }}
             {% endif %}
+        {% elif doc.modulename == "mitmproxy.websocket" %}
+            {% if doc.qualname != "WebSocketMessage.type" %}
+                {{ default_is_public(doc) }}
+            {% endif %}
         {% else %}
             {{ default_is_public(doc) }}
         {% endif %}

mitmproxy/websocket.py

@@ -20,18 +20,16 @@ class WebSocketMessage(serializable.Serializable):
     """
     A single WebSocket message sent from one peer to the other.
 
-    Fragmented WebSocket messages are reassembled by mitmproxy and the
+    Fragmented WebSocket messages are reassembled by mitmproxy and then
     represented as a single instance of this class.
 
     The [WebSocket RFC](https://tools.ietf.org/html/rfc6455) specifies both
     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:
+    mitmproxy stores all message contents as `bytes`. If you need a `str`, you can access the `text` property
+    on text messages:
 
-    >>> from wsproto.frame_protocol import Opcode
+    >>> if message.is_text:
-    >>> if message.type == Opcode.TEXT:
+    >>>     text = message.text
-    >>>     text = message.content.decode()
-
-    Per the WebSocket spec, text messages always use UTF-8 encoding.
     """
 
     from_client: bool
@@ -40,8 +38,7 @@ class WebSocketMessage(serializable.Serializable):
     """
     The message type, as per RFC 6455's [opcode](https://tools.ietf.org/html/rfc6455#section-5.2).
 
-    Note that mitmproxy will always store the message contents as *bytes*.
+    Mitmproxy currently only exposes messages assembled from `TEXT` and `BINARY` frames.
-    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."""
@@ -81,10 +78,39 @@ class WebSocketMessage(serializable.Serializable):
         else:
             return repr(self.content)
 
+    @property
+    def is_text(self) -> bool:
+        """
+        `True` if this message is assembled from WebSocket `TEXT` frames,
+        `False` if it is assembled from `BINARY` frames.
+        """
+        return self.type == Opcode.TEXT
+
     def kill(self):
         # Likely to be replaced with .drop() in the future, see https://github.com/mitmproxy/mitmproxy/pull/4486
         self.killed = True
 
+    @property
+    def text(self) -> str:
+        """
+        The message content as text.
+
+        This attribute is only available if `WebSocketMessage.is_text` is `True`.
+
+        *See also:* `WebSocketMessage.content`
+        """
+        if self.type != Opcode.TEXT:
+            raise AttributeError(f"{self.type.name.title()} WebSocket frames do not have a 'text' attribute.")
+
+        return self.content.decode()
+
+    @text.setter
+    def text(self, value: str) -> None:
+        if self.type != Opcode.TEXT:
+            raise AttributeError(f"{self.type.name.title()} WebSocket frames do not have a 'text' attribute.")
+
+        self.content = value.encode()
+
 
 class WebSocketData(stateobject.StateObject):
     """
@@ -97,9 +123,9 @@ class WebSocketData(stateobject.StateObject):
 
     closed_by_client: Optional[bool] = None
     """
-    True if the client closed the connection,
+    `True` if the client closed the connection,
-    False if the server closed the connection,
+    `False` if the server closed the connection,
-    None if the connection is active.
+    `None` if the connection is active.
     """
     close_code: Optional[int] = None
     """[Close Code](https://tools.ietf.org/html/rfc6455#section-7.1.5)"""

test/mitmproxy/test_websocket.py

@@ -1,3 +1,5 @@
+import pytest
+
 from mitmproxy import http
 from mitmproxy import websocket
 from mitmproxy.test import tflow
@@ -26,3 +28,18 @@ class TestWebSocketMessage:
         assert not m.killed
         m.kill()
         assert m.killed
+
+    def test_text(self):
+        txt = websocket.WebSocketMessage(Opcode.TEXT, True, b"foo")
+        bin = websocket.WebSocketMessage(Opcode.BINARY, True, b"foo")
+
+        assert txt.is_text
+        assert txt.text
+        txt.text = "bar"
+        assert txt.content == b"bar"
+
+        assert not bin.is_text
+        with pytest.raises(AttributeError, match="do not have a 'text' attribute."):
+            _ = bin.text
+        with pytest.raises(AttributeError, match="do not have a 'text' attribute."):
+            bin.text = "bar"