AVRO-4070: Optimize Check Max Collection Length for New Collections (a…

…pache#3195)
belugabehr · Oct 8, 2024 · 9304cb7 · 9304cb7
1 parent 2ac5a46
commit 9304cb7
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 2 deletions.
diff --git a/lang/java/avro/src/main/java/org/apache/avro/SystemLimitException.java b/lang/java/avro/src/main/java/org/apache/avro/SystemLimitException.java
@@ -153,6 +153,35 @@ public static int checkMaxCollectionLength(long existing, long items) {
     return (int) length;
   }
 
+  /**
+   * Check to ensure that reading the specified number of items remains within the
+   * specified limits.
+   *
+   * @param items The next number of items to read. In normal usage, this is
+   *              always a positive, permitted value. Negative and zero values
+   *              have a special meaning in Avro decoding.
+   * @return The total number of items in the collection if and only if it is
+   *         within the limit and non-negative.
+   * @throws UnsupportedOperationException if reading the items would allocate a
+   *                                       collection that the Java VM would be
+   *                                       unable to handle
+   * @throws SystemLimitException          if the decoding should fail because it
+   *                                       would otherwise result in an allocation
+   *                                       exceeding the set limit
+   * @throws AvroRuntimeException          if the length is negative
+   */
+  public static int checkMaxCollectionLength(long items) {
+    if (items > MAX_ARRAY_VM_LIMIT) {
+      throw new UnsupportedOperationException(
+          "Cannot read collections larger than " + MAX_ARRAY_VM_LIMIT + " items in Java library");
+    }
+    if (items > maxCollectionLength) {
+      throw new SystemLimitException(
+          "Collection length " + items + " exceeds the maximum allowed of " + maxCollectionLength);
+    }
+    return (int) items;
+  }
+
   /**
    * Check to ensure that reading the string size is within the specified limits.
    *

diff --git a/lang/java/avro/src/main/java/org/apache/avro/io/BinaryDecoder.java b/lang/java/avro/src/main/java/org/apache/avro/io/BinaryDecoder.java
@@ -428,7 +428,7 @@ private long doSkipItems() throws IOException {
 
   @Override
   public long readArrayStart() throws IOException {
-    collectionCount = SystemLimitException.checkMaxCollectionLength(0L, doReadItemCount());
+    collectionCount = SystemLimitException.checkMaxCollectionLength(doReadItemCount());
     return collectionCount;
   }
 
@@ -446,7 +446,7 @@ public long skipArray() throws IOException {
 
   @Override
   public long readMapStart() throws IOException {
-    collectionCount = SystemLimitException.checkMaxCollectionLength(0L, doReadItemCount());
+    collectionCount = SystemLimitException.checkMaxCollectionLength(doReadItemCount());
     return collectionCount;
   }