From d0d827b0074e4d78cdc8284d18c1d535db0bf39f Mon Sep 17 00:00:00 2001 From: liufeng Date: Sat, 1 Aug 2026 17:27:03 +0800 Subject: [PATCH 01/12] 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 02/12] 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 03/12] 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 04/12] 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 05/12] 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(" Date: Thu, 13 Aug 2026 14:28:36 +0800 Subject: [PATCH 06/12] Make max_message_size required and fix decompress spy compatibility - Make max_message_size a required parameter in decompress() and remove the unbounded None branch, ensuring the decompression limit is always applied. - Fix the test_compression_commands decompress spy to accept the max_message_size argument, resolving the TypeError that caused widespread CI failures after merging main. - Add tracemalloc-based peak memory test verifying the max_length bound actually limits allocation (2 MB vs 215 MB without the fix). - Add exact-boundary and snappy over-limit regression tests. --- pymongo/compression_support.py | 20 ++++------- test/asynchronous/test_client.py | 8 +++-- test/test_client.py | 8 +++-- test/test_compression_support.py | 57 ++++++++++++++++++++++++-------- 4 files changed, 62 insertions(+), 31 deletions(-) diff --git a/pymongo/compression_support.py b/pymongo/compression_support.py index 63263de085..78118949d2 100644 --- a/pymongo/compression_support.py +++ b/pymongo/compression_support.py @@ -164,9 +164,7 @@ def compress(data: bytes) -> bytes: return zstd.compress(data) -def decompress( - data: bytes | memoryview, compressor_id: int, max_message_size: int | None = None -) -> bytes: +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 @@ -178,25 +176,19 @@ def decompress( elif compressor_id == ZlibContext.compressor_id: import zlib - 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) + # 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 - if max_message_size is None: - result = zstd.decompress(data) - else: - result = zstd.ZstdDecompressor().decompress(data, max_message_size + 1) + 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: + if len(result) > max_message_size: from pymongo.errors import ProtocolError raise ProtocolError( diff --git a/test/asynchronous/test_client.py b/test/asynchronous/test_client.py index 5606c7f8c1..e76e229bb0 100644 --- a/test/asynchronous/test_client.py +++ b/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. diff --git a/test/test_client.py b/test/test_client.py index 9df392e64f..0fd55fe205 100644 --- a/test/test_client.py +++ b/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. diff --git a/test/test_compression_support.py b/test/test_compression_support.py index 8e8dceaf2d..7f18c83642 100644 --- a/test/test_compression_support.py +++ b/test/test_compression_support.py @@ -174,13 +174,15 @@ def test_compress_and_decompress_roundtrip(self): 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(): @@ -201,26 +203,55 @@ 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_decompressed_size_exceeds_max_raises(self): + def test_decompression_peak_memory_bounded(self): + import tracemalloc import zlib # High expansion ratio payload (repeated zeros, ~100000:1 ratio) - payload = zlib.compress(b"\x00" * 10_000_000) - original_len = len(zlib.decompress(payload)) - self.assertGreater(original_len, 1000) # high expansion ratio - # Raise when decompressed size exceeds small limit + payload = zlib.compress(b"\x00" * 100_000_000) + max_size = 1_000_000 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=20_000_000) - self.assertEqual(result, b"\x00" * original_len) + 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): + import zlib + + # 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) + from pymongo.errors import ProtocolError + + 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") + from pymongo.errors import ProtocolError + + data = b"\x00" * 100_000 + payload = SnappyContext.compress(data) + with self.assertRaises(ProtocolError): + decompress(payload, SnappyContext.compressor_id, max_message_size=1000) if __name__ == "__main__": From c57a7b171ff53e5c76500b919e99bc855552429f Mon Sep 17 00:00:00 2001 From: liufeng Date: Thu, 27 Aug 2026 10:57:16 +0800 Subject: [PATCH 07/12] PYTHON-5983 Validate snappy declared uncompressed length before decompressing --- pymongo/compression_support.py | 23 +++++++++- test/test_compression_support.py | 72 ++++++++++++++++++++++++++++++-- 2 files changed, 90 insertions(+), 5 deletions(-) diff --git a/pymongo/compression_support.py b/pymongo/compression_support.py index 78118949d2..501125f024 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,12 +165,32 @@ def compress(data: bytes) -> bytes: return zstd.compress(data) +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 result = snappy.uncompress(bytes(data)) @@ -189,8 +210,6 @@ def decompress(data: bytes | memoryview, compressor_id: int, max_message_size: i 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})" diff --git a/test/test_compression_support.py b/test/test_compression_support.py index 7f18c83642..7c452a2e87 100644 --- a/test/test_compression_support.py +++ b/test/test_compression_support.py @@ -29,6 +29,7 @@ _have_snappy, _have_zlib, _have_zstd, + _snappy_uncompressed_length, decompress, validate_compressors, validate_zlib_compression_level, @@ -171,6 +172,26 @@ def test_compress_and_decompress_roundtrip(self): 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): + from pymongo.errors import ProtocolError + + with self.assertRaises(ProtocolError): + _snappy_uncompressed_length(b"\xff") + + def test_overlong_varint(self): + from pymongo.errors import ProtocolError + + 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: @@ -180,9 +201,7 @@ def test_unknown_compressor_id_raises(self): 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, max_message_size=2**20), data - ) + self.assertEqual(decompress(payload, compressor_id, max_message_size=2**20), data) def test_zlib_roundtrip(self): if not _have_zlib(): @@ -253,6 +272,53 @@ def test_snappy_exceeds_max_rejected(self): 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") + import tracemalloc + + from pymongo.errors import ProtocolError + + 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") + import tracemalloc + + from pymongo.errors import ProtocolError + + 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): + from pymongo.errors import ProtocolError + + # 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) + if __name__ == "__main__": unittest.main() From 06da304c637abaee1893e6ffe8376d0005748722 Mon Sep 17 00:00:00 2001 From: liufeng Date: Tue, 1 Sep 2026 11:36:19 +0800 Subject: [PATCH 08/12] PYTHON-5983 Reject truncated and trailing-garbage compressed messages The incremental decompressors (zlib decompressobj / zstd ZstdDecompressor) silently accept truncated streams and bytes trailing the compressed frame, whereas the previous one-shot APIs raised. Check eof and unused_data after bounded decompression and raise ProtocolError for truncated/trailing data. Also account for the 16-byte standard message header when validating the uncompressed_size from the OP_COMPRESSED sub-header, so the reconstructed message can no longer exceed max_message_size by up to 16 bytes. --- pymongo/compression_support.py | 24 ++++++++--- pymongo/network_layer.py | 8 ++-- test/asynchronous/test_async_network_layer.py | 42 ++++++++++++++++--- test/test_compression_support.py | 41 +++++++++++++----- 4 files changed, 89 insertions(+), 26 deletions(-) diff --git a/pymongo/compression_support.py b/pymongo/compression_support.py index 501125f024..1795e0900e 100644 --- a/pymongo/compression_support.py +++ b/pymongo/compression_support.py @@ -181,10 +181,6 @@ def _snappy_uncompressed_length(data: bytes | memoryview) -> int: 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( @@ -193,20 +189,36 @@ def decompress(data: bytes | memoryview, compressor_id: int, max_message_size: i ) 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. result = snappy.uncompress(bytes(data)) elif compressor_id == ZlibContext.compressor_id: import zlib + dc = zlib.decompressobj() # 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) + 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 - result = zstd.ZstdDecompressor().decompress(data, max_message_size + 1) + 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: diff --git a/pymongo/network_layer.py b/pymongo/network_layer.py index b4408a5740..d270365b8d 100644 --- a/pymongo/network_layer.py +++ b/pymongo/network_layer.py @@ -609,10 +609,10 @@ def buffer_updated(self, nbytes: int) -> None: uncompressed_size, self._compressor_id, ) = self.process_compression_header() - if uncompressed_size > self._max_message_size: + if uncompressed_size + 16 > self._max_message_size: self.close( ProtocolError( - f"Uncompressed message size ({uncompressed_size!r}) " + f"Uncompressed message size ({uncompressed_size + 16!r}) " f"is larger than server max message size " f"({self._max_message_size!r})" ) @@ -797,9 +797,9 @@ def receive_message( op_code, uncompressed_size, compressor_id = _UNPACK_COMPRESSION_HEADER( receive_data(conn, 9, deadline) ) - if uncompressed_size > max_message_size: + if uncompressed_size + 16 > max_message_size: raise ProtocolError( - f"Uncompressed message size ({uncompressed_size!r}) is larger " + f"Uncompressed message size ({uncompressed_size + 16!r}) is larger " f"than server max message size ({max_message_size!r})" ) data = decompress( diff --git a/test/asynchronous/test_async_network_layer.py b/test/asynchronous/test_async_network_layer.py index 64bba148a6..c4db486234 100644 --- a/test/asynchronous/test_async_network_layer.py +++ b/test/asynchronous/test_async_network_layer.py @@ -91,9 +91,7 @@ def test_length_exceeds_max_raises(self): def test_process_compression_header_returns_uncompressed_size(self): self.protocol._compression_header[:] = struct.pack(" Date: Wed, 2 Sep 2026 17:47:38 +0800 Subject: [PATCH 09/12] PYTHON-5983 Align decompress bound with message size, reject nonpositive uncompressed_size Pass max_message_size - 16 into decompress() at both call sites so the decompression bound agrees with the network-layer pre-check: the reconstructed message (16-byte header plus body) never exceeds max_message_size even when the header-size field understates the body. Also reject uncompressed_size <= 0 (signed int32) as malformed instead of letting it slip past the upper-bound check. Add regression tests for both paths at the asyncio-protocol and sync receive_message levels. --- pymongo/network_layer.py | 14 ++++---- test/asynchronous/test_async_network_layer.py | 36 +++++++++++++++++++ 2 files changed, 43 insertions(+), 7 deletions(-) diff --git a/pymongo/network_layer.py b/pymongo/network_layer.py index d270365b8d..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, self._max_message_size) + data = decompress(data, compressor_id, self._max_message_size - 16) return data, op_code raise OSError("connection closed") @@ -609,11 +609,11 @@ def buffer_updated(self, nbytes: int) -> None: uncompressed_size, self._compressor_id, ) = self.process_compression_header() - if uncompressed_size + 16 > self._max_message_size: + if uncompressed_size <= 0 or uncompressed_size + 16 > self._max_message_size: self.close( ProtocolError( - f"Uncompressed message size ({uncompressed_size + 16!r}) " - f"is larger than server max message size " + f"Uncompressed message size ({uncompressed_size!r}) is invalid or larger " + f"than server max message size " f"({self._max_message_size!r})" ) ) @@ -797,13 +797,13 @@ def receive_message( op_code, uncompressed_size, compressor_id = _UNPACK_COMPRESSION_HEADER( receive_data(conn, 9, deadline) ) - if uncompressed_size + 16 > max_message_size: + if uncompressed_size <= 0 or uncompressed_size + 16 > max_message_size: raise ProtocolError( - f"Uncompressed message size ({uncompressed_size + 16!r}) is larger " + 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 + 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 c4db486234..2d5d0feb94 100644 --- a/test/asynchronous/test_async_network_layer.py +++ b/test/asynchronous/test_async_network_layer.py @@ -210,6 +210,27 @@ async def test_uncompressed_size_equal_max_closes_connection(self): 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(" Date: Wed, 2 Sep 2026 06:01:39 -0500 Subject: [PATCH 10/12] PYTHON-5983 Fix max message size wording, satisfy lint and typing decompress() now reports the value it actually enforces as the maximum allowed payload size rather than the server max message size, since callers pass the body allowance (max_message_size - 16) rather than the raw negotiated limit. Also adds missing type: ignore annotations for the _FakeConn test double and removes stray blank lines left behind by an earlier import cleanup. --- pymongo/compression_support.py | 4 ++-- test/asynchronous/test_async_network_layer.py | 6 +++--- test/test_compression_support.py | 3 --- 3 files changed, 5 insertions(+), 8 deletions(-) diff --git a/pymongo/compression_support.py b/pymongo/compression_support.py index 1795e0900e..39f0834a8c 100644 --- a/pymongo/compression_support.py +++ b/pymongo/compression_support.py @@ -185,7 +185,7 @@ def decompress(data: bytes | memoryview, compressor_id: int, max_message_size: i 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})" + f"maximum allowed payload size ({max_message_size!r})" ) import snappy @@ -224,6 +224,6 @@ def decompress(data: bytes | memoryview, compressor_id: int, max_message_size: i 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})" + f"maximum allowed payload size ({max_message_size!r})" ) return result diff --git a/test/asynchronous/test_async_network_layer.py b/test/asynchronous/test_async_network_layer.py index 2d5d0feb94..7d36711ef1 100644 --- a/test/asynchronous/test_async_network_layer.py +++ b/test/asynchronous/test_async_network_layer.py @@ -285,7 +285,7 @@ def test_oversized_uncompressed_size_rejected(self): sub_header = struct.pack(" Date: Wed, 2 Sep 2026 06:11:47 -0500 Subject: [PATCH 11/12] PYTHON-5983 Hoist remaining inline imports to module scope in tests Moves the repeated inline "import zlib", "import tracemalloc", and "from pymongo.network_layer import receive_message" imports to the top of their respective test files. --- test/asynchronous/test_async_network_layer.py | 8 +------- test/test_compression_support.py | 19 ++----------------- 2 files changed, 3 insertions(+), 24 deletions(-) diff --git a/test/asynchronous/test_async_network_layer.py b/test/asynchronous/test_async_network_layer.py index 7d36711ef1..59638debb7 100644 --- a/test/asynchronous/test_async_network_layer.py +++ b/test/asynchronous/test_async_network_layer.py @@ -25,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 @@ -276,8 +276,6 @@ def set_conn_timeout(self, t): 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) @@ -288,8 +286,6 @@ def test_oversized_uncompressed_size_rejected(self): receive_message(conn, request_id=99, max_message_size=1024) # type: ignore[arg-type] def test_uncompressed_size_equal_max_rejected(self): - from pymongo.network_layer import receive_message - # uncompressed_size == max_message_size; the reconstructed message # (uncompressed_size + 16-byte header) must exceed the limit. compressed = b"x" * 10 @@ -301,8 +297,6 @@ def test_uncompressed_size_equal_max_rejected(self): receive_message(conn, request_id=99, max_message_size=1024) # type: ignore[arg-type] def test_nonpositive_uncompressed_size_rejected(self): - from pymongo.network_layer import receive_message - # 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): diff --git a/test/test_compression_support.py b/test/test_compression_support.py index cf363de32f..0a1be53d94 100644 --- a/test/test_compression_support.py +++ b/test/test_compression_support.py @@ -17,6 +17,8 @@ from __future__ import annotations import sys +import tracemalloc +import zlib from unittest.mock import patch sys.path[0:0] = [""] @@ -165,8 +167,6 @@ 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) @@ -203,8 +203,6 @@ def _assert_roundtrip(self, compressed, compressor_id, 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) @@ -225,9 +223,6 @@ def test_zstd_roundtrip(self): class TestDecompressSizeLimit(unittest.TestCase): def test_decompression_peak_memory_bounded(self): - import tracemalloc - import zlib - # High expansion ratio payload (repeated zeros, ~100000:1 ratio) payload = zlib.compress(b"\x00" * 100_000_000) max_size = 1_000_000 @@ -243,8 +238,6 @@ def test_decompression_peak_memory_bounded(self): self.assertLess(peak, 10 * max_size) def test_zlib_exact_boundary(self): - import zlib - # Data that decompresses to exactly max_message_size must be accepted. data = b"\x00" * 1000 payload = zlib.compress(data) @@ -269,8 +262,6 @@ def test_snappy_exceeds_max_rejected(self): def test_snappy_peak_memory_bounded(self): if not _have_snappy(): self.skipTest("python-snappy not installed") - import tracemalloc - payload = SnappyContext.compress(b"\x00" * 100_000_000) max_size = 1_000_000 tracemalloc.start() @@ -286,8 +277,6 @@ def test_snappy_peak_memory_bounded(self): def test_zstd_peak_memory_bounded(self): if not _have_zstd(): self.skipTest("zstd not available") - import tracemalloc - payload = ZstdContext.compress(b"\x00" * 100_000_000) max_size = 1_000_000 tracemalloc.start() @@ -308,8 +297,6 @@ def test_snappy_declared_size_exceeds_max_rejected(self): decompress(payload, SnappyContext.compressor_id, max_message_size=1000) def test_zlib_truncated_rejected(self): - import zlib - payload = zlib.compress(b"\x00" * 1000)[:-1] with self.assertRaises(ProtocolError): decompress(payload, ZlibContext.compressor_id, max_message_size=10_000) @@ -323,8 +310,6 @@ def test_zstd_truncated_rejected(self): decompress(payload, ZstdContext.compressor_id, max_message_size=10_000) def test_zlib_trailing_data_rejected(self): - import zlib - payload = zlib.compress(b"\x00" * 1000) + b"GARBAGE" with self.assertRaises(ProtocolError): decompress(payload, ZlibContext.compressor_id, max_message_size=10_000) From aad4d0b205e4e82bfafe35e72e9ba6c4bfb1fe7e Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Wed, 2 Sep 2026 06:49:19 -0500 Subject: [PATCH 12/12] PYTHON-5983 Fix CI failures on Windows/PyPy and PyPy tracemalloc import receive_data() takes a separate wait_for_read() path on Windows and PyPy that reaches through conn.conn.sock, conn.socket_checker, and conn.cancel_context before ever calling recv_into(). _FakeSocket and _FakeConn only supported the POSIX/CPython fast path, so TestReceiveMessage failed on Windows CI with AttributeError: '_FakeSocket' object has no attribute 'sock'. Separately, PyPy does not ship the _tracemalloc C extension, so hoisting "import tracemalloc" to module scope broke collection of the whole file on PyPy CI. Gate it behind a _have_tracemalloc() check in the same style as the file's existing _have_zlib/_have_snappy/ _have_zstd helpers, and skip the three tests that need it when unavailable. --- test/asynchronous/test_async_network_layer.py | 22 +++++++++++++++++- test/test_compression_support.py | 23 ++++++++++++++++++- 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/test/asynchronous/test_async_network_layer.py b/test/asynchronous/test_async_network_layer.py index 59638debb7..b1da68e5b1 100644 --- a/test/asynchronous/test_async_network_layer.py +++ b/test/asynchronous/test_async_network_layer.py @@ -245,15 +245,24 @@ async def test_raises_on_connection_closed(self): class _FakeSocket: - """Feeds a byte buffer, simulating a socket.""" + """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: @@ -263,9 +272,20 @@ def recv_into(self, buf): 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 diff --git a/test/test_compression_support.py b/test/test_compression_support.py index 0a1be53d94..e1b903f32b 100644 --- a/test/test_compression_support.py +++ b/test/test_compression_support.py @@ -17,7 +17,6 @@ from __future__ import annotations import sys -import tracemalloc import zlib from unittest.mock import patch @@ -40,6 +39,16 @@ from test import unittest +def _have_tracemalloc() -> 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): @@ -223,6 +232,10 @@ def test_zstd_roundtrip(self): 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 @@ -262,6 +275,10 @@ def test_snappy_exceeds_max_rejected(self): 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() @@ -277,6 +294,10 @@ def test_snappy_peak_memory_bounded(self): 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()