PYTHON-5983 Validate uncompressed size in OP_COMPRESSED messages

doc/changelog.rst

@@ -20,6 +20,9 @@ PyMongo 4.18 brings a number of changes including:
   attempts, so consumers can correlate a retried operation's events. As a
   result, ``operation_id`` is no longer equal to the per-attempt ``request_id``
   for these operations.
+- Added validation of OP_COMPRESSED decompressed message size against
+  ``max_message_size`` to prevent memory exhaustion from maliciously crafted
+  compressed server responses.
 - Fixed a potential out-of-bounds read in the C extension when decoding an
   array of BSON documents. An embedded document whose declared length exceeds
   the bytes remaining in the array now raises

pymongo/compression_support.py

@@ -18,6 +18,7 @@
 from collections.abc import Iterable
 from typing import Any, Optional, Union
 
+from pymongo.errors import ProtocolError
 from pymongo.hello import HelloCompat
 from pymongo.helpers_shared import _SENSITIVE_COMMANDS
 
@@ -164,25 +165,53 @@ def compress(data: bytes) -> bytes:
         return zstd.compress(data)
 
 
-def decompress(data: bytes | memoryview, compressor_id: int) -> bytes:
+def _snappy_uncompressed_length(data: bytes | memoryview) -> int:
+    """Read the varint-encoded uncompressed length from a raw snappy block."""
+    result = shift = 0
+    for i in range(5):
+        if i >= len(data):
+            raise ProtocolError("Truncated snappy payload")
+        byte = data[i]
+        result |= (byte & 0x7F) << shift
+        if not byte & 0x80:
+            return result
+        shift += 7
+    raise ProtocolError("Invalid snappy uncompressed length header")
+
+
+def decompress(data: bytes | memoryview, compressor_id: int, max_message_size: int) -> bytes:
     if compressor_id == SnappyContext.compressor_id:
         # python-snappy doesn't support the buffer interface.
         # https://github.com/andrix/python-snappy/issues/65
         # This only matters when data is a memoryview since
         # id(bytes(data)) == id(data) when data is a bytes.
+        declared = _snappy_uncompressed_length(data)
+        if declared > max_message_size:
+            raise ProtocolError(
+                f"Decompressed message size ({declared!r}) is larger than "
+                f"server max message size ({max_message_size!r})"
+            )
         import snappy
 
-        return snappy.uncompress(bytes(data))
+        result = snappy.uncompress(bytes(data))
     elif compressor_id == ZlibContext.compressor_id:
         import zlib
 
-        return zlib.decompress(data)
+        # Bound the decompressed output during decompression to avoid
+        # allocating a huge buffer before the size check runs.
+        result = zlib.decompressobj().decompress(data, max_message_size + 1)
     elif compressor_id == ZstdContext.compressor_id:
         if sys.version_info >= (3, 14):
             from compression import zstd
         else:
             from backports import zstd
 
-        return zstd.decompress(data)
+        result = zstd.ZstdDecompressor().decompress(data, max_message_size + 1)
     else:
         raise ValueError(f"Unknown compressorId {compressor_id}")
+    if len(result) > max_message_size:
+        raise ProtocolError(
+            f"Decompressed message size ({len(result)!r}) is larger than "
+            f"server max message size ({max_message_size!r})"
+        )
+    return result

pymongo/network_layer.py

@@ -551,7 +551,7 @@ async def read(self, request_id: Optional[int], max_message_size: int) -> tuple[
                         f"Got response id {response_to!r} but expected {request_id!r}"
                     )
             if compressor_id is not None:
-                data = decompress(data, compressor_id)
+                data = decompress(data, compressor_id, self._max_message_size)
             return data, op_code
         raise OSError("connection closed")
 
@@ -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)
@@ -779,8 +794,17 @@ def receive_message(
             raise ProtocolError(
                 f"Message length ({length!r}) not longer than standard OP_COMPRESSED message header size (25)"
             )
-        op_code, _, compressor_id = _UNPACK_COMPRESSION_HEADER(receive_data(conn, 9, deadline))
-        data = decompress(receive_data(conn, length - 25, deadline), compressor_id)
+        op_code, uncompressed_size, compressor_id = _UNPACK_COMPRESSION_HEADER(
+            receive_data(conn, 9, deadline)
+        )
+        if uncompressed_size > max_message_size:
+            raise ProtocolError(
+                f"Uncompressed message size ({uncompressed_size!r}) is larger "
+                f"than server max message size ({max_message_size!r})"
+            )
+        data = decompress(
+            receive_data(conn, length - 25, deadline), compressor_id, max_message_size
+        )
     else:
         data = receive_data(conn, length - 16, deadline)
 

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,15 @@ def test_length_exceeds_max_raises(self):
         with self.assertRaisesRegex(ProtocolError, "larger than server max"):
             self.protocol.process_header()
 
+    def test_process_compression_header_returns_uncompressed_size(self):
+        self.protocol._compression_header[:] = struct.pack("<iiB", 2013, 9999, 2)
+        op_code, uncompressed_size, compressor_id = (
+            self.protocol.process_compression_header()
+        )
+        self.assertEqual(op_code, 2013)
+        self.assertEqual(uncompressed_size, 9999)
+        self.assertEqual(compressor_id, 2)
+
 
 class TestClose(AsyncUnitTest):
     async def asyncSetUp(self):
@@ -160,6 +170,29 @@ async def test_resolves_pending_read(self):
         _data, op_code = await read_task
         self.assertEqual(op_code, 2013)
 
+    async def test_oversized_uncompressed_size_closes_connection(self):
+        self.protocol._max_message_size = 1024
+        read_task = asyncio.create_task(
+            self.protocol.read(request_id=None, max_message_size=1024)
+        )
+        await asyncio.sleep(0)
+
+        # Feed OP_COMPRESSED header (length = 16 + 9 + 1 = 26).
+        header = pack_msg_header(length=26, request_id=1, response_to=99, op_code=2012)
+        buf = self.protocol.get_buffer(16)
+        buf[:16] = header
+        self.protocol.buffer_updated(16)
+        self.assertTrue(self.protocol._expecting_compression)
+
+        # Feed compression sub-header with uncompressed_size > max (1024).
+        buf = self.protocol.get_buffer(9)
+        buf[:9] = struct.pack("<iiB", 2013, 9999, 2)
+        self.protocol.buffer_updated(9)
+
+        self.assertTrue(self.protocol.transport.abort.called)
+        with self.assertRaisesRegex(ProtocolError, "Uncompressed message size"):
+            await read_task
+
 
 class TestAsyncSocketReceive(AsyncUnitTest):
     async def test_raises_on_connection_closed(self):
@@ -173,5 +206,49 @@ async def test_raises_on_connection_closed(self):
                 await _async_socket_receive(mock_socket, 10, loop)
 
 
+class _FakeSocket:
+    """Feeds a byte buffer, simulating a socket."""
+
+    def __init__(self, data: bytes):
+        self.data = data
+        self.pos = 0
+
+    def gettimeout(self):
+        return None
+
+    def recv_into(self, buf):
+        n = min(len(buf), len(self.data) - self.pos)
+        if n <= 0:
+            return 0
+        buf[:n] = self.data[self.pos : self.pos + n]
+        self.pos += n
+        return n
+
+
+class _FakeConn:
+    def __init__(self, data: bytes):
+        self.conn = _FakeSocket(data)
+
+    def gettimeout(self):
+        return None
+
+    def set_conn_timeout(self, t):
+        pass
+
+
+class TestReceiveMessage(unittest.TestCase):
+    def test_oversized_uncompressed_size_rejected(self):
+        from pymongo.network_layer import receive_message
+
+        # Build OP_COMPRESSED with uncompressed_size > max_message_size.
+        compressed = b"x" * 10
+        total_len = 16 + 9 + len(compressed)
+        header = struct.pack("<iiii", total_len, 1, 99, 2012)
+        sub_header = struct.pack("<iiB", 2013, 9999, 2)
+        conn = _FakeConn(header + sub_header + compressed)
+        with self.assertRaisesRegex(ProtocolError, "Uncompressed message size"):
+            receive_message(conn, request_id=99, max_message_size=1024)
+
+
 if __name__ == "__main__":
     unittest.main()

test/asynchronous/test_client.py

@@ -1950,10 +1950,14 @@ def spy(data, _original=original, _recorded=compressed):
                     original_decompress = network_layer.decompress
 
                     def decompress_spy(
-                        data, compressor_id, _original=original_decompress, _recorded=decompressed
+                        data,
+                        compressor_id,
+                        max_message_size=None,
+                        _original=original_decompress,
+                        _recorded=decompressed,
                     ):
                         _recorded.append(compressor_id)
-                        return _original(data, compressor_id)
+                        return _original(data, compressor_id, max_message_size)
 
                     # Round-trip a command. Every non-sensitive command is
                     # compressed.

test/test_client.py

@@ -1903,10 +1903,14 @@ def spy(data, _original=original, _recorded=compressed):
                     original_decompress = network_layer.decompress
 
                     def decompress_spy(
-                        data, compressor_id, _original=original_decompress, _recorded=decompressed
+                        data,
+                        compressor_id,
+                        max_message_size=None,
+                        _original=original_decompress,
+                        _recorded=decompressed,
                     ):
                         _recorded.append(compressor_id)
-                        return _original(data, compressor_id)
+                        return _original(data, compressor_id, max_message_size)
 
                     # Round-trip a command. Every non-sensitive command is
                     # compressed.