Add missing bounds checks

airlift · Feb 13, 2024 · d01ecb7 · d01ecb7
1 parent 2cea90a
commit d01ecb7
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 3 deletions.
diff --git a/src/main/java/io/airlift/compress/lz4/Lz4RawDecompressor.java b/src/main/java/io/airlift/compress/lz4/Lz4RawDecompressor.java
@@ -62,6 +62,9 @@ public static int decompress(
             // decode literal length
             int literalLength = token >>> 4; // top-most 4 bits of token
             if (literalLength == 0xF) {
+                if (input >= inputLimit) {
+                    throw new MalformedInputException(input - inputAddress);
+                }
                 int value;
                 do {
                     value = UNSAFE.getByte(inputBase, input++) & 0xFF;

diff --git a/src/main/java/io/airlift/compress/lzo/LzoRawDecompressor.java b/src/main/java/io/airlift/compress/lzo/LzoRawDecompressor.java
@@ -325,7 +325,7 @@ else if ((command & 0b1100_0000) != 0) {
                 }
                 long literalOutputLimit = output + literalLength;
                 if (literalOutputLimit > fastOutputLimit || input + literalLength > inputLimit - SIZE_OF_LONG) {
-                    if (literalOutputLimit > outputLimit) {
+                    if (literalOutputLimit > outputLimit || input + literalLength > inputLimit) {
                         throw new MalformedInputException(input - inputAddress);
                     }
 

diff --git a/src/main/java/io/airlift/compress/zstd/ZstdFrameDecompressor.java b/src/main/java/io/airlift/compress/zstd/ZstdFrameDecompressor.java
@@ -197,6 +197,7 @@ public int decompress(
 
                 long hash = XxHash64.hash(0, outputBase, outputStart, decodedFrameSize);
 
+                verify(input + SIZE_OF_INT <= inputLimit, input, "Not enough input bytes");
                 int checksum = UNSAFE.getInt(inputBase, input);
                 if (checksum != (int) hash) {
                     throw new MalformedInputException(input, format("Bad checksum. Expected: %s, actual: %s", Integer.toHexString(checksum), Integer.toHexString((int) hash)));
@@ -510,14 +511,15 @@ else if (sequenceCount > 127) {
         }
 
         // last literal segment
-        output = copyLastLiteral(outputBase, literalsBase, literalsLimit, output, literalsInput);
+        output = copyLastLiteral(input, literalsBase, literalsInput, literalsLimit, outputBase, output, outputLimit);
 
         return (int) (output - outputAddress);
     }
 
-    private static long copyLastLiteral(Object outputBase, Object literalsBase, long literalsLimit, long output, long literalsInput)
+    private static long copyLastLiteral(long input, Object literalsBase, long literalsInput, long literalsLimit, Object outputBase, long output, long outputLimit)
     {
         long lastLiteralsSize = literalsLimit - literalsInput;
+        verify(output + lastLiteralsSize <= outputLimit, input, "Output buffer too small");
         UNSAFE.copyMemory(literalsBase, literalsInput, outputBase, output, lastLiteralsSize);
         output += lastLiteralsSize;
         return output;