Cap array allocations in ProtocolUtils readers

proxy/src/main/java/com/velocitypowered/proxy/protocol/ProtocolUtils.java

@@ -357,8 +357,21 @@ public static void writeMinimalKey(ByteBuf buf, Key key) {
    * @return the decoded key array
    */
   public static Key[] readKeyArray(ByteBuf buf) {
+    return readKeyArray(buf, DEFAULT_MAX_STRING_SIZE);
+  }
+
+  /**
+   * Reads a standard Mojang Text namespaced:key array from the buffer, refusing
+   * to allocate more than {@code cap} entries.
+   *
+   * @param buf the buffer to read from
+   * @param cap the maximum number of entries to accept
+   * @return the decoded key array
+   */
+  public static Key[] readKeyArray(ByteBuf buf, int cap) {
     int length = readVarInt(buf);
     checkFrame(length >= 0, "Got a negative-length array (%s)", length);
+    checkFrame(length <= cap, "Bad array size (got %s, maximum is %s)", length, cap);
     checkFrame(buf.isReadable(length),
         "Trying to read an array that is too long (wanted %s, only have %s)", length,
         buf.readableBytes());
@@ -419,8 +432,21 @@ public static void writeByteArray(ByteBuf buf, byte[] array) {
    * @return an array of integers
    */
   public static int[] readIntegerArray(ByteBuf buf) {
+    return readIntegerArray(buf, DEFAULT_MAX_STRING_SIZE);
+  }
+
+  /**
+   * Reads an VarInt-prefixed array of VarInt integers from the {@code buf},
+   * refusing to allocate more than {@code cap} entries.
+   *
+   * @param buf the buffer to read from
+   * @param cap the maximum number of entries to accept
+   * @return an array of integers
+   */
+  public static int[] readIntegerArray(ByteBuf buf, int cap) {
     int len = readVarInt(buf);
     checkFrame(len >= 0, "Got a negative-length integer array (%s)", len);
+    checkFrame(len <= cap, "Bad integer array size (got %s, maximum is %s)", len, cap);
     checkFrame(buf.isReadable(len),
         "Trying to read an array that is too long (wanted %s, only have %s)", len,
         buf.readableBytes());
@@ -542,8 +568,21 @@ public static <T extends BinaryTag> void writeBinaryTag(ByteBuf buf, ProtocolVer
    * @return the String array from the buffer
    */
   public static String[] readStringArray(ByteBuf buf) {
+    return readStringArray(buf, DEFAULT_MAX_STRING_SIZE);
+  }
+
+  /**
+   * Reads a String array from the {@code buf}, refusing to allocate more than
+   * {@code cap} entries.
+   *
+   * @param buf the buffer to read from
+   * @param cap the maximum number of entries to accept
+   * @return the String array from the buffer
+   */
+  public static String[] readStringArray(ByteBuf buf, int cap) {
     int length = readVarInt(buf);
     checkFrame(length >= 0, "Got a negative-length array (%s)", length);
+    checkFrame(length <= cap, "Bad array size (got %s, maximum is %s)", length, cap);
     checkFrame(buf.isReadable(length),
         "Trying to read an array that is too long (wanted %s, only have %s)", length,
         buf.readableBytes());
@@ -574,8 +613,21 @@ public static void writeStringArray(ByteBuf buf, String[] stringArray) {
    * @return the Integer array from the buffer
    */
   public static int[] readVarIntArray(ByteBuf buf) {
+    return readVarIntArray(buf, DEFAULT_MAX_STRING_SIZE);
+  }
+
+  /**
+   * Reads a VarInt-prefixed VarInt array from the {@code buf}, refusing to
+   * allocate more than {@code cap} entries.
+   *
+   * @param buf the buffer to read from
+   * @param cap the maximum number of entries to accept
+   * @return the Integer array from the buffer
+   */
+  public static int[] readVarIntArray(ByteBuf buf, int cap) {
     int length = readVarInt(buf);
     checkFrame(length >= 0, "Got a negative-length array (%s)", length);
+    checkFrame(length <= cap, "Bad array size (got %s, maximum is %s)", length, cap);
     checkFrame(buf.isReadable(length),
         "Trying to read an array that is too long (wanted %s, only have %s)", length,
         buf.readableBytes());

proxy/src/test/java/com/velocitypowered/proxy/protocol/ProtocolUtilsTest.java

@@ -20,7 +20,9 @@
 import static com.velocitypowered.proxy.protocol.ProtocolUtils.encode21BitVarInt;
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 
+import com.velocitypowered.proxy.util.except.QuietDecoderException;
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.ByteBufUtil;
 import io.netty.buffer.Unpooled;
@@ -155,6 +157,65 @@ private void writeVarIntOld(ByteBuf buf, int value) {
     }
   }
 
+  @Test
+  void readIntegerArrayRejectsCountAboveCap() {
+    // length 4 stays inside the per-buffer "isReadable(length)" check (we put
+    // 5 bytes in the buffer below) but is above the cap of 2, so the new
+    // check is the one that must trip.
+    ByteBuf buf = Unpooled.buffer();
+    ProtocolUtils.writeVarInt(buf, 4);
+    for (int i = 0; i < 4; i++) {
+      ProtocolUtils.writeVarInt(buf, 0);
+    }
+    assertThrows(QuietDecoderException.class,
+        () -> ProtocolUtils.readIntegerArray(buf, 2));
+  }
+
+  @Test
+  void readVarIntArrayRejectsCountAboveCap() {
+    ByteBuf buf = Unpooled.buffer();
+    ProtocolUtils.writeVarInt(buf, 4);
+    for (int i = 0; i < 4; i++) {
+      ProtocolUtils.writeVarInt(buf, 0);
+    }
+    assertThrows(QuietDecoderException.class,
+        () -> ProtocolUtils.readVarIntArray(buf, 2));
+  }
+
+  @Test
+  void readStringArrayRejectsCountAboveCap() {
+    ByteBuf buf = Unpooled.buffer();
+    ProtocolUtils.writeVarInt(buf, 4);
+    for (int i = 0; i < 4; i++) {
+      ProtocolUtils.writeString(buf, "");
+    }
+    assertThrows(QuietDecoderException.class,
+        () -> ProtocolUtils.readStringArray(buf, 2));
+  }
+
+  @Test
+  void readIntegerArrayDefaultCapRejectsAmplificationAttempt() {
+    // PoC for the bug the cap is fixing. Pre-fix readIntegerArray only
+    // gated on isReadable(length), and length is a count of entries while
+    // isReadable counts bytes. A count of 256k packed as one-byte varints
+    // is 256 KiB on the wire but drives a 1 MiB int[] (4x amplification).
+    // The default cap of DEFAULT_MAX_STRING_SIZE = 65536 now rejects this.
+    int len = 256 * 1024;
+    ByteBuf buf = Unpooled.buffer();
+    ProtocolUtils.writeVarInt(buf, len);
+    for (int i = 0; i < len; i++) {
+      buf.writeByte(0);
+    }
+    int reader = buf.readerIndex();
+    assertThrows(QuietDecoderException.class,
+        () -> ProtocolUtils.readIntegerArray(buf));
+    // Same input is accepted when the cap is effectively disabled, which
+    // is the pre-fix behaviour: a 256k-element int[] gets allocated.
+    buf.readerIndex(reader);
+    int[] result = ProtocolUtils.readIntegerArray(buf, Integer.MAX_VALUE);
+    assertEquals(len, result.length);
+  }
+
   private int conventionalWrittenBytes(int value) {
     int wouldBeWritten = 0;
     while (true) {