From 9304cb793aee4d5b1da56b1dced3b7e65329779f Mon Sep 17 00:00:00 2001
From: belugabehr <12578579+belugabehr@users.noreply.github.com>
Date: Tue, 8 Oct 2024 19:33:29 -0400
Subject: [PATCH] AVRO-4070: Optimize Check Max Collection Length for New
 Collections (#3195)

---
 .../org/apache/avro/SystemLimitException.java | 29 +++++++++++++++++++
 .../org/apache/avro/io/BinaryDecoder.java     |  4 +--
 2 files changed, 31 insertions(+), 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
index a96f812d84d..b4e46834599 100644
--- 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
index 7217be3addd..865bc9c06c2 100644
--- 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;
   }