AVRO-4282: [python] Bound collection size when decoding arrays and maps

lang/py/avro/errors.py

@@ -83,6 +83,15 @@ def __init__(self, *args):
         return super().__init__(f"The exponent of {datum}, {exponent}, is too large for the schema scale of {scale}")
 
 
+class AvroCollectionSizeException(AvroException):
+    """Raised when an array or map declares more items than the configured maximum.
+
+    This guards against unbounded memory allocation when the block count of an
+    array or map read from a decoder is very large, for example because the
+    input is malformed or truncated.
+    """
+
+
 class SchemaResolutionException(AvroException):
     def __init__(self, fail_msg, writers_schema=None, readers_schema=None, *args):
         writers_message = f"\nWriter's Schema: {_safe_pretty(writers_schema)}" if writers_schema else ""

lang/py/avro/io.py

@@ -87,6 +87,7 @@
 import collections
 import datetime
 import decimal
+import os
 import struct
 import warnings
 from typing import IO, Generator, Iterable, List, Mapping, Optional, Sequence, Union
@@ -103,6 +104,56 @@
 STRUCT_SIGNED_INT = struct.Struct(">i")  # big-endian signed int
 STRUCT_SIGNED_LONG = struct.Struct(">q")  # big-endian signed long
 
+# Name of the environment variable used to override the default maximum number
+# of items permitted in a single decoded array or map.
+MAX_COLLECTION_ITEMS_ENV = "AVRO_MAX_COLLECTION_ITEMS"
+
+# Default upper bound on the number of items in a single array or map decoded
+# from a stream. When decoding, the block count of an array or map is read from
+# the (potentially untrusted or truncated) input and drives the allocation of
+# the resulting collection. Reading a collection that declares more items than
+# this limit raises an :class:`avro.errors.AvroCollectionSizeException` instead
+# of attempting a potentially huge allocation. This mirrors the Java SDK's
+# ``org.apache.avro.limits.collectionItems.maxLength`` limit. The default may be
+# overridden with the ``AVRO_MAX_COLLECTION_ITEMS`` environment variable.
+DEFAULT_MAX_COLLECTION_ITEMS = (1 << 31) - 8
+
+
+def _default_max_collection_items() -> int:
+    """Return the default collection item limit, honoring the environment override."""
+    value = os.environ.get(MAX_COLLECTION_ITEMS_ENV)
+    if value is None:
+        return DEFAULT_MAX_COLLECTION_ITEMS
+    try:
+        parsed = int(value)
+    except ValueError:
+        warnings.warn(avro.errors.AvroWarning(f"Ignoring invalid {MAX_COLLECTION_ITEMS_ENV} value: {value!r}"))
+        return DEFAULT_MAX_COLLECTION_ITEMS
+    if parsed < 0:
+        warnings.warn(avro.errors.AvroWarning(f"Ignoring negative {MAX_COLLECTION_ITEMS_ENV} value: {value!r}"))
+        return DEFAULT_MAX_COLLECTION_ITEMS
+    return parsed
+
+
+def check_max_collection_items(existing: int, items: int, limit: int) -> None:
+    """Guard against unbounded allocation when decoding an array or map.
+
+    :param existing: the number of items already read for the collection.
+    :param items: the number of items in the next block to be read. Callers pass
+                  the normalized (non-negative) block count -- in the Avro
+                  encoding a negative count merely signals that a block-size long
+                  follows and its absolute value is the count.
+    :param limit: the maximum total number of items permitted.
+    :raises avro.errors.AvroCollectionSizeException: if the running total would
+        exceed ``limit`` (or, defensively, if ``items`` is negative, which would
+        indicate an un-normalized count was passed).
+    """
+    if items < 0:
+        raise avro.errors.AvroCollectionSizeException(f"Block count must be normalized before checking, got: {items}")
+    if existing + items > limit:
+        raise avro.errors.AvroCollectionSizeException(f"Cannot read collections larger than {limit} items")
+
+
 ValidationNode = collections.namedtuple("ValidationNode", ["schema", "datum", "name"])
 ValidationNodeGeneratorType = Generator[ValidationNode, None, None]
 JsonScalarFieldType = Union[None, bool, str, int, float]
@@ -613,6 +664,10 @@ def __init__(self, writers_schema: Optional[avro.schema.Schema] = None, readers_
         """
         self._writers_schema = writers_schema
         self._readers_schema = readers_schema
+        # Maximum number of items permitted in a single decoded array or map.
+        # Guards against unbounded allocation from a large block count read from
+        # untrusted or truncated input. See DEFAULT_MAX_COLLECTION_ITEMS.
+        self.max_collection_items: int = _default_max_collection_items()
 
     @property
     def writers_schema(self) -> Optional[avro.schema.Schema]:
@@ -795,26 +850,35 @@ def read_array(self, writers_schema: avro.schema.ArraySchema, readers_schema: av
         The actual count in this case
         is the absolute value of the count written.
         """
-        read_items = []
+        read_items: List[object] = []
         block_count = decoder.read_long()
         while block_count != 0:
             if block_count < 0:
                 block_count = -block_count
                 decoder.skip_long()
+            check_max_collection_items(len(read_items), block_count, self.max_collection_items)
             for i in range(block_count):
                 read_items.append(self.read_data(writers_schema.items, readers_schema.items, decoder))
             block_count = decoder.read_long()
         return read_items
 
     def skip_array(self, writers_schema: avro.schema.ArraySchema, decoder: BinaryDecoder) -> None:
         block_count = decoder.read_long()
+        items_skipped = 0
         while block_count != 0:
             if block_count < 0:
                 block_size = decoder.read_long()
+                if block_size < 0:
+                    raise avro.errors.InvalidAvroBinaryEncoding(f"Array block size must be non-negative, got: {block_size}")
                 decoder.skip(block_size)
             else:
+                # A negative block is skipped by byte size above; only the
+                # positive per-item skip loop can be driven unbounded by a large
+                # untrusted block count, so apply the same cumulative cap here.
+                check_max_collection_items(items_skipped, block_count, self.max_collection_items)
                 for i in range(block_count):
                     self.skip_data(writers_schema.items, decoder)
+                items_skipped += block_count
             block_count = decoder.read_long()
 
     def read_map(self, writers_schema: avro.schema.MapSchema, readers_schema: avro.schema.MapSchema, decoder: BinaryDecoder) -> Mapping[str, object]:
@@ -834,26 +898,39 @@ def read_map(self, writers_schema: avro.schema.MapSchema, readers_schema: avro.s
         """
         read_items = {}
         block_count = decoder.read_long()
+        # Track the number of pairs decoded rather than len(read_items): repeated
+        # keys collapse in the dict and would otherwise let the cumulative check
+        # be bypassed by a stream that keeps rewriting the same key.
+        items_read = 0
         while block_count != 0:
             if block_count < 0:
                 block_count = -block_count
                 decoder.skip_long()
+            check_max_collection_items(items_read, block_count, self.max_collection_items)
             for i in range(block_count):
                 key = decoder.read_utf8()
                 read_items[key] = self.read_data(writers_schema.values, readers_schema.values, decoder)
+            items_read += block_count
             block_count = decoder.read_long()
         return read_items
 
     def skip_map(self, writers_schema: avro.schema.MapSchema, decoder: BinaryDecoder) -> None:
         block_count = decoder.read_long()
+        items_skipped = 0
         while block_count != 0:
             if block_count < 0:
                 block_size = decoder.read_long()
+                if block_size < 0:
+                    raise avro.errors.InvalidAvroBinaryEncoding(f"Map block size must be non-negative, got: {block_size}")
                 decoder.skip(block_size)
             else:
+                # As in skip_array, only the positive per-item skip loop can be
+                # driven unbounded, so bound the cumulative count here too.
+                check_max_collection_items(items_skipped, block_count, self.max_collection_items)
                 for i in range(block_count):
                     decoder.skip_utf8()
                     self.skip_data(writers_schema.values, decoder)
+                items_skipped += block_count
             block_count = decoder.read_long()
 
     def read_union(self, writers_schema: avro.schema.UnionSchema, readers_schema: avro.schema.UnionSchema, decoder: BinaryDecoder) -> object:

lang/py/avro/test/test_io.py

@@ -23,11 +23,13 @@
 import io
 import itertools
 import json
+import os
 import unittest
 import uuid
 import warnings
 from typing import BinaryIO, Collection, Dict, List, Optional, Tuple, Union, cast
 
+import avro.errors
 import avro.io
 import avro.schema
 import avro.timezones
@@ -715,6 +717,167 @@ def test_type_exception_record(self) -> None:
             write_datum(datum_to_write, writers_schema)
 
 
+class TestCollectionSizeLimit(unittest.TestCase):
+    """Decoding arrays and maps must bound the block count read from the input.
+
+    The block count of an array or map is read from a potentially untrusted or
+    truncated stream and drives allocation of the resulting collection. These
+    tests ensure a pathological block count raises an error instead of
+    attempting an unbounded allocation.
+    """
+
+    def _encode_longs(self, *values: int) -> io.BytesIO:
+        buffer = io.BytesIO()
+        encoder = avro.io.BinaryEncoder(buffer)
+        for value in values:
+            encoder.write_long(value)
+        return buffer
+
+    def _decode(self, buffer: io.BytesIO, schema: avro.schema.Schema, max_items: Optional[int] = None) -> object:
+        reader = io.BytesIO(buffer.getvalue())
+        decoder = avro.io.BinaryDecoder(reader)
+        datum_reader = avro.io.DatumReader(schema)
+        if max_items is not None:
+            datum_reader.max_collection_items = max_items
+        return datum_reader.read(decoder)
+
+    def test_array_block_count_exceeds_default_limit(self) -> None:
+        """A tiny input declaring more items than the default limit is rejected."""
+        schema = avro.schema.parse('{"type": "array", "items": "null"}')
+        buffer = self._encode_longs(avro.io.DEFAULT_MAX_COLLECTION_ITEMS + 1)
+        with self.assertRaises(avro.errors.AvroCollectionSizeException):
+            self._decode(buffer, schema)
+
+    def test_map_block_count_exceeds_default_limit(self) -> None:
+        schema = avro.schema.parse('{"type": "map", "values": "null"}')
+        buffer = self._encode_longs(avro.io.DEFAULT_MAX_COLLECTION_ITEMS + 1)
+        with self.assertRaises(avro.errors.AvroCollectionSizeException):
+            self._decode(buffer, schema)
+
+    def test_array_respects_configured_limit(self) -> None:
+        schema = avro.schema.parse('{"type": "array", "items": "null"}')
+        buffer = self._encode_longs(11, 0)  # 11 items in one block, then end-of-array
+        with self.assertRaises(avro.errors.AvroCollectionSizeException):
+            self._decode(buffer, schema, max_items=10)
+
+    def test_map_respects_configured_limit(self) -> None:
+        schema = avro.schema.parse('{"type": "map", "values": "null"}')
+        buffer = self._encode_longs(11, 0)
+        with self.assertRaises(avro.errors.AvroCollectionSizeException):
+            self._decode(buffer, schema, max_items=10)
+
+    def test_array_cumulative_limit_across_blocks(self) -> None:
+        """Two blocks that individually fit but together exceed the limit are rejected."""
+        schema = avro.schema.parse('{"type": "array", "items": "null"}')
+        buffer = self._encode_longs(6, 6, 0)
+        with self.assertRaises(avro.errors.AvroCollectionSizeException):
+            self._decode(buffer, schema, max_items=10)
+
+    def test_map_cumulative_limit_with_repeated_keys(self) -> None:
+        """Repeated keys must not let the cumulative pair count bypass the limit."""
+        schema = avro.schema.parse('{"type": "map", "values": "null"}')
+        buffer = io.BytesIO()
+        encoder = avro.io.BinaryEncoder(buffer)
+        encoder.write_long(6)  # first block: 6 pairs (all the same key)
+        for _ in range(6):
+            encoder.write_utf8("a")  # identical key; a null value occupies no bytes
+        encoder.write_long(6)  # second block would push the total to 12 > 10
+        with self.assertRaises(avro.errors.AvroCollectionSizeException):
+            self._decode(buffer, schema, max_items=10)
+
+    def test_negative_block_count_is_bounded(self) -> None:
+        """A negative block count (with block size) uses its absolute value and is bounded."""
+        schema = avro.schema.parse('{"type": "array", "items": "null"}')
+        buffer = self._encode_longs(-11, 0)  # count=-11 -> 11 items, followed by block size 0
+        with self.assertRaises(avro.errors.AvroCollectionSizeException):
+            self._decode(buffer, schema, max_items=10)
+
+    def test_array_within_limit_still_reads(self) -> None:
+        schema = avro.schema.parse('{"type": "array", "items": "null"}')
+        buffer = self._encode_longs(3, 0)  # 3 null items, then end-of-array
+        self.assertEqual([None, None, None], self._decode(buffer, schema, max_items=10))
+
+    def _reader(self, buffer: io.BytesIO, schema: avro.schema.Schema, max_items: int) -> "tuple[avro.io.DatumReader, avro.io.BinaryDecoder]":
+        decoder = avro.io.BinaryDecoder(io.BytesIO(buffer.getvalue()))
+        datum_reader = avro.io.DatumReader(schema)
+        datum_reader.max_collection_items = max_items
+        return datum_reader, decoder
+
+    def test_skip_array_respects_limit(self) -> None:
+        """Skipping an array must bound the per-item loop like reading does."""
+        schema = avro.schema.parse('{"type": "array", "items": "null"}')
+        buffer = self._encode_longs(11, 0)  # 11 items in one block
+        datum_reader, decoder = self._reader(buffer, schema, max_items=10)
+        with self.assertRaises(avro.errors.AvroCollectionSizeException):
+            datum_reader.skip_array(cast(avro.schema.ArraySchema, schema), decoder)
+
+    def test_skip_map_respects_limit(self) -> None:
+        schema = avro.schema.parse('{"type": "map", "values": "null"}')
+        buffer = self._encode_longs(11, 0)
+        datum_reader, decoder = self._reader(buffer, schema, max_items=10)
+        with self.assertRaises(avro.errors.AvroCollectionSizeException):
+            datum_reader.skip_map(cast(avro.schema.MapSchema, schema), decoder)
+
+    def test_skip_array_rejects_negative_block_size(self) -> None:
+        """A negative block size must not seek the decoder backwards."""
+        schema = avro.schema.parse('{"type": "array", "items": "null"}')
+        buffer = self._encode_longs(-1, -1)  # negative block count, then negative block size
+        datum_reader, decoder = self._reader(buffer, schema, max_items=10)
+        with self.assertRaises(avro.errors.InvalidAvroBinaryEncoding):
+            datum_reader.skip_array(cast(avro.schema.ArraySchema, schema), decoder)
+
+    def test_skip_map_rejects_negative_block_size(self) -> None:
+        schema = avro.schema.parse('{"type": "map", "values": "null"}')
+        buffer = self._encode_longs(-1, -1)
+        datum_reader, decoder = self._reader(buffer, schema, max_items=10)
+        with self.assertRaises(avro.errors.InvalidAvroBinaryEncoding):
+            datum_reader.skip_map(cast(avro.schema.MapSchema, schema), decoder)
+
+    def _default_limit_with_env(self, value: Optional[str]) -> int:
+        """Return the default collection-item limit with the env var set to `value`."""
+        previous = os.environ.get(avro.io.MAX_COLLECTION_ITEMS_ENV)
+        try:
+            if value is None:
+                os.environ.pop(avro.io.MAX_COLLECTION_ITEMS_ENV, None)
+            else:
+                os.environ[avro.io.MAX_COLLECTION_ITEMS_ENV] = value
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore", avro.errors.AvroWarning)
+                return avro.io._default_max_collection_items()
+        finally:
+            if previous is None:
+                os.environ.pop(avro.io.MAX_COLLECTION_ITEMS_ENV, None)
+            else:
+                os.environ[avro.io.MAX_COLLECTION_ITEMS_ENV] = previous
+
+    def test_env_override_is_honored(self) -> None:
+        self.assertEqual(5, self._default_limit_with_env("5"))
+
+    def test_env_override_invalid_falls_back_to_default(self) -> None:
+        for bad in ("not-a-number", "-1", ""):
+            self.assertEqual(avro.io.DEFAULT_MAX_COLLECTION_ITEMS, self._default_limit_with_env(bad))
+
+    def test_env_override_unset_uses_default(self) -> None:
+        self.assertEqual(avro.io.DEFAULT_MAX_COLLECTION_ITEMS, self._default_limit_with_env(None))
+
+    def test_env_override_rejects_over_limit_collection(self) -> None:
+        """A reader constructed with the env-set limit enforces it end to end."""
+        previous = os.environ.get(avro.io.MAX_COLLECTION_ITEMS_ENV)
+        try:
+            os.environ[avro.io.MAX_COLLECTION_ITEMS_ENV] = "5"
+            schema = avro.schema.parse('{"type": "array", "items": "null"}')
+            datum_reader = avro.io.DatumReader(schema)  # reads the env at construction
+            self.assertEqual(5, datum_reader.max_collection_items)
+            decoder = avro.io.BinaryDecoder(io.BytesIO(self._encode_longs(6, 0).getvalue()))
+            with self.assertRaises(avro.errors.AvroCollectionSizeException):
+                datum_reader.read(decoder)
+        finally:
+            if previous is None:
+                os.environ.pop(avro.io.MAX_COLLECTION_ITEMS_ENV, None)
+            else:
+                os.environ[avro.io.MAX_COLLECTION_ITEMS_ENV] = previous
+
+
 def load_tests(loader: unittest.TestLoader, default_tests: None, pattern: None) -> unittest.TestSuite:
     """Generate test cases across many test schema."""
     suite = unittest.TestSuite()
@@ -729,6 +892,7 @@ def load_tests(loader: unittest.TestLoader, default_tests: None, pattern: None)
     )
     suite.addTests(DefaultValueTestCase(field_type, default) for field_type, default in DEFAULT_VALUE_EXAMPLES)
     suite.addTests(loader.loadTestsFromTestCase(TestIncompatibleSchemaReading))
+    suite.addTests(loader.loadTestsFromTestCase(TestCollectionSizeLimit))
     return suite