diff --git a/lang/py/avro/io.py b/lang/py/avro/io.py index a67230f0734..f5063a6686f 100644 --- a/lang/py/avro/io.py +++ b/lang/py/avro/io.py @@ -686,14 +686,16 @@ def write_timestamp_micros_long(self, datum: datetime.datetime) -> None: # elements in any collection, allocated from a single decode. MAX_COLLECTION_ITEMS_ENV = "AVRO_MAX_COLLECTION_ITEMS" -# Default maximum number of zero-byte-encoded collection elements to allocate. -# Elements whose schema encodes to zero bytes (``null``, a zero-length ``fixed``, -# or a record with only zero-byte fields) consume no input, so the bytes-remaining -# check cannot bound their count; without a cap a tiny payload can declare a huge -# block count and exhaust memory. A legitimate collection of zero-byte elements is -# small, so this default is generous while still rejecting pathological input. It -# can be raised (or lowered) with the ``AVRO_MAX_COLLECTION_ITEMS`` environment -# variable. +# Default maximum number of zero-byte-encoded collection elements to allocate +# across a single decoded datum. Elements whose schema encodes to zero bytes +# (``null``, a zero-length ``fixed``, or a record with only zero-byte fields) +# consume no input, so the bytes-remaining check cannot bound their count; without +# a cap a tiny payload can declare a huge block count and exhaust memory. The cap +# is cumulative over the whole datum rather than per collection, because a record's +# schema can declare many such collection fields, each individually under the limit +# but jointly unbounded. A legitimate datum of zero-byte elements is small, so this +# default is generous while still rejecting pathological input. It can be raised (or +# lowered) with the ``AVRO_MAX_COLLECTION_ITEMS`` environment variable. DEFAULT_MAX_COLLECTION_ITEMS = 10_000_000 # Default structural cap on the number of elements in any array or map (a @@ -805,6 +807,15 @@ def __init__(self, writers_schema: Optional[avro.schema.Schema] = None, readers_ """ self._writers_schema = writers_schema self._readers_schema = readers_schema + # Cumulative number of zero-byte-encoded collection elements (e.g. an + # array of nulls) allocated while decoding the *current* datum. Reset at + # the start of each top-level read(). Because such elements consume no + # input bytes, the bytes-remaining check cannot bound them; capping the + # count per collection is not enough either, since a single record's + # schema can declare many collection fields, each individually under the + # limit but together unbounded. The cap is therefore applied across the + # whole datum. See _ensure_collection_available. + self._zero_byte_items_read = 0 @property def writers_schema(self) -> Optional[avro.schema.Schema]: @@ -828,6 +839,12 @@ def read(self, decoder: "BinaryDecoder") -> object: reader_schema = self.readers_schema if reader_schema is None: reader_schema = self.writers_schema + # Start a fresh zero-byte-element budget for this datum. The cap bounds + # the *cumulative* number of zero-byte collection elements decoded across + # every collection in this datum, not per collection, so a record made of + # many small collection fields cannot bypass it (see + # _ensure_collection_available). + self._zero_byte_items_read = 0 return self.read_data(self.writers_schema, reader_schema, decoder) def read_data(self, writers_schema: avro.schema.Schema, readers_schema: avro.schema.Schema, decoder: "BinaryDecoder") -> object: @@ -982,8 +999,8 @@ def read_enum(self, writers_schema: avro.schema.EnumSchema, readers_schema: avro def skip_enum(self, writers_schema: avro.schema.EnumSchema, decoder: BinaryDecoder) -> None: return decoder.skip_int() - @staticmethod def _ensure_collection_available( + self, decoder: BinaryDecoder, existing: int, count: int, @@ -997,10 +1014,17 @@ def _ensure_collection_available( For elements with a positive minimum on-wire size, the declared count is checked against the bytes actually remaining and against a structural - limit (an overflow/defense-in-depth cap). For zero-byte elements (e.g. an - array of nulls), which consume no input and so cannot be bounded by the - bytes remaining, the cumulative count is checked against the (tighter) - zero-byte limit. + limit (an overflow/defense-in-depth cap); both are naturally bounded per + collection because decoding consumes input, so ``existing`` is the count + already read in *this* collection. + + For zero-byte elements (e.g. an array of nulls), which consume no input, + neither the bytes remaining nor a per-collection count can bound them: a + single datum's schema may declare many such collections (one per record + field), each individually small but jointly unbounded. Their count is + therefore accumulated on the reader across the whole datum + (``self._zero_byte_items_read``, reset per top-level ``read()``) and + checked against the (tighter) zero-byte limit. """ if count <= 0: return @@ -1024,10 +1048,14 @@ def _ensure_collection_available( f"Cannot read a collection of more than {structural_limit} elements " f"(declared {existing + count}); raise the {MAX_COLLECTION_ITEMS_ENV} limit if this is legitimate." ) - elif existing + count > zero_byte_limit: + return + # Zero-byte element type: bound the cumulative count across the datum. + self._zero_byte_items_read += count + if self._zero_byte_items_read > zero_byte_limit: raise avro.errors.AvroCollectionSizeException( - f"Cannot read a collection of more than {zero_byte_limit} zero-byte elements " - f"(declared {existing + count}); raise the {MAX_COLLECTION_ITEMS_ENV} limit if this is legitimate." + f"Cannot read more than {zero_byte_limit} zero-byte collection elements " + f"in a single datum (reached {self._zero_byte_items_read}); raise the " + f"{MAX_COLLECTION_ITEMS_ENV} limit if this is legitimate." ) def read_array(self, writers_schema: avro.schema.ArraySchema, readers_schema: avro.schema.ArraySchema, decoder: BinaryDecoder) -> List[object]: diff --git a/lang/py/avro/test/test_io.py b/lang/py/avro/test/test_io.py index 9ee39b1937b..1a71e0598e5 100644 --- a/lang/py/avro/test/test_io.py +++ b/lang/py/avro/test/test_io.py @@ -1142,6 +1142,52 @@ def test_array_of_null_cumulative_across_blocks(self) -> None: buf.getvalue(), ) + def test_record_of_array_of_null_fields_cumulative_across_datum(self) -> None: + # AVRO-4296 follow-up: the cap is per decoded datum, not per collection. + # A record whose schema declares several array fields, each block + # individually under the limit, must still be rejected once their combined + # count exceeds it. Here two fields of 600 nulls each (1200 > 1000) are + # rejected on the second field, mirroring the multi-field container-file + # amplification (many small collection fields, unbounded in aggregate). + schema_json = json.dumps( + { + "type": "record", + "name": "R", + "fields": [ + {"name": "a", "type": {"type": "array", "items": "null"}}, + {"name": "b", "type": {"type": "array", "items": "null"}}, + ], + } + ) + payload = self._array_block(600) + self._array_block(600) + with unittest.mock.patch.dict(os.environ, {"AVRO_MAX_COLLECTION_ITEMS": "1000"}): + self.assertRaises(avro.errors.AvroCollectionSizeException, self._decode, schema_json, payload) + + def test_record_of_array_of_null_fields_within_datum_limit_reads(self) -> None: + # The complement of the amplification test: two array fields whose + # combined count stays under the per-datum cap decode normally, and a + # fresh read() resets the budget so a subsequent datum is not penalized. + schema_json = json.dumps( + { + "type": "record", + "name": "R", + "fields": [ + {"name": "a", "type": {"type": "array", "items": "null"}}, + {"name": "b", "type": {"type": "array", "items": "null"}}, + ], + } + ) + payload = self._array_block(400) + self._array_block(400) + schema = avro.schema.parse(schema_json) + reader = avro.io.DatumReader(schema) + with unittest.mock.patch.dict(os.environ, {"AVRO_MAX_COLLECTION_ITEMS": "1000"}): + with io.BytesIO(payload) as bio: + self.assertEqual(reader.read(avro.io.BinaryDecoder(bio)), {"a": [None] * 400, "b": [None] * 400}) + # The budget resets per top-level read(): decoding the same datum + # again on the same reader must not accumulate across datums. + with io.BytesIO(payload) as bio: + self.assertEqual(reader.read(avro.io.BinaryDecoder(bio)), {"a": [None] * 400, "b": [None] * 400}) + def test_map_duplicate_keys_counted_cumulatively(self) -> None: # Two blocks of 600 pairs that repeat the SAME key: len(read_items) would # be 1, so a separate decoded-pair counter is needed to reject 1200 > 1000.