AVRO-4290: [python] Enforce a maximum decompressed block size#3850
AVRO-4290: [python] Enforce a maximum decompressed block size#3850iemejia wants to merge 8 commits into
Conversation
When reading a data file, each block is decompressed according to the file's codec. A block with a very high compression ratio (or a malformed block) could expand to far more memory than its compressed size. Enforce a configurable maximum decompressed size across the deflate, bzip2, snappy and zstandard codecs, mirroring the Java SDK's decompression limit (AVRO-4247). The limit defaults to 200 MiB and can be overridden with the AVRO_MAX_DECOMPRESS_LENGTH environment variable; exceeding it raises AvroDecompressionSizeException. Assisted-by: GitHub Copilot:claude-opus-4.8
There was a problem hiding this comment.
Pull request overview
This PR hardens the Python Avro data-file reader against decompression bombs by enforcing a configurable maximum decompressed block size (default 200 MiB, overrideable via AVRO_MAX_DECOMPRESS_LENGTH) and raising avro.errors.AvroDecompressionSizeException when the limit is exceeded.
Changes:
- Add
AvroDecompressionSizeExceptionand enforce a maximum decompressed block size across deflate, bzip2, snappy, and zstandard codecs. - Introduce
AVRO_MAX_DECOMPRESS_LENGTHenv override and a default limit of 200 MiB inavro.codecs. - Add unit tests covering over-limit rejection and within-limit decoding for available codecs.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| lang/py/avro/codecs.py | Adds max-decompressed-size enforcement and env-configurable limit across supported codecs. |
| lang/py/avro/errors.py | Introduces the new exception type raised when decompressed output exceeds the configured cap. |
| lang/py/avro/test/test_codecs.py | Adds coverage ensuring blocks over the limit are rejected and normal blocks still decode correctly. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| uncompressed = decompressor.decompress(data, limit + 1) | ||
| if len(uncompressed) > limit: | ||
| _raise_decompression_too_large(limit) | ||
| uncompressed += decompressor.flush() | ||
| if len(uncompressed) > limit: | ||
| _raise_decompression_too_large(limit) |
There was a problem hiding this comment.
Fixed in 638aac5 — DeflateCodec.decompress now accumulates into a bytearray, so flush() output is appended in place instead of creating an extra full-size copy of the decompressed data.
| uncompressed.extend(chunk) | ||
| if len(uncompressed) > limit: | ||
| _raise_decompression_too_large(limit) |
There was a problem hiding this comment.
Fixed in 638aac5 — the zstandard path now checks len(uncompressed) + len(chunk) > limit before extending, so the buffer never grows past the limit.
…fore-extend - DeflateCodec.decompress accumulates into a bytearray so the flush() output is appended in place instead of creating an extra full-size copy of the already decompressed data. - ZstandardCodec.decompress checks len(uncompressed) + len(chunk) before extending, so the buffer never grows past the configured limit. Assisted-by: GitHub Copilot:claude-opus-4.8
| decompressor = bz2.BZ2Decompressor() | ||
| uncompressed = decompressor.decompress(data, limit + 1) | ||
| if len(uncompressed) > limit: | ||
| _raise_decompression_too_large(limit) | ||
| return avro.io.BinaryDecoder(io.BytesIO(uncompressed)) |
There was a problem hiding this comment.
Good catch. Reworked BZip2Codec.decompress to verify the decompressor reached end-of-stream and to drain any trailing data, looping over concatenated bzip2 streams so a multi-stream block is fully decoded. Added a test for a truncated block. Pushed in 9f3c49d.
BZip2Codec.decompress now loops the BZ2Decompressor: it drains all buffered output (bounded by the limit), verifies the stream reached EOF (rejecting a truncated/corrupt block with InvalidAvroBinaryEncoding), and handles concatenated bzip2 streams as bz2.decompress does. Add a truncated-block test. Assisted-by: GitHub Copilot:claude-opus-4.8
| limit = _max_decompress_length() | ||
| decompressor = zlib.decompressobj(-15) | ||
| # Request at most limit + 1 bytes so that an over-large block is detected | ||
| # without allocating the whole (potentially huge) output. Accumulate into | ||
| # a bytearray so the flush() output is appended in place rather than | ||
| # creating an extra full-size copy of the already-decompressed data. | ||
| uncompressed = bytearray(decompressor.decompress(data, limit + 1)) | ||
| if len(uncompressed) > limit: | ||
| _raise_decompression_too_large(limit) | ||
| uncompressed += decompressor.flush() | ||
| if len(uncompressed) > limit: | ||
| _raise_decompression_too_large(limit) |
There was a problem hiding this comment.
Good catch. Added a decompressor.eof check after flush() — a truncated/incomplete deflate block now raises InvalidAvroBinaryEncoding, restoring the reject-truncated behavior that zlib.decompress had.
| # Compressed data includes a 4-byte CRC32 checksum | ||
| length = readers_decoder.read_long() | ||
| data = readers_decoder.read(length - 4) | ||
| limit = _max_decompress_length() |
There was a problem hiding this comment.
Fixed — SnappyCodec.decompress now validates length >= 4 and raises a codec-specific InvalidAvroBinaryEncoding before reading length - 4 bytes.
Address review feedback: - DeflateCodec: after flushing the decompressor, verify decompressor.eof so a truncated/incomplete raw-deflate block is rejected with InvalidAvroBinaryEncoding (zlib.decompress used to raise for this; the decompressobj-based size cap otherwise silently accepted partial output). - SnappyCodec: validate the block length is >= 4 before reading length - 4 bytes, raising a codec-specific InvalidAvroBinaryEncoding instead of falling through to the generic decoder error. Assisted-by: GitHub Copilot:claude-opus-4.8
| if not decompressor.eof: | ||
| # The end-of-stream marker was not reached: the block is truncated or | ||
| # corrupt. zlib.decompress() used to raise for this; preserve that. | ||
| raise avro.errors.InvalidAvroBinaryEncoding("Truncated or corrupt deflate block") |
There was a problem hiding this comment.
Added test_deflate_truncated_block_rejected — it truncates a valid deflate block and asserts DeflateCodec.decompress() raises InvalidAvroBinaryEncoding (mirrors the existing bzip2 truncation test).
Cover the new DeflateCodec behavior (raising InvalidAvroBinaryEncoding when the end-of-stream marker isn't reached) with a test that truncates a valid deflate block and asserts decompress() rejects it, mirroring the bzip2 truncation test. Assisted-by: GitHub Copilot:claude-opus-4.8
| limit = _max_decompress_length() | ||
| decompressor = zlib.decompressobj(-15) | ||
| # Request at most limit + 1 bytes so that an over-large block is detected | ||
| # without allocating the whole (potentially huge) output. Accumulate into | ||
| # a bytearray so the flush() output is appended in place rather than | ||
| # creating an extra full-size copy of the already-decompressed data. | ||
| uncompressed = bytearray(decompressor.decompress(data, limit + 1)) | ||
| if len(uncompressed) > limit: | ||
| _raise_decompression_too_large(limit) | ||
| uncompressed += decompressor.flush() | ||
| if len(uncompressed) > limit: | ||
| _raise_decompression_too_large(limit) | ||
| if not decompressor.eof: | ||
| # The end-of-stream marker was not reached: the block is truncated or | ||
| # corrupt. zlib.decompress() used to raise for this; preserve that. | ||
| raise avro.errors.InvalidAvroBinaryEncoding("Truncated or corrupt deflate block") |
There was a problem hiding this comment.
Good catch — decompress(data, limit+1) leaves unconsumed input, and the subsequent flush() was unbounded. Reworked to drain unconsumed_tail in a loop with a per-call max_length and to pass a bounded length to flush(), so total output can't exceed limit+1 before being rejected.
There was a problem hiding this comment.
Already bounded — the final flush is decompressor.flush(limit + 1 - len(uncompressed)) (codecs.py:194) followed by the size check, so flush() can't materialize past the cap. The eof check runs after.
The previous fix bounded decompress(data, limit+1) but then called decompressor.flush() with no limit, so the residual output (from the unconsumed input left by max_length) could still expand unbounded. Drain the unconsumed_tail in a loop with a per-call max_length and pass a bounded length to flush(), so the accumulated output can never exceed the limit by more than one byte before being rejected. Assisted-by: GitHub Copilot:claude-opus-4.8
| def _max_decompress_length() -> int: | ||
| """Return the maximum decompressed block size, honoring the environment override.""" | ||
| value = os.environ.get(MAX_DECOMPRESS_LENGTH_ENV) | ||
| if value is None: | ||
| return DEFAULT_MAX_DECOMPRESS_LENGTH | ||
| try: | ||
| parsed = int(value) | ||
| except ValueError: | ||
| return DEFAULT_MAX_DECOMPRESS_LENGTH | ||
| return parsed if parsed > 0 else DEFAULT_MAX_DECOMPRESS_LENGTH |
There was a problem hiding this comment.
Fixed — _max_decompress_length() now returns min(parsed, sys.maxsize) so an oversized override stays a valid Py_ssize_t and won't raise OverflowError inside zlib/bz2 decompress().
_max_decompress_length() returned the raw parsed int. An absurdly large override is passed as max_length to zlib/bz2 decompress(), which raise OverflowError when it exceeds Py_ssize_t. Clamp to sys.maxsize so an oversized override is honored as "effectively unbounded" instead of producing a confusing OverflowError. Assisted-by: GitHub Copilot:claude-opus-4.8
| uncompressed = bytearray() | ||
| pending = data | ||
| while True: | ||
| want = limit + 1 - len(uncompressed) | ||
| uncompressed += decompressor.decompress(pending, want) | ||
| if len(uncompressed) > limit: | ||
| _raise_decompression_too_large(limit) | ||
| pending = decompressor.unconsumed_tail | ||
| if not pending: | ||
| break | ||
| uncompressed += decompressor.flush(limit + 1 - len(uncompressed)) |
There was a problem hiding this comment.
Fixed — added _decompress_read_ceiling(limit) which returns limit + 1 normally but sys.maxsize when limit == sys.maxsize, so the deflate max_length never overflows Py_ssize_t.
| limit = _max_decompress_length() | ||
| uncompressed = bytearray() | ||
| decompressor = bz2.BZ2Decompressor() | ||
| input_data = data | ||
| while True: | ||
| # Request enough to detect exceeding the limit without allocating | ||
| # the full (potentially huge) output. | ||
| want = limit + 1 - len(uncompressed) | ||
| if want <= 0: | ||
| want = 1 | ||
| uncompressed += decompressor.decompress(input_data, want) |
There was a problem hiding this comment.
Fixed — the bz2 path uses the same _decompress_read_ceiling() so want can't become sys.maxsize + 1.
After clamping AVRO_MAX_DECOMPRESS_LENGTH to sys.maxsize, the deflate/bz2 paths still computed `limit + 1` for the decompress() max_length, which overflows Py_ssize_t (and raises OverflowError) when limit == sys.maxsize. Add _decompress_read_ceiling(), which returns limit + 1 normally but sys.maxsize when the limit is already sys.maxsize (no realistic block can exceed it), and use it for the bounded read/flush in both paths. Assisted-by: GitHub Copilot:claude-opus-4.8
What is the purpose of the change
The deflate and bzip2 codecs are inflated incrementally and rejected once the output would exceed the limit; snappy rejects an over-large declared length up front; zstandard caps its streaming loop. Exceeding the limit raises
avro.errors.AvroDecompressionSizeException.When reading a data file, each block is decompressed according to the file's codec. A block with a very high compression ratio (or a malformed block) could expand to far more memory than its compressed size (a decompression bomb). This enforces a configurable maximum decompressed size while reading each block, mirroring the Java SDK's decompression limit (AVRO-4247). The limit defaults to 200 MiB and can be overridden with the
AVRO_MAX_DECOMPRESS_LENGTHenvironment variable.This is part of the umbrella issue AVRO-4283.
Verifying this change
This change added tests and can be verified as follows:
avro/test/test_codecs.py(TestDecompressionSizeLimit) covering deflate and bzip2 over-limit rejection and within-limit decoding (snappy/zstandard covered when those optional deps are installed).cd lang/py && python3 -m unittest avro.test.test_codecsDocumentation
AVRO_MAX_DECOMPRESS_LENGTHenvironment variable is documented in code comments)