CVE-2021-37137

@@ -45,13 +45,19 @@ private enum ChunkType {
     }
 
     private static final int SNAPPY_IDENTIFIER_LEN = 6;
+    // See https://github.com/google/snappy/blob/1.1.9/framing_format.txt#L95
     private static final int MAX_UNCOMPRESSED_DATA_SIZE = 65536 + 4;
+    // See https://github.com/google/snappy/blob/1.1.9/framing_format.txt#L82
+    private static final int MAX_DECOMPRESSED_DATA_SIZE = 65536;
+    // See https://github.com/google/snappy/blob/1.1.9/framing_format.txt#L82
+    private static final int MAX_COMPRESSED_CHUNK_SIZE = 16777216 - 1;
 
     private final Snappy snappy = new Snappy();
     private final boolean validateChecksums;
 
     private boolean started;
     private boolean corrupted;
+    private int numBytesToSkip;
 
     /**
      * Creates a new snappy-framed decoder with validation of checksums
@@ -82,6 +88,16 @@ protected void decode(ChannelHandlerContext ctx, ByteBuf in, List<Object> out) t
             return;
         }
 
+        if (numBytesToSkip != 0) {
+            // The last chunkType we detected was RESERVED_SKIPPABLE and we still have some bytes to skip.
+            int skipBytes = Math.min(numBytesToSkip, in.readableBytes());
+            in.skipBytes(skipBytes);
+            numBytesToSkip -= skipBytes;
+
+            // Let's return and try again.
+            return;
+        }
+
         try {
             int idx = in.readerIndex();
             final int inSize = in.readableBytes();
@@ -123,12 +139,15 @@ protected void decode(ChannelHandlerContext ctx, ByteBuf in, List<Object> out) t
                         throw new DecompressionException("Received RESERVED_SKIPPABLE tag before STREAM_IDENTIFIER");
                     }
 
-                    if (inSize < 4 + chunkLength) {
-                        // TODO: Don't keep skippable bytes
-                        return;
-                    }
+                    in.skipBytes(4);
 
-                    in.skipBytes(4 + chunkLength);
+                    int skipBytes = Math.min(chunkLength, in.readableBytes());
+                    in.skipBytes(skipBytes);
+                    if (skipBytes != chunkLength) {
+                        // We could skip all bytes, let's store the remaining so we can do so once we receive more
+                        // data.
+                        numBytesToSkip = chunkLength - skipBytes;
+                    }
                     break;
                 case RESERVED_UNSKIPPABLE:
                     // The spec mandates that reserved unskippable chunks must immediately
@@ -141,7 +160,8 @@ protected void decode(ChannelHandlerContext ctx, ByteBuf in, List<Object> out) t
                         throw new DecompressionException("Received UNCOMPRESSED_DATA tag before STREAM_IDENTIFIER");
                     }
                     if (chunkLength > MAX_UNCOMPRESSED_DATA_SIZE) {
-                        throw new DecompressionException("Received UNCOMPRESSED_DATA larger than 65540 bytes");
+                        throw new DecompressionException("Received UNCOMPRESSED_DATA larger than " +
+                                MAX_UNCOMPRESSED_DATA_SIZE + " bytes");
                     }
 
                     if (inSize < 4 + chunkLength) {
@@ -162,13 +182,25 @@ protected void decode(ChannelHandlerContext ctx, ByteBuf in, List<Object> out) t
                         throw new DecompressionException("Received COMPRESSED_DATA tag before STREAM_IDENTIFIER");
                     }
 
+                    if (chunkLength > MAX_COMPRESSED_CHUNK_SIZE) {
+                        throw new DecompressionException("Received COMPRESSED_DATA that contains" +
+                                " chunk that exceeds " + MAX_COMPRESSED_CHUNK_SIZE + " bytes");
+                    }
+
                     if (inSize < 4 + chunkLength) {
                         return;
                     }
 
                     in.skipBytes(4);
                     int checksum = in.readIntLE();
-                    ByteBuf uncompressed = ctx.alloc().buffer();
+
+                    int uncompressedSize = snappy.getPreamble(in);
+                    if (uncompressedSize > MAX_DECOMPRESSED_DATA_SIZE) {
+                        throw new DecompressionException("Received COMPRESSED_DATA that contains" +
+                                " uncompressed data that exceeds " + MAX_DECOMPRESSED_DATA_SIZE + " bytes");
+                    }
+
+                    ByteBuf uncompressed = ctx.alloc().buffer(uncompressedSize, MAX_DECOMPRESSED_DATA_SIZE);
                     try {
                         if (validateChecksums) {
                             int oldWriterIndex = in.writerIndex();

@@ -38,20 +38,19 @@ public final class Snappy {
     private static final int COPY_2_BYTE_OFFSET = 2;
     private static final int COPY_4_BYTE_OFFSET = 3;
 
-    private State state = State.READY;
+    private State state = State.READING_PREAMBLE;
     private byte tag;
     private int written;
 
     private enum State {
-        READY,
         READING_PREAMBLE,
         READING_TAG,
         READING_LITERAL,
         READING_COPY
     }
 
     public void reset() {
-        state = State.READY;
+        state = State.READING_PREAMBLE;
         tag = 0;
         written = 0;
     }
@@ -270,9 +269,6 @@ private static void encodeCopy(ByteBuf out, int offset, int length) {
     public void decode(ByteBuf in, ByteBuf out) {
         while (in.isReadable()) {
             switch (state) {
-            case READY:
-                state = State.READING_PREAMBLE;
-                // fall through
             case READING_PREAMBLE:
                 int uncompressedLength = readPreamble(in);
                 if (uncompressedLength == PREAMBLE_NOT_FULL) {
@@ -281,7 +277,6 @@ public void decode(ByteBuf in, ByteBuf out) {
                 }
                 if (uncompressedLength == 0) {
                     // Should never happen, but it does mean we have nothing further to do
-                    state = State.READY;
                     return;
                 }
                 out.ensureWritable(uncompressedLength);
@@ -378,6 +373,27 @@ private static int readPreamble(ByteBuf in) {
         return 0;
     }
 
+    /**
+     * Get the length varint (a series of bytes, where the lower 7 bits
+     * are data and the upper bit is a flag to indicate more bytes to be
+     * read).
+     *
+     * @param in The input buffer to get the preamble from
+     * @return The calculated length based on the input buffer, or 0 if
+     *   no preamble is able to be calculated
+     */
+    int getPreamble(ByteBuf in) {
+        if (state == State.READING_PREAMBLE) {
+            int readerIndex = in.readerIndex();
+            try {
+                return readPreamble(in);
+            } finally {
+                in.readerIndex(readerIndex);
+            }
+        }
+        return 0;
+    }
+
     /**
      * Reads a literal from the input buffer directly to the output buffer.
      * A "literal" is an uncompressed segment of data stored directly in the