Fix WebSocket reader with fragmented masked messages (#10764)

Author: bdraco

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.17

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(b"".join(self._payload_fragments))
                         websocket_mask(self._frame_mask, payload_bytearray)
                         payload = payload_bytearray
                     else:

tests/test_websocket_parser.py

@@ -1,13 +1,21 @@
 import asyncio
 import pickle
+import random
 import struct
 from typing import Optional, Union
 from unittest import mock
 
 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
@@ -352,6 +371,51 @@ def test_fragmentation_header(
     assert res == WSMessageText(data="a", size=1, extra="")
 
 
+def test_large_message(
+    out: WebSocketDataQueue, parser: PatchableWebSocketReader
+) -> None:
+    large_payload = b"b" * 131072
+    data = build_frame(large_payload, WSMsgType.BINARY)
+    parser._feed_data(data)
+
+    res = out._buffer[0]
+    assert res == WSMessageBinary(data=large_payload, size=131072, extra="")
+
+
+def test_large_masked_message(
+    out: WebSocketDataQueue, parser: PatchableWebSocketReader
+) -> None:
+    large_payload = b"b" * 131072
+    data = build_frame(large_payload, WSMsgType.BINARY, mask=True)
+    parser._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: