Repository: mitmproxy/mitmproxy

Issue:
Unfragmented WebSocket messages getting fragmented
https://github.com/mitmproxy/mitmproxy/blob/e270399a3eba7b3212bf7325beaa50e159a7ca0a/mitmproxy/proxy/layers/websocket.py#L153-L175

While handling WebSocket events, mitmproxy doesn't distinguish between `message_finished` and `frame_finished` ([Message class](https://python-hyper.org/projects/wsproto/en/stable/api.html#wsproto.events.Message)) which in my case led to continuation frames being sent while there were none in the initial WebSocket message.

This is because the wsproto API doesn't always emit complete frames (I guess this is caused by the TCP fragmentation?), they could be chunked while the original WebSocket message has no fragmentation and I think even WebSocket messages using fragmentation with large continuation frames could be emitted as chunks themselves?

To avoid this behavior each `frame_buf` entry must be a complete frame, here is my fix suggestion:

```python
        for ws_event in src_ws.events():
            if isinstance(ws_event, wsproto.events.Message):
                is_text = isinstance(ws_event.data, str)

                # Add the data variable to avoid multiple conditions
                if is_text:
                    typ = Opcode.TEXT
                    data = ws_event.data.encode()
                else:
                    typ = Opcode.BINARY
                    data = ws_event.data

                # Make each frame one entry to frame_buf, append if empty to avoid IndexError
                if src_ws.frame_buf:
                    src_ws.frame_buf[-1] += data
                else:
                    src_ws.frame_buf.append(data)

                if ws_event.message_finished:
                    content = b"".join(src_ws.frame_buf)

                    fragmentizer = Fragmentizer(src_ws.frame_buf, is_text)
                    src_ws.frame_buf.clear()

                    message = websocket.WebSocketMessage(typ, from_client, content)
                    self.flow.websocket.messages.append(message)
                    yield WebsocketMessageHook(self.flow)

                    if not message.dropped:
                        for msg in fragmentizer(message.content):
                            yield dst_ws.send2(msg)

                # Initialize next frame entry
                elif ws_event.frame_finished:
                    src_ws.frame_buf.append(b"")
```

It works for me but I didn't test it with in an environment using WebSocket continuation frames. Also this only works for unmodified WebSocket messages, injected or modified messages are still concerned by the issue because the `Fragmentizer` class compares the lengths so they will always fall into the else condition (unless you made the modified message keep its original length):

https://github.com/mitmproxy/mitmproxy/blob/e270399a3eba7b3212bf7325beaa50e159a7ca0a/mitmproxy/proxy/layers/websocket.py#L229-L243

For my use case I didn't make a proper fix for this, I just made the first condition always true, maybe a boolean variable can be added to the `WebSocketMessage` and `Fragmentizer` classes or something like that?


Discussion:

Resolving patch:
diff --git a/mitmproxy/proxy/layers/websocket.py b/mitmproxy/proxy/layers/websocket.py
--- a/mitmproxy/proxy/layers/websocket.py
+++ b/mitmproxy/proxy/layers/websocket.py
@@ -65,7 +65,7 @@ class WebsocketConnection(wsproto.Connection):
     def __init__(self, *args, conn: connection.Connection, **kwargs):
         super(WebsocketConnection, self).__init__(*args, **kwargs)
         self.conn = conn
-        self.frame_buf = []
+        self.frame_buf = [b""]
 
     def send2(self, event: wsproto.events.Event) -> commands.SendData:
         data = self.send(event)
@@ -155,10 +155,10 @@ def relay_messages(self, event: events.Event) -> layer.CommandGenerator[None]:
                 is_text = isinstance(ws_event.data, str)
                 if is_text:
                     typ = Opcode.TEXT
-                    src_ws.frame_buf.append(ws_event.data.encode())
+                    src_ws.frame_buf[-1] += ws_event.data.encode()
                 else:
                     typ = Opcode.BINARY
-                    src_ws.frame_buf.append(ws_event.data)
+                    src_ws.frame_buf[-1] += ws_event.data
 
                 if ws_event.message_finished:
                     content = b"".join(src_ws.frame_buf)
@@ -174,6 +174,9 @@ def relay_messages(self, event: events.Event) -> layer.CommandGenerator[None]:
                         for msg in fragmentizer(message.content):
                             yield dst_ws.send2(msg)
 
+                elif ws_event.frame_finished:
+                    src_ws.frame_buf.append(b"")
+
             elif isinstance(ws_event, (wsproto.events.Ping, wsproto.events.Pong)):
                 yield commands.Log(
                     f"Received WebSocket {ws_event.__class__.__name__.lower()} from {from_str} "
@@ -209,8 +212,11 @@ class Fragmentizer:
 
     Practice:
         Some WebSocket servers reject large payload sizes.
+        Other WebSocket servers reject CONTINUATION frames.
 
     As a workaround, we either retain the original chunking or, if the payload has been modified, use ~4kB chunks.
+    If one deals with web servers that do not support CONTINUATION frames, addons need to monkeypatch FRAGMENT_SIZE
+    if they need to modify the message.
     """
     # A bit less than 4kb to accommodate for headers.
     FRAGMENT_SIZE = 4000


Test patch:
diff --git a/test/mitmproxy/proxy/layers/test_websocket.py b/test/mitmproxy/proxy/layers/test_websocket.py
--- a/test/mitmproxy/proxy/layers/test_websocket.py
+++ b/test/mitmproxy/proxy/layers/test_websocket.py
@@ -193,6 +193,26 @@ def test_fragmented(ws_testdata):
     assert flow.websocket.messages[-1].content == b"foobar"
 
 
+def test_unfragmented(ws_testdata):
+    tctx, playbook, flow = ws_testdata
+    assert (
+            playbook
+            << websocket.WebsocketStartHook(flow)
+            >> reply()
+            >> DataReceived(tctx.server, b"\x81\x06foo")
+    )
+    # This already triggers wsproto to emit a wsproto.events.Message, see
+    # https://github.com/mitmproxy/mitmproxy/issues/4701
+    assert(
+            playbook
+            >> DataReceived(tctx.server, b"bar")
+            << websocket.WebsocketMessageHook(flow)
+            >> reply()
+            << SendData(tctx.client, b"\x81\x06foobar")
+    )
+    assert flow.websocket.messages[-1].content == b"foobar"
+
+
 def test_protocol_error(ws_testdata):
     tctx, playbook, flow = ws_testdata
     assert (