fix(avro): bound the decompressed size of a block

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:
 

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

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

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:

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)