PYTHON-5983 Validate uncompressed size in OP_COMPRESSED messages by dfengliu · Pull Request #2976 · mongodb/mongo-python-driver

dfengliu · 2026-08-01T09:27:22Z

Summary

Validate the uncompressed_size field from the OP_COMPRESSED wire protocol compression sub-header against max_message_size.

Details

The process_compression_header method in network_layer.py previously unpacked the compression sub-header and discarded the uncompressed_size field. A malicious or compromised MongoDB server could send a small compressed envelope (passing the envelope size check) that decompresses to a very large payload, causing memory exhaustion.

Changes

process_compression_header now returns uncompressed_size in addition to op_code and compressor_id
The caller validates uncompressed_size against self._max_message_size and raises ProtocolError if it exceeds the limit
Added unit test test_compression_uncompressed_size_exceeds_max_closes

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.

Copilot

Pull request overview

This PR hardens PyMongo’s handling of OP_COMPRESSED messages by validating the compression sub-header’s uncompressed_size against the configured max_message_size, preventing decompression-driven memory exhaustion.

Changes:

Extend process_compression_header() to return uncompressed_size along with op_code and compressor_id.
Add a uncompressed_size > self._max_message_size guard in the asyncio protocol receive path (buffer_updated()), closing the connection with ProtocolError when exceeded.
Add a unit test intended to assert the connection closes when uncompressed_size exceeds the max.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
`pymongo/network_layer.py`	Adds `uncompressed_size` extraction and validation in the asyncio protocol path for OP_COMPRESSED messages.
`test/asynchronous/test_async_network_layer.py`	Adds a new test meant to cover oversized `uncompressed_size` behavior.

+    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(
+            "<iiB", 2013, 9999, 2
+        )
+        self.protocol._compression_index = 9
+        self.protocol.buffer_updated(0)
+        self.protocol.transport.abort.assert_called()


+                (
+                    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


blink1073 · 2026-08-03T16:01:42Z

Hi @dfengliu, thanks for the PR! Could you move the bound into decompress() in
pymongo/compression_support.py instead? That is the right layer for it, and it covers both
the async and sync read paths in one place.

dfengliu · 2026-08-04T03:37:10Z

@blink1073 Thank you for the suggestion. I have moved the size validation into _decompress() in compression_support.py, which covers both the async and sync read paths. The original process_compression_header has been reverted to its previous behavior.

Copilot

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

Suppressed comments (4)

pymongo/network_layer.py:783

The OP_COMPRESSED sub-header’s uncompressed_size field is still unpacked and discarded (_). This means a server can advertise an extremely large uncompressed size and the client will proceed to decompress, which is the memory-exhaustion vector described in the PR. Validate uncompressed_size against max_message_size before calling _decompress.

        op_code, _, compressor_id = _UNPACK_COMPRESSION_HEADER(receive_data(conn, 9, deadline))
        data = _decompress(receive_data(conn, length - 25, deadline), compressor_id, max_message_size)

test/asynchronous/test_async_network_layer.py:104

This new test exercises _decompress’s post-decompression length check, but it does not cover validating the OP_COMPRESSED uncompressed_size field against max_message_size (the core requirement described in the PR). Add a regression test that feeds an OP_COMPRESSED header + compression sub-header with uncompressed_size > MAX_MESSAGE_SIZE into the protocol/receive path and asserts a ProtocolError (and that the connection is closed/aborted, if applicable).

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)

pymongo/network_layer.py:37

decompress is imported but no longer used in this module (all call sites were switched to _decompress). This will fail linting (unused import).

from pymongo.compression_support import _decompress, decompress

test/asynchronous/test_async_network_layer.py:22

struct is imported but not used in this test module, which will fail linting (unused import).

import asyncio
import struct
import sys

            if compressor_id is not None:
-                data = decompress(data, compressor_id)
+                data = _decompress(data, compressor_id, self._max_message_size)
            return data, op_code


+    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


blink1073

Thanks for moving this forward. One thing still needs to change: the check runs after decompression completes, so it measures the result rather than bounding it. Both zlib and zstd accept a max length on their incremental decompressors, which would apply the limit during decompression instead.

Smaller items:

The 2**31 - 1 default on decompress() is effectively unbounded, and nothing in the
driver calls it now. Worth dropping it or defaulting to MAX_MESSAGE_SIZE.
Please restore the snappy comment explaining the bytes(data) conversion.
The test payload is small enough that it only exercises the comparison. A large expansion case would be more useful, and test/test_compression_support.py is a better home for it.
just lint-manual picks up a few things, and this needs a doc/changelog.rst entry.

…ompression 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.

- 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

dfengliu · 2026-08-05T02:29:14Z

@blink1073 Thank you for the detailed review. I have addressed the items you raised:

Restored decompress() as the original standalone function, removing the 2**31 - 1 default.
Restored the snappy bytes(data) comment that was lost during refactoring.
Moved the size-limit test to test/test_compression_support.py with a high-expansion-ratio payload.
Added a changelog entry.

Regarding the incremental decompressor suggestion: the current approach uses pre-validation (checking uncompressed_size from the OP_COMPRESSED sub-header before decompression in both async and sync paths) combined with the post-decompression size check in _decompress() as defense-in-depth. The pre-validation blocks oversized messages at zero memory cost, while the post-validation catches servers that misreport the uncompressed size. This provides more complete coverage than relying on the per-call max_length parameter in zlib's incremental API alone.

Would this approach be acceptable?

blink1073 · 2026-08-05T11:12:39Z

The layered structure is fine to keep, but the max length arguments need to go in as well.
Without them the trailing check can only run once the buffer already exists, so it reports
the size rather than capping it.

     elif compressor_id == ZlibContext.compressor_id:
         import zlib

-        result = zlib.decompress(data)
+        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)
+        result = zstd.ZstdDecompressor().decompress(data, max_message_size + 1)

Snappy has no equivalent, so it keeps relying on the trailing check.

Rest of the commit:

decompress() has no callers in the driver now, and it's a near-copy of _decompress.
Probably worth collapsing to one. The snappy comment is currently only on the unused one.
The changelog entry landed inside the operation_id bullet, between "As a" and
"result, ...". It needs its own bullet.
The test's payload is 100 KB against a 1000 byte cap, so it passes either way. Something
large enough to be obviously wrong unbounded would catch a regression here.

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.

dfengliu · 2026-08-06T03:46:08Z

@blink1073 Thank you for the detailed feedback. I have implemented the changes:

Added max_length bounding during decompression for both zlib (zlib.decompressobj().decompress(data, max_message_size + 1)) and zstd (zstd.ZstdDecompressor().decompress(data, max_message_size + 1)), so memory is bounded before the size check runs. Snappy has no equivalent API and keeps relying on the trailing check.
Collapsed decompress() and _decompress() into a single function with an optional max_message_size parameter, and restored the snappy bytes(data) comment.
Fixed the changelog entry to be its own bullet.
Enlarged the test payload to 10 MB with a high expansion ratio, moved it to test/test_compression_support.py.
Added regression tests for the async/sync pre-validation of uncompressed_size and the decompression bomb path.

All 39 unit tests and the end-to-end scenarios pass.

blink1073 · 2026-08-12T11:05:01Z

Thank you! We're getting close. I have two more requests:

Make max_message_size required rather than defaulting to None. Both call sites
already pass it, and the None branch still reaches the unbounded zlib.decompress
and zstd.decompress calls, which is the path this PR exists to remove.
TestDecompressSizeLimit still passes if you take the max_length arguments back
out, so it's covering the size comparison rather than the bound itself. It needs to
fail without the fix. The property to assert is that peak allocation stays near the
limit rather than near the payload size, which tracemalloc can measure:
```
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()
self.assertLess(peak, 10 * max_size)
```
With the bound that peak lands around 2 MB. Without it, around 215 MB. That is a
wide enough gap that the threshold won't be flaky across platforms.

- 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.

Copilot

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.

Suppressed comments (5)

pymongo/network_layer.py:800

This compares the header-excluding uncompressed_size with a limit that applies to the complete wire message, allowing the reconstructed message to exceed the negotiated maximum by 16 bytes. Include the standard header in this check and bound the decompressed body to max_message_size - 16.

        if uncompressed_size > max_message_size:

pymongo/compression_support.py:188

ZstdDecompressor.decompress is incremental as well, so this can silently return partial output for a truncated frame when the result stays under the size limit. Check eof to retain the complete-frame validation previously provided by the one-shot zstd.decompress call.

        result = zstd.ZstdDecompressor().decompress(data, max_message_size + 1)

pymongo/network_layer.py:612

uncompressed_size is the original message body size and excludes the 16-byte standard header (pymongo/message.py:262 packs len(data)). Comparing it directly with the total wire-message limit allows a reconstructed message 16 bytes over max_message_size; the later decompress call has the same mismatch. Validate uncompressed_size + 16 here and pass self._max_message_size - 16 as the payload bound to decompress.

This issue also appears on line 800 of the same file.

                if uncompressed_size > self._max_message_size:

test/asynchronous/test_async_network_layer.py:240

This test exercises the synchronous receive_message path inside a module explicitly documented as async-only (test/asynchronous/test_async_network_layer.py:15-19). Move it to test/test_network_layer.py, whose module documentation identifies receive_message as sync-only and which already provides connection/receive helpers; that also avoids maintaining the duplicate _FakeSocket and _FakeConn fixtures added above.

class TestReceiveMessage(unittest.TestCase):
    def test_oversized_uncompressed_size_rejected(self):

pymongo/compression_support.py:181

Switching from one-shot zlib.decompress to the incremental API removes end-of-stream validation: truncated zlib data can return partial bytes with eof == False instead of raising as before. Preserve the bounded output while rejecting an incomplete stream.

        result = zlib.decompressobj().decompress(data, max_message_size + 1)

        import snappy

-        return snappy.uncompress(bytes(data))
+        result = snappy.uncompress(bytes(data))


dfengliu · 2026-08-13T11:27:13Z

@blink1073 Thank you for raising these points. Both requests are now addressed:

max_message_size is now a required parameter in decompress() — the None default and the unbounded zlib.decompress/zstd.decompress branches have been removed. Both call sites already pass it.
Added a tracemalloc-based peak memory test (test_decompression_peak_memory_bounded): it decompresses a 100 MB payload with a 1 MB limit and asserts that peak allocation stays below 10 MB. With the max_length bound it lands around 2 MB; without it, around 215 MB — so the test fails without the fix.
Also added exact-boundary (== max accepted, max+1 rejected) and snappy over-limit regression tests.

Additionally, I fixed a compatibility issue in the test_compression_commands spy that was causing CI failures after merging main — the spy now accepts and forwards the max_message_size argument.

CI is currently waiting on approval for both GitHub Actions and Evergreen. Could you help approve the runs when you have a moment? Thanks!

dfengliu · 2026-08-13T11:29:22Z

Thank you! We're getting close. I have two more requests:
Make max_message_size required rather than defaulting to None. Both call sites
already pass it, and the None branch still reaches the unbounded zlib.decompress
and zstd.decompress calls, which is the path this PR exists to remove.
TestDecompressSizeLimit still passes if you take the max_length arguments back
out, so it's covering the size comparison rather than the bound itself. It needs to
fail without the fix. The property to assert is that peak allocation stays near the
limit rather than near the payload size, which tracemalloc can measure:
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()
self.assertLess(peak, 10 * max_size)
With the bound that peak lands around 2 MB. Without it, around 215 MB. That is a
wide enough gap that the threshold won't be flaky across platforms.

@blink1073 Thank you for raising these points. Both requests are now addressed:

max_message_size is now a required parameter in decompress() — the None default and the unbounded zlib.decompress/zstd.decompress branches have been removed. Both call sites already pass it.
Added a tracemalloc-based peak memory test (test_decompression_peak_memory_bounded): it decompresses a 100 MB payload with a 1 MB limit and asserts that peak allocation stays below 10 MB. With the max_length bound it lands around 2 MB; without it, around 215 MB — so the test fails without the fix.
Also added exact-boundary (== max accepted, max+1 rejected) and snappy over-limit regression tests.

Additionally, I fixed a compatibility issue in the test_compression_commands spy that was causing CI failures after merging main — the spy now accepts and forwards the max_message_size argument.

CI is currently waiting on approval for both GitHub Actions and Evergreen. Could you help approve the runs when you have a moment? Thanks!

blink1073 · 2026-08-13T11:40:05Z

I agree with the copilot comment about snappy. We should check the length the raw block declares before calling
uncompress. For a valid stream that value is the real output length, and cramjam rejects streams where it disagrees, so the check is sufficient:

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")

Then before snappy.uncompress(bytes(data)):

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})"
    )

cramjam.snappy.decompress_raw_len(data) does the same in one line, but cramjam is
python-snappy's implementation detail rather than one of ours, and the requirement
has no upper bound, so we should parse it ourselves.

I think it is worth extending test_decompression_peak_memory_bounded to cover snappy on the same
pattern once this is in.

dfengliu requested a review from a team as a code owner August 1, 2026 09:27

dfengliu requested a review from blink1073 August 1, 2026 09:27

blink1073 changed the title ~~Validate uncompressed size in OP_COMPRESSED messages~~ PYTHON-XXXX Validate uncompressed size in OP_COMPRESSED messages Aug 3, 2026

blink1073 marked this pull request as draft August 3, 2026 13:48

blink1073 marked this pull request as ready for review August 3, 2026 13:48

Copilot AI review requested due to automatic review settings August 3, 2026 13:48

Copilot started reviewing on behalf of blink1073 August 3, 2026 13:49 View session

blink1073 changed the title ~~PYTHON-XXXX Validate uncompressed size in OP_COMPRESSED messages~~ PYTHON-5983 Validate uncompressed size in OP_COMPRESSED messages Aug 3, 2026

Copilot AI reviewed Aug 3, 2026

View reviewed changes

Move decompression size validation to _decompress in compression_support

670d454

blink1073 requested a review from Copilot August 4, 2026 11:47

Copilot started reviewing on behalf of blink1073 August 4, 2026 11:48 View session

Copilot AI reviewed Aug 4, 2026

View reviewed changes

blink1073 requested changes Aug 4, 2026

View reviewed changes

dfengliu added 2 commits August 5, 2026 09:21

aclark4life and others added 2 commits August 12, 2026 13:03

Merge branch 'main' into fix/op-compressed-size-validation

7c1fe73

blink1073 requested a balanced review from Copilot August 13, 2026 11:04

Copilot started reviewing on behalf of blink1073 August 13, 2026 11:05 View session

Copilot AI reviewed Aug 13, 2026

View reviewed changes

Comment thread pymongo/compression_support.py

import snappy

return snappy.uncompress(bytes(data))

result = snappy.uncompress(bytes(data))

dfengliu closed this Aug 13, 2026

dfengliu reopened this Aug 13, 2026

Navigation Menu

Uh oh!

PYTHON-5983 Validate uncompressed size in OP_COMPRESSED messages - #2976

PYTHON-5983 Validate uncompressed size in OP_COMPRESSED messages#2976
dfengliu wants to merge 7 commits into
mongodb:mainfrom
dfengliu:fix/op-compressed-size-validation

dfengliu commented Aug 1, 2026 •

edited by blink1073

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

blink1073 commented Aug 3, 2026

Uh oh!

dfengliu commented Aug 4, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

blink1073 left a comment

Uh oh!

dfengliu commented Aug 5, 2026

Uh oh!

blink1073 commented Aug 5, 2026

Uh oh!

dfengliu commented Aug 6, 2026

Uh oh!

blink1073 commented Aug 12, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

dfengliu commented Aug 13, 2026

Uh oh!

dfengliu commented Aug 13, 2026

Uh oh!

blink1073 commented Aug 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

dfengliu commented Aug 1, 2026 • edited by blink1073 Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Details

Changes

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

blink1073 commented Aug 3, 2026

Uh oh!

dfengliu commented Aug 4, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

blink1073 left a comment

Choose a reason for hiding this comment

Uh oh!

dfengliu commented Aug 5, 2026

Uh oh!

blink1073 commented Aug 5, 2026

Uh oh!

dfengliu commented Aug 6, 2026

Uh oh!

blink1073 commented Aug 12, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

dfengliu commented Aug 13, 2026

Uh oh!

dfengliu commented Aug 13, 2026

Uh oh!

blink1073 commented Aug 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

dfengliu commented Aug 1, 2026 •

edited by blink1073

Loading