[pull] master from aio-libs:master

aiohttp/_websocket/reader_c.pxd

@@ -8,6 +8,8 @@ cdef unsigned int READ_PAYLOAD_LENGTH
 cdef unsigned int READ_PAYLOAD_MASK
 cdef unsigned int READ_PAYLOAD
 
+cdef unsigned long long MAX_PAYLOAD_LEN
+
 cdef int OP_CODE_NOT_SET
 cdef int OP_CODE_CONTINUATION
 cdef int OP_CODE_TEXT
@@ -115,6 +117,7 @@ cdef class WebSocketReader:
         fin=bint,
         had_fragments=Py_ssize_t,
         partial_len=Py_ssize_t,
+        frame_len="unsigned long long",
         payload_bytearray=bytearray,
     )
     cpdef void _feed_data(self, bytes data) except *

aiohttp/_websocket/reader_py.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import builtins
+import sys
 from collections import deque
 from typing import Final
 
@@ -33,6 +34,12 @@
 READ_PAYLOAD_MASK = 3
 READ_PAYLOAD = 4
 
+# Largest declared payload length the reader can represent: the compiled
+# reader stores it in a Py_ssize_t, which holds 2**31-1 on the 32-bit builds
+# (the win32 and armv7l wheels) and 2**63-1 everywhere else.
+# TODO: Remove when we drop 32 bit support (and from reader_c.pxd).
+MAX_PAYLOAD_LEN = sys.maxsize
+
 WS_MSG_TYPE_BINARY = WSMsgType.BINARY
 WS_MSG_TYPE_TEXT = WSMsgType.TEXT
 
@@ -453,7 +460,16 @@ def _feed_data(self, data: bytes) -> None:
                 elif len_flag > 126:
                     if data_len - start_pos < 8:
                         break
-                    self._payload_bytes_to_read = UNPACK_LEN3(data, start_pos)[0]
+                    # The declared length is an unsigned 64-bit integer that
+                    # does not necessarily fit _payload_bytes_to_read.
+                    frame_len = UNPACK_LEN3(data, start_pos)[0]
+                    if frame_len > MAX_PAYLOAD_LEN:
+                        raise WebSocketError(
+                            WSCloseCode.MESSAGE_TOO_BIG,
+                            f"Message size {int(frame_len) + len(self._partial)} "
+                            f"exceeds limit {self._max_msg_size or MAX_PAYLOAD_LEN}",
+                        )
+                    self._payload_bytes_to_read = frame_len
                     start_pos += 8
                 else:
                     self._payload_bytes_to_read = len_flag
@@ -468,7 +484,7 @@ def _feed_data(self, data: bytes) -> None:
                 }:
                     # partial_len declared in reader_c.pxd to keep it in C.
                     partial_len = len(self._partial)
-                    # payload_bytes_to_read is a signed 64-bit C value,
+                    # payload_bytes_to_read is a signed Py_ssize_t C value,
                     # use subtraction here to avoid an integer overflow.
                     if self._payload_bytes_to_read >= self._max_msg_size - partial_len:
                         raise WebSocketError(

tests/test_benchmarks_multipart.py

@@ -3,11 +3,11 @@
 import asyncio
 import base64
 from typing import TYPE_CHECKING
-from unittest import mock
 
 import pytest
 from multidict import CIMultiDict
 
+from aiohttp.base_protocol import BaseProtocol
 from aiohttp.hdrs import CONTENT_TRANSFER_ENCODING
 from aiohttp.helpers import DEFAULT_CHUNK_SIZE, HeadersDictProxy
 from aiohttp.multipart import BodyPartReader
@@ -23,10 +23,18 @@
 BASE64_HEADERS: CIMultiDict[str] = CIMultiDict({CONTENT_TRANSFER_ENCODING: "base64"})
 
 
+class _NoFlowControlProtocol(BaseProtocol):
+    """The whole body is fed up front, so there is nothing to pause or resume."""
+
+    def pause_reading(self) -> None:
+        """Swallow pause."""
+
+    def resume_reading(self, resume_parser: bool = True) -> None:
+        """Swallow resume."""
+
+
 def _part(body: bytes, loop: asyncio.AbstractEventLoop) -> BodyPartReader:
-    stream = StreamReader(
-        mock.Mock(_reading_paused=False), DEFAULT_CHUNK_SIZE, loop=loop
-    )
+    stream = StreamReader(_NoFlowControlProtocol(loop), DEFAULT_CHUNK_SIZE, loop=loop)
     stream.feed_data(body)
     stream.feed_eof()
     return BodyPartReader(

tests/test_websocket_parser.py

@@ -2,6 +2,7 @@
 import pickle
 import random
 import struct
+import sys
 import zlib
 from unittest import mock
 
@@ -833,19 +834,19 @@ def test_msg_too_large_declared_length_near_ssize_t_max(
 ) -> None:
     # Regression test: the header-time size check compares
     # `payload_bytes_to_read + len(partial)` against max_msg_size. In the
-    # compiled Cython reader payload_bytes_to_read is a signed 64-bit C
-    # value, and RFC 6455 allows a declared length up to 2**63-1, so a
-    # naive addition can wrap around to a negative number and bypass the
-    # limit entirely once anything is already buffered in `partial`.
+    # compiled Cython reader payload_bytes_to_read is a signed Py_ssize_t C
+    # value, and RFC 6455 allows a declared length up to 2**63-1, so a naive
+    # addition can wrap around to a negative number and bypass the limit
+    # entirely once anything is already buffered in `partial`.
     max_msg_size = 4 * 1024 * 1024
     parser = WebSocketReader(out, max_msg_size, compress=False, decode_text=True)
 
     # Buffer one byte in `partial` via a non-fin fragment.
     first = build_frame(b"a", WSMsgType.TEXT, is_fin=False)
     parser._feed_data(first)
 
-    # Continuation header alone declares the maximum length a signed
-    # 64-bit payload_bytes_to_read can hold: 2**63-1.
+    # Continuation header alone declares the largest length RFC 6455
+    # allows: 2**63-1.
     header = PACK_LEN3(WSMsgType.CONTINUATION, 127, 2**63 - 1)
     with pytest.raises(
         WebSocketError,
@@ -855,6 +856,54 @@ def test_msg_too_large_declared_length_near_ssize_t_max(
     assert ctx.value.code == WSCloseCode.MESSAGE_TOO_BIG
 
 
+@pytest.mark.parametrize(
+    "declared",
+    (sys.maxsize + 1, 2**64 - 1),
+    ids=("ssize_t_max_plus_one", "uint64_max"),
+)
+@pytest.mark.parametrize(
+    "max_msg_size", (4 * 1024 * 1024, 0), ids=("bounded", "unbounded")
+)
+def test_msg_too_large_declared_length_above_ssize_t_max(
+    out: WebSocketDataQueue, max_msg_size: int, declared: int
+) -> None:
+    # RFC 6455 allows an unsigned 64-bit declared length, which does not fit
+    # the signed Py_ssize_t payload_bytes_to_read of the compiled reader:
+    # anything above sys.maxsize (2**31-1 on the 32-bit builds) has to be
+    # rejected before it is stored there, or storing it raises OverflowError
+    # and escapes as an abnormal closure instead of MESSAGE_TOO_BIG. Such a
+    # frame can never be buffered, so an unlimited max_msg_size is checked
+    # too.
+    parser = WebSocketReader(out, max_msg_size, compress=False, decode_text=True)
+
+    header = PACK_LEN3(0x80 | WSMsgType.TEXT, 127, declared)
+    with pytest.raises(
+        WebSocketError,
+        match=rf"^Message size {declared} exceeds limit {max_msg_size or sys.maxsize}$",
+    ) as ctx:
+        parser._feed_data(header)
+    assert ctx.value.code == WSCloseCode.MESSAGE_TOO_BIG
+
+
+def test_msg_too_large_above_ssize_t_max_reports_buffered_fragments(
+    out: WebSocketDataQueue,
+) -> None:
+    # The reported size adds what is already buffered, so the message reads the
+    # same whether the declared length is caught by the Py_ssize_t range check
+    # or by the max_msg_size comparison below it.
+    max_msg_size = 4 * 1024 * 1024
+    parser = WebSocketReader(out, max_msg_size, compress=False, decode_text=True)
+
+    parser._feed_data(build_frame(b"a", WSMsgType.TEXT, is_fin=False))
+
+    header = PACK_LEN3(WSMsgType.CONTINUATION, 127, 2**64 - 1)
+    with pytest.raises(
+        WebSocketError, match=rf"^Message size {2**64} exceeds limit {max_msg_size}$"
+    ) as ctx:
+        parser._feed_data(header)
+    assert ctx.value.code == WSCloseCode.MESSAGE_TOO_BIG
+
+
 @pytest.mark.parametrize(
     "opcode",
     (0x3, 0x4, 0x5, 0x6, 0x7, 0xB, 0xC, 0xD, 0xE, 0xF),