Support separate limits for messages and fragments.

Fix #1602.

Author: aaugustin

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

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

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

src/websockets/asyncio/connection.py

@@ -64,8 +64,9 @@ def __init__(
         self.ping_timeout = ping_timeout
         self.close_timeout = close_timeout
         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

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

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__(

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):

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

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__(

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

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

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,