From 4b35f6cce8001e31d9703084d0c2bcfefa7b6549 Mon Sep 17 00:00:00 2001 From: ghoshp83 Date: Thu, 17 Sep 2026 20:35:19 +0100 Subject: [PATCH] fix(avro): bound the decompressed size of a block The bzip2, deflate and zstandard codecs decompressed a block without a cap, so the compressed block alone decided how much was allocated. Bound each path to MAX_DECOMPRESSED_BLOCK_SIZE and raise once a block decodes past it. --- pyiceberg/avro/codecs/bzip2.py | 8 +++- pyiceberg/avro/codecs/codec.py | 4 ++ pyiceberg/avro/codecs/deflate.py | 8 +++- pyiceberg/avro/codecs/zstandard_codec.py | 4 +- tests/avro/test_codecs.py | 50 ++++++++++++++++++++++++ 5 files changed, 69 insertions(+), 5 deletions(-) create mode 100644 tests/avro/test_codecs.py diff --git a/pyiceberg/avro/codecs/bzip2.py b/pyiceberg/avro/codecs/bzip2.py index 24786681f2..019dd1341c 100644 --- a/pyiceberg/avro/codecs/bzip2.py +++ b/pyiceberg/avro/codecs/bzip2.py @@ -16,7 +16,7 @@ # under the License. from __future__ import annotations -from pyiceberg.avro.codecs.codec import Codec +from pyiceberg.avro.codecs.codec import MAX_DECOMPRESSED_BLOCK_SIZE, Codec try: import bz2 @@ -29,7 +29,11 @@ def compress(data: bytes) -> tuple[bytes, int]: @staticmethod def decompress(data: bytes) -> bytes: - return bz2.decompress(data) + decompressor = bz2.BZ2Decompressor() + uncompressed = decompressor.decompress(data, max_length=MAX_DECOMPRESSED_BLOCK_SIZE) + if not decompressor.eof: + raise ValueError(f"Decompressed block exceeds the maximum of {MAX_DECOMPRESSED_BLOCK_SIZE} bytes") + return uncompressed except ImportError: diff --git a/pyiceberg/avro/codecs/codec.py b/pyiceberg/avro/codecs/codec.py index ab97ab6990..5e56028176 100644 --- a/pyiceberg/avro/codecs/codec.py +++ b/pyiceberg/avro/codecs/codec.py @@ -18,6 +18,10 @@ from abc import ABC, abstractmethod +# The compressed block alone determines how much a codec decodes, so decompression +# is bounded to keep a small block from expanding without limit. +MAX_DECOMPRESSED_BLOCK_SIZE = 1 << 30 + class Codec(ABC): """Abstract base class for all Avro codec classes.""" diff --git a/pyiceberg/avro/codecs/deflate.py b/pyiceberg/avro/codecs/deflate.py index 33fc11cd43..a1a62a58a6 100644 --- a/pyiceberg/avro/codecs/deflate.py +++ b/pyiceberg/avro/codecs/deflate.py @@ -18,7 +18,7 @@ import zlib -from pyiceberg.avro.codecs.codec import Codec +from pyiceberg.avro.codecs.codec import MAX_DECOMPRESSED_BLOCK_SIZE, Codec class DeflateCodec(Codec): @@ -33,4 +33,8 @@ def compress(data: bytes) -> tuple[bytes, int]: def decompress(data: bytes) -> bytes: # -15 is the log of the window size; negative indicates # "raw" (no zlib headers) decompression. See zlib.h. - return zlib.decompress(data, -15) + decompressor = zlib.decompressobj(-15) + uncompressed = decompressor.decompress(data, MAX_DECOMPRESSED_BLOCK_SIZE) + if decompressor.unconsumed_tail: + raise ValueError(f"Decompressed block exceeds the maximum of {MAX_DECOMPRESSED_BLOCK_SIZE} bytes") + return uncompressed diff --git a/pyiceberg/avro/codecs/zstandard_codec.py b/pyiceberg/avro/codecs/zstandard_codec.py index 4cc815214f..ab55999866 100644 --- a/pyiceberg/avro/codecs/zstandard_codec.py +++ b/pyiceberg/avro/codecs/zstandard_codec.py @@ -18,7 +18,7 @@ from io import BytesIO -from pyiceberg.avro.codecs.codec import Codec +from pyiceberg.avro.codecs.codec import MAX_DECOMPRESSED_BLOCK_SIZE, Codec try: from zstandard import ZstdCompressor, ZstdDecompressor @@ -39,6 +39,8 @@ def decompress(data: bytes) -> bytes: if not chunk: break uncompressed.extend(chunk) + if len(uncompressed) > MAX_DECOMPRESSED_BLOCK_SIZE: + raise ValueError(f"Decompressed block exceeds the maximum of {MAX_DECOMPRESSED_BLOCK_SIZE} bytes") return bytes(uncompressed) except ImportError: diff --git a/tests/avro/test_codecs.py b/tests/avro/test_codecs.py new file mode 100644 index 0000000000..0f5c937933 --- /dev/null +++ b/tests/avro/test_codecs.py @@ -0,0 +1,50 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from types import ModuleType + +import pytest + +from pyiceberg.avro.codecs import bzip2, deflate, zstandard_codec +from pyiceberg.avro.codecs.codec import Codec + +CODEC_MODULES = [ + (bzip2, bzip2.BZip2Codec), + (deflate, deflate.DeflateCodec), + (zstandard_codec, zstandard_codec.ZStandardCodec), +] + + +@pytest.mark.parametrize("module, codec", CODEC_MODULES) +def test_roundtrip(module: ModuleType, codec: type[Codec]) -> None: + data = b"aaaaaaaaaa" * 1000 + + compressed, _ = codec.compress(data) + + assert codec.decompress(compressed) == data + + +@pytest.mark.parametrize("module, codec", CODEC_MODULES) +def test_decompress_stops_at_the_limit(module: ModuleType, codec: type[Codec], monkeypatch: pytest.MonkeyPatch) -> None: + # A highly compressible block expands far beyond its compressed size, so the + # decoder must refuse it rather than let the block decide how much it allocates. + compressed, compressed_size = codec.compress(b"\x00" * 1_000_000) + monkeypatch.setattr(module, "MAX_DECOMPRESSED_BLOCK_SIZE", 1024) + + assert compressed_size < 1024 + + with pytest.raises(ValueError, match="Decompressed block exceeds the maximum of 1024 bytes"): + codec.decompress(compressed)