Report size correctly in PayloadTooBig.

Previously, it was reported incorrectly for fragmented messages.

Fix #1522.

Author: aaugustin

docs/project/changelog.rst

@@ -67,15 +67,6 @@ Backwards-incompatible changes
     Aliases for deprecated API were removed from ``__all__``. As a consequence,
     they cannot be imported e.g. with ``from websockets import *`` anymore.
 
-.. admonition:: :attr:`Frame.data <frames.Frame.data>` is now a bytes-like object.
-    :class: note
-
-    In addition to :class:`bytes`, it may be a :class:`bytearray` or a
-    :class:`memoryview`.
-
-    If you wrote an :class:`extension <extensions.Extension>` that relies on
-    methods not provided by these new types, you may need to update your code.
-
 .. admonition:: Several API raise :exc:`ValueError` instead of :exc:`TypeError`
     on invalid arguments.
     :class: note
@@ -88,6 +79,26 @@ Backwards-incompatible changes
     raise :exc:`ValueError` when a required argument isn't provided or an
     argument that is incompatible with others is provided.
 
+.. admonition:: :attr:`Frame.data <frames.Frame.data>` is now a bytes-like object.
+    :class: note
+
+    In addition to :class:`bytes`, it may be a :class:`bytearray` or a
+    :class:`memoryview`.
+
+    If you wrote an :class:`extension <extensions.Extension>` that relies on
+    methods not provided by these new types, you may need to update your code.
+
+.. admonition:: The signature of :exc:`~exceptions.PayloadTooBig` changed.
+    :class: note
+
+    If you wrote an extension that raises :exc:`~exceptions.PayloadTooBig` in
+    :meth:`~extensions.Extension.decode`, for example, you must replace::
+
+        PayloadTooBig(f"over size limit ({size} > {max_size} bytes)")
+
+    with::
+
+        PayloadTooBig(size, max_size)
 
 New features
 ............
@@ -105,6 +116,8 @@ Improvements
 
 * Sending or receiving large compressed messages is now faster.
 
+* Errors when a fragmented message is too large are clearer.
+
 .. _13.1:
 
 13.1

src/websockets/exceptions.py

@@ -334,6 +334,47 @@ class PayloadTooBig(WebSocketException):
 
     """
 
+    def __init__(
+        self,
+        size_or_message: int | None | str,
+        max_size: int | None = None,
+        cur_size: int | None = None,
+    ) -> None:
+        if isinstance(size_or_message, str):
+            assert max_size is None
+            assert cur_size is None
+            warnings.warn(  # deprecated in 14.0
+                "PayloadTooBig(message) is deprecated; "
+                "change to PayloadTooBig(size, max_size)",
+                DeprecationWarning,
+            )
+            self.message: str | None = size_or_message
+        else:
+            self.message = None
+            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)
+
+    def __str__(self) -> str:
+        if self.message is not None:
+            return self.message
+        else:
+            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 "
+            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
+
 
 class InvalidState(WebSocketException, AssertionError):
     """

src/websockets/extensions/permessage_deflate.py

@@ -139,7 +139,8 @@ def decode(
         try:
             data = self.decoder.decompress(data, max_length)
             if self.decoder.unconsumed_tail:
-                raise PayloadTooBig(f"over size limit (? > {max_size} bytes)")
+                assert max_size is not None  # help mypy
+                raise PayloadTooBig(None, max_size)
             if frame.fin and len(frame.data) >= 2044:
                 # This cannot generate additional data.
                 self.decoder.decompress(_EMPTY_UNCOMPRESSED_BLOCK)

src/websockets/frames.py

@@ -252,7 +252,7 @@ def parse(
             data = yield from read_exact(8)
             (length,) = struct.unpack("!Q", data)
         if max_size is not None and length > max_size:
-            raise PayloadTooBig(f"over size limit ({length} > {max_size} bytes)")
+            raise PayloadTooBig(length, max_size)
         if mask:
             mask_bytes = yield from read_exact(4)
 

src/websockets/legacy/framing.py

@@ -93,7 +93,7 @@ async def read(
             data = await reader(8)
             (length,) = struct.unpack("!Q", data)
         if max_size is not None and length > max_size:
-            raise PayloadTooBig(f"over size limit ({length} > {max_size} bytes)")
+            raise PayloadTooBig(length, max_size)
         if mask:
             mask_bits = await reader(4)
 

src/websockets/protocol.py

@@ -587,6 +587,7 @@ def parse(self) -> Generator[None]:
             self.parser_exc = exc
 
         except PayloadTooBig as exc:
+            exc.set_current_size(self.cur_size)
             self.fail(CloseCode.MESSAGE_TOO_BIG, str(exc))
             self.parser_exc = exc
 
@@ -639,9 +640,7 @@ 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:
                 raise ProtocolError("expected a continuation frame")
-            if frame.fin:
-                self.cur_size = None
-            else:
+            if not frame.fin:
                 self.cur_size = len(frame.data)
 
         elif frame.opcode is OP_CONT:

tests/test_exceptions.py

@@ -160,8 +160,16 @@ def test_str(self):
                 "invalid opcode: 7",
             ),
             (
-                PayloadTooBig("payload length exceeds limit: 2 > 1 bytes"),
-                "payload length exceeds limit: 2 > 1 bytes",
+                PayloadTooBig(None, 4),
+                "frame exceeds limit of 4 bytes",
+            ),
+            (
+                PayloadTooBig(8, 4),
+                "frame with 8 bytes exceeds limit of 4 bytes",
+            ),
+            (
+                PayloadTooBig(8, 4, 12),
+                "frame with 8 bytes after reading 12 bytes exceeds limit of 16 bytes",
             ),
             (
                 InvalidState("WebSocket connection isn't established yet"),
@@ -202,3 +210,11 @@ def test_connection_closed_attributes_deprecation_defaults(self):
             "use Protocol.close_reason or ConnectionClosed.rcvd.reason"
         ):
             self.assertEqual(exception.reason, "")
+
+    def test_payload_too_big_with_message(self):
+        with self.assertDeprecationWarning(
+            "PayloadTooBig(message) is deprecated; "
+            "change to PayloadTooBig(size, max_size)",
+        ):
+            exc = PayloadTooBig("payload length exceeds limit: 2 > 1 bytes")
+        self.assertEqual(str(exc), "payload length exceeds limit: 2 > 1 bytes")

tests/test_protocol.py

@@ -265,18 +265,28 @@ def test_client_receives_text_over_size_limit(self):
         client = Protocol(CLIENT, max_size=3)
         client.receive_data(b"\x81\x04\xf0\x9f\x98\x80")
         self.assertIsInstance(client.parser_exc, PayloadTooBig)
-        self.assertEqual(str(client.parser_exc), "over size limit (4 > 3 bytes)")
+        self.assertEqual(
+            str(client.parser_exc),
+            "frame with 4 bytes exceeds limit of 3 bytes",
+        )
         self.assertConnectionFailing(
-            client, CloseCode.MESSAGE_TOO_BIG, "over size limit (4 > 3 bytes)"
+            client,
+            CloseCode.MESSAGE_TOO_BIG,
+            "frame with 4 bytes exceeds limit of 3 bytes",
         )
 
     def test_server_receives_text_over_size_limit(self):
         server = Protocol(SERVER, max_size=3)
         server.receive_data(b"\x81\x84\x00\x00\x00\x00\xf0\x9f\x98\x80")
         self.assertIsInstance(server.parser_exc, PayloadTooBig)
-        self.assertEqual(str(server.parser_exc), "over size limit (4 > 3 bytes)")
+        self.assertEqual(
+            str(server.parser_exc),
+            "frame with 4 bytes exceeds limit of 3 bytes",
+        )
         self.assertConnectionFailing(
-            server, CloseCode.MESSAGE_TOO_BIG, "over size limit (4 > 3 bytes)"
+            server,
+            CloseCode.MESSAGE_TOO_BIG,
+            "frame with 4 bytes exceeds limit of 3 bytes",
         )
 
     def test_client_receives_text_without_size_limit(self):
@@ -363,9 +373,14 @@ def test_client_receives_fragmented_text_over_size_limit(self):
         )
         client.receive_data(b"\x80\x02\x98\x80")
         self.assertIsInstance(client.parser_exc, PayloadTooBig)
-        self.assertEqual(str(client.parser_exc), "over size limit (2 > 1 bytes)")
+        self.assertEqual(
+            str(client.parser_exc),
+            "frame with 2 bytes after reading 2 bytes exceeds limit of 3 bytes",
+        )
         self.assertConnectionFailing(
-            client, CloseCode.MESSAGE_TOO_BIG, "over size limit (2 > 1 bytes)"
+            client,
+            CloseCode.MESSAGE_TOO_BIG,
+            "frame with 2 bytes after reading 2 bytes exceeds limit of 3 bytes",
         )
 
     def test_server_receives_fragmented_text_over_size_limit(self):
@@ -377,9 +392,14 @@ def test_server_receives_fragmented_text_over_size_limit(self):
         )
         server.receive_data(b"\x80\x82\x00\x00\x00\x00\x98\x80")
         self.assertIsInstance(server.parser_exc, PayloadTooBig)
-        self.assertEqual(str(server.parser_exc), "over size limit (2 > 1 bytes)")
+        self.assertEqual(
+            str(server.parser_exc),
+            "frame with 2 bytes after reading 2 bytes exceeds limit of 3 bytes",
+        )
         self.assertConnectionFailing(
-            server, CloseCode.MESSAGE_TOO_BIG, "over size limit (2 > 1 bytes)"
+            server,
+            CloseCode.MESSAGE_TOO_BIG,
+            "frame with 2 bytes after reading 2 bytes exceeds limit of 3 bytes",
         )
 
     def test_client_receives_fragmented_text_without_size_limit(self):
@@ -533,18 +553,28 @@ def test_client_receives_binary_over_size_limit(self):
         client = Protocol(CLIENT, max_size=3)
         client.receive_data(b"\x82\x04\x01\x02\xfe\xff")
         self.assertIsInstance(client.parser_exc, PayloadTooBig)
-        self.assertEqual(str(client.parser_exc), "over size limit (4 > 3 bytes)")
+        self.assertEqual(
+            str(client.parser_exc),
+            "frame with 4 bytes exceeds limit of 3 bytes",
+        )
         self.assertConnectionFailing(
-            client, CloseCode.MESSAGE_TOO_BIG, "over size limit (4 > 3 bytes)"
+            client,
+            CloseCode.MESSAGE_TOO_BIG,
+            "frame with 4 bytes exceeds limit of 3 bytes",
         )
 
     def test_server_receives_binary_over_size_limit(self):
         server = Protocol(SERVER, max_size=3)
         server.receive_data(b"\x82\x84\x00\x00\x00\x00\x01\x02\xfe\xff")
         self.assertIsInstance(server.parser_exc, PayloadTooBig)
-        self.assertEqual(str(server.parser_exc), "over size limit (4 > 3 bytes)")
+        self.assertEqual(
+            str(server.parser_exc),
+            "frame with 4 bytes exceeds limit of 3 bytes",
+        )
         self.assertConnectionFailing(
-            server, CloseCode.MESSAGE_TOO_BIG, "over size limit (4 > 3 bytes)"
+            server,
+            CloseCode.MESSAGE_TOO_BIG,
+            "frame with 4 bytes exceeds limit of 3 bytes",
         )
 
     def test_client_sends_fragmented_binary(self):
@@ -615,9 +645,14 @@ def test_client_receives_fragmented_binary_over_size_limit(self):
         )
         client.receive_data(b"\x80\x02\xfe\xff")
         self.assertIsInstance(client.parser_exc, PayloadTooBig)
-        self.assertEqual(str(client.parser_exc), "over size limit (2 > 1 bytes)")
+        self.assertEqual(
+            str(client.parser_exc),
+            "frame with 2 bytes after reading 2 bytes exceeds limit of 3 bytes",
+        )
         self.assertConnectionFailing(
-            client, CloseCode.MESSAGE_TOO_BIG, "over size limit (2 > 1 bytes)"
+            client,
+            CloseCode.MESSAGE_TOO_BIG,
+            "frame with 2 bytes after reading 2 bytes exceeds limit of 3 bytes",
         )
 
     def test_server_receives_fragmented_binary_over_size_limit(self):
@@ -629,9 +664,14 @@ def test_server_receives_fragmented_binary_over_size_limit(self):
         )
         server.receive_data(b"\x80\x82\x00\x00\x00\x00\xfe\xff")
         self.assertIsInstance(server.parser_exc, PayloadTooBig)
-        self.assertEqual(str(server.parser_exc), "over size limit (2 > 1 bytes)")
+        self.assertEqual(
+            str(server.parser_exc),
+            "frame with 2 bytes after reading 2 bytes exceeds limit of 3 bytes",
+        )
         self.assertConnectionFailing(
-            server, CloseCode.MESSAGE_TOO_BIG, "over size limit (2 > 1 bytes)"
+            server,
+            CloseCode.MESSAGE_TOO_BIG,
+            "frame with 2 bytes after reading 2 bytes exceeds limit of 3 bytes",
         )
 
     def test_client_sends_unexpected_binary(self):