PYTHON-5983 Validate uncompressed size in OP_COMPRESSED messages

doc/changelog.rst

@@ -31,6 +31,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.
 - Improved the performance and memory usage of decoding large documents to
   :class:`~bson.raw_bson.RawBSONDocument`. Documents and subdocuments that are 4KB or greater
   and decoded from an immutable buffer are now exposed as read-only :class:`memoryview`

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,65 @@ 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:
+        declared = _snappy_uncompressed_length(data)
+        if declared > max_message_size:
+            raise ProtocolError(
+                f"Decompressed message size ({declared!r}) is larger than "
+                f"maximum allowed payload size ({max_message_size!r})"
+            )
+        import snappy
+
         # 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.
-        import snappy
-
-        return snappy.uncompress(bytes(data))
+        result = snappy.uncompress(bytes(data))
     elif compressor_id == ZlibContext.compressor_id:
         import zlib
 
-        return zlib.decompress(data)
+        dc = zlib.decompressobj()
+        # Bound the decompressed output during decompression to avoid
+        # allocating a huge buffer before the size check runs.
+        result = dc.decompress(data, max_message_size + 1)
+        if len(result) <= max_message_size:
+            if not dc.eof:
+                raise ProtocolError("Truncated zlib-compressed message")
+            if dc.unused_data:
+                raise ProtocolError("Trailing data after zlib-compressed message")
     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)
+        zdc = zstd.ZstdDecompressor()
+        result = zdc.decompress(data, max_message_size + 1)
+        if len(result) <= max_message_size:
+            if not zdc.eof:
+                raise ProtocolError("Truncated zstd-compressed message")
+            if zdc.unused_data:
+                raise ProtocolError("Trailing data after zstd-compressed message")
     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"maximum allowed payload 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 - 16)
             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 <= 0 or uncompressed_size + 16 > self._max_message_size:
+                    self.close(
+                        ProtocolError(
+                            f"Uncompressed message size ({uncompressed_size!r}) is invalid or larger "
+                            f"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 <= 0 or uncompressed_size + 16 > max_message_size:
+            raise ProtocolError(
+                f"Uncompressed message size ({uncompressed_size!r}) is invalid or larger "
+                f"than server max message size ({max_message_size!r})"
+            )
+        data = decompress(
+            receive_data(conn, length - 25, deadline), compressor_id, max_message_size - 16
+        )
     else:
         data = receive_data(conn, length - 16, deadline)
 

test/asynchronous/test_async_network_layer.py

@@ -17,14 +17,15 @@
 from __future__ import annotations
 
 import asyncio
+import struct
 import sys
 from unittest.mock import AsyncMock, MagicMock, patch
 
 sys.path[0:0] = [""]
 
 from pymongo.common import MAX_MESSAGE_SIZE
 from pymongo.errors import ProtocolError
-from pymongo.network_layer import PyMongoProtocol, _async_socket_receive
+from pymongo.network_layer import PyMongoProtocol, _async_socket_receive, receive_message
 from test.asynchronous import AsyncUnitTest, unittest
 from test.utils_shared import pack_msg_header
 
@@ -88,6 +89,13 @@ 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 +168,69 @@ 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
+
+    async def test_uncompressed_size_equal_max_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)
+
+        # uncompressed_size == max_message_size: reconstructed message
+        # (uncompressed_size + 16-byte header) must exceed the limit.
+        buf = self.protocol.get_buffer(9)
+        buf[:9] = struct.pack("<iiB", 2013, 1024, 2)
+        self.protocol.buffer_updated(9)
+
+        self.assertTrue(self.protocol.transport.abort.called)
+        with self.assertRaisesRegex(ProtocolError, "Uncompressed message size"):
+            await read_task
+
+    async def test_nonpositive_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)
+
+        # uncompressed_size is a signed int32; zero and negative values must be
+        # rejected as malformed rather than accepted by the upper-bound check.
+        buf = self.protocol.get_buffer(9)
+        buf[:9] = struct.pack("<iiB", 2013, 0, 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 +244,91 @@ 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.
+
+    On Windows and PyPy, receive_data() calls wait_for_read() before
+    recv_into(), which reaches through conn.conn.sock, so this also has to
+    look like a socket to that code path.
+    """
+
+    def __init__(self, data: bytes):
+        self.data = data
+        self.pos = 0
+        self.sock = self
+
+    def gettimeout(self):
+        return None
+
+    def fileno(self):
+        return 1
+
+    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 _FakeSocketChecker:
+    def select(self, sock, read=False, write=False, timeout=None):
+        return True
+
+
+class _FakeCancelContext:
+    cancelled = False
+
+
+class _FakeConn:
+    def __init__(self, data: bytes):
+        self.conn = _FakeSocket(data)
+        self.socket_checker = _FakeSocketChecker()
+        self.cancel_context = _FakeCancelContext()
+
+    def gettimeout(self):
+        return None
+
+    def set_conn_timeout(self, t):
+        pass
+
+
+class TestReceiveMessage(unittest.TestCase):
+    def test_oversized_uncompressed_size_rejected(self):
+        # 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)  # type: ignore[arg-type]
+
+    def test_uncompressed_size_equal_max_rejected(self):
+        # uncompressed_size == max_message_size; the reconstructed message
+        # (uncompressed_size + 16-byte header) must exceed the limit.
+        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, 1024, 2)
+        conn = _FakeConn(header + sub_header + compressed)
+        with self.assertRaisesRegex(ProtocolError, "Uncompressed message size"):
+            receive_message(conn, request_id=99, max_message_size=1024)  # type: ignore[arg-type]
+
+    def test_nonpositive_uncompressed_size_rejected(self):
+        # uncompressed_size is a signed int32; zero and negative values must be
+        # rejected as malformed rather than accepted by the upper-bound check.
+        for size in (0, -1):
+            with self.subTest(uncompressed_size=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, size, 2)
+                conn = _FakeConn(header + sub_header + compressed)
+                with self.assertRaisesRegex(ProtocolError, "Uncompressed message size"):
+                    receive_message(conn, request_id=99, max_message_size=1024)  # type: ignore[arg-type]
+
+
 if __name__ == "__main__":
     unittest.main()

test/asynchronous/test_client.py

@@ -1949,10 +1949,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

@@ -1902,10 +1902,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.