iemejia commented on code in PR #3861:
URL: https://github.com/apache/avro/pull/3861#discussion_r3567551373
##########
lang/py/avro/io.py:
##########
@@ -599,6 +649,90 @@ def write_timestamp_micros_long(self, datum:
datetime.datetime) -> None:
#
# DatumReader/Writer
#
+#: Environment variable overriding the maximum number of zero-byte-encoded
+#: collection elements (e.g. an array of nulls) to allocate 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_MAX_COLLECTION_ITEMS = 10_000_000
+
+#: Default structural cap on the number of elements in any array or map (a
+#: collection larger than this is treated as malformed regardless of element
+#: type). Matches the historical ``Integer.MAX_VALUE - 8`` limit. Elements
with a
+#: positive on-wire size are also bounded by the bytes remaining; this cap is
an
+#: additional overflow/defense-in-depth guard. ``AVRO_MAX_COLLECTION_ITEMS``,
when
+#: set, overrides both this and :data:`DEFAULT_MAX_COLLECTION_ITEMS`.
+DEFAULT_MAX_COLLECTION_STRUCTURAL = (1 << 31) - 1 - 8 # Integer.MAX_VALUE - 8
+
+
+def _collection_limits() -> Tuple[int, int]:
+ """Return ``(zero_byte_limit, structural_limit)``.
+
+ ``AVRO_MAX_COLLECTION_ITEMS``, when set to a non-negative integer, caps
both
+ zero-byte-element collections and all other collections at that value.
+ Otherwise zero-byte elements use the tighter
:data:`DEFAULT_MAX_COLLECTION_ITEMS`
+ and all collections use :data:`DEFAULT_MAX_COLLECTION_STRUCTURAL`.
+ """
+ value = os.environ.get(MAX_COLLECTION_ITEMS_ENV)
+ if value is None:
+ return DEFAULT_MAX_COLLECTION_ITEMS, DEFAULT_MAX_COLLECTION_STRUCTURAL
+ 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, DEFAULT_MAX_COLLECTION_STRUCTURAL
+ if parsed < 0:
+ warnings.warn(avro.errors.AvroWarning(f"Ignoring negative
{MAX_COLLECTION_ITEMS_ENV} value: {value!r}"))
+ return DEFAULT_MAX_COLLECTION_ITEMS, DEFAULT_MAX_COLLECTION_STRUCTURAL
+ return parsed, parsed
+
+
+def _max_collection_items() -> int:
+ """Return the configured zero-byte-element collection limit."""
+ return _collection_limits()[0]
+
+
+def _min_bytes_per_element(schema: avro.schema.Schema, visited:
Optional[Set[int]] = None) -> int:
+ """
+ Return the minimum number of bytes a single value of ``schema`` can occupy
+ on the wire. Used to reject an array/map block count that could not
possibly
+ be backed by the bytes remaining. A type that can encode to zero bytes
+ (``null``) returns 0, which disables the collection check for it (avoiding
a
+ false positive on, e.g., an array of nulls).
+ """
+ if visited is None:
+ visited = set()
+ schema_type = schema.type
+ if schema_type == "null":
+ return 0
+ if schema_type == "float":
+ return 4
+ if schema_type == "double":
+ return 8
+ if schema_type == "fixed":
+ return getattr(schema, "size", 1)
+ if schema_type in ("record", "error"):
+ # Guard against self-referencing records (recursion would not
terminate).
+ if id(schema) in visited:
+ return 0
+ visited.add(id(schema))
+ total = 0
+ for field in getattr(schema, "fields", []):
+ total += _min_bytes_per_element(field.type, visited)
+ visited.discard(id(schema))
+ return total
+ # boolean, int, long, bytes, string, enum, union, array, map: all >= 1 byte
+ # (a union encodes at least a 1-byte branch index).
+ return 1
Review Comment:
Fixed — _min_bytes_per_element now returns `1 + min(branch minimums)` for a
union (the branch index plus the smallest branch payload), so a union whose
branches all have a positive minimum is no longer underestimated as 1. It stays
1 when any branch is null/zero-byte.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
