diff --git a/aiohttp/_websocket/reader_c.pxd b/aiohttp/_websocket/reader_c.pxd index 33587831d0c..918fea048ad 100644 --- a/aiohttp/_websocket/reader_c.pxd +++ b/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 * diff --git a/aiohttp/_websocket/reader_py.py b/aiohttp/_websocket/reader_py.py index 40dda00d039..b22884cec62 100644 --- a/aiohttp/_websocket/reader_py.py +++ b/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( diff --git a/tests/test_benchmarks_multipart.py b/tests/test_benchmarks_multipart.py index 4596f00ae7f..6c82aacbcb1 100644 --- a/tests/test_benchmarks_multipart.py +++ b/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( diff --git a/tests/test_websocket_parser.py b/tests/test_websocket_parser.py index a2438d461e2..2cd383cdb77 100644 --- a/tests/test_websocket_parser.py +++ b/tests/test_websocket_parser.py @@ -2,6 +2,7 @@ import pickle import random import struct +import sys import zlib from unittest import mock @@ -833,10 +834,10 @@ 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) @@ -844,8 +845,8 @@ def test_msg_too_large_declared_length_near_ssize_t_max( 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),