Allow a limit to be set on the decompressed buffer size for ZlibDecod…

…ers (#9924) Motivation: It is impossible to know in advance how much memory will be needed to decompress a stream of bytes that was compressed using the DEFLATE algorithm. In theory, up to 1032 times the compressed size could be needed. For untrusted input, an attacker could exploit this to exhaust the memory pool. Modifications: ZlibDecoder and its subclasses now support an optional limit on the size of the decompressed buffer. By default, if the limit is reached, decompression stops and a DecompressionException is thrown. Behavior upon reaching the limit is modifiable by subclasses in case they desire something else. Result: The decompressed buffer can now be limited to a configurable size, thus mitigating the possibility of memory pool exhaustion.
netty · Jan 31, 2020 · 1543218 · 1543218
1 parent 6e5f229
commit 1543218
Show file tree

Hide file tree

Showing 8 changed files with 245 additions and 20 deletions.
diff --git a/codec/src/main/java/io/netty/handler/codec/compression/JZlibDecoder.java b/codec/src/main/java/io/netty/handler/codec/compression/JZlibDecoder.java
@@ -18,6 +18,7 @@
 import com.jcraft.jzlib.Inflater;
 import com.jcraft.jzlib.JZlib;
 import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufAllocator;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.util.internal.ObjectUtil;
 
@@ -35,7 +36,21 @@ public class JZlibDecoder extends ZlibDecoder {
      * @throws DecompressionException if failed to initialize zlib
      */
     public JZlibDecoder() {
-        this(ZlibWrapper.ZLIB);
+        this(ZlibWrapper.ZLIB, 0);
+    }
+
+    /**
+     * Creates a new instance with the default wrapper ({@link ZlibWrapper#ZLIB})
+     * and specified maximum buffer allocation.
+     *
+     * @param maxAllocation
+     *          Maximum size of the decompression buffer. Must be &gt;= 0.
+     *          If zero, maximum size is decided by the {@link ByteBufAllocator}.
+     *
+     * @throws DecompressionException if failed to initialize zlib
+     */
+    public JZlibDecoder(int maxAllocation) {
+        this(ZlibWrapper.ZLIB, maxAllocation);
     }
 
     /**
@@ -44,6 +59,21 @@ public JZlibDecoder() {
      * @throws DecompressionException if failed to initialize zlib
      */
     public JZlibDecoder(ZlibWrapper wrapper) {
+        this(wrapper, 0);
+    }
+
+    /**
+     * Creates a new instance with the specified wrapper and maximum buffer allocation.
+     *
+     * @param maxAllocation
+     *          Maximum size of the decompression buffer. Must be &gt;= 0.
+     *          If zero, maximum size is decided by the {@link ByteBufAllocator}.
+     *
+     * @throws DecompressionException if failed to initialize zlib
+     */
+    public JZlibDecoder(ZlibWrapper wrapper, int maxAllocation) {
+        super(maxAllocation);
+
         ObjectUtil.checkNotNull(wrapper, "wrapper");
 
         int resultCode = z.init(ZlibUtil.convertWrapperType(wrapper));
@@ -60,6 +90,22 @@ public JZlibDecoder(ZlibWrapper wrapper) {
      * @throws DecompressionException if failed to initialize zlib
      */
     public JZlibDecoder(byte[] dictionary) {
+        this(dictionary, 0);
+    }
+
+    /**
+     * Creates a new instance with the specified preset dictionary and maximum buffer allocation.
+     * The wrapper is always {@link ZlibWrapper#ZLIB} because it is the only format that
+     * supports the preset dictionary.
+     *
+     * @param maxAllocation
+     *          Maximum size of the decompression buffer. Must be &gt;= 0.
+     *          If zero, maximum size is decided by the {@link ByteBufAllocator}.
+     *
+     * @throws DecompressionException if failed to initialize zlib
+     */
+    public JZlibDecoder(byte[] dictionary, int maxAllocation) {
+        super(maxAllocation);
         this.dictionary = ObjectUtil.checkNotNull(dictionary, "dictionary");
         int resultCode;
         resultCode = z.inflateInit(JZlib.W_ZLIB);
@@ -105,11 +151,11 @@ protected void decode(ChannelHandlerContext ctx, ByteBuf in, List<Object> out) t
             final int oldNextInIndex = z.next_in_index;
 
             // Configure output.
-            ByteBuf decompressed = ctx.alloc().heapBuffer(inputLength << 1);
+            ByteBuf decompressed = prepareDecompressBuffer(ctx, null, inputLength << 1);
 
             try {
                 loop: for (;;) {
-                    decompressed.ensureWritable(z.avail_in << 1);
+                    decompressed = prepareDecompressBuffer(ctx, decompressed, z.avail_in << 1);
                     z.avail_out = decompressed.writableBytes();
                     z.next_out = decompressed.array();
                     z.next_out_index = decompressed.arrayOffset() + decompressed.writerIndex();
@@ -165,4 +211,9 @@ protected void decode(ChannelHandlerContext ctx, ByteBuf in, List<Object> out) t
             z.next_out = null;
         }
     }
+
+    @Override
+    protected void decompressionBufferExhausted(ByteBuf buffer) {
+        finished = true;
+    }
 }
diff --git a/codec/src/main/java/io/netty/handler/codec/compression/JdkZlibDecoder.java b/codec/src/main/java/io/netty/handler/codec/compression/JdkZlibDecoder.java
@@ -16,6 +16,7 @@
 package io.netty.handler.codec.compression;
 
 import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufAllocator;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.util.internal.ObjectUtil;
 
@@ -65,7 +66,19 @@ private enum GzipState {
      * Creates a new instance with the default wrapper ({@link ZlibWrapper#ZLIB}).
      */
     public JdkZlibDecoder() {
-        this(ZlibWrapper.ZLIB, null, false);
+        this(ZlibWrapper.ZLIB, null, false, 0);
+    }
+
+    /**
+     * Creates a new instance with the default wrapper ({@link ZlibWrapper#ZLIB})
+     * and the specified maximum buffer allocation.
+     *
+     * @param maxAllocation
+     *          Maximum size of the decompression buffer. Must be &gt;= 0.
+     *          If zero, maximum size is decided by the {@link ByteBufAllocator}.
+     */
+    public JdkZlibDecoder(int maxAllocation) {
+        this(ZlibWrapper.ZLIB, null, false, maxAllocation);
     }
 
     /**
@@ -74,7 +87,20 @@ public JdkZlibDecoder() {
      * supports the preset dictionary.
      */
     public JdkZlibDecoder(byte[] dictionary) {
-        this(ZlibWrapper.ZLIB, dictionary, false);
+        this(ZlibWrapper.ZLIB, dictionary, false, 0);
+    }
+
+    /**
+     * Creates a new instance with the specified preset dictionary and maximum buffer allocation.
+     * The wrapper is always {@link ZlibWrapper#ZLIB} because it is the only format that
+     * supports the preset dictionary.
+     *
+     * @param maxAllocation
+     *          Maximum size of the decompression buffer. Must be &gt;= 0.
+     *          If zero, maximum size is decided by the {@link ByteBufAllocator}.
+     */
+    public JdkZlibDecoder(byte[] dictionary, int maxAllocation) {
+        this(ZlibWrapper.ZLIB, dictionary, false, maxAllocation);
     }
 
     /**
@@ -83,18 +109,41 @@ public JdkZlibDecoder(byte[] dictionary) {
      * supported atm.
      */
     public JdkZlibDecoder(ZlibWrapper wrapper) {
-        this(wrapper, null, false);
+        this(wrapper, null, false, 0);
+    }
+
+    /**
+     * Creates a new instance with the specified wrapper and maximum buffer allocation.
+     * Be aware that only {@link ZlibWrapper#GZIP}, {@link ZlibWrapper#ZLIB} and {@link ZlibWrapper#NONE} are
+     * supported atm.
+     *
+     * @param maxAllocation
+     *          Maximum size of the decompression buffer. Must be &gt;= 0.
+     *          If zero, maximum size is decided by the {@link ByteBufAllocator}.
+     */
+    public JdkZlibDecoder(ZlibWrapper wrapper, int maxAllocation) {
+        this(wrapper, null, false, maxAllocation);
     }
 
     public JdkZlibDecoder(ZlibWrapper wrapper, boolean decompressConcatenated) {
-        this(wrapper, null, decompressConcatenated);
+        this(wrapper, null, decompressConcatenated, 0);
+    }
+
+    public JdkZlibDecoder(ZlibWrapper wrapper, boolean decompressConcatenated, int maxAllocation) {
+        this(wrapper, null, decompressConcatenated, maxAllocation);
     }
 
     public JdkZlibDecoder(boolean decompressConcatenated) {
-        this(ZlibWrapper.GZIP, null, decompressConcatenated);
+        this(ZlibWrapper.GZIP, null, decompressConcatenated, 0);
     }
 
-    private JdkZlibDecoder(ZlibWrapper wrapper, byte[] dictionary, boolean decompressConcatenated) {
+    public JdkZlibDecoder(boolean decompressConcatenated, int maxAllocation) {
+        this(ZlibWrapper.GZIP, null, decompressConcatenated, maxAllocation);
+    }
+
+    private JdkZlibDecoder(ZlibWrapper wrapper, byte[] dictionary, boolean decompressConcatenated, int maxAllocation) {
+        super(maxAllocation);
+
         ObjectUtil.checkNotNull(wrapper, "wrapper");
 
         this.decompressConcatenated = decompressConcatenated;
@@ -177,7 +226,7 @@ protected void decode(ChannelHandlerContext ctx, ByteBuf in, List<Object> out) t
             inflater.setInput(array);
         }
 
-        ByteBuf decompressed = ctx.alloc().heapBuffer(inflater.getRemaining() << 1);
+        ByteBuf decompressed = prepareDecompressBuffer(ctx, null, inflater.getRemaining() << 1);
         try {
             boolean readFooter = false;
             while (!inflater.needsInput()) {
@@ -208,7 +257,7 @@ protected void decode(ChannelHandlerContext ctx, ByteBuf in, List<Object> out) t
                     }
                     break;
                 } else {
-                    decompressed.ensureWritable(inflater.getRemaining() << 1);
+                    decompressed = prepareDecompressBuffer(ctx, decompressed, inflater.getRemaining() << 1);
                 }
             }
 
@@ -238,6 +287,11 @@ protected void decode(ChannelHandlerContext ctx, ByteBuf in, List<Object> out) t
         }
     }
 
+    @Override
+    protected void decompressionBufferExhausted(ByteBuf buffer) {
+        finished = true;
+    }
+
     @Override
     protected void handlerRemoved0(ChannelHandlerContext ctx) throws Exception {
         super.handlerRemoved0(ctx);

diff --git a/codec/src/main/java/io/netty/handler/codec/compression/ZlibDecoder.java b/codec/src/main/java/io/netty/handler/codec/compression/ZlibDecoder.java
@@ -16,16 +16,81 @@
 package io.netty.handler.codec.compression;
 
 import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufAllocator;
+import io.netty.channel.ChannelHandlerContext;
 import io.netty.handler.codec.ByteToMessageDecoder;
 
 /**
  * Decompresses a {@link ByteBuf} using the deflate algorithm.
  */
 public abstract class ZlibDecoder extends ByteToMessageDecoder {
 
+    /**
+     * Maximum allowed size of the decompression buffer.
+     */
+    protected final int maxAllocation;
+
+    /**
+     * Same as {@link #ZlibDecoder(int)} with maxAllocation = 0.
+     */
+    public ZlibDecoder() {
+        this(0);
+    }
+
+    /**
+     * Construct a new ZlibDecoder.
+     * @param maxAllocation
+     *          Maximum size of the decompression buffer. Must be &gt;= 0.
+     *          If zero, maximum size is decided by the {@link ByteBufAllocator}.
+     */
+    public ZlibDecoder(int maxAllocation) {
+        if (maxAllocation < 0) {
+            throw new IllegalArgumentException("maxAllocation must be >= 0");
+        }
+        this.maxAllocation = maxAllocation;
+    }
+
     /**
      * Returns {@code true} if and only if the end of the compressed stream
      * has been reached.
      */
     public abstract boolean isClosed();
+
+    /**
+     * Allocate or expand the decompression buffer, without exceeding the maximum allocation.
+     * Calls {@link #decompressionBufferExhausted(ByteBuf)} if the buffer is full and cannot be expanded further.
+     */
+    protected ByteBuf prepareDecompressBuffer(ChannelHandlerContext ctx, ByteBuf buffer, int preferredSize) {
+        if (buffer == null) {
+            if (maxAllocation == 0) {
+                return ctx.alloc().heapBuffer(preferredSize);
+            }
+
+            return ctx.alloc().heapBuffer(Math.min(preferredSize, maxAllocation), maxAllocation);
+        }
+
+        // this always expands the buffer if possible, even if the expansion is less than preferredSize
+        // we throw the exception only if the buffer could not be expanded at all
+        // this means that one final attempt to deserialize will always be made with the buffer at maxAllocation
+        if (buffer.ensureWritable(preferredSize, true) == 1) {
+            // buffer must be consumed so subclasses don't add it to output
+            // we therefore duplicate it when calling decompressionBufferExhausted() to guarantee non-interference
+            // but wait until after to consume it so the subclass can tell how much output is really in the buffer
+            decompressionBufferExhausted(buffer.duplicate());
+            buffer.skipBytes(buffer.readableBytes());
+            throw new DecompressionException("Decompression buffer has reached maximum size: " + buffer.maxCapacity());
+        }
+
+        return buffer;
+    }
+
+    /**
+     * Called when the decompression buffer cannot be expanded further.
+     * Default implementation is a no-op, but subclasses can override in case they want to
+     * do something before the {@link DecompressionException} is thrown, such as log the
+     * data that was decompressed so far.
+     */
+    protected void decompressionBufferExhausted(ByteBuf buffer) {
+    }
+
 }
diff --git a/codec/src/test/java/io/netty/handler/codec/compression/JZlibTest.java b/codec/src/test/java/io/netty/handler/codec/compression/JZlibTest.java
@@ -23,7 +23,7 @@ protected ZlibEncoder createEncoder(ZlibWrapper wrapper) {
     }
 
     @Override
-    protected ZlibDecoder createDecoder(ZlibWrapper wrapper) {
-        return new JZlibDecoder(wrapper);
+    protected ZlibDecoder createDecoder(ZlibWrapper wrapper, int maxAllocation) {
+        return new JZlibDecoder(wrapper, maxAllocation);
     }
 }
diff --git a/codec/src/test/java/io/netty/handler/codec/compression/JdkZlibTest.java b/codec/src/test/java/io/netty/handler/codec/compression/JdkZlibTest.java
@@ -38,8 +38,8 @@ protected ZlibEncoder createEncoder(ZlibWrapper wrapper) {
     }
 
     @Override
-    protected ZlibDecoder createDecoder(ZlibWrapper wrapper) {
-        return new JdkZlibDecoder(wrapper);
+    protected ZlibDecoder createDecoder(ZlibWrapper wrapper, int maxAllocation) {
+        return new JdkZlibDecoder(wrapper, maxAllocation);
     }
 
     @Test(expected = DecompressionException.class)

diff --git a/codec/src/test/java/io/netty/handler/codec/compression/ZlibCrossTest1.java b/codec/src/test/java/io/netty/handler/codec/compression/ZlibCrossTest1.java
@@ -23,7 +23,7 @@ protected ZlibEncoder createEncoder(ZlibWrapper wrapper) {
     }
 
     @Override
-    protected ZlibDecoder createDecoder(ZlibWrapper wrapper) {
-        return new JZlibDecoder(wrapper);
+    protected ZlibDecoder createDecoder(ZlibWrapper wrapper, int maxAllocation) {
+        return new JZlibDecoder(wrapper, maxAllocation);
     }
 }
diff --git a/codec/src/test/java/io/netty/handler/codec/compression/ZlibCrossTest2.java b/codec/src/test/java/io/netty/handler/codec/compression/ZlibCrossTest2.java
@@ -25,8 +25,8 @@ protected ZlibEncoder createEncoder(ZlibWrapper wrapper) {
     }
 
     @Override
-    protected ZlibDecoder createDecoder(ZlibWrapper wrapper) {
-        return new JdkZlibDecoder(wrapper);
+    protected ZlibDecoder createDecoder(ZlibWrapper wrapper, int maxAllocation) {
+        return new JdkZlibDecoder(wrapper, maxAllocation);
     }
 
     @Test(expected = DecompressionException.class)