diff --git a/docs/faq/misc.rst b/docs/faq/misc.rst index 3b5106006..634827186 100644 --- a/docs/faq/misc.rst +++ b/docs/faq/misc.rst @@ -22,6 +22,7 @@ you must disable: * Compression: set ``compression=None`` * Keepalive: set ``ping_interval=None`` +* Limits: set ``max_size=None`` * UTF-8 decoding: send ``bytes`` rather than ``str`` Then, please consider whether websockets is the bottleneck of the performance diff --git a/docs/project/changelog.rst b/docs/project/changelog.rst index 12fc8c32e..5821cbc54 100644 --- a/docs/project/changelog.rst +++ b/docs/project/changelog.rst @@ -35,6 +35,7 @@ notice. Improvements ............ +* Allowed setting separate limits for messages and fragments with ``max_size``. * Added support for HTTP/1.0 proxies. 15.0.1 diff --git a/src/websockets/asyncio/client.py b/src/websockets/asyncio/client.py index 0e4eba329..63cd2be2e 100644 --- a/src/websockets/asyncio/client.py +++ b/src/websockets/asyncio/client.py @@ -236,7 +236,9 @@ class connect: close_timeout: Timeout for closing the connection in seconds. :obj:`None` disables the timeout. max_size: Maximum size of incoming messages in bytes. - :obj:`None` disables the limit. + :obj:`None` disables the limit. You may pass a ``(max_message_size, + max_fragment_size)`` tuple to set different limits for messages and + fragments when you expect long messages sent in short fragments. max_queue: High-water mark of the buffer where frames are received. It defaults to 16 frames. The low-water mark defaults to ``max_queue // 4``. You may pass a ``(high, low)`` tuple to set the high-water @@ -314,7 +316,7 @@ def __init__( ping_timeout: float | None = 20, close_timeout: float | None = 10, # Limits - max_size: int | None = 2**20, + max_size: int | None | tuple[int | None, int | None] = 2**20, max_queue: int | None | tuple[int | None, int | None] = 16, write_limit: int | tuple[int, int | None] = 2**15, # Logging diff --git a/src/websockets/asyncio/connection.py b/src/websockets/asyncio/connection.py index 1b51e4791..61c300d63 100644 --- a/src/websockets/asyncio/connection.py +++ b/src/websockets/asyncio/connection.py @@ -63,9 +63,11 @@ def __init__( self.ping_interval = ping_interval self.ping_timeout = ping_timeout self.close_timeout = close_timeout + self.max_queue: tuple[int | None, int | None] if isinstance(max_queue, int) or max_queue is None: - max_queue = (max_queue, None) - self.max_queue = max_queue + self.max_queue = (max_queue, None) + else: + self.max_queue = max_queue if isinstance(write_limit, int): write_limit = (write_limit, None) self.write_limit = write_limit diff --git a/src/websockets/asyncio/server.py b/src/websockets/asyncio/server.py index ec7fc4383..745869299 100644 --- a/src/websockets/asyncio/server.py +++ b/src/websockets/asyncio/server.py @@ -644,7 +644,9 @@ def handler(websocket): close_timeout: Timeout for closing connections in seconds. :obj:`None` disables the timeout. max_size: Maximum size of incoming messages in bytes. - :obj:`None` disables the limit. + :obj:`None` disables the limit. You may pass a ``(max_message_size, + max_fragment_size)`` tuple to set different limits for messages and + fragments when you expect long messages sent in short fragments. max_queue: High-water mark of the buffer where frames are received. It defaults to 16 frames. The low-water mark defaults to ``max_queue // 4``. You may pass a ``(high, low)`` tuple to set the high-water @@ -719,7 +721,7 @@ def __init__( ping_timeout: float | None = 20, close_timeout: float | None = 10, # Limits - max_size: int | None = 2**20, + max_size: int | None | tuple[int | None, int | None] = 2**20, max_queue: int | None | tuple[int | None, int | None] = 16, write_limit: int | tuple[int, int | None] = 2**15, # Logging diff --git a/src/websockets/client.py b/src/websockets/client.py index 9ea21c39c..0fbcda60c 100644 --- a/src/websockets/client.py +++ b/src/websockets/client.py @@ -60,8 +60,10 @@ class ClientProtocol(Protocol): subprotocols: List of supported subprotocols, in order of decreasing preference. state: Initial state of the WebSocket connection. - max_size: Maximum size of incoming messages in bytes; - :obj:`None` disables the limit. + max_size: Maximum size of incoming messages in bytes. + :obj:`None` disables the limit. You may pass a ``(max_message_size, + max_fragment_size)`` tuple to set different limits for messages and + fragments when you expect long messages sent in short fragments. logger: Logger for this connection; defaults to ``logging.getLogger("websockets.client")``; see the :doc:`logging guide <../../topics/logging>` for details. @@ -76,7 +78,7 @@ def __init__( extensions: Sequence[ClientExtensionFactory] | None = None, subprotocols: Sequence[Subprotocol] | None = None, state: State = CONNECTING, - max_size: int | None = 2**20, + max_size: int | None | tuple[int | None, int | None] = 2**20, logger: LoggerLike | None = None, ) -> None: super().__init__( diff --git a/src/websockets/exceptions.py b/src/websockets/exceptions.py index ab1a15ca8..a88deaa66 100644 --- a/src/websockets/exceptions.py +++ b/src/websockets/exceptions.py @@ -394,11 +394,11 @@ def __init__( self, size_or_message: int | None | str, max_size: int | None = None, - cur_size: int | None = None, + current_size: int | None = None, ) -> None: if isinstance(size_or_message, str): assert max_size is None - assert cur_size is None + assert current_size is None warnings.warn( # deprecated in 14.0 - 2024-11-09 "PayloadTooBig(message) is deprecated; " "change to PayloadTooBig(size, max_size)", @@ -410,8 +410,8 @@ def __init__( self.size: int | None = size_or_message assert max_size is not None self.max_size: int = max_size - self.cur_size: int | None = None - self.set_current_size(cur_size) + self.current_size: int | None = None + self.set_current_size(current_size) def __str__(self) -> str: if self.message is not None: @@ -420,16 +420,16 @@ def __str__(self) -> str: message = "frame " if self.size is not None: message += f"with {self.size} bytes " - if self.cur_size is not None: - message += f"after reading {self.cur_size} bytes " + if self.current_size is not None: + message += f"after reading {self.current_size} bytes " message += f"exceeds limit of {self.max_size} bytes" return message - def set_current_size(self, cur_size: int | None) -> None: - assert self.cur_size is None - if cur_size is not None: - self.max_size += cur_size - self.cur_size = cur_size + def set_current_size(self, current_size: int | None) -> None: + assert self.current_size is None + if current_size is not None: + self.max_size += current_size + self.current_size = current_size class InvalidState(WebSocketException, AssertionError): diff --git a/src/websockets/protocol.py b/src/websockets/protocol.py index bc64a216a..c2522b8a5 100644 --- a/src/websockets/protocol.py +++ b/src/websockets/protocol.py @@ -77,8 +77,10 @@ class Protocol: Args: side: :attr:`~Side.CLIENT` or :attr:`~Side.SERVER`. state: Initial state of the WebSocket connection. - max_size: Maximum size of incoming messages in bytes; - :obj:`None` disables the limit. + max_size: Maximum size of incoming messages in bytes. + :obj:`None` disables the limit. You may pass a ``(max_message_size, + max_fragment_size)`` tuple to set different limits for messages and + fragments when you expect long messages sent in short fragments. logger: Logger for this connection; depending on ``side``, defaults to ``logging.getLogger("websockets.client")`` or ``logging.getLogger("websockets.server")``; @@ -91,7 +93,7 @@ def __init__( side: Side, *, state: State = OPEN, - max_size: int | None = 2**20, + max_size: tuple[int | None, int | None] | int | None = 2**20, logger: LoggerLike | None = None, ) -> None: # Unique identifier. For logs. @@ -114,11 +116,14 @@ def __init__( self.state = state # Maximum size of incoming messages in bytes. - self.max_size = max_size + if isinstance(max_size, int) or max_size is None: + self.max_message_size, self.max_fragment_size = max_size, None + else: + self.max_message_size, self.max_fragment_size = max_size # Current size of incoming message in bytes. Only set while reading a # fragmented message i.e. a data frames with the FIN bit not set. - self.cur_size: int | None = None + self.current_size: int | None = None # True while sending a fragmented message i.e. a data frames with the # FIN bit not set. @@ -578,12 +583,19 @@ def parse(self) -> Generator[None]: # connection isn't closed cleanly. raise EOFError("unexpected end of stream") - if self.max_size is None: - max_size = None - elif self.cur_size is None: - max_size = self.max_size - else: - max_size = self.max_size - self.cur_size + max_size = None + + if self.max_message_size is not None: + if self.current_size is None: + max_size = self.max_message_size + else: + max_size = self.max_message_size - self.current_size + + if self.max_fragment_size is not None: + if max_size is None: + max_size = self.max_fragment_size + else: + max_size = min(max_size, self.max_fragment_size) # During a normal closure, execution ends here on the next # iteration of the loop after receiving a close frame. At @@ -613,7 +625,7 @@ def parse(self) -> Generator[None]: self.parser_exc = exc except PayloadTooBig as exc: - exc.set_current_size(self.cur_size) + exc.set_current_size(self.current_size) self.fail(CloseCode.MESSAGE_TOO_BIG, str(exc)) self.parser_exc = exc @@ -664,18 +676,18 @@ def recv_frame(self, frame: Frame) -> None: """ if frame.opcode is OP_TEXT or frame.opcode is OP_BINARY: - if self.cur_size is not None: + if self.current_size is not None: raise ProtocolError("expected a continuation frame") if not frame.fin: - self.cur_size = len(frame.data) + self.current_size = len(frame.data) elif frame.opcode is OP_CONT: - if self.cur_size is None: + if self.current_size is None: raise ProtocolError("unexpected continuation frame") if frame.fin: - self.cur_size = None + self.current_size = None else: - self.cur_size += len(frame.data) + self.current_size += len(frame.data) elif frame.opcode is OP_PING: # 5.5.2. Ping: "Upon receipt of a Ping frame, an endpoint MUST @@ -696,7 +708,7 @@ def recv_frame(self, frame: Frame) -> None: assert self.close_sent is not None self.close_rcvd_then_sent = False - if self.cur_size is not None: + if self.current_size is not None: raise ProtocolError("incomplete fragmented message") # 5.5.1 Close: "If an endpoint receives a Close frame and did diff --git a/src/websockets/server.py b/src/websockets/server.py index 174441203..de2c63548 100644 --- a/src/websockets/server.py +++ b/src/websockets/server.py @@ -65,8 +65,10 @@ class ServerProtocol(Protocol): signature as the :meth:`select_subprotocol` method, including a :class:`ServerProtocol` instance as first argument. state: Initial state of the WebSocket connection. - max_size: Maximum size of incoming messages in bytes; - :obj:`None` disables the limit. + max_size: Maximum size of incoming messages in bytes. + :obj:`None` disables the limit. You may pass a ``(max_message_size, + max_fragment_size)`` tuple to set different limits for messages and + fragments when you expect long messages sent in short fragments. logger: Logger for this connection; defaults to ``logging.getLogger("websockets.server")``; see the :doc:`logging guide <../../topics/logging>` for details. @@ -87,7 +89,7 @@ def __init__( | None ) = None, state: State = CONNECTING, - max_size: int | None = 2**20, + max_size: int | None | tuple[int | None, int | None] = 2**20, logger: LoggerLike | None = None, ) -> None: super().__init__( diff --git a/src/websockets/sync/client.py b/src/websockets/sync/client.py index fd0ccb6d1..bfaef61eb 100644 --- a/src/websockets/sync/client.py +++ b/src/websockets/sync/client.py @@ -150,7 +150,7 @@ def connect( ping_timeout: float | None = 20, close_timeout: float | None = 10, # Limits - max_size: int | None = 2**20, + max_size: int | None | tuple[int | None, int | None] = 2**20, max_queue: int | None | tuple[int | None, int | None] = 16, # Logging logger: LoggerLike | None = None, @@ -210,7 +210,9 @@ def connect( close_timeout: Timeout for closing the connection in seconds. :obj:`None` disables the timeout. max_size: Maximum size of incoming messages in bytes. - :obj:`None` disables the limit. + :obj:`None` disables the limit. You may pass a ``(max_message_size, + max_fragment_size)`` tuple to set different limits for messages and + fragments when you expect long messages sent in short fragments. max_queue: High-water mark of the buffer where frames are received. It defaults to 16 frames. The low-water mark defaults to ``max_queue // 4``. You may pass a ``(high, low)`` tuple to set the high-water diff --git a/src/websockets/sync/connection.py b/src/websockets/sync/connection.py index 8b9e06257..bedbf4def 100644 --- a/src/websockets/sync/connection.py +++ b/src/websockets/sync/connection.py @@ -59,9 +59,11 @@ def __init__( self.ping_interval = ping_interval self.ping_timeout = ping_timeout self.close_timeout = close_timeout + self.max_queue: tuple[int | None, int | None] if isinstance(max_queue, int) or max_queue is None: - max_queue = (max_queue, None) - self.max_queue = max_queue + self.max_queue = (max_queue, None) + else: + self.max_queue = max_queue # Inject reference to this instance in the protocol's logger. self.protocol.logger = logging.LoggerAdapter( diff --git a/src/websockets/sync/server.py b/src/websockets/sync/server.py index efb40a7f4..cb554c21f 100644 --- a/src/websockets/sync/server.py +++ b/src/websockets/sync/server.py @@ -365,7 +365,7 @@ def serve( ping_timeout: float | None = 20, close_timeout: float | None = 10, # Limits - max_size: int | None = 2**20, + max_size: int | None | tuple[int | None, int | None] = 2**20, max_queue: int | None | tuple[int | None, int | None] = 16, # Logging logger: LoggerLike | None = None, @@ -450,7 +450,9 @@ def handler(websocket): close_timeout: Timeout for closing connections in seconds. :obj:`None` disables the timeout. max_size: Maximum size of incoming messages in bytes. - :obj:`None` disables the limit. + :obj:`None` disables the limit. You may pass a ``(max_message_size, + max_fragment_size)`` tuple to set different limits for messages and + fragments when you expect long messages sent in short fragments. max_queue: High-water mark of the buffer where frames are received. It defaults to 16 frames. The low-water mark defaults to ``max_queue // 4``. You may pass a ``(high, low)`` tuple to set the high-water diff --git a/tests/test_protocol.py b/tests/test_protocol.py index 9e2d65041..3cad8b1f2 100644 --- a/tests/test_protocol.py +++ b/tests/test_protocol.py @@ -365,8 +365,26 @@ def test_server_receives_fragmented_text(self): Frame(OP_CONT, "😀".encode()[2:]), ) - def test_client_receives_fragmented_text_over_size_limit(self): - client = Protocol(CLIENT, max_size=3) + def test_client_receives_fragmented_text_over_fragment_size_limit(self): + client = Protocol(CLIENT, max_size=(None, 3)) + client.receive_data(b"\x01\x04\xf0\x9f\x98\x80") + self.assertIsInstance(client.parser_exc, PayloadTooBig) + self.assertEqual( + str(client.parser_exc), + "frame with 4 bytes exceeds limit of 3 bytes", + ) + + def test_server_receives_fragmented_text_over_fragment_size_limit(self): + server = Protocol(SERVER, max_size=(None, 3)) + server.receive_data(b"\x01\x84\x00\x00\x00\x00\xf0\x9f\x98\x80") + self.assertIsInstance(server.parser_exc, PayloadTooBig) + self.assertEqual( + str(server.parser_exc), + "frame with 4 bytes exceeds limit of 3 bytes", + ) + + def test_client_receives_fragmented_text_over_message_size_limit(self): + client = Protocol(CLIENT, max_size=(3, 2)) client.receive_data(b"\x01\x02\xf0\x9f") self.assertFrameReceived( client, @@ -384,8 +402,8 @@ def test_client_receives_fragmented_text_over_size_limit(self): "frame with 2 bytes after reading 2 bytes exceeds limit of 3 bytes", ) - def test_server_receives_fragmented_text_over_size_limit(self): - server = Protocol(SERVER, max_size=3) + def test_server_receives_fragmented_text_over_message_size_limit(self): + server = Protocol(SERVER, max_size=(3, 2)) server.receive_data(b"\x01\x82\x00\x00\x00\x00\xf0\x9f") self.assertFrameReceived( server,