PYTHON-5983 Validate uncompressed size in OP_COMPRESSED messages

pymongo/network_layer.py

@@ -604,7 +604,20 @@ def buffer_updated(self, nbytes: int) -> None:
             self._compression_index += nbytes
             if self._compression_index >= 9:
                 self._expecting_compression = False
-                self._op_code, self._compressor_id = self.process_compression_header()
+                (
+                    self._op_code,
+                    uncompressed_size,
+                    self._compressor_id,
+                ) = self.process_compression_header()
+                if uncompressed_size > self._max_message_size:
+                    self.close(
+                        ProtocolError(
+                            f"Uncompressed message size ({uncompressed_size!r}) "
+                            f"is larger than server max message size "
+                            f"({self._max_message_size!r})"
+                        )
+                    )
+                    return
             return
 
         self._message_index += nbytes
@@ -658,10 +671,12 @@ def process_header(self) -> tuple[int, int, int, bool]:
 
         return length - 16, op_code, response_to, expecting_compression
 
-    def process_compression_header(self) -> tuple[int, int]:
+    def process_compression_header(self) -> tuple[int, int, int]:
         """Unpack a MongoDB Wire Protocol compression header."""
-        op_code, _, compressor_id = _UNPACK_COMPRESSION_HEADER(self._compression_header)
-        return op_code, compressor_id
+        op_code, uncompressed_size, compressor_id = _UNPACK_COMPRESSION_HEADER(
+            self._compression_header
+        )
+        return op_code, uncompressed_size, compressor_id
 
     def _resolve_pending_messages(self, exc: Optional[Exception] = None) -> None:
         pending = list(self._pending_messages)

test/asynchronous/test_async_network_layer.py

@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import asyncio
+import struct
 import sys
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -88,6 +89,24 @@ def test_length_exceeds_max_raises(self):
         with self.assertRaisesRegex(ProtocolError, "larger than server max"):
             self.protocol.process_header()
 
+    def test_compression_uncompressed_size_exceeds_max_closes(self):
+        self.protocol._max_message_size = 1024
+        self.protocol._header = memoryview(
+            bytearray(
+                pack_msg_header(
+                    length=35, request_id=1, response_to=0, op_code=2012
+                )
+            )
+        )
+        self.protocol.process_header()
+        # Now feed compression sub-header with uncompressed_size > max
+        self.protocol._compression_header[:] = struct.pack(
+            "<iiB", 2013, 9999, 2
+        )
+        self.protocol._compression_index = 9
+        self.protocol.buffer_updated(0)
+        self.protocol.transport.abort.assert_called()
+
 
 class TestClose(AsyncUnitTest):
     async def asyncSetUp(self):