We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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?

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `mitmproxy/proxy/layers/websocket.py`
Content:
```
1 import time
2 from dataclasses import dataclass
3 from typing import Iterator, List
4 
5 import wsproto
6 import wsproto.extensions
7 import wsproto.frame_protocol
8 import wsproto.utilities
9 from mitmproxy import connection, http, websocket
10 from mitmproxy.proxy import commands, events, layer
11 from mitmproxy.proxy.commands import StartHook
12 from mitmproxy.proxy.context import Context
13 from mitmproxy.proxy.events import MessageInjected
14 from mitmproxy.proxy.utils import expect
15 from wsproto import ConnectionState
16 from wsproto.frame_protocol import Opcode
17 
18 
19 @dataclass
20 class WebsocketStartHook(StartHook):
21     """
22     A WebSocket connection has commenced.
23     """
24     flow: http.HTTPFlow
25 
26 
27 @dataclass
28 class WebsocketMessageHook(StartHook):
29     """
30     Called when a WebSocket message is received from the client or
31     server. The most recent message will be flow.messages[-1]. The
32     message is user-modifiable. Currently there are two types of
33     messages, corresponding to the BINARY and TEXT frame types.
34     """
35     flow: http.HTTPFlow
36 
37 
38 @dataclass
39 class WebsocketEndHook(StartHook):
40     """
41     A WebSocket connection has ended.
42     You can check `flow.websocket.close_code` to determine why it ended.
43     """
44 
45     flow: http.HTTPFlow
46 
47 
48 class WebSocketMessageInjected(MessageInjected[websocket.WebSocketMessage]):
49     """
50     The user has injected a custom WebSocket message.
51     """
52 
53 
54 class WebsocketConnection(wsproto.Connection):
55     """
56     A very thin wrapper around wsproto.Connection:
57 
58      - we keep the underlying connection as an attribute for easy access.
59      - we add a framebuffer for incomplete messages
60      - we wrap .send() so that we can directly yield it.
61     """
62     conn: connection.Connection
63     frame_buf: List[bytes]
64 
65     def __init__(self, *args, conn: connection.Connection, **kwargs):
66         super(WebsocketConnection, self).__init__(*args, **kwargs)
67         self.conn = conn
68         self.frame_buf = []
69 
70     def send2(self, event: wsproto.events.Event) -> commands.SendData:
71         data = self.send(event)
72         return commands.SendData(self.conn, data)
73 
74     def __repr__(self):
75         return f"WebsocketConnection<{self.state.name}, {self.conn}>"
76 
77 
78 class WebsocketLayer(layer.Layer):
79     """
80     WebSocket layer that intercepts and relays messages.
81     """
82     flow: http.HTTPFlow
83     client_ws: WebsocketConnection
84     server_ws: WebsocketConnection
85 
86     def __init__(self, context: Context, flow: http.HTTPFlow):
87         super().__init__(context)
88         self.flow = flow
89         assert context.server.connected
90 
91     @expect(events.Start)
92     def start(self, _) -> layer.CommandGenerator[None]:
93 
94         client_extensions = []
95         server_extensions = []
96 
97         # Parse extension headers. We only support deflate at the moment and ignore everything else.
98         assert self.flow.response  # satisfy type checker
99         ext_header = self.flow.response.headers.get("Sec-WebSocket-Extensions", "")
100         if ext_header:
101             for ext in wsproto.utilities.split_comma_header(ext_header.encode("ascii", "replace")):
102                 ext_name = ext.split(";", 1)[0].strip()
103                 if ext_name == wsproto.extensions.PerMessageDeflate.name:
104                     client_deflate = wsproto.extensions.PerMessageDeflate()
105                     client_deflate.finalize(ext)
106                     client_extensions.append(client_deflate)
107                     server_deflate = wsproto.extensions.PerMessageDeflate()
108                     server_deflate.finalize(ext)
109                     server_extensions.append(server_deflate)
110                 else:
111                     yield commands.Log(f"Ignoring unknown WebSocket extension {ext_name!r}.")
112 
113         self.client_ws = WebsocketConnection(wsproto.ConnectionType.SERVER, client_extensions, conn=self.context.client)
114         self.server_ws = WebsocketConnection(wsproto.ConnectionType.CLIENT, server_extensions, conn=self.context.server)
115 
116         yield WebsocketStartHook(self.flow)
117 
118         self._handle_event = self.relay_messages
119 
120     _handle_event = start
121 
122     @expect(events.DataReceived, events.ConnectionClosed, WebSocketMessageInjected)
123     def relay_messages(self, event: events.Event) -> layer.CommandGenerator[None]:
124         assert self.flow.websocket  # satisfy type checker
125 
126         if isinstance(event, events.ConnectionEvent):
127             from_client = event.connection == self.context.client
128         elif isinstance(event, WebSocketMessageInjected):
129             from_client = event.message.from_client
130         else:
131             raise AssertionError(f"Unexpected event: {event}")
132 
133         from_str = 'client' if from_client else 'server'
134         if from_client:
135             src_ws = self.client_ws
136             dst_ws = self.server_ws
137         else:
138             src_ws = self.server_ws
139             dst_ws = self.client_ws
140 
141         if isinstance(event, events.DataReceived):
142             src_ws.receive_data(event.data)
143         elif isinstance(event, events.ConnectionClosed):
144             src_ws.receive_data(None)
145         elif isinstance(event, WebSocketMessageInjected):
146             fragmentizer = Fragmentizer([], event.message.type == Opcode.TEXT)
147             src_ws._events.extend(
148                 fragmentizer(event.message.content)
149             )
150         else:  # pragma: no cover
151             raise AssertionError(f"Unexpected event: {event}")
152 
153         for ws_event in src_ws.events():
154             if isinstance(ws_event, wsproto.events.Message):
155                 is_text = isinstance(ws_event.data, str)
156                 if is_text:
157                     typ = Opcode.TEXT
158                     src_ws.frame_buf.append(ws_event.data.encode())
159                 else:
160                     typ = Opcode.BINARY
161                     src_ws.frame_buf.append(ws_event.data)
162 
163                 if ws_event.message_finished:
164                     content = b"".join(src_ws.frame_buf)
165 
166                     fragmentizer = Fragmentizer(src_ws.frame_buf, is_text)
167                     src_ws.frame_buf.clear()
168 
169                     message = websocket.WebSocketMessage(typ, from_client, content)
170                     self.flow.websocket.messages.append(message)
171                     yield WebsocketMessageHook(self.flow)
172 
173                     if not message.dropped:
174                         for msg in fragmentizer(message.content):
175                             yield dst_ws.send2(msg)
176 
177             elif isinstance(ws_event, (wsproto.events.Ping, wsproto.events.Pong)):
178                 yield commands.Log(
179                     f"Received WebSocket {ws_event.__class__.__name__.lower()} from {from_str} "
180                     f"(payload: {bytes(ws_event.payload)!r})"
181                 )
182                 yield dst_ws.send2(ws_event)
183             elif isinstance(ws_event, wsproto.events.CloseConnection):
184                 self.flow.websocket.timestamp_end = time.time()
185                 self.flow.websocket.closed_by_client = from_client
186                 self.flow.websocket.close_code = ws_event.code
187                 self.flow.websocket.close_reason = ws_event.reason
188 
189                 for ws in [self.server_ws, self.client_ws]:
190                     if ws.state in {ConnectionState.OPEN, ConnectionState.REMOTE_CLOSING}:
191                         # response == original event, so no need to differentiate here.
192                         yield ws.send2(ws_event)
193                     yield commands.CloseConnection(ws.conn)
194                 yield WebsocketEndHook(self.flow)
195                 self._handle_event = self.done
196             else:  # pragma: no cover
197                 raise AssertionError(f"Unexpected WebSocket event: {ws_event}")
198 
199     @expect(events.DataReceived, events.ConnectionClosed, WebSocketMessageInjected)
200     def done(self, _) -> layer.CommandGenerator[None]:
201         yield from ()
202 
203 
204 class Fragmentizer:
205     """
206     Theory (RFC 6455):
207        Unless specified otherwise by an extension, frames have no semantic
208        meaning.  An intermediary might coalesce and/or split frames, [...]
209 
210     Practice:
211         Some WebSocket servers reject large payload sizes.
212 
213     As a workaround, we either retain the original chunking or, if the payload has been modified, use ~4kB chunks.
214     """
215     # A bit less than 4kb to accommodate for headers.
216     FRAGMENT_SIZE = 4000
217 
218     def __init__(self, fragments: List[bytes], is_text: bool):
219         self.fragment_lengths = [len(x) for x in fragments]
220         self.is_text = is_text
221 
222     def msg(self, data: bytes, message_finished: bool):
223         if self.is_text:
224             data_str = data.decode(errors="replace")
225             return wsproto.events.TextMessage(data_str, message_finished=message_finished)
226         else:
227             return wsproto.events.BytesMessage(data, message_finished=message_finished)
228 
229     def __call__(self, content: bytes) -> Iterator[wsproto.events.Message]:
230         if len(content) == sum(self.fragment_lengths):
231             # message has the same length, we can reuse the same sizes
232             offset = 0
233             for fl in self.fragment_lengths[:-1]:
234                 yield self.msg(content[offset:offset + fl], False)
235                 offset += fl
236             yield self.msg(content[offset:], True)
237         else:
238             offset = 0
239             total = len(content) - self.FRAGMENT_SIZE
240             while offset < total:
241                 yield self.msg(content[offset:offset + self.FRAGMENT_SIZE], False)
242                 offset += self.FRAGMENT_SIZE
243             yield self.msg(content[offset:], True)
244
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
     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 @@
                 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 @@
                         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 @@
 
     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

{"golden_diff": "diff --git a/mitmproxy/proxy/layers/websocket.py b/mitmproxy/proxy/layers/websocket.py\n--- a/mitmproxy/proxy/layers/websocket.py\n+++ b/mitmproxy/proxy/layers/websocket.py\n@@ -65,7 +65,7 @@\n     def __init__(self, *args, conn: connection.Connection, **kwargs):\n         super(WebsocketConnection, self).__init__(*args, **kwargs)\n         self.conn = conn\n-        self.frame_buf = []\n+        self.frame_buf = [b\"\"]\n \n     def send2(self, event: wsproto.events.Event) -> commands.SendData:\n         data = self.send(event)\n@@ -155,10 +155,10 @@\n                 is_text = isinstance(ws_event.data, str)\n                 if is_text:\n                     typ = Opcode.TEXT\n-                    src_ws.frame_buf.append(ws_event.data.encode())\n+                    src_ws.frame_buf[-1] += ws_event.data.encode()\n                 else:\n                     typ = Opcode.BINARY\n-                    src_ws.frame_buf.append(ws_event.data)\n+                    src_ws.frame_buf[-1] += ws_event.data\n \n                 if ws_event.message_finished:\n                     content = b\"\".join(src_ws.frame_buf)\n@@ -174,6 +174,9 @@\n                         for msg in fragmentizer(message.content):\n                             yield dst_ws.send2(msg)\n \n+                elif ws_event.frame_finished:\n+                    src_ws.frame_buf.append(b\"\")\n+\n             elif isinstance(ws_event, (wsproto.events.Ping, wsproto.events.Pong)):\n                 yield commands.Log(\n                     f\"Received WebSocket {ws_event.__class__.__name__.lower()} from {from_str} \"\n@@ -209,8 +212,11 @@\n \n     Practice:\n         Some WebSocket servers reject large payload sizes.\n+        Other WebSocket servers reject CONTINUATION frames.\n \n     As a workaround, we either retain the original chunking or, if the payload has been modified, use ~4kB chunks.\n+    If one deals with web servers that do not support CONTINUATION frames, addons need to monkeypatch FRAGMENT_SIZE\n+    if they need to modify the message.\n     \"\"\"\n     # A bit less than 4kb to accommodate for headers.\n     FRAGMENT_SIZE = 4000\n", "issue": "Unfragmented WebSocket messages getting fragmented\nhttps://github.com/mitmproxy/mitmproxy/blob/e270399a3eba7b3212bf7325beaa50e159a7ca0a/mitmproxy/proxy/layers/websocket.py#L153-L175\r\n\r\nWhile 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.\r\n\r\nThis 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?\r\n\r\nTo avoid this behavior each `frame_buf` entry must be a complete frame, here is my fix suggestion:\r\n\r\n```python\r\n        for ws_event in src_ws.events():\r\n            if isinstance(ws_event, wsproto.events.Message):\r\n                is_text = isinstance(ws_event.data, str)\r\n\r\n                # Add the data variable to avoid multiple conditions\r\n                if is_text:\r\n                    typ = Opcode.TEXT\r\n                    data = ws_event.data.encode()\r\n                else:\r\n                    typ = Opcode.BINARY\r\n                    data = ws_event.data\r\n\r\n                # Make each frame one entry to frame_buf, append if empty to avoid IndexError\r\n                if src_ws.frame_buf:\r\n                    src_ws.frame_buf[-1] += data\r\n                else:\r\n                    src_ws.frame_buf.append(data)\r\n\r\n                if ws_event.message_finished:\r\n                    content = b\"\".join(src_ws.frame_buf)\r\n\r\n                    fragmentizer = Fragmentizer(src_ws.frame_buf, is_text)\r\n                    src_ws.frame_buf.clear()\r\n\r\n                    message = websocket.WebSocketMessage(typ, from_client, content)\r\n                    self.flow.websocket.messages.append(message)\r\n                    yield WebsocketMessageHook(self.flow)\r\n\r\n                    if not message.dropped:\r\n                        for msg in fragmentizer(message.content):\r\n                            yield dst_ws.send2(msg)\r\n\r\n                # Initialize next frame entry\r\n                elif ws_event.frame_finished:\r\n                    src_ws.frame_buf.append(b\"\")\r\n```\r\n\r\nIt 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):\r\n\r\nhttps://github.com/mitmproxy/mitmproxy/blob/e270399a3eba7b3212bf7325beaa50e159a7ca0a/mitmproxy/proxy/layers/websocket.py#L229-L243\r\n\r\nFor 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?\n", "before_files": [{"content": "import time\nfrom dataclasses import dataclass\nfrom typing import Iterator, List\n\nimport wsproto\nimport wsproto.extensions\nimport wsproto.frame_protocol\nimport wsproto.utilities\nfrom mitmproxy import connection, http, websocket\nfrom mitmproxy.proxy import commands, events, layer\nfrom mitmproxy.proxy.commands import StartHook\nfrom mitmproxy.proxy.context import Context\nfrom mitmproxy.proxy.events import MessageInjected\nfrom mitmproxy.proxy.utils import expect\nfrom wsproto import ConnectionState\nfrom wsproto.frame_protocol import Opcode\n\n\n@dataclass\nclass WebsocketStartHook(StartHook):\n    \"\"\"\n    A WebSocket connection has commenced.\n    \"\"\"\n    flow: http.HTTPFlow\n\n\n@dataclass\nclass WebsocketMessageHook(StartHook):\n    \"\"\"\n    Called when a WebSocket message is received from the client or\n    server. The most recent message will be flow.messages[-1]. The\n    message is user-modifiable. Currently there are two types of\n    messages, corresponding to the BINARY and TEXT frame types.\n    \"\"\"\n    flow: http.HTTPFlow\n\n\n@dataclass\nclass WebsocketEndHook(StartHook):\n    \"\"\"\n    A WebSocket connection has ended.\n    You can check `flow.websocket.close_code` to determine why it ended.\n    \"\"\"\n\n    flow: http.HTTPFlow\n\n\nclass WebSocketMessageInjected(MessageInjected[websocket.WebSocketMessage]):\n    \"\"\"\n    The user has injected a custom WebSocket message.\n    \"\"\"\n\n\nclass WebsocketConnection(wsproto.Connection):\n    \"\"\"\n    A very thin wrapper around wsproto.Connection:\n\n     - we keep the underlying connection as an attribute for easy access.\n     - we add a framebuffer for incomplete messages\n     - we wrap .send() so that we can directly yield it.\n    \"\"\"\n    conn: connection.Connection\n    frame_buf: List[bytes]\n\n    def __init__(self, *args, conn: connection.Connection, **kwargs):\n        super(WebsocketConnection, self).__init__(*args, **kwargs)\n        self.conn = conn\n        self.frame_buf = []\n\n    def send2(self, event: wsproto.events.Event) -> commands.SendData:\n        data = self.send(event)\n        return commands.SendData(self.conn, data)\n\n    def __repr__(self):\n        return f\"WebsocketConnection<{self.state.name}, {self.conn}>\"\n\n\nclass WebsocketLayer(layer.Layer):\n    \"\"\"\n    WebSocket layer that intercepts and relays messages.\n    \"\"\"\n    flow: http.HTTPFlow\n    client_ws: WebsocketConnection\n    server_ws: WebsocketConnection\n\n    def __init__(self, context: Context, flow: http.HTTPFlow):\n        super().__init__(context)\n        self.flow = flow\n        assert context.server.connected\n\n    @expect(events.Start)\n    def start(self, _) -> layer.CommandGenerator[None]:\n\n        client_extensions = []\n        server_extensions = []\n\n        # Parse extension headers. We only support deflate at the moment and ignore everything else.\n        assert self.flow.response  # satisfy type checker\n        ext_header = self.flow.response.headers.get(\"Sec-WebSocket-Extensions\", \"\")\n        if ext_header:\n            for ext in wsproto.utilities.split_comma_header(ext_header.encode(\"ascii\", \"replace\")):\n                ext_name = ext.split(\";\", 1)[0].strip()\n                if ext_name == wsproto.extensions.PerMessageDeflate.name:\n                    client_deflate = wsproto.extensions.PerMessageDeflate()\n                    client_deflate.finalize(ext)\n                    client_extensions.append(client_deflate)\n                    server_deflate = wsproto.extensions.PerMessageDeflate()\n                    server_deflate.finalize(ext)\n                    server_extensions.append(server_deflate)\n                else:\n                    yield commands.Log(f\"Ignoring unknown WebSocket extension {ext_name!r}.\")\n\n        self.client_ws = WebsocketConnection(wsproto.ConnectionType.SERVER, client_extensions, conn=self.context.client)\n        self.server_ws = WebsocketConnection(wsproto.ConnectionType.CLIENT, server_extensions, conn=self.context.server)\n\n        yield WebsocketStartHook(self.flow)\n\n        self._handle_event = self.relay_messages\n\n    _handle_event = start\n\n    @expect(events.DataReceived, events.ConnectionClosed, WebSocketMessageInjected)\n    def relay_messages(self, event: events.Event) -> layer.CommandGenerator[None]:\n        assert self.flow.websocket  # satisfy type checker\n\n        if isinstance(event, events.ConnectionEvent):\n            from_client = event.connection == self.context.client\n        elif isinstance(event, WebSocketMessageInjected):\n            from_client = event.message.from_client\n        else:\n            raise AssertionError(f\"Unexpected event: {event}\")\n\n        from_str = 'client' if from_client else 'server'\n        if from_client:\n            src_ws = self.client_ws\n            dst_ws = self.server_ws\n        else:\n            src_ws = self.server_ws\n            dst_ws = self.client_ws\n\n        if isinstance(event, events.DataReceived):\n            src_ws.receive_data(event.data)\n        elif isinstance(event, events.ConnectionClosed):\n            src_ws.receive_data(None)\n        elif isinstance(event, WebSocketMessageInjected):\n            fragmentizer = Fragmentizer([], event.message.type == Opcode.TEXT)\n            src_ws._events.extend(\n                fragmentizer(event.message.content)\n            )\n        else:  # pragma: no cover\n            raise AssertionError(f\"Unexpected event: {event}\")\n\n        for ws_event in src_ws.events():\n            if isinstance(ws_event, wsproto.events.Message):\n                is_text = isinstance(ws_event.data, str)\n                if is_text:\n                    typ = Opcode.TEXT\n                    src_ws.frame_buf.append(ws_event.data.encode())\n                else:\n                    typ = Opcode.BINARY\n                    src_ws.frame_buf.append(ws_event.data)\n\n                if ws_event.message_finished:\n                    content = b\"\".join(src_ws.frame_buf)\n\n                    fragmentizer = Fragmentizer(src_ws.frame_buf, is_text)\n                    src_ws.frame_buf.clear()\n\n                    message = websocket.WebSocketMessage(typ, from_client, content)\n                    self.flow.websocket.messages.append(message)\n                    yield WebsocketMessageHook(self.flow)\n\n                    if not message.dropped:\n                        for msg in fragmentizer(message.content):\n                            yield dst_ws.send2(msg)\n\n            elif isinstance(ws_event, (wsproto.events.Ping, wsproto.events.Pong)):\n                yield commands.Log(\n                    f\"Received WebSocket {ws_event.__class__.__name__.lower()} from {from_str} \"\n                    f\"(payload: {bytes(ws_event.payload)!r})\"\n                )\n                yield dst_ws.send2(ws_event)\n            elif isinstance(ws_event, wsproto.events.CloseConnection):\n                self.flow.websocket.timestamp_end = time.time()\n                self.flow.websocket.closed_by_client = from_client\n                self.flow.websocket.close_code = ws_event.code\n                self.flow.websocket.close_reason = ws_event.reason\n\n                for ws in [self.server_ws, self.client_ws]:\n                    if ws.state in {ConnectionState.OPEN, ConnectionState.REMOTE_CLOSING}:\n                        # response == original event, so no need to differentiate here.\n                        yield ws.send2(ws_event)\n                    yield commands.CloseConnection(ws.conn)\n                yield WebsocketEndHook(self.flow)\n                self._handle_event = self.done\n            else:  # pragma: no cover\n                raise AssertionError(f\"Unexpected WebSocket event: {ws_event}\")\n\n    @expect(events.DataReceived, events.ConnectionClosed, WebSocketMessageInjected)\n    def done(self, _) -> layer.CommandGenerator[None]:\n        yield from ()\n\n\nclass Fragmentizer:\n    \"\"\"\n    Theory (RFC 6455):\n       Unless specified otherwise by an extension, frames have no semantic\n       meaning.  An intermediary might coalesce and/or split frames, [...]\n\n    Practice:\n        Some WebSocket servers reject large payload sizes.\n\n    As a workaround, we either retain the original chunking or, if the payload has been modified, use ~4kB chunks.\n    \"\"\"\n    # A bit less than 4kb to accommodate for headers.\n    FRAGMENT_SIZE = 4000\n\n    def __init__(self, fragments: List[bytes], is_text: bool):\n        self.fragment_lengths = [len(x) for x in fragments]\n        self.is_text = is_text\n\n    def msg(self, data: bytes, message_finished: bool):\n        if self.is_text:\n            data_str = data.decode(errors=\"replace\")\n            return wsproto.events.TextMessage(data_str, message_finished=message_finished)\n        else:\n            return wsproto.events.BytesMessage(data, message_finished=message_finished)\n\n    def __call__(self, content: bytes) -> Iterator[wsproto.events.Message]:\n        if len(content) == sum(self.fragment_lengths):\n            # message has the same length, we can reuse the same sizes\n            offset = 0\n            for fl in self.fragment_lengths[:-1]:\n                yield self.msg(content[offset:offset + fl], False)\n                offset += fl\n            yield self.msg(content[offset:], True)\n        else:\n            offset = 0\n            total = len(content) - self.FRAGMENT_SIZE\n            while offset < total:\n                yield self.msg(content[offset:offset + self.FRAGMENT_SIZE], False)\n                offset += self.FRAGMENT_SIZE\n            yield self.msg(content[offset:], True)\n", "path": "mitmproxy/proxy/layers/websocket.py"}], "after_files": [{"content": "import time\nfrom dataclasses import dataclass\nfrom typing import Iterator, List\n\nimport wsproto\nimport wsproto.extensions\nimport wsproto.frame_protocol\nimport wsproto.utilities\nfrom mitmproxy import connection, http, websocket\nfrom mitmproxy.proxy import commands, events, layer\nfrom mitmproxy.proxy.commands import StartHook\nfrom mitmproxy.proxy.context import Context\nfrom mitmproxy.proxy.events import MessageInjected\nfrom mitmproxy.proxy.utils import expect\nfrom wsproto import ConnectionState\nfrom wsproto.frame_protocol import Opcode\n\n\n@dataclass\nclass WebsocketStartHook(StartHook):\n    \"\"\"\n    A WebSocket connection has commenced.\n    \"\"\"\n    flow: http.HTTPFlow\n\n\n@dataclass\nclass WebsocketMessageHook(StartHook):\n    \"\"\"\n    Called when a WebSocket message is received from the client or\n    server. The most recent message will be flow.messages[-1]. The\n    message is user-modifiable. Currently there are two types of\n    messages, corresponding to the BINARY and TEXT frame types.\n    \"\"\"\n    flow: http.HTTPFlow\n\n\n@dataclass\nclass WebsocketEndHook(StartHook):\n    \"\"\"\n    A WebSocket connection has ended.\n    You can check `flow.websocket.close_code` to determine why it ended.\n    \"\"\"\n\n    flow: http.HTTPFlow\n\n\nclass WebSocketMessageInjected(MessageInjected[websocket.WebSocketMessage]):\n    \"\"\"\n    The user has injected a custom WebSocket message.\n    \"\"\"\n\n\nclass WebsocketConnection(wsproto.Connection):\n    \"\"\"\n    A very thin wrapper around wsproto.Connection:\n\n     - we keep the underlying connection as an attribute for easy access.\n     - we add a framebuffer for incomplete messages\n     - we wrap .send() so that we can directly yield it.\n    \"\"\"\n    conn: connection.Connection\n    frame_buf: List[bytes]\n\n    def __init__(self, *args, conn: connection.Connection, **kwargs):\n        super(WebsocketConnection, self).__init__(*args, **kwargs)\n        self.conn = conn\n        self.frame_buf = [b\"\"]\n\n    def send2(self, event: wsproto.events.Event) -> commands.SendData:\n        data = self.send(event)\n        return commands.SendData(self.conn, data)\n\n    def __repr__(self):\n        return f\"WebsocketConnection<{self.state.name}, {self.conn}>\"\n\n\nclass WebsocketLayer(layer.Layer):\n    \"\"\"\n    WebSocket layer that intercepts and relays messages.\n    \"\"\"\n    flow: http.HTTPFlow\n    client_ws: WebsocketConnection\n    server_ws: WebsocketConnection\n\n    def __init__(self, context: Context, flow: http.HTTPFlow):\n        super().__init__(context)\n        self.flow = flow\n        assert context.server.connected\n\n    @expect(events.Start)\n    def start(self, _) -> layer.CommandGenerator[None]:\n\n        client_extensions = []\n        server_extensions = []\n\n        # Parse extension headers. We only support deflate at the moment and ignore everything else.\n        assert self.flow.response  # satisfy type checker\n        ext_header = self.flow.response.headers.get(\"Sec-WebSocket-Extensions\", \"\")\n        if ext_header:\n            for ext in wsproto.utilities.split_comma_header(ext_header.encode(\"ascii\", \"replace\")):\n                ext_name = ext.split(\";\", 1)[0].strip()\n                if ext_name == wsproto.extensions.PerMessageDeflate.name:\n                    client_deflate = wsproto.extensions.PerMessageDeflate()\n                    client_deflate.finalize(ext)\n                    client_extensions.append(client_deflate)\n                    server_deflate = wsproto.extensions.PerMessageDeflate()\n                    server_deflate.finalize(ext)\n                    server_extensions.append(server_deflate)\n                else:\n                    yield commands.Log(f\"Ignoring unknown WebSocket extension {ext_name!r}.\")\n\n        self.client_ws = WebsocketConnection(wsproto.ConnectionType.SERVER, client_extensions, conn=self.context.client)\n        self.server_ws = WebsocketConnection(wsproto.ConnectionType.CLIENT, server_extensions, conn=self.context.server)\n\n        yield WebsocketStartHook(self.flow)\n\n        self._handle_event = self.relay_messages\n\n    _handle_event = start\n\n    @expect(events.DataReceived, events.ConnectionClosed, WebSocketMessageInjected)\n    def relay_messages(self, event: events.Event) -> layer.CommandGenerator[None]:\n        assert self.flow.websocket  # satisfy type checker\n\n        if isinstance(event, events.ConnectionEvent):\n            from_client = event.connection == self.context.client\n        elif isinstance(event, WebSocketMessageInjected):\n            from_client = event.message.from_client\n        else:\n            raise AssertionError(f\"Unexpected event: {event}\")\n\n        from_str = 'client' if from_client else 'server'\n        if from_client:\n            src_ws = self.client_ws\n            dst_ws = self.server_ws\n        else:\n            src_ws = self.server_ws\n            dst_ws = self.client_ws\n\n        if isinstance(event, events.DataReceived):\n            src_ws.receive_data(event.data)\n        elif isinstance(event, events.ConnectionClosed):\n            src_ws.receive_data(None)\n        elif isinstance(event, WebSocketMessageInjected):\n            fragmentizer = Fragmentizer([], event.message.type == Opcode.TEXT)\n            src_ws._events.extend(\n                fragmentizer(event.message.content)\n            )\n        else:  # pragma: no cover\n            raise AssertionError(f\"Unexpected event: {event}\")\n\n        for ws_event in src_ws.events():\n            if isinstance(ws_event, wsproto.events.Message):\n                is_text = isinstance(ws_event.data, str)\n                if is_text:\n                    typ = Opcode.TEXT\n                    src_ws.frame_buf[-1] += ws_event.data.encode()\n                else:\n                    typ = Opcode.BINARY\n                    src_ws.frame_buf[-1] += ws_event.data\n\n                if ws_event.message_finished:\n                    content = b\"\".join(src_ws.frame_buf)\n\n                    fragmentizer = Fragmentizer(src_ws.frame_buf, is_text)\n                    src_ws.frame_buf.clear()\n\n                    message = websocket.WebSocketMessage(typ, from_client, content)\n                    self.flow.websocket.messages.append(message)\n                    yield WebsocketMessageHook(self.flow)\n\n                    if not message.dropped:\n                        for msg in fragmentizer(message.content):\n                            yield dst_ws.send2(msg)\n\n                elif ws_event.frame_finished:\n                    src_ws.frame_buf.append(b\"\")\n\n            elif isinstance(ws_event, (wsproto.events.Ping, wsproto.events.Pong)):\n                yield commands.Log(\n                    f\"Received WebSocket {ws_event.__class__.__name__.lower()} from {from_str} \"\n                    f\"(payload: {bytes(ws_event.payload)!r})\"\n                )\n                yield dst_ws.send2(ws_event)\n            elif isinstance(ws_event, wsproto.events.CloseConnection):\n                self.flow.websocket.timestamp_end = time.time()\n                self.flow.websocket.closed_by_client = from_client\n                self.flow.websocket.close_code = ws_event.code\n                self.flow.websocket.close_reason = ws_event.reason\n\n                for ws in [self.server_ws, self.client_ws]:\n                    if ws.state in {ConnectionState.OPEN, ConnectionState.REMOTE_CLOSING}:\n                        # response == original event, so no need to differentiate here.\n                        yield ws.send2(ws_event)\n                    yield commands.CloseConnection(ws.conn)\n                yield WebsocketEndHook(self.flow)\n                self._handle_event = self.done\n            else:  # pragma: no cover\n                raise AssertionError(f\"Unexpected WebSocket event: {ws_event}\")\n\n    @expect(events.DataReceived, events.ConnectionClosed, WebSocketMessageInjected)\n    def done(self, _) -> layer.CommandGenerator[None]:\n        yield from ()\n\n\nclass Fragmentizer:\n    \"\"\"\n    Theory (RFC 6455):\n       Unless specified otherwise by an extension, frames have no semantic\n       meaning.  An intermediary might coalesce and/or split frames, [...]\n\n    Practice:\n        Some WebSocket servers reject large payload sizes.\n        Other WebSocket servers reject CONTINUATION frames.\n\n    As a workaround, we either retain the original chunking or, if the payload has been modified, use ~4kB chunks.\n    If one deals with web servers that do not support CONTINUATION frames, addons need to monkeypatch FRAGMENT_SIZE\n    if they need to modify the message.\n    \"\"\"\n    # A bit less than 4kb to accommodate for headers.\n    FRAGMENT_SIZE = 4000\n\n    def __init__(self, fragments: List[bytes], is_text: bool):\n        self.fragment_lengths = [len(x) for x in fragments]\n        self.is_text = is_text\n\n    def msg(self, data: bytes, message_finished: bool):\n        if self.is_text:\n            data_str = data.decode(errors=\"replace\")\n            return wsproto.events.TextMessage(data_str, message_finished=message_finished)\n        else:\n            return wsproto.events.BytesMessage(data, message_finished=message_finished)\n\n    def __call__(self, content: bytes) -> Iterator[wsproto.events.Message]:\n        if len(content) == sum(self.fragment_lengths):\n            # message has the same length, we can reuse the same sizes\n            offset = 0\n            for fl in self.fragment_lengths[:-1]:\n                yield self.msg(content[offset:offset + fl], False)\n                offset += fl\n            yield self.msg(content[offset:], True)\n        else:\n            offset = 0\n            total = len(content) - self.FRAGMENT_SIZE\n            while offset < total:\n                yield self.msg(content[offset:offset + self.FRAGMENT_SIZE], False)\n                offset += self.FRAGMENT_SIZE\n            yield self.msg(content[offset:], True)\n", "path": "mitmproxy/proxy/layers/websocket.py"}]}