Compressor cleanups

native/src/main/c/jni_zlib_inflate.c

@@ -34,8 +34,7 @@ Java_com_velocitypowered_natives_compression_NativeZlibInflate_process(JNIEnv *e
     jlong sourceAddress,
     jint sourceLength,
     jlong destinationAddress,
-    jint destinationLength,
+    jint destinationLength)
-    jlong maximumSize)
 {
     struct libdeflate_decompressor *decompress = (struct libdeflate_decompressor *) ctx;
     enum libdeflate_result result = libdeflate_zlib_decompress(decompress, (void *) sourceAddress,

native/src/main/java/com/velocitypowered/natives/compression/JavaVelocityCompressor.java

@@ -56,24 +56,45 @@ public class JavaVelocityCompressor implements VelocityCompressor {
     final int origIdx = source.readerIndex();
     inflater.setInput(source.nioBuffer());
 
+    int totalProduced = 0;
     try {
       final int readable = source.readableBytes();
       while (!inflater.finished() && inflater.getBytesRead() < readable) {
+        if (totalProduced >= uncompressedSize) {
+          throw new DataFormatException("Decompressed data exceeds the claimed uncompressed size "
+              + "of " + uncompressedSize + " bytes");
+        }
+
+        final int remaining = uncompressedSize - totalProduced;
         if (!destination.isWritable()) {
-          destination.ensureWritable(ZLIB_BUFFER_SIZE);
+          destination.ensureWritable(Math.min(ZLIB_BUFFER_SIZE, remaining));
         }
 
         ByteBuffer destNioBuf = destination.nioBuffer(destination.writerIndex(),
             destination.writableBytes());
+
+        // Never let a single inflate step write past the claimed size
+        if (destNioBuf.remaining() > remaining) {
+          destNioBuf.limit(destNioBuf.position() + remaining);
+        }
+
         int produced = inflater.inflate(destNioBuf);
+        if (produced == 0 && !inflater.finished()) {
+          // Output space was available yet the inflater made no progress: the stream is truncated
+          // or corrupt (this also covers a peer that over-reported the uncompressed size).
+          throw new DataFormatException("Received a truncated or malformed deflate stream, "
+              + "expected " + uncompressedSize + " bytes");
+        }
+        totalProduced += produced;
         destination.writerIndex(destination.writerIndex() + produced);
       }
 
       if (!inflater.finished()) {
-        throw new DataFormatException("Received a deflate stream that was too large, wanted "
+        throw new DataFormatException("Received a truncated or malformed deflate stream, expected "
-            + uncompressedSize);
+            + uncompressedSize + " bytes");
       }
-      source.readerIndex(origIdx + inflater.getTotalIn());
+
+      source.readerIndex(origIdx + (int) inflater.getBytesRead());
     } finally {
       inflater.reset();
     }
@@ -102,7 +123,7 @@ public class JavaVelocityCompressor implements VelocityCompressor {
       destination.writerIndex(destination.writerIndex() + produced);
     }
 
-    source.readerIndex(origIdx + deflater.getTotalIn());
+    source.readerIndex(origIdx + (int) deflater.getBytesRead());
     deflater.reset();
   }
 

native/src/test/java/com/velocitypowered/natives/compression/VelocityCompressorTest.java

@@ -17,6 +17,7 @@
 
 package com.velocitypowered.natives.compression;
 
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
@@ -77,6 +78,83 @@ class VelocityCompressorTest {
     check(compressor, () -> Unpooled.buffer(TEST_DATA.length + 32));
   }
 
+  private static final int BOMB_ACTUAL_SIZE = 1 << 20;
+  private static final int BOMB_LYING_CLAIM = 1024;
+
+  @Test
+  void javaRejectsUnderReportedUncompressedSize() throws DataFormatException {
+    VelocityCompressor compressor = JavaVelocityCompressor.FACTORY
+        .create(Deflater.DEFAULT_COMPRESSION);
+    DataFormatException ex = assertRejectsDecompressionBomb(compressor);
+    // The Java compressor's size guard names the claimed size in its message, so operators can
+    // tell an over-size rejection apart from a genuinely corrupt stream.
+    assertTrue(ex.getMessage().contains(String.valueOf(BOMB_LYING_CLAIM)),
+        "rejection must originate from the uncompressed-size guard, got: " + ex.getMessage());
+  }
+
+  @Test
+  @EnabledOnOs({LINUX})
+  void nativeRejectsUnderReportedUncompressedSize() throws DataFormatException {
+    VelocityCompressor compressor = Natives.compress.get().create(Deflater.DEFAULT_COMPRESSION);
+    if (compressor.preferredBufferType() != BufferPreference.DIRECT_REQUIRED) {
+      compressor.close();
+      fail("Loaded regular compressor");
+    }
+    // libdeflate rejects with its own native-origin message ("uncompressed size is inaccurate"),
+    // so we only assert the behavioural guarantee here, not the message text.
+    assertRejectsDecompressionBomb(compressor);
+  }
+
+  /**
+   * Asserts that a compressor refuses a decompression bomb: a small, valid deflate stream whose
+   * claimed uncompressed size is far smaller than what it actually inflates to. Verifies the same
+   * stream round-trips when the claimed size is honest (proving the rejection is caused by the
+   * under-reported size, not corrupt input) and that no output is written past the claimed size.
+   * Closes the compressor before returning the exception thrown by the rejected inflate.
+   */
+  private DataFormatException assertRejectsDecompressionBomb(VelocityCompressor compressor)
+      throws DataFormatException {
+    // Direct buffers so this works for the native compressor, which requires them.
+    ByteBuf source = Unpooled.directBuffer(BOMB_ACTUAL_SIZE);
+    ByteBuf compressed = Unpooled.directBuffer();
+    try {
+      source.writeZero(BOMB_ACTUAL_SIZE);
+      compressor.deflate(source, compressed);
+      final int compressedSize = compressed.readableBytes();
+      assertTrue(compressedSize < BOMB_ACTUAL_SIZE / 100,
+          "sanity: payload really is a decompression bomb (" + compressedSize + " -> "
+              + BOMB_ACTUAL_SIZE + ")");
+
+      // Positive control: the compressed stream is perfectly valid and round-trips when the peer
+      // tells the truth about its uncompressed size.
+      ByteBuf honest = Unpooled.directBuffer();
+      try {
+        compressor.inflate(compressed.duplicate(), honest, BOMB_ACTUAL_SIZE);
+        assertEquals(BOMB_ACTUAL_SIZE, honest.readableBytes(),
+            "valid stream must fully decompress when the claimed size is honest");
+      } finally {
+        honest.release();
+      }
+
+      // Attack: same valid stream, but a tiny claimed size. inflate must refuse rather than grow
+      // the destination without bound.
+      ByteBuf decompressed = Unpooled.directBuffer();
+      try {
+        DataFormatException ex = assertThrows(DataFormatException.class,
+            () -> compressor.inflate(compressed.duplicate(), decompressed, BOMB_LYING_CLAIM));
+        assertTrue(decompressed.writerIndex() <= BOMB_LYING_CLAIM,
+            "inflate must not write past the claimed uncompressed size");
+        return ex;
+      } finally {
+        decompressed.release();
+      }
+    } finally {
+      source.release();
+      compressed.release();
+      compressor.close();
+    }
+  }
+
   private void check(VelocityCompressor compressor, Supplier<ByteBuf> bufSupplier)
       throws DataFormatException {
     ByteBuf source = bufSupplier.get();