diff --git a/doc/changelog.rst b/doc/changelog.rst index f90025963f..f657b819a6 100644 --- a/doc/changelog.rst +++ b/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` diff --git a/pymongo/compression_support.py b/pymongo/compression_support.py index d669e02b75..39f0834a8c 100644 --- a/pymongo/compression_support.py +++ b/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 diff --git a/pymongo/network_layer.py b/pymongo/network_layer.py index 102f560d65..dbb3e9f903 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 - 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) diff --git a/test/asynchronous/test_async_network_layer.py b/test/asynchronous/test_async_network_layer.py index 5adb7aaeac..b1da68e5b1 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 @@ -24,7 +25,7 @@ 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(" 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(" bool: + try: + import tracemalloc + + return True + except ImportError: + # PyPy does not ship the _tracemalloc C extension. + return False + + class TestValidateCompressors(unittest.TestCase): def test_string_input_single(self): with patch("pymongo.compression_support._have_zlib", return_value=True): @@ -163,30 +176,42 @@ def setUp(self): self.skipTest("zlib not available") def test_compress_and_decompress_roundtrip(self): - import zlib - ctx = ZlibContext(level=-1) data = b"hello world" * 100 compressed = ctx.compress(data) self.assertEqual(zlib.decompress(compressed), data) +class TestSnappyUncompressedLength(unittest.TestCase): + def test_single_byte(self): + self.assertEqual(_snappy_uncompressed_length(b"\x03"), 3) + + def test_multi_byte(self): + self.assertEqual(_snappy_uncompressed_length(b"\xac\x02"), 300) + + def test_truncated(self): + with self.assertRaises(ProtocolError): + _snappy_uncompressed_length(b"\xff") + + def test_overlong_varint(self): + with self.assertRaises(ProtocolError): + _snappy_uncompressed_length(b"\xff" * 5) + + class TestDecompress(unittest.TestCase): def test_unknown_compressor_id_raises(self): with self.assertRaises(ValueError) as ctx: - decompress(b"data", 99) + decompress(b"data", 99, max_message_size=2**20) self.assertIn("Unknown compressorId 99", str(ctx.exception)) def _assert_roundtrip(self, compressed, compressor_id, data): for payload in (compressed, memoryview(compressed)): with self.subTest(type=type(payload).__name__): - self.assertEqual(decompress(payload, compressor_id), data) + self.assertEqual(decompress(payload, compressor_id, max_message_size=2**20), data) def test_zlib_roundtrip(self): if not _have_zlib(): self.skipTest("zlib not available") - import zlib - data = b"hello world" self._assert_roundtrip(zlib.compress(data), ZlibContext.compressor_id, data) @@ -201,8 +226,122 @@ def test_zstd_roundtrip(self): self.skipTest("zstd not available") data = b"hello world" * 50 compressed = ZstdContext.compress(data) - result = decompress(compressed, ZstdContext.compressor_id) + result = decompress(compressed, ZstdContext.compressor_id, max_message_size=2**20) + self.assertEqual(result, data) + + +class TestDecompressSizeLimit(unittest.TestCase): + def test_decompression_peak_memory_bounded(self): + if not _have_tracemalloc(): + self.skipTest("tracemalloc not available") + import tracemalloc + + # High expansion ratio payload (repeated zeros, ~100000:1 ratio) + payload = zlib.compress(b"\x00" * 100_000_000) + max_size = 1_000_000 + + tracemalloc.start() + try: + with self.assertRaises(ProtocolError): + decompress(payload, ZlibContext.compressor_id, max_message_size=max_size) + _current, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + # Peak allocation should stay near the bound rather than the payload size. + self.assertLess(peak, 10 * max_size) + + def test_zlib_exact_boundary(self): + # Data that decompresses to exactly max_message_size must be accepted. + data = b"\x00" * 1000 + payload = zlib.compress(data) + result = decompress(payload, ZlibContext.compressor_id, max_message_size=1000) self.assertEqual(result, data) + # One byte over the limit must be rejected. + data_over = b"\x00" * 1001 + payload_over = zlib.compress(data_over) + + with self.assertRaises(ProtocolError): + decompress(payload_over, ZlibContext.compressor_id, max_message_size=1000) + + def test_snappy_exceeds_max_rejected(self): + if not _have_snappy(): + self.skipTest("python-snappy not installed") + + data = b"\x00" * 100_000 + payload = SnappyContext.compress(data) + with self.assertRaises(ProtocolError): + decompress(payload, SnappyContext.compressor_id, max_message_size=1000) + + def test_snappy_peak_memory_bounded(self): + if not _have_snappy(): + self.skipTest("python-snappy not installed") + if not _have_tracemalloc(): + self.skipTest("tracemalloc not available") + import tracemalloc + + payload = SnappyContext.compress(b"\x00" * 100_000_000) + max_size = 1_000_000 + tracemalloc.start() + try: + with self.assertRaises(ProtocolError): + decompress(payload, SnappyContext.compressor_id, max_message_size=max_size) + _current, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + # Peak allocation should stay near the bound rather than the payload size. + self.assertLess(peak, 10 * max_size) + + def test_zstd_peak_memory_bounded(self): + if not _have_zstd(): + self.skipTest("zstd not available") + if not _have_tracemalloc(): + self.skipTest("tracemalloc not available") + import tracemalloc + + payload = ZstdContext.compress(b"\x00" * 100_000_000) + max_size = 1_000_000 + tracemalloc.start() + try: + with self.assertRaises(ProtocolError): + decompress(payload, ZstdContext.compressor_id, max_message_size=max_size) + _current, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + # Peak allocation should stay near the bound rather than the payload size. + self.assertLess(peak, 10 * max_size) + + def test_snappy_declared_size_exceeds_max_rejected(self): + # Varint declaring 2^31 uncompressed bytes; rejected by the declared + # size pre-check before python-snappy is imported. + payload = b"\x80\x80\x80\x80\x08" + with self.assertRaises(ProtocolError): + decompress(payload, SnappyContext.compressor_id, max_message_size=1000) + + def test_zlib_truncated_rejected(self): + payload = zlib.compress(b"\x00" * 1000)[:-1] + with self.assertRaises(ProtocolError): + decompress(payload, ZlibContext.compressor_id, max_message_size=10_000) + + def test_zstd_truncated_rejected(self): + if not _have_zstd(): + self.skipTest("zstd not available") + + payload = ZstdContext.compress(b"\x00" * 1000)[:-1] + with self.assertRaises(ProtocolError): + decompress(payload, ZstdContext.compressor_id, max_message_size=10_000) + + def test_zlib_trailing_data_rejected(self): + payload = zlib.compress(b"\x00" * 1000) + b"GARBAGE" + with self.assertRaises(ProtocolError): + decompress(payload, ZlibContext.compressor_id, max_message_size=10_000) + + def test_zstd_trailing_data_rejected(self): + if not _have_zstd(): + self.skipTest("zstd not available") + + payload = ZstdContext.compress(b"\x00" * 1000) + b"GARBAGE" + with self.assertRaises(ProtocolError): + decompress(payload, ZstdContext.compressor_id, max_message_size=10_000) if __name__ == "__main__":