diff --git a/lang/py/avro/errors.py b/lang/py/avro/errors.py index b961a04ae43..138c0feefb5 100644 --- a/lang/py/avro/errors.py +++ b/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 "" diff --git a/lang/py/avro/io.py b/lang/py/avro/io.py index 4be3d8e0574..9b246db44f9 100644 --- a/lang/py/avro/io.py +++ b/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,12 +850,13 @@ 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() @@ -808,13 +864,21 @@ def read_array(self, writers_schema: avro.schema.ArraySchema, readers_schema: av 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: diff --git a/lang/py/avro/test/test_io.py b/lang/py/avro/test/test_io.py index 41a0e366bad..ad4ff23ca78 100644 --- a/lang/py/avro/test/test_io.py +++ b/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