AVRO-4241: [Java] Bound zero-byte collection elements per datum, not per collection

lang/java/avro/src/main/java/org/apache/avro/SystemLimitException.java

@@ -119,6 +119,24 @@ public class SystemLimitException extends AvroRuntimeException {
    */
   private static long maxCollectionAllocation = defaultMaxCollectionAllocation();
 
+  /**
+   * Per-thread cumulative accounting of zero-byte collection elements allocated
+   * while decoding a single datum. The {@link #maxCollectionAllocation} cap on
+   * such elements must apply across the whole datum, not per collection: a
+   * container file carries its own schema, so an attacker can declare a record
+   * with many collection fields, each block individually under the limit but
+   * jointly unbounded. A depth counter marks the outermost decode scope so the
+   * running total is reset only there and accumulates across every (possibly
+   * nested) collection in between.
+   */
+  private static final class CollectionAllocationScope {
+    private int depth;
+    private long allocated;
+  }
+
+  private static final ThreadLocal<CollectionAllocationScope> COLLECTION_ALLOCATION_SCOPE = ThreadLocal
+      .withInitial(CollectionAllocationScope::new);
+
   static {
     resetLimits();
   }
@@ -334,6 +352,68 @@ public static long checkMaxCollectionAllocation(long existing, long items) {
     return total;
   }
 
+  /**
+   * Begin an outermost decode scope for cumulative zero-byte collection
+   * allocation accounting. Must be paired with
+   * {@link #endCollectionAllocationScope()} in a {@code finally} block. Scopes
+   * nest: only the outermost one resets the running total, so the cap applies
+   * across the whole datum rather than per collection. See
+   * {@link #checkMaxCollectionAllocation(long)}.
+   */
+  public static void beginCollectionAllocationScope() {
+    CollectionAllocationScope scope = COLLECTION_ALLOCATION_SCOPE.get();
+    if (scope.depth == 0) {
+      scope.allocated = 0;
+    }
+    scope.depth++;
+  }
+
+  /**
+   * End a decode scope opened by {@link #beginCollectionAllocationScope()}. When
+   * the outermost scope closes the running total is cleared so it never leaks
+   * into an unrelated later decode on the same thread.
+   */
+  public static void endCollectionAllocationScope() {
+    CollectionAllocationScope scope = COLLECTION_ALLOCATION_SCOPE.get();
+    if (scope.depth > 0) {
+      scope.depth--;
+      if (scope.depth == 0) {
+        scope.allocated = 0;
+      }
+    }
+  }
+
+  /**
+   * Accumulate {@code items} zero-byte-minimum collection elements into the
+   * current decode scope and verify the running total stays within
+   * {@link #MAX_COLLECTION_ALLOCATION_PROPERTY the allocation limit}.
+   * <p>
+   * Unlike {@link #checkMaxCollectionAllocation(long, long)}, which bounds a
+   * single collection, this bounds the cumulative count across every collection
+   * decoded within the enclosing {@link #beginCollectionAllocationScope() scope}
+   * (one datum), so a record made of many small zero-byte collection fields
+   * cannot bypass the cap in aggregate. When called outside any scope it falls
+   * back to a stateless single-collection check, preserving the previous
+   * behaviour for callers that do not delimit a datum.
+   *
+   * @param items The next number of zero-byte elements to allocate.
+   * @return The cumulative element count if and only if it is within the limit.
+   * @throws SystemLimitException if the cumulative allocation would exceed the
+   *                              limit.
+   * @throws AvroRuntimeException if {@code items} is negative.
+   */
+  public static long checkMaxCollectionAllocation(long items) {
+    CollectionAllocationScope scope = COLLECTION_ALLOCATION_SCOPE.get();
+    if (scope.depth == 0) {
+      // Not inside a delimited datum: behave as a per-collection check so this
+      // path is never stricter than before for callers that do not open a scope.
+      return checkMaxCollectionAllocation(0L, items);
+    }
+    long total = checkMaxCollectionAllocation(scope.allocated, items);
+    scope.allocated = total;
+    return total;
+  }
+
   /**
    * Check to ensure that reading the string size is within the specified limits.
    *

lang/java/avro/src/main/java/org/apache/avro/generic/GenericDatumReader.java

@@ -168,18 +168,29 @@ protected final ResolvingDecoder getResolver(Schema actual, Schema expected) thr
   @Override
   @SuppressWarnings("unchecked")
   public D read(D reuse, Decoder in) throws IOException {
-    if (data.isFastReaderEnabled()) {
-      if (this.fastDatumReader == null) {
-        this.fastDatumReader = data.getFastReaderBuilder().createDatumReader(actual, expected);
+    // Open a decode scope so the zero-byte collection-element allocation cap is
+    // enforced cumulatively across this datum (see SystemLimitException): a
+    // record with many small array<null>-style fields, each individually under
+    // the limit, must not be able to over-allocate in aggregate. Nested scopes
+    // (e.g. the delegated fast reader, or skipped writer fields) accumulate into
+    // this one; only the outermost resets the running total.
+    SystemLimitException.beginCollectionAllocationScope();
+    try {
+      if (data.isFastReaderEnabled()) {
+        if (this.fastDatumReader == null) {
+          this.fastDatumReader = data.getFastReaderBuilder().createDatumReader(actual, expected);
+        }
+        return fastDatumReader.read(reuse, in);
       }
-      return fastDatumReader.read(reuse, in);
-    }
 
-    ResolvingDecoder resolver = getResolver(actual, expected);
-    resolver.configure(in);
-    D result = (D) read(reuse, expected, resolver);
-    resolver.drain();
-    return result;
+      ResolvingDecoder resolver = getResolver(actual, expected);
+      resolver.configure(in);
+      D result = (D) read(reuse, expected, resolver);
+      resolver.drain();
+      return result;
+    } finally {
+      SystemLimitException.endCollectionAllocationScope();
+    }
   }
 
   /** Called to read data. */
@@ -326,7 +337,7 @@ protected Object readArray(Object old, Schema expected, ResolvingDecoder in) thr
       // backing-array allocation.
       boolean zeroByteElements = isZeroByteSchema(expectedType);
       if (zeroByteElements) {
-        SystemLimitException.checkMaxCollectionAllocation(base, l);
+        SystemLimitException.checkMaxCollectionAllocation(l);
       }
       LogicalType logicalType = expectedType.getLogicalType();
       Conversion<?> conversion = getData().getConversionFor(logicalType);
@@ -345,7 +356,7 @@ protected Object readArray(Object old, Schema expected, ResolvingDecoder in) thr
         base += l;
         l = arrayNext(in, expectedType);
         if (zeroByteElements && l > 0) {
-          SystemLimitException.checkMaxCollectionAllocation(base, l);
+          SystemLimitException.checkMaxCollectionAllocation(l);
         }
       } while (l > 0);
       return pruneArray(array);
@@ -792,10 +803,24 @@ protected Object createBytes(byte[] value) {
 
   /** Skip an instance of a schema. */
   public static void skip(Schema schema, Decoder in) throws IOException {
+    // Delimit a decode scope so a huge count of zero-byte elements split across
+    // fields/blocks is bounded cumulatively (see SystemLimitException). Scopes
+    // nest, so a skip invoked mid-read (e.g. an unused writer field) accumulates
+    // into the enclosing datum budget instead of resetting it, while a top-level
+    // skip (e.g. from BinaryData.compare) is bounded per invocation.
+    SystemLimitException.beginCollectionAllocationScope();
+    try {
+      skipInternal(schema, in);
+    } finally {
+      SystemLimitException.endCollectionAllocationScope();
+    }
+  }
+
+  private static void skipInternal(Schema schema, Decoder in) throws IOException {
     switch (schema.getType()) {
     case RECORD:
       for (Field field : schema.getFields())
-        skip(field.schema(), in);
+        skipInternal(field.schema(), in);
       break;
     case ENUM:
       in.readEnum();
@@ -816,11 +841,11 @@ public static void skip(Schema schema, Decoder in) throws IOException {
         // cannot drive an unbounded skip loop.
         SystemLimitException.checkMaxCollectionLength(arrayTotal, l);
         if (zeroByteElements) {
-          SystemLimitException.checkMaxCollectionAllocation(arrayTotal, l);
+          SystemLimitException.checkMaxCollectionAllocation(l);
         }
         arrayTotal += l;
         for (long i = 0; i < l; i++) {
-          skip(elementType, in);
+          skipInternal(elementType, in);
         }
       }
       break;
@@ -833,12 +858,12 @@ public static void skip(Schema schema, Decoder in) throws IOException {
         mapTotal += l;
         for (long i = 0; i < l; i++) {
           in.skipString();
-          skip(value, in);
+          skipInternal(value, in);
         }
       }
       break;
     case UNION:
-      skip(schema.getTypes().get(in.readIndex()), in);
+      skipInternal(schema.getTypes().get(in.readIndex()), in);
       break;
     case FIXED:
       in.skipFixed(schema.getFixedSize());

lang/java/avro/src/main/java/org/apache/avro/io/FastReaderBuilder.java

@@ -478,38 +478,48 @@ private FieldReader createArrayReader(Schema readerSchema, Container action) thr
     boolean zeroByteElements = GenericDatumReader.isZeroByteSchema(elementType);
 
     return reusingReader((reuse, decoder) -> {
-      if (reuse instanceof GenericArray) {
-        GenericArray<Object> reuseArray = (GenericArray<Object>) reuse;
-        long l = decoder.readArrayStart();
-        long total = 0;
-        checkArrayBlock(decoder, elementType, zeroByteElements, total, l);
-        reuseArray.clear();
-
-        while (l > 0) {
-          for (long i = 0; i < l; i++) {
-            reuseArray.add(elementReader.read(reuseArray.peek(), decoder));
+      // Open a decode scope so the zero-byte element allocation cap is cumulative
+      // across every block of this array even when the fast reader is used
+      // standalone (i.e. without GenericDatumReader.read opening the outer datum
+      // scope); otherwise a huge array split into many small blocks would bypass
+      // the cap. The scope nests: when a datum scope is already open this simply
+      // accumulates into it, and only the outermost scope resets the running
+      // total (see SystemLimitException). The try/finally guarantees the scope is
+      // always closed so ThreadLocal state cannot leak into later decodes on the
+      // same thread.
+      SystemLimitException.beginCollectionAllocationScope();
+      try {
+        if (reuse instanceof GenericArray) {
+          GenericArray<Object> reuseArray = (GenericArray<Object>) reuse;
+          long l = decoder.readArrayStart();
+          checkArrayBlock(decoder, elementType, zeroByteElements, l);
+          reuseArray.clear();
+
+          while (l > 0) {
+            for (long i = 0; i < l; i++) {
+              reuseArray.add(elementReader.read(reuseArray.peek(), decoder));
+            }
+            l = decoder.arrayNext();
+            checkArrayBlock(decoder, elementType, zeroByteElements, l);
           }
-          total += l;
-          l = decoder.arrayNext();
-          checkArrayBlock(decoder, elementType, zeroByteElements, total, l);
-        }
-        return reuseArray;
-      } else {
-        long l = decoder.readArrayStart();
-        long total = 0;
-        checkArrayBlock(decoder, elementType, zeroByteElements, total, l);
-        List<Object> array = (reuse instanceof List) ? (List<Object>) reuse
-            : new GenericData.Array<>(GenericDatumReader.initialCollectionCapacity(l), readerSchema);
-        array.clear();
-        while (l > 0) {
-          for (long i = 0; i < l; i++) {
-            array.add(elementReader.read(null, decoder));
+          return reuseArray;
+        } else {
+          long l = decoder.readArrayStart();
+          checkArrayBlock(decoder, elementType, zeroByteElements, l);
+          List<Object> array = (reuse instanceof List) ? (List<Object>) reuse
+              : new GenericData.Array<>(GenericDatumReader.initialCollectionCapacity(l), readerSchema);
+          array.clear();
+          while (l > 0) {
+            for (long i = 0; i < l; i++) {
+              array.add(elementReader.read(null, decoder));
+            }
+            l = decoder.arrayNext();
+            checkArrayBlock(decoder, elementType, zeroByteElements, l);
           }
-          total += l;
-          l = decoder.arrayNext();
-          checkArrayBlock(decoder, elementType, zeroByteElements, total, l);
+          return array;
         }
-        return array;
+      } finally {
+        SystemLimitException.endCollectionAllocationScope();
       }
     });
   }
@@ -521,16 +531,18 @@ private FieldReader createArrayReader(Schema readerSchema, Container action) thr
    * heap-aware allocation cap for zero-byte elements (which the bytes check
    * cannot bound).
    */
-  private static void checkArrayBlock(Decoder decoder, Schema elementType, boolean zeroByteElements, long total,
-      long count) throws IOException {
+  private static void checkArrayBlock(Decoder decoder, Schema elementType, boolean zeroByteElements, long count)
+      throws IOException {
     if (count <= 0) {
       return;
     }
     if (zeroByteElements) {
       // The bytes-remaining check cannot bound zero-byte elements (minBytes is
       // 0, so ensureAvailableCollectionBytes would no-op after recomputing it);
-      // apply the heap-aware allocation cap instead.
-      SystemLimitException.checkMaxCollectionAllocation(total, count);
+      // apply the heap-aware allocation cap instead. The cap is cumulative across
+      // the enclosing datum scope (see SystemLimitException), so a record of many
+      // small array<null>-style fields cannot over-allocate in aggregate.
+      SystemLimitException.checkMaxCollectionAllocation(count);
     } else {
       GenericDatumReader.ensureAvailableCollectionBytes(decoder, count, elementType);
     }

lang/java/avro/src/main/java/org/apache/avro/reflect/ReflectDatumReader.java

@@ -153,7 +153,7 @@ protected Object readArray(Object old, Schema expected, ResolvingDecoder in) thr
     // eager allocation before any element is read.
     ensureAvailableCollectionBytes(in, l, expectedType);
     if (isZeroByteSchema(expectedType)) {
-      SystemLimitException.checkMaxCollectionAllocation(0, l);
+      SystemLimitException.checkMaxCollectionAllocation(l);
     }
     Object array = newArray(old, (int) l, expected);
     if (array instanceof Collection) {
@@ -209,7 +209,7 @@ private Object readObjectArray(Object[] array, Schema expectedType, long l, Reso
           array[index] = element;
           index++;
         }
-      } while ((l = nextArrayBlock(in, expectedType, index, zeroByte)) > 0);
+      } while ((l = nextArrayBlock(in, expectedType, zeroByte)) > 0);
     } else {
       do {
         int limit = index + (int) l;
@@ -218,7 +218,7 @@ private Object readObjectArray(Object[] array, Schema expectedType, long l, Reso
           array[index] = element;
           index++;
         }
-      } while ((l = nextArrayBlock(in, expectedType, index, zeroByte)) > 0);
+      } while ((l = nextArrayBlock(in, expectedType, zeroByte)) > 0);
     }
     return array;
   }
@@ -228,23 +228,20 @@ private Object readCollection(Collection<Object> c, Schema expectedType, long l,
     LogicalType logicalType = expectedType.getLogicalType();
     Conversion<?> conversion = getData().getConversionFor(logicalType);
     boolean zeroByte = isZeroByteSchema(expectedType);
-    long count = 0;
     if (logicalType != null && conversion != null) {
       do {
         for (int i = 0; i < l; i++) {
           Object element = readWithConversion(null, expectedType, logicalType, conversion, in);
           c.add(element);
         }
-        count += l;
-      } while ((l = nextArrayBlock(in, expectedType, count, zeroByte)) > 0);
+      } while ((l = nextArrayBlock(in, expectedType, zeroByte)) > 0);
     } else {
       do {
         for (int i = 0; i < l; i++) {
           Object element = readWithoutConversion(null, expectedType, in);
           c.add(element);
         }
-        count += l;
-      } while ((l = nextArrayBlock(in, expectedType, count, zeroByte)) > 0);
+      } while ((l = nextArrayBlock(in, expectedType, zeroByte)) > 0);
     }
     return c;
   }
@@ -254,22 +251,22 @@ private Object readCollection(Collection<Object> c, Schema expectedType, long l,
    * {@link org.apache.avro.generic.GenericDatumReader#readArray}: bound the
    * declared count against the bytes remaining, and for element types whose
    * minimum encoded size is zero bound the cumulative allocation (which the
-   * bytes-remaining check cannot). This closes the gap where a large logical
-   * array split across multiple blocks would otherwise pass only the first
-   * block's guard.
+   * bytes-remaining check cannot). The zero-byte allocation cap is cumulative
+   * across the enclosing datum scope (see
+   * {@link org.apache.avro.SystemLimitException}), closing the gap where a large
+   * logical array split across multiple blocks would otherwise pass only the
+   * first block's guard.
    *
    * @param in           the decoder
    * @param expectedType the array element schema
-   * @param existing     the number of elements already read
    * @param zeroByte     whether the element type's minimum encoded size is zero
    * @return the validated next block count
    */
-  private long nextArrayBlock(ResolvingDecoder in, Schema expectedType, long existing, boolean zeroByte)
-      throws IOException {
+  private long nextArrayBlock(ResolvingDecoder in, Schema expectedType, boolean zeroByte) throws IOException {
     long l = in.arrayNext();
     ensureAvailableCollectionBytes(in, l, expectedType);
     if (zeroByte && l > 0) {
-      SystemLimitException.checkMaxCollectionAllocation(existing, l);
+      SystemLimitException.checkMaxCollectionAllocation(l);
     }
     return l;
   }