[CORE][MINOR] Validate array counts against received bytes in shuffle fetch decoders

FetchShuffleBlocks and FetchShuffleBlockChunks decode an outer element
count and eagerly allocate a 2D array of that size. Check the count
against the bytes actually received before allocating, mirroring the
bounds checks Encoders already applies: each inner group carries its own
4-byte length prefix, so a valid count never exceeds readableBytes() / 4.

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-Authored-By: Holden Karau <holden@pigscanfly.ca>
(cherry picked from commit 9687793)

Author: 3 people

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/protocol/FetchShuffleBlockChunks.java

@@ -18,6 +18,7 @@
 package org.apache.spark.network.shuffle.protocol;
 
 import java.util.Arrays;
+import java.util.Objects;
 
 import io.netty.buffer.ByteBuf;
 
@@ -130,6 +131,9 @@ public static FetchShuffleBlockChunks decode(ByteBuf buf) {
     int shuffleMergeId = buf.readInt();
     int[] reduceIds = Encoders.IntArrays.decode(buf);
     int chunkIdsLen = buf.readInt();
+    // The divisor 4 is the minimum on-wire size of one element, since each chunk-id group
+    // is prefixed with its own 4-byte length.
+    Objects.checkFromIndexSize(0, chunkIdsLen, buf.readableBytes() / 4);
     int[][] chunkIds = new int[chunkIdsLen][];
     for (int i = 0; i < chunkIdsLen; i++) {
       chunkIds[i] = Encoders.IntArrays.decode(buf);

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/protocol/FetchShuffleBlocks.java

@@ -18,6 +18,7 @@
 package org.apache.spark.network.shuffle.protocol;
 
 import java.util.Arrays;
+import java.util.Objects;
 
 import io.netty.buffer.ByteBuf;
 
@@ -130,6 +131,9 @@ public static FetchShuffleBlocks decode(ByteBuf buf) {
     int shuffleId = buf.readInt();
     long[] mapIds = Encoders.LongArrays.decode(buf);
     int reduceIdsSize = buf.readInt();
+    // The divisor 4 is the minimum on-wire size of one element, since each reduce-id group
+    // is prefixed with its own 4-byte length.
+    Objects.checkFromIndexSize(0, reduceIdsSize, buf.readableBytes() / 4);
     int[][] reduceIds = new int[reduceIdsSize][];
     for (int i = 0; i < reduceIdsSize; i++) {
       reduceIds[i] = Encoders.IntArrays.decode(buf);

common/network-shuffle/src/test/java/org/apache/spark/network/shuffle/protocol/FetchShuffleBlockChunksSuite.java

@@ -24,6 +24,8 @@
 
 import static org.junit.jupiter.api.Assertions.*;
 
+import org.apache.spark.network.protocol.Encoders;
+
 public class FetchShuffleBlockChunksSuite {
 
   @Test
@@ -39,4 +41,17 @@ public void testFetchShuffleBlockChunksEncodeDecode() {
     FetchShuffleBlockChunks decoded = FetchShuffleBlockChunks.decode(buf);
     assertEquals(shuffleBlockChunks, decoded);
   }
+
+  @Test
+  public void testDecodeRejectsCountExceedingReceivedBytes() {
+    // Decode must reject the count rather than allocate an array of the claimed size.
+    ByteBuf buf = Unpooled.buffer(64);
+    Encoders.Strings.encode(buf, "app0");
+    Encoders.Strings.encode(buf, "exec1");
+    buf.writeInt(0); // shuffleId
+    buf.writeInt(0); // shuffleMergeId
+    Encoders.IntArrays.encode(buf, new int[] {0});
+    buf.writeInt(Integer.MAX_VALUE); // chunkIds count
+    assertThrows(IndexOutOfBoundsException.class, () -> FetchShuffleBlockChunks.decode(buf));
+  }
 }

common/network-shuffle/src/test/java/org/apache/spark/network/shuffle/protocol/FetchShuffleBlocksSuite.java

@@ -24,6 +24,8 @@
 
 import static org.junit.jupiter.api.Assertions.*;
 
+import org.apache.spark.network.protocol.Encoders;
+
 public class FetchShuffleBlocksSuite {
 
   @Test
@@ -39,4 +41,16 @@ public void testFetchShuffleBlockEncodeDecode() {
     FetchShuffleBlocks decoded = FetchShuffleBlocks.decode(buf);
     assertEquals(fetchShuffleBlocks, decoded);
   }
+
+  @Test
+  public void testDecodeRejectsCountExceedingReceivedBytes() {
+    // Decode must reject the count rather than allocate an array of the claimed size.
+    ByteBuf buf = Unpooled.buffer(64);
+    Encoders.Strings.encode(buf, "app0");
+    Encoders.Strings.encode(buf, "exec1");
+    buf.writeInt(0); // shuffleId
+    Encoders.LongArrays.encode(buf, new long[] {0});
+    buf.writeInt(Integer.MAX_VALUE); // reduceIds count
+    assertThrows(IndexOutOfBoundsException.class, () -> FetchShuffleBlocks.decode(buf));
+  }
 }