AVRO-4300: [java] Bound array/map allocation and skipping for zero-byte elements and on the fast reader path

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

@@ -28,13 +28,20 @@
  * The following system properties can be set to limit the size of bytes,
  * strings and collection types to be allocated:
  * <ul>
- * <li><tt>org.apache.avro.limits.byte.maxLength</tt></li> limits the maximum
- * size of <tt>byte</tt> types.</li>
- * <li><tt>org.apache.avro.limits.collectionItems.maxLength</tt></li> limits the
- * maximum number of <tt>map</tt> and <tt>list</tt> items that can be read at
- * once single sequence.</li>
- * <li><tt>org.apache.avro.limits.string.maxLength</tt></li> limits the maximum
- * size of <tt>string</tt> types.</li>
+ * <li><tt>org.apache.avro.limits.bytes.maxLength</tt> limits the maximum size
+ * of <tt>bytes</tt> types.</li>
+ * <li><tt>org.apache.avro.limits.collectionItems.maxLength</tt> limits the
+ * maximum number of <tt>map</tt> and <tt>list</tt> items that can be read in a
+ * single sequence.</li>
+ * <li><tt>org.apache.avro.limits.string.maxLength</tt> limits the maximum size
+ * of <tt>string</tt> types.</li>
+ * <li><tt>org.apache.avro.limits.collectionItems.maxAllocation</tt> limits the
+ * number of <tt>array</tt> elements whose minimum encoded size is zero (such as
+ * <tt>null</tt>, a zero-length <tt>fixed</tt>, a record whose fields are all
+ * zero-byte, or a recursive schema whose cycle is conservatively broken with a
+ * 0 minimum) that may be allocated at once. Unlike other element types, these
+ * cannot be bounded by the number of bytes remaining in the stream, so the
+ * limit defaults to a fraction of the maximum heap.</li>
  * </ul>
  *
  * The default is to permit sizes up to {@link #MAX_ARRAY_VM_LIMIT}.
@@ -83,6 +90,35 @@ public class SystemLimitException extends AvroRuntimeException {
   public static final long MAX_DECOMPRESS_LENGTH = getLongLimitFromProperty(MAX_DECOMPRESS_LENGTH_PROPERTY,
       defaultMaxDecompressLength());
 
+  /**
+   * System property declaring the maximum number of array elements whose minimum
+   * encoded size is zero (e.g. {@code null}, a zero-length fixed, a record whose
+   * fields are all zero-byte, or a recursive schema conservatively treated as a 0
+   * minimum) to allocate at once: {@value}.
+   */
+  public static final String MAX_COLLECTION_ALLOCATION_PROPERTY = "org.apache.avro.limits.collectionItems.maxAllocation";
+
+  /**
+   * Fraction of the maximum heap a single decoded collection of elements whose
+   * minimum encoded size is zero may occupy by default. Keeps the backing
+   * allocation below the heap so a small payload declaring a huge block count
+   * cannot exhaust the JVM.
+   */
+  private static final long DEFAULT_MAX_COLLECTION_ALLOCATION_HEAP_FRACTION = 4;
+
+  /**
+   * Estimated bytes retained per pre-allocated collection slot (a single object
+   * reference), used to translate the heap budget into an element count.
+   */
+  private static final long BYTES_PER_COLLECTION_SLOT = 8;
+
+  /**
+   * Maximum number of array elements whose minimum encoded size is zero to
+   * allocate at once. Recomputed from the system property (or the heap) by
+   * {@link #resetLimits()}.
+   */
+  private static long maxCollectionAllocation = defaultMaxCollectionAllocation();
+
   static {
     resetLimits();
   }
@@ -153,6 +189,20 @@ private static long defaultMaxDecompressLength() {
         Math.max(1L, Runtime.getRuntime().maxMemory() / DEFAULT_MAX_DECOMPRESS_HEAP_FRACTION));
   }
 
+  /**
+   * Calculate the default maximum number of array elements whose minimum encoded
+   * size is zero to allocate at once, as a fraction of the maximum heap. Such
+   * elements consume no guaranteed input bytes, so the usual "bytes remaining"
+   * bound does not apply and the allocation must instead be capped relative to
+   * the available memory.
+   *
+   * @return the calculated default max element count.
+   */
+  private static long defaultMaxCollectionAllocation() {
+    long heapBudget = Math.max(1L, Runtime.getRuntime().maxMemory() / DEFAULT_MAX_COLLECTION_ALLOCATION_HEAP_FRACTION);
+    return Math.max(1L, heapBudget / BYTES_PER_COLLECTION_SLOT);
+  }
+
   /**
    * Check to ensure that reading the bytes is within the specified limits.
    *
@@ -246,6 +296,44 @@ public static int checkMaxCollectionLength(long items) {
     return (int) items;
   }
 
+  /**
+   * Check to ensure that allocating storage for the specified number of array
+   * elements whose minimum encoded size is zero remains within the heap-aware
+   * limit.
+   * <p>
+   * Elements whose minimum encoded size is zero (e.g. {@code null}, a zero-length
+   * fixed, a record whose fields are all zero-byte, or a recursive schema whose
+   * cycle is conservatively broken with a 0 minimum) consume no guaranteed input
+   * bytes, so the number that may be declared is not bounded by the bytes
+   * remaining in the stream. Without a cap, a tiny payload can declare an
+   * enormous block count and drive an unbounded backing-array allocation. This
+   * limit is derived from the maximum heap (see
+   * {@link #MAX_COLLECTION_ALLOCATION_PROPERTY}).
+   *
+   * @param existing The number of elements already allocated for the collection.
+   * @param items    The next number of 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 either argument is negative.
+   */
+  public static long checkMaxCollectionAllocation(long existing, long items) {
+    if (existing < 0) {
+      throw new AvroRuntimeException("Malformed data. Length is negative: " + existing);
+    }
+    if (items < 0) {
+      throw new AvroRuntimeException("Malformed data. Length is negative: " + items);
+    }
+    long total = existing + items;
+    if (total < existing || total > maxCollectionAllocation) {
+      throw new SystemLimitException("Cannot allocate " + (total < existing ? "more than Long.MAX_VALUE" : total)
+          + " collection elements whose minimum encoded size is zero: exceeds the maximum allowed of "
+          + maxCollectionAllocation + " (configure with the system property " + MAX_COLLECTION_ALLOCATION_PROPERTY
+          + ")");
+    }
+    return total;
+  }
+
   /**
    * Check to ensure that reading the string size is within the specified limits.
    *
@@ -294,5 +382,11 @@ static void resetLimits() {
     maxBytesLength = getLimitFromProperty(MAX_BYTES_LENGTH_PROPERTY, MAX_ARRAY_VM_LIMIT);
     maxCollectionLength = getLimitFromProperty(MAX_COLLECTION_LENGTH_PROPERTY, MAX_ARRAY_VM_LIMIT);
     maxStringLength = getLimitFromProperty(MAX_STRING_LENGTH_PROPERTY, MAX_ARRAY_VM_LIMIT);
+    maxCollectionAllocation = getLongLimitFromProperty(MAX_COLLECTION_ALLOCATION_PROPERTY,
+        defaultMaxCollectionAllocation());
+    // A collection cannot hold more than MAX_ARRAY_VM_LIMIT elements, so keep the
+    // zero-byte allocation cap consistent with the other collection limits even
+    // when it is configured (or derived from a very large heap) above that.
+    maxCollectionAllocation = Math.min(maxCollectionAllocation, MAX_ARRAY_VM_LIMIT);
   }
 }

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

@@ -37,6 +37,7 @@
 import org.apache.avro.LogicalType;
 import org.apache.avro.Schema;
 import org.apache.avro.Schema.Field;
+import org.apache.avro.SystemLimitException;
 import org.apache.avro.io.DatumReader;
 import org.apache.avro.io.Decoder;
 import org.apache.avro.io.DecoderFactory;
@@ -114,6 +115,26 @@ public void setExpected(Schema reader) {
   private static final ThreadLocal<Map<Schema, Map<Schema, ResolvingDecoder>>> RESOLVER_CACHE = ThreadLocalWithInitial
       .of(WeakIdentityHashMap::new);
 
+  /**
+   * Upper bound on the initial capacity eagerly allocated for a collection from
+   * its declared block count. The backing array/map grows on demand as elements
+   * are read, so this is only a starting hint: it prevents a large declared count
+   * from driving a huge up-front allocation before any element is decoded. This
+   * matters most for stream sources, where the decoder cannot know how many bytes
+   * remain and so cannot otherwise bound the declared count against the input.
+   */
+  private static final int MAX_COLLECTION_PREALLOC = 1024;
+
+  /**
+   * Clamp a declared collection block count to a safe initial allocation size.
+   *
+   * @param count the declared (already limit-checked) block count
+   * @return {@code count} capped at {@link #MAX_COLLECTION_PREALLOC}
+   */
+  public static int initialCollectionCapacity(long count) {
+    return (int) Math.min(count, MAX_COLLECTION_PREALLOC);
+  }
+
   /**
    * Gets a resolving decoder for use by this GenericDatumReader. Unstable API.
    * Currently uses a thread local cache to prevent constructing the resolvers too
@@ -296,9 +317,20 @@ protected Object readArray(Object old, Schema expected, ResolvingDecoder in) thr
     long base = 0;
     if (l > 0) {
       ensureAvailableCollectionBytes(in, l, expectedType);
+      // Elements whose minimum encoded size is zero (null, a zero-length fixed, a
+      // record whose fields are all zero-byte, or a recursive schema where the
+      // cycle is broken with a 0 minimum) consume no guaranteed input, so
+      // ensureAvailableCollectionBytes cannot bound their count from the bytes
+      // remaining. Cap such collections against a heap-aware limit so a tiny
+      // payload cannot declare a huge block count and drive an unbounded
+      // backing-array allocation.
+      boolean zeroByteElements = minBytesPerElement(expectedType) == 0;
+      if (zeroByteElements) {
+        SystemLimitException.checkMaxCollectionAllocation(base, l);
+      }
       LogicalType logicalType = expectedType.getLogicalType();
       Conversion<?> conversion = getData().getConversionFor(logicalType);
-      Object array = newArray(old, (int) l, expected);
+      Object array = newArray(old, initialCollectionCapacity(l), expected);
       do {
         if (logicalType != null && conversion != null) {
           for (long i = 0; i < l; i++) {
@@ -311,7 +343,11 @@ protected Object readArray(Object old, Schema expected, ResolvingDecoder in) thr
           }
         }
         base += l;
-      } while ((l = arrayNext(in, expectedType)) > 0);
+        l = arrayNext(in, expectedType);
+        if (zeroByteElements && l > 0) {
+          SystemLimitException.checkMaxCollectionAllocation(base, l);
+        }
+      } while (l > 0);
       return pruneArray(array);
     } else {
       return pruneArray(newArray(old, 0, expected));
@@ -364,7 +400,7 @@ protected Object readMap(Object old, Schema expected, ResolvingDecoder in) throw
     LogicalType logicalType = eValue.getLogicalType();
     Conversion<?> conversion = getData().getConversionFor(logicalType);
     ensureAvailableMapBytes(in, l, eValue);
-    Object map = newMap(old, (int) l);
+    Object map = newMap(old, initialCollectionCapacity(l));
     if (l > 0) {
       do {
         if (logicalType != null && conversion != null) {
@@ -441,6 +477,23 @@ static int minBytesPerElement(Schema schema) {
     return minBytesPerElement(schema, Collections.newSetFromMap(new IdentityHashMap<>()));
   }
 
+  /**
+   * Whether the minimum encoded size of the given schema is zero, i.e.
+   * {@link #minBytesPerElement(Schema)} is {@code 0}. This is true for values
+   * that always encode to zero bytes (e.g. {@code null}, a zero-length
+   * {@code fixed}, or a record whose fields are all zero-byte), and
+   * conservatively for recursive schemas, where the cycle is broken by returning
+   * a 0 minimum. Such elements cannot be bounded by the number of bytes remaining
+   * in the stream, so a collection of them must be bounded by a heap-aware
+   * allocation limit instead.
+   *
+   * @param schema the element (or map value) schema
+   * @return {@code true} if the schema's minimum encoded size is zero
+   */
+  public static boolean isZeroByteSchema(Schema schema) {
+    return minBytesPerElement(schema) == 0;
+  }
+
   private static int minBytesPerElement(Schema schema, Set<Schema> visited) {
     switch (schema.getType()) {
     case NULL:
@@ -481,9 +534,11 @@ private static int minBytesPerElement(Schema schema, Set<Schema> visited) {
    * reports fewer remaining bytes than required.
    * <p>
    * This check prevents out-of-memory errors from pre-allocating huge backing
-   * arrays when the source data is truncated or malicious.
+   * arrays when the source data is truncated or malicious. It is exposed so the
+   * fast reader ({@code FastReaderBuilder}) can apply the same guard as this
+   * classic reader.
    */
-  private static void ensureAvailableCollectionBytes(Decoder decoder, long count, Schema elementSchema)
+  public static void ensureAvailableCollectionBytes(Decoder decoder, long count, Schema elementSchema)
       throws EOFException {
     if (count <= 0) {
       return;
@@ -747,15 +802,35 @@ public static void skip(Schema schema, Decoder in) throws IOException {
       break;
     case ARRAY:
       Schema elementType = schema.getElementType();
+      // Bound the cumulative element count: skipping a huge block of elements
+      // whose minimum encoded size is zero (e.g. null) would otherwise loop
+      // unboundedly (a CPU exhaustion) even though it reads nothing. Such
+      // elements use the heap-aware allocation cap; others the structural
+      // collection cap.
+      boolean zeroByteElements = isZeroByteSchema(elementType);
+      long arrayTotal = 0;
       for (long l = in.skipArray(); l > 0; l = in.skipArray()) {
+        // Always enforce the cumulative structural cap, then additionally the
+        // heap-aware allocation cap for zero-byte elements (which the structural
+        // cap alone does not bound tightly), so a huge count split across blocks
+        // cannot drive an unbounded skip loop.
+        SystemLimitException.checkMaxCollectionLength(arrayTotal, l);
+        if (zeroByteElements) {
+          SystemLimitException.checkMaxCollectionAllocation(arrayTotal, l);
+        }
+        arrayTotal += l;
         for (long i = 0; i < l; i++) {
           skip(elementType, in);
         }
       }
       break;
     case MAP:
       Schema value = schema.getValueType();
+      // Map entries always carry a >= 1 byte key, so the structural cap applies.
+      long mapTotal = 0;
       for (long l = in.skipMap(); l > 0; l = in.skipMap()) {
+        SystemLimitException.checkMaxCollectionLength(mapTotal, l);
+        mapTotal += l;
         for (long i = 0; i < l; i++) {
           in.skipString();
           skip(value, in);

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

@@ -408,8 +408,21 @@ protected void doReadBytes(byte[] bytes, int start, int length) throws IOExcepti
   protected long doReadItemCount() throws IOException {
     long result = readLong();
     if (result < 0L) {
-      // Consume byte-count if present
-      readLong();
+      // A negative block count is followed by a block byte-size; consume it.
+      final long bytecount = readLong();
+      if (result == Long.MIN_VALUE) {
+        // Long.MIN_VALUE cannot be negated (-Long.MIN_VALUE overflows back to
+        // Long.MIN_VALUE), so it is not a valid block count. Reject it rather
+        // than letting it fall through as a negative "count" that would later be
+        // truncated to 0 and silently terminate the collection without consuming
+        // the end marker, desynchronizing decoding of subsequent fields.
+        throw new AvroRuntimeException("Malformed data. Block count is invalid: " + result);
+      }
+      if (bytecount < 0L) {
+        // The block byte-size is a byte count and must be non-negative, matching
+        // doSkipItems().
+        throw new AvroRuntimeException("Malformed data. Block byte-size is negative: " + bytecount);
+      }
       result = -result;
     }
     return result;
@@ -436,7 +449,17 @@ protected long doReadItemCount() throws IOException {
   private long doSkipItems() throws IOException {
     long result = readLong();
     while (result < 0L) {
+      if (result == Long.MIN_VALUE) {
+        // Consistent with doReadItemCount: Long.MIN_VALUE is not a valid block
+        // count (it cannot be negated), so reject it rather than treating it as
+        // a byte-sized block and continuing to skip.
+        readLong();
+        throw new AvroRuntimeException("Malformed data. Block count is invalid: " + result);
+      }
       final long bytecount = readLong();
+      if (bytecount < 0L) {
+        throw new AvroRuntimeException("Malformed data. Block byte-size is negative: " + bytecount);
+      }
       doSkipBytes(bytecount);
       result = readLong();
     }
@@ -458,7 +481,7 @@ public long arrayNext() throws IOException {
 
   @Override
   public long skipArray() throws IOException {
-    return doSkipItems();
+    return SystemLimitException.checkMaxCollectionLength(doSkipItems());
   }
 
   @Override
@@ -476,7 +499,7 @@ public long mapNext() throws IOException {
 
   @Override
   public long skipMap() throws IOException {
-    return doSkipItems();
+    return SystemLimitException.checkMaxCollectionLength(doSkipItems());
   }
 
   @Override