AVRO-4324: [Java] Align ReflectDatumReader.readArray with GenericDatumReader eager-allocation guards

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

@@ -32,6 +32,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.generic.IndexedRecord;
 import org.apache.avro.io.Decoder;
 import org.apache.avro.io.ResolvingDecoder;
@@ -143,6 +144,17 @@ protected Object readArray(Object old, Schema expected, ResolvingDecoder in) thr
     if (l <= 0) {
       return newArray(old, 0, expected);
     }
+    // Match GenericDatumReader.readArray: before eagerly allocating the backing
+    // array for the declared block count, verify the input could plausibly hold
+    // that many elements (guarding against a malformed or truncated payload),
+    // and separately cap element types whose minimum encoded size is zero, which
+    // the bytes-remaining check cannot bound. Without this a small malformed
+    // record mapped to a Java array field (e.g. long[]) could drive a very large
+    // eager allocation before any element is read.
+    ensureAvailableCollectionBytes(in, l, expectedType);
+    if (isZeroByteSchema(expectedType)) {
+      SystemLimitException.checkMaxCollectionAllocation(0, l);
+    }
     Object array = newArray(old, (int) l, expected);
     if (array instanceof Collection) {
       @SuppressWarnings("unchecked")
@@ -187,6 +199,7 @@ private Object readPrimitiveArray(Object array, Class<?> c, long l, ResolvingDec
   private Object readObjectArray(Object[] array, Schema expectedType, long l, ResolvingDecoder in) throws IOException {
     LogicalType logicalType = expectedType.getLogicalType();
     Conversion<?> conversion = getData().getConversionFor(logicalType);
+    boolean zeroByte = isZeroByteSchema(expectedType);
     int index = 0;
     if (logicalType != null && conversion != null) {
       do {
@@ -196,7 +209,7 @@ private Object readObjectArray(Object[] array, Schema expectedType, long l, Reso
           array[index] = element;
           index++;
         }
-      } while ((l = in.arrayNext()) > 0);
+      } while ((l = nextArrayBlock(in, expectedType, index, zeroByte)) > 0);
     } else {
       do {
         int limit = index + (int) l;
@@ -205,7 +218,7 @@ private Object readObjectArray(Object[] array, Schema expectedType, long l, Reso
           array[index] = element;
           index++;
         }
-      } while ((l = in.arrayNext()) > 0);
+      } while ((l = nextArrayBlock(in, expectedType, index, zeroByte)) > 0);
     }
     return array;
   }
@@ -214,24 +227,53 @@ private Object readCollection(Collection<Object> c, Schema expectedType, long l,
       throws IOException {
     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);
         }
-      } while ((l = in.arrayNext()) > 0);
+        count += l;
+      } while ((l = nextArrayBlock(in, expectedType, count, zeroByte)) > 0);
     } else {
       do {
         for (int i = 0; i < l; i++) {
           Object element = readWithoutConversion(null, expectedType, in);
           c.add(element);
         }
-      } while ((l = in.arrayNext()) > 0);
+        count += l;
+      } while ((l = nextArrayBlock(in, expectedType, count, zeroByte)) > 0);
     }
     return c;
   }
 
+  /**
+   * Read and validate the next array block count, mirroring
+   * {@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.
+   *
+   * @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 {
+    long l = in.arrayNext();
+    ensureAvailableCollectionBytes(in, l, expectedType);
+    if (zeroByte && l > 0) {
+      SystemLimitException.checkMaxCollectionAllocation(existing, l);
+    }
+    return l;
+  }
+
   @Override
   protected Object readString(Object old, Decoder in) throws IOException {
     return super.readString(null, in).toString();

lang/java/avro/src/test/java/org/apache/avro/reflect/TestReflectDatumReader.java

@@ -22,6 +22,7 @@
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import java.io.ByteArrayOutputStream;
+import java.io.EOFException;
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.HashSet;
@@ -32,6 +33,7 @@
 import java.util.Optional;
 
 import org.apache.avro.Schema;
+import org.apache.avro.SystemLimitException;
 import org.apache.avro.io.Decoder;
 import org.apache.avro.io.DecoderFactory;
 import org.apache.avro.io.Encoder;
@@ -102,6 +104,73 @@ void read_PojoWithArray() throws IOException {
     assertEquals(pojoWithArray, deserialized);
   }
 
+  /**
+   * A malformed or truncated record can declare an array block count far larger
+   * than the data that follows. The reader must reject it before eagerly
+   * allocating the backing Java array, the same way GenericDatumReader does.
+   */
+  @Test
+  void read_PojoWithArray_rejectsOversizedArrayCount() throws IOException {
+    ByteArrayOutputStream out = new ByteArrayOutputStream();
+    Encoder encoder = EncoderFactory.get().binaryEncoder(out, null);
+    encoder.writeInt(42); // record field "id"
+    encoder.writeLong(2_000_000_000L); // array block count for "relatedIds", with no items following
+    encoder.flush();
+
+    byte[] malformed = out.toByteArray();
+
+    Decoder decoder = DecoderFactory.get().binaryDecoder(malformed, null);
+    ReflectDatumReader<PojoWithArray> reflectDatumReader = new ReflectDatumReader<>(PojoWithArray.class);
+
+    assertThrows(EOFException.class, () -> reflectDatumReader.read(new PojoWithArray(), decoder));
+  }
+
+  /**
+   * Elements whose minimum encoded size is zero (here an empty record) consume no
+   * input, so a large logical array can be split across blocks that each pass the
+   * first-block guard. The cumulative allocation must still be bounded across
+   * blocks, mirroring GenericDatumReader.
+   */
+  @Test
+  void read_PojoWithZeroByteList_rejectsCumulativeCountAcrossBlocks() throws IOException {
+    System.setProperty(SystemLimitException.MAX_COLLECTION_ALLOCATION_PROPERTY, "1000");
+    org.apache.avro.TestSystemLimitException.resetLimits();
+    try {
+      // Two blocks of 600 zero-byte records each (1200 > 1000): the first block
+      // passes, the cumulative count must be rejected on the second block.
+      ByteArrayOutputStream out = new ByteArrayOutputStream();
+      Encoder encoder = EncoderFactory.get().binaryEncoder(out, null);
+      encoder.writeLong(600L);
+      encoder.writeLong(600L);
+      encoder.writeLong(0L); // array terminator (not reached)
+      encoder.flush();
+
+      byte[] malformed = out.toByteArray();
+      Decoder decoder = DecoderFactory.get().binaryDecoder(malformed, null);
+      ReflectDatumReader<PojoWithZeroByteList> reader = new ReflectDatumReader<>(PojoWithZeroByteList.class);
+
+      assertThrows(SystemLimitException.class, () -> reader.read(new PojoWithZeroByteList(), decoder));
+    } finally {
+      System.clearProperty(SystemLimitException.MAX_COLLECTION_ALLOCATION_PROPERTY);
+      org.apache.avro.TestSystemLimitException.resetLimits();
+    }
+  }
+
+  /**
+   * An empty record encodes to zero bytes, making it a zero-byte element type.
+   */
+  public static class EmptyRecord {
+    public EmptyRecord() {
+    }
+  }
+
+  public static class PojoWithZeroByteList {
+    public List<EmptyRecord> items;
+
+    public PojoWithZeroByteList() {
+    }
+  }
+
   @Test
   public void testRead_PojoWithSet() throws IOException {
     PojoWithSet pojoWithSet = new PojoWithSet();