Skip to content

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

Open
dfengliu wants to merge 5 commits into
mongodb:mainfrom
dfengliu:fix/op-compressed-size-validation
Open

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

Conversation

@dfengliu

@dfengliu dfengliu commented Aug 1, 2026

Copy link
Copy Markdown

PYTHON-5983

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.
@dfengliu
dfengliu requested a review from a team as a code owner August 1, 2026 09:27
@dfengliu
dfengliu requested a review from blink1073 August 1, 2026 09:27
@blink1073 blink1073 changed the title Validate uncompressed size in OP_COMPRESSED messages PYTHON-XXXX Validate uncompressed size in OP_COMPRESSED messages Aug 3, 2026
@blink1073
blink1073 marked this pull request as draft August 3, 2026 13:48
@blink1073
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
@blink1073 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 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment on lines +92 to +108
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()
Comment thread pymongo/network_layer.py
Comment on lines +607 to +620
(
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

Copy link
Copy Markdown
Member

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

dfengliu commented Aug 4, 2026

Copy link
Copy Markdown
Author

@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 AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment thread pymongo/network_layer.py
Comment on lines 553 to 555
if compressor_id is not None:
data = decompress(data, compressor_id)
data = _decompress(data, compressor_id, self._max_message_size)
return data, op_code
Comment thread pymongo/compression_support.py Outdated
Comment on lines +189 to +196
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 blink1073 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

dfengliu commented Aug 5, 2026

Copy link
Copy Markdown
Author

@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

Copy link
Copy Markdown
Member

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

dfengliu commented Aug 6, 2026

Copy link
Copy Markdown
Author

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants