Repository: python-websockets/websockets

Issue:
Respecting max_size on fragments but not on message
I don't *think* this is an asyncio issue/question as I think it's above that level.
What I've got is a websocket receiving and consuming fragments from its peer just as it's indicated to in the docs:
```python
  async for frag in web_skt.recv_streaming():
    consume_fragment(frag)
```
What I don't see (please correct me if I'm missing this) is a way to limit the size of the fragments (so that no peer can use up all of my memory) without also putting a limit on the total message size.
That is, I can certainly set the `max_size` arg in `websockets.asyncio.server.serve()` but when the sum total of the size of the fragments (that I have already consumed) reaches that limit, then an exception is raised. Since it's a "streaming" mode, seems like the total size of a message could be large.
Setting `max_size=None` certainly bypasses that issue, but then I have no way to limit the size of fragments before they (potentially) exhaust my memory.
Or, am I missing something?

Discussion:
@aaugustin: `max_size` limits the sum of the size of all fragments belonging to the same message.

This is tested here: https://github.com/python-websockets/websockets/blob/6304d108a15d79763a57ecb073860a624e649878/tests/test_protocol.py#L387-L404
@aaugustin: There's no built-in way to limit the size of each individual fragment. I'm not convinced that this would be greatly useful in general. If you want it, you can simply do:


```python
async for frag in web_skt.recv_streaming():
    if len(frag) > 100_000:
        raise ValueError("fragment too large")
    consume_fragment(frag)
```@aaugustin: I'm skeptical because RFC 6455 makes it clear that any intermediary can split or reassemble fragments in any way they want. Enforcing limits at that level seems sketchy.@aaugustin: Here's a more structured recap of the situation.

**Legacy implementation**

Since version 3.2 (exactly 9 years ago!) websockets has had `max_size` and `max_queue` parameters to limit the size of the buffer of parsed messages. It provides these two parameters because some servers receive very small messages at a high rate and can benefit from a larger buffer e.g. [crypto](https://websockets.readthedocs.io/en/stable/project/support.html#cryptocurrency-users) trackers... while others receive larger messages at a slower rate; I've talked to someone building a remote desktop over websockets, where each message contained a video frame!

The legacy implementation contained a buffer of messages after reassembling fragmented messages. The maximum size of the buffer was literally `4 * max_size * max_queue`. There's a x4 because each UTF-8 encoded byte can expand into a character that Python stores on 4 bytes in memory and we test the length of the UTF-8 encoded message — we'd have to decode it to tell its length but we cannot process before applying security limits.

The legacy implementation didn't let you access frames with `recv_streaming` — this was one of the drivers for rewriting it — nor did it provide any control at this level.

**Current implementations**

The new implementations kept the logic of limiting the size of frames with `max_size`. When fragmentation is enabled, they keep track of how much room remains and check that incoming fragments don't exceed that limit. This behavior is necessary because fragmentation is uncommon and most users aren't aware that it exists. I don't even expect most users to take 5 mins to think about the messages they expect to receive and configure a limit properly...

However, they changed the logic of limiting the size of the buffer with `max_queue`: indeed, in order to support `recv_streaming`, the buffer is now a buffer of frames instead of buffer of messages. This change leaked into the public API. As a consequence:

- without fragmentation, the limit is still `max_size * max_queue`
- with fragmentation, the limit is effectively lower — it's roughly `max_size * max_queue / fragments_per_message`

**Challenge**

We have a conflict here between two goals:

1. Keeping the semantics of `max_size` simple to increase the chances that many users configure a reasonable value.
2. Controlling memory usage accurately — the knobs don't match the actual memory usage anymore.

Changing `max_size` to be the size of a fragment would break the first goal because users could receive an arbitrarily large message in `recv()` if an attacker sent a large message in small fragments. That would be unacceptable.

The current implementation fails the second goal in the case of messages with many small fragments, as described in this issue. When I wrote it, I treated fragmentation as a transport-level concern. I considered that the receiver should always be able to handle the full message.

Upon further thought, this seems incorrect. For example, if you want to transfer a large file in a WebSocket message (why not?) you can send it chunk-by-chunk at one end and write it to disk chunk-by-chunk at the other end. That should work even if it's a 10GB file and you don't want to hold it in memory.

**Potential solution**

I'm considering supporting `max_size=(max_message_size, max_fragment_size)`. This would provide more control for those who want it without increasing the API surface. TBD if this can be implemented with reasonable complexity vs. the benefits.@aaugustin: Hello, would you be able to test if #1622 does what you want?@gmabey: I was initially unhappy with the concept of overloading `max_size` with an tuple/iterable type, but I have to admit your rationale is very good; the frequency of this use case warrants a more arcane form.
I hope to test #1622 later this week.@aaugustin: Yes, I know, I've been hesitating between overloading adding a new argument... I went for overloading because:

- there are too many arguments already...
- having `max_size` + `max_fragment_size` would look weird and I didn't want to rename `max_size` to `max_message_size`.

Solution patch:
diff --git a/docs/faq/misc.rst b/docs/faq/misc.rst
index 3b51060..6348271 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 12fc8c3..5821cbc 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 0e4eba3..63cd2be 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 @@ class 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,
         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 1b51e47..61c300d 100644
--- a/src/websockets/asyncio/connection.py
+++ b/src/websockets/asyncio/connection.py
@@ -63,9 +63,11 @@ class Connection(asyncio.Protocol):
         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 ec7fc43..7458692 100644
--- a/src/websockets/asyncio/server.py
+++ b/src/websockets/asyncio/server.py
@@ -644,7 +644,9 @@ class serve:
         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 @@ class 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,
         write_limit: int | tuple[int, int | None] = 2**15,
         # Logging
diff --git a/src/websockets/client.py b/src/websockets/client.py
index 9ea21c3..0fbcda6 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 @@ class ClientProtocol(Protocol):
         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 ab1a15c..a88deaa 100644
--- a/src/websockets/exceptions.py
+++ b/src/websockets/exceptions.py
@@ -394,11 +394,11 @@ class PayloadTooBig(WebSocketException):
         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 @@ class PayloadTooBig(WebSocketException):
             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 @@ class PayloadTooBig(WebSocketException):
             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 bc64a21..c2522b8 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 @@ class Protocol:
         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 @@ class Protocol:
         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 @@ class Protocol:
                     # 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 @@ class Protocol:
             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 @@ class Protocol:
 
         """
         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 @@ class Protocol:
                 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 1744412..de2c635 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 @@ class ServerProtocol(Protocol):
             | 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 fd0ccb6..bfaef61 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 8b9e062..bedbf4d 100644
--- a/src/websockets/sync/connection.py
+++ b/src/websockets/sync/connection.py
@@ -59,9 +59,11 @@ class Connection:
         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 efb40a7..cb554c2 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 serve(
         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


Test patch:
diff --git a/tests/test_protocol.py b/tests/test_protocol.py
index 9e2d650..3cad8b1 100644
--- a/tests/test_protocol.py
+++ b/tests/test_protocol.py
@@ -365,8 +365,26 @@ class TextTests(ProtocolTestCase):
             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 @@ class TextTests(ProtocolTestCase):
             "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,