Repository: aio-libs/aiohttp
Pull request: Fix WebSocket reader with fragmented masked messages
Fix WebSocket reader with fragmented masked messages

We didn't have proper coverage for these cases. Add mask and large payload coverage for the reported cases since they would come through as 0 instead of the correct size.

fixes #10763
regressed in #10747



Problem statement:
When the client sends a relatively large bytes data packet to the server, the server will receive zero length message.
### Describe the bug

If the client sends a 20M bytes message, the server will receive zero length message.

The 3.11.16 version works fine.

### To Reproduce

1. Start a WebSocket server, when a request is received, create a WebSocket object "ws = web.WebSocketResponse(max_msg_size = 0)"
2. await ws.prepare(request)
3. await ws.receive_bytes()
4. At this point, if the client sends a 20M bytes message, the server will get zero length message. Setting max_msg_size to 100M yields the same result.

### Expected behavior

Receive correct message.


### Logs/tracebacks

```python-traceback
No exception, but error message.
```

### Python Version

```console
$ python --version
3.11.2
```

### aiohttp Version

```console
$ python -m pip show aiohttp
3.11.17
```

### multidict Version

```console
$ python -m pip show multidict
6.0.4
```

### propcache Version

```console
$ python -m pip show propcache
0.2.0
```

### yarl Version

```console
$ python -m pip show yarl
1.18.3
```

### OS

Linux, debian 12

### Related component

Server

### Additional context

```python
from aiohttp import web

async def Handler(request: web.Request):
    try:
        ws = web.WebSocketResponse(max_msg_size=100*1024*1024)
        await ws.prepare(request)
        while True:
            rawMsg = await ws.receive_bytes()
            print(len(rawMsg))
            await ws.send_bytes(rawMsg)
    except BaseException as e:
        print(e)


if __name__ == "__main__":
    app = web.Application()
    app.add_routes([web.get("/get_task", Handler)])
    web.run_app(app, path="./test_websocket.socket", handle_signals=False)
```

Client

### Additional context
```python
import asyncio
import aiohttp
import pickle

async def Main():
    async with aiohttp.ClientSession(
        connector=aiohttp.UnixConnector("./test_websocket.socket")
    ) as client:
        while True:
            try:
                async with client.ws_connect("ws://xxxxx/get_task", max_msg_size=0) as ws:

                    async def RecvTask():
                        while True:
                            rawMsg = await ws.receive_bytes()
                            print(len(rawMsg))

                    async def SendTask():
                        while True:
                            await ws.send_bytes(pickle.dumps([{"aa": 123}]*10000000))
                            await asyncio.sleep(5)

                    await asyncio.gather(RecvTask(), SendTask())
            except BaseException as e:
                print(f"{e!r}")
                await asyncio.sleep(10)


if __name__ == "__main__":
    asyncio.run(Main())
```

### Code of Conduct

- [x] I agree to follow the aio-libs Code of Conduct

Hints from issue discussion:
Thanks for the report.  I'll take a look in the morning when I get up 
20022767 byte message
confirmed the reproducer. yanked 3.11.17

Will fix it in the morning. I have a good idea whats wrong

Commit message: lint

Patch to review:
diff --git a/CHANGES/10764.bugfix.rst b/CHANGES/10764.bugfix.rst
new file mode 100644
index 00000000000..43a34295353
--- /dev/null
+++ b/CHANGES/10764.bugfix.rst
@@ -0,0 +1,3 @@
+Fixed reading fragmented WebSocket messages when the payload was masked -- by :user:`bdraco`.
+
+The problem first appeared in 3.11.16
diff --git a/aiohttp/_websocket/reader_py.py b/aiohttp/_websocket/reader_py.py
index df322a436dd..257b6d9e1ac 100644
--- a/aiohttp/_websocket/reader_py.py
+++ b/aiohttp/_websocket/reader_py.py
@@ -458,8 +458,7 @@ def _feed_data(self, data: bytes) -> None:
                     self._payload_fragments.append(data_cstr[f_start_pos:f_end_pos])
                     if self._has_mask:
                         assert self._frame_mask is not None
-                        payload_bytearray = bytearray()
-                        payload_bytearray.join(self._payload_fragments)
+                        payload_bytearray = bytearray().join(self._payload_fragments)
                         websocket_mask(self._frame_mask, payload_bytearray)
                         payload = payload_bytearray
                     else:
diff --git a/tests/test_websocket_parser.py b/tests/test_websocket_parser.py
index 6199abae359..c909fe2c080 100644
--- a/tests/test_websocket_parser.py
+++ b/tests/test_websocket_parser.py
@@ -1,5 +1,6 @@
 import asyncio
 import pickle
+import random
 import struct
 from typing import Optional, Union
 from unittest import mock
@@ -7,7 +8,14 @@
 import pytest
 
 from aiohttp._websocket import helpers as _websocket_helpers
-from aiohttp._websocket.helpers import PACK_CLOSE_CODE, PACK_LEN1, PACK_LEN2
+from aiohttp._websocket.helpers import (
+    PACK_CLOSE_CODE,
+    PACK_LEN1,
+    PACK_LEN2,
+    PACK_LEN3,
+    PACK_RANDBITS,
+    websocket_mask,
+)
 from aiohttp._websocket.models import WS_DEFLATE_TRAILING
 from aiohttp._websocket.reader import WebSocketDataQueue
 from aiohttp.base_protocol import BaseProtocol
@@ -52,6 +60,7 @@ def build_frame(
     noheader: bool = False,
     is_fin: bool = True,
     ZLibBackend: Optional[ZLibBackendWrapper] = None,
+    mask: bool = False,
 ) -> bytes:
     # Send a frame over the websocket with message as its payload.
     compress = False
@@ -72,11 +81,21 @@ def build_frame(
     if compress:
         header_first_byte |= 0x40
 
+    mask_bit = 0x80 if mask else 0
+
     if msg_length < 126:
-        header = PACK_LEN1(header_first_byte, msg_length)
+        header = PACK_LEN1(header_first_byte, msg_length | mask_bit)
+    elif msg_length < 65536:
+        header = PACK_LEN2(header_first_byte, 126 | mask_bit, msg_length)
     else:
-        assert msg_length < (1 << 16)
-        header = PACK_LEN2(header_first_byte, 126, msg_length)
+        header = PACK_LEN3(header_first_byte, 127 | mask_bit, msg_length)
+
+    if mask:
+        assert not noheader
+        mask_bytes = PACK_RANDBITS(random.getrandbits(32))
+        message_arr = bytearray(message)
+        websocket_mask(mask_bytes, message_arr)
+        return header + mask_bytes + message_arr
 
     if noheader:
         return message
@@ -125,6 +144,11 @@ def parser(out: WebSocketDataQueue) -> PatchableWebSocketReader:
     return PatchableWebSocketReader(out, 4 * 1024 * 1024)
 
 
+@pytest.fixture()
+def parser_with_large_limit(out: WebSocketDataQueue) -> PatchableWebSocketReader:
+    return PatchableWebSocketReader(out, 1024 * 1024 * 1024)
+
+
 def test_feed_data_remembers_exception(parser: WebSocketReader) -> None:
     """Verify that feed_data remembers an exception was already raised internally."""
     error, data = parser.feed_data(struct.pack("!BB", 0b01100000, 0b00000000))
@@ -352,6 +376,51 @@ def test_fragmentation_header(
     assert res == WSMessageText(data="a", size=1, extra="")
 
 
+def test_large_message(
+    out: WebSocketDataQueue, parser_with_large_limit: PatchableWebSocketReader
+) -> None:
+    large_payload = b"b" * 131072
+    data = build_frame(large_payload, WSMsgType.BINARY)
+    parser_with_large_limit._feed_data(data)
+
+    res = out._buffer[0]
+    assert res == WSMessageBinary(data=large_payload, size=131072, extra="")
+
+
+def test_large_masked_message(
+    out: WebSocketDataQueue, parser_with_large_limit: PatchableWebSocketReader
+) -> None:
+    large_payload = b"b" * 131072
+    data = build_frame(large_payload, WSMsgType.BINARY, mask=True)
+    parser_with_large_limit._feed_data(data)
+
+    res = out._buffer[0]
+    assert res == WSMessageBinary(data=large_payload, size=131072, extra="")
+
+
+def test_fragmented_masked_message(
+    out: WebSocketDataQueue, parser: PatchableWebSocketReader
+) -> None:
+    large_payload = b"b" * 100
+    data = build_frame(large_payload, WSMsgType.BINARY, mask=True)
+    for i in range(len(data)):
+        parser._feed_data(data[i : i + 1])
+
+    res = out._buffer[0]
+    assert res == WSMessageBinary(data=large_payload, size=100, extra="")
+
+
+def test_large_fragmented_masked_message(
+    out: WebSocketDataQueue, parser: PatchableWebSocketReader
+) -> None:
+    large_payload = b"b" * 131072
+    data = build_frame(large_payload, WSMsgType.BINARY, mask=True)
+    for i in range(0, len(data), 16384):
+        parser._feed_data(data[i : i + 16384])
+    res = out._buffer[0]
+    assert res == WSMessageBinary(data=large_payload, size=131072, extra="")
+
+
 def test_continuation(
     out: WebSocketDataQueue, parser: PatchableWebSocketReader
 ) -> None:


Code review comments:
diff_hunk: @@ -458,8 +458,7 @@ def _feed_data(self, data: bytes) -> None:
                     self._payload_fragments.append(data_cstr[f_start_pos:f_end_pos])
                     if self._has_mask:
                         assert self._frame_mask is not None
-                        payload_bytearray = bytearray()
-                        payload_bytearray.join(self._payload_fragments)
+                        payload_bytearray = bytearray().join(self._payload_fragments)
path: aiohttp/_websocket/reader_py.py
text: @author:
Sadly this was an simple refactoring error which I didn't notice because the coverage here wasn't actually testing this case. We now have coverage for this.