diff --git a/doc/changelog.rst b/doc/changelog.rst index cda288c575..3e626ac2b6 100644 --- a/doc/changelog.rst +++ b/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 diff --git a/pymongo/compression_support.py b/pymongo/compression_support.py index d669e02b75..63263de085 100644 --- a/pymongo/compression_support.py +++ b/pymongo/compression_support.py @@ -164,7 +164,9 @@ def compress(data: bytes) -> bytes: return zstd.compress(data) -def decompress(data: bytes | memoryview, compressor_id: int) -> bytes: +def decompress( + data: bytes | memoryview, compressor_id: int, max_message_size: int | None = None +) -> bytes: if compressor_id == SnappyContext.compressor_id: # python-snappy doesn't support the buffer interface. # https://github.com/andrix/python-snappy/issues/65 @@ -172,17 +174,33 @@ def decompress(data: bytes | memoryview, compressor_id: int) -> bytes: # 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) + if max_message_size is None: + result = zlib.decompress(data) + else: + # 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) + if max_message_size is None: + result = zstd.decompress(data) + else: + result = zstd.ZstdDecompressor().decompress(data, max_message_size + 1) else: raise ValueError(f"Unknown compressorId {compressor_id}") + if max_message_size is not None and len(result) > max_message_size: + from pymongo.errors import ProtocolError + + raise ProtocolError( + f"Decompressed message size ({len(result)!r}) is larger than " + f"server max message size ({max_message_size!r})" + ) + return result diff --git a/pymongo/network_layer.py b/pymongo/network_layer.py index 102f560d65..b4408a5740 100644 --- a/pymongo/network_layer.py +++ b/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) diff --git a/test/asynchronous/test_async_network_layer.py b/test/asynchronous/test_async_network_layer.py index 5adb7aaeac..64bba148a6 100644 --- a/test/asynchronous/test_async_network_layer.py +++ b/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(" max (1024). + buf = self.protocol.get_buffer(9) + buf[:9] = struct.pack(" max_message_size. + compressed = b"x" * 10 + total_len = 16 + 9 + len(compressed) + header = struct.pack("