From d0d827b0074e4d78cdc8284d18c1d535db0bf39f Mon Sep 17 00:00:00 2001 From: liufeng Date: Sat, 1 Aug 2026 17:27:03 +0800 Subject: [PATCH 1/5] Validate uncompressed size in OP_COMPRESSED messages The process_compression_header method previously discarded the uncompressed_size field from the compression sub-header. A malicious or compromised server could send a small compressed envelope (passing the max_message_size check) that decompresses to a very large payload, causing memory exhaustion. This change returns the uncompressed_size from the compression header and validates it against max_message_size before accepting the compressed payload. --- pymongo/network_layer.py | 23 +++++++++++++++---- test/asynchronous/test_async_network_layer.py | 19 +++++++++++++++ 2 files changed, 38 insertions(+), 4 deletions(-) diff --git a/pymongo/network_layer.py b/pymongo/network_layer.py index 102f560d65..87bcd18455 100644 --- a/pymongo/network_layer.py +++ b/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) diff --git a/test/asynchronous/test_async_network_layer.py b/test/asynchronous/test_async_network_layer.py index 5adb7aaeac..31b7812bad 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,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( + " Date: Tue, 4 Aug 2026 11:13:21 +0800 Subject: [PATCH 2/5] Move decompression size validation to _decompress in compression_support --- pymongo/compression_support.py | 22 +++++++++----- pymongo/network_layer.py | 29 +++++------------- test/asynchronous/test_async_network_layer.py | 30 ++++++++----------- 3 files changed, 35 insertions(+), 46 deletions(-) diff --git a/pymongo/compression_support.py b/pymongo/compression_support.py index d669e02b75..08d172521a 100644 --- a/pymongo/compression_support.py +++ b/pymongo/compression_support.py @@ -165,24 +165,32 @@ def compress(data: bytes) -> bytes: def decompress(data: bytes | memoryview, compressor_id: int) -> bytes: + return _decompress(data, compressor_id, max_message_size=2**31 - 1) + + +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. import snappy - return snappy.uncompress(bytes(data)) + result = snappy.uncompress(bytes(data)) elif compressor_id == ZlibContext.compressor_id: import zlib - return zlib.decompress(data) + result = zlib.decompress(data) 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.decompress(data) else: raise ValueError(f"Unknown compressorId {compressor_id}") + if 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 87bcd18455..a0744d9c92 100644 --- a/pymongo/network_layer.py +++ b/pymongo/network_layer.py @@ -34,7 +34,7 @@ from pymongo import _csot, ssl_support from pymongo._asyncio_task import create_task from pymongo.common import MAX_MESSAGE_SIZE -from pymongo.compression_support import decompress +from pymongo.compression_support import _decompress, decompress from pymongo.errors import ProtocolError, _OperationCancelled from pymongo.message import _UNPACK_REPLY, _OpMsg from pymongo.socket_checker import _errno_from_exception @@ -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,20 +604,7 @@ def buffer_updated(self, nbytes: int) -> None: self._compression_index += nbytes if self._compression_index >= 9: self._expecting_compression = False - ( - 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 + self._op_code, self._compressor_id = self.process_compression_header() return self._message_index += nbytes @@ -671,12 +658,10 @@ 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, int]: + def process_compression_header(self) -> tuple[int, int]: """Unpack a MongoDB Wire Protocol compression header.""" - op_code, uncompressed_size, compressor_id = _UNPACK_COMPRESSION_HEADER( - self._compression_header - ) - return op_code, uncompressed_size, compressor_id + op_code, _, compressor_id = _UNPACK_COMPRESSION_HEADER(self._compression_header) + return op_code, compressor_id def _resolve_pending_messages(self, exc: Optional[Exception] = None) -> None: pending = list(self._pending_messages) @@ -795,7 +780,7 @@ def receive_message( 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) + 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 31b7812bad..43792739de 100644 --- a/test/asynchronous/test_async_network_layer.py +++ b/test/asynchronous/test_async_network_layer.py @@ -89,23 +89,19 @@ 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( - " Date: Wed, 5 Aug 2026 09:21:39 +0800 Subject: [PATCH 3/5] Add pre-decompression uncompressed_size validation alongside post-decompression check Validate uncompressed_size from the OP_COMPRESSED sub-header against max_message_size before calling _decompress, in both async and sync receive paths. The internal _decompress function retains a post-decompression length check as defense-in-depth against servers that misreport the uncompressed size. --- pymongo/network_layer.py | 32 ++++++++++++++++--- test/asynchronous/test_async_network_layer.py | 10 ++++++ 2 files changed, 37 insertions(+), 5 deletions(-) diff --git a/pymongo/network_layer.py b/pymongo/network_layer.py index a0744d9c92..81e4d86a04 100644 --- a/pymongo/network_layer.py +++ b/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) @@ -779,7 +794,14 @@ 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)) + 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 43792739de..96d545d733 100644 --- a/test/asynchronous/test_async_network_layer.py +++ b/test/asynchronous/test_async_network_layer.py @@ -89,6 +89,16 @@ 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(" Date: Wed, 5 Aug 2026 10:25:53 +0800 Subject: [PATCH 4/5] Restore decompress as standalone function, improve test coverage - Restore public decompress() as the original function without wrapper - Keep _decompress() with required max_message_size for internal validation - Restore snappy bytes(data) comment that was lost during refactoring - Move decompress size-limit test to test_compression_support.py with high expansion ratio payload - Add changelog entry --- doc/changelog.rst | 3 +++ pymongo/compression_support.py | 22 ++++++++++++++++++- pymongo/network_layer.py | 2 +- test/asynchronous/test_async_network_layer.py | 15 ------------- test/test_compression_support.py | 19 ++++++++++++++++ 5 files changed, 44 insertions(+), 17 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index cda288c575..796576d2a9 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -18,6 +18,9 @@ PyMongo 4.18 brings a number of changes including: - Command monitoring events and command log messages for a single logical operation now share one stable ``operation_id`` across all of its retry attempts, so consumers can correlate a retried operation's events. As a +- Added validation of OP_COMPRESSED decompressed message size against + ``max_message_size`` to prevent memory exhaustion from maliciously crafted + compressed server responses. result, ``operation_id`` is no longer equal to the per-attempt ``request_id`` for these operations. - Fixed a potential out-of-bounds read in the C extension when decoding an diff --git a/pymongo/compression_support.py b/pymongo/compression_support.py index 08d172521a..9c48394e80 100644 --- a/pymongo/compression_support.py +++ b/pymongo/compression_support.py @@ -165,7 +165,27 @@ def compress(data: bytes) -> bytes: def decompress(data: bytes | memoryview, compressor_id: int) -> bytes: - return _decompress(data, compressor_id, max_message_size=2**31 - 1) + 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. + import snappy + + return snappy.uncompress(bytes(data)) + elif compressor_id == ZlibContext.compressor_id: + import zlib + + return zlib.decompress(data) + 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) + else: + raise ValueError(f"Unknown compressorId {compressor_id}") def _decompress(data: bytes | memoryview, compressor_id: int, max_message_size: int) -> bytes: diff --git a/pymongo/network_layer.py b/pymongo/network_layer.py index 81e4d86a04..2c1266dbfc 100644 --- a/pymongo/network_layer.py +++ b/pymongo/network_layer.py @@ -34,7 +34,7 @@ from pymongo import _csot, ssl_support from pymongo._asyncio_task import create_task from pymongo.common import MAX_MESSAGE_SIZE -from pymongo.compression_support import _decompress, decompress +from pymongo.compression_support import _decompress from pymongo.errors import ProtocolError, _OperationCancelled from pymongo.message import _UNPACK_REPLY, _OpMsg from pymongo.socket_checker import _errno_from_exception diff --git a/test/asynchronous/test_async_network_layer.py b/test/asynchronous/test_async_network_layer.py index 96d545d733..6846df2d81 100644 --- a/test/asynchronous/test_async_network_layer.py +++ b/test/asynchronous/test_async_network_layer.py @@ -99,21 +99,6 @@ def test_process_compression_header_returns_uncompressed_size(self): self.assertEqual(compressor_id, 2) -class TestDecompress(unittest.TestCase): - def test_decompressed_size_exceeds_max_raises(self): - from pymongo.compression_support import _decompress - - import zlib - - # Compress a small payload that decompresses larger than max - payload = zlib.compress(b"x" * 100) - with self.assertRaisesRegex(ProtocolError, "Decompressed message size"): - _decompress(payload, 2, max_message_size=5) - # Normal decompression still works - result = _decompress(payload, 2, max_message_size=1024) - self.assertEqual(result, b"x" * 100) - - class TestClose(AsyncUnitTest): async def asyncSetUp(self): self.protocol = _make_protocol() diff --git a/test/test_compression_support.py b/test/test_compression_support.py index 0c37627f21..ca802071ab 100644 --- a/test/test_compression_support.py +++ b/test/test_compression_support.py @@ -26,6 +26,7 @@ SnappyContext, ZlibContext, ZstdContext, + _decompress, _have_snappy, _have_zlib, _have_zstd, @@ -205,5 +206,23 @@ def test_zstd_roundtrip(self): self.assertEqual(result, data) +class TestDecompressSizeLimit(unittest.TestCase): + def test_decompressed_size_exceeds_max_raises(self): + import zlib + + # High expansion ratio payload (repeated zeros, ~1000:1 ratio) + payload = zlib.compress(b"\x00" * 100_000) + original_len = len(zlib.decompress(payload)) + self.assertGreater(original_len, 1000) # high expansion ratio + # Raise when decompressed size exceeds small limit + from pymongo.errors import ProtocolError + + with self.assertRaisesRegex(ProtocolError, "Decompressed message size"): + _decompress(payload, ZlibContext.compressor_id, max_message_size=1000) + # Normal decompression with adequate limit + result = _decompress(payload, ZlibContext.compressor_id, max_message_size=1_000_000) + self.assertEqual(result, b"\x00" * original_len) + + if __name__ == "__main__": unittest.main() From 0d4c748e7586d31ee838060d7d65628772e26e3f Mon Sep 17 00:00:00 2001 From: liufeng Date: Thu, 6 Aug 2026 11:44:18 +0800 Subject: [PATCH 5/5] Bound decompression output size during decompression Apply the max_message_size limit during decompression for zlib and zstd using their incremental decompressor max_length parameter, so memory is bounded before the size check runs. Snappy has no such API and continues to rely on the post-decompression check. Collapse decompress into a single function with an optional max_message_size parameter, and restore the snappy bytes(data) comment. Add pre-validation of the OP_COMPRESSED sub-header's uncompressed_size in both async and sync receive paths, plus regression tests covering oversized declarations and decompression bombs. --- doc/changelog.rst | 4 +- pymongo/compression_support.py | 38 ++++------- pymongo/network_layer.py | 8 ++- test/asynchronous/test_async_network_layer.py | 67 +++++++++++++++++++ test/test_compression_support.py | 9 ++- 5 files changed, 92 insertions(+), 34 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 796576d2a9..3e626ac2b6 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -18,11 +18,11 @@ PyMongo 4.18 brings a number of changes including: - Command monitoring events and command log messages for a single logical operation now share one stable ``operation_id`` across all of its retry 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. - result, ``operation_id`` is no longer equal to the per-attempt ``request_id`` - for these operations. - 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 9c48394e80..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,41 +174,29 @@ 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)) - elif compressor_id == ZlibContext.compressor_id: - import zlib - - return zlib.decompress(data) - 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) - else: - raise ValueError(f"Unknown compressorId {compressor_id}") - - -def _decompress(data: bytes | memoryview, compressor_id: int, max_message_size: int) -> bytes: - if compressor_id == SnappyContext.compressor_id: - import snappy - result = snappy.uncompress(bytes(data)) elif compressor_id == ZlibContext.compressor_id: import zlib - result = 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 - result = 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 len(result) > max_message_size: + if max_message_size is not None and len(result) > max_message_size: from pymongo.errors import ProtocolError raise ProtocolError( diff --git a/pymongo/network_layer.py b/pymongo/network_layer.py index 2c1266dbfc..b4408a5740 100644 --- a/pymongo/network_layer.py +++ b/pymongo/network_layer.py @@ -34,7 +34,7 @@ from pymongo import _csot, ssl_support from pymongo._asyncio_task import create_task from pymongo.common import MAX_MESSAGE_SIZE -from pymongo.compression_support import _decompress +from pymongo.compression_support import decompress from pymongo.errors import ProtocolError, _OperationCancelled from pymongo.message import _UNPACK_REPLY, _OpMsg from pymongo.socket_checker import _errno_from_exception @@ -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, self._max_message_size) + data = decompress(data, compressor_id, self._max_message_size) return data, op_code raise OSError("connection closed") @@ -802,7 +802,9 @@ def receive_message( 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) + 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 6846df2d81..64bba148a6 100644 --- a/test/asynchronous/test_async_network_layer.py +++ b/test/asynchronous/test_async_network_layer.py @@ -170,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(" max_message_size. + compressed = b"x" * 10 + total_len = 16 + 9 + len(compressed) + header = struct.pack("