From 1543218d3e7afcb33a90b728b14370395a3deca0 Mon Sep 17 00:00:00 2001
From: Rich DiCroce <Rich.DiCroce@scientificgames.com>
Date: Fri, 31 Jan 2020 06:11:06 -0500
Subject: [PATCH] Allow a limit to be set on the decompressed buffer size for
 ZlibDecoders (#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.
---
 .../codec/compression/JZlibDecoder.java       | 57 ++++++++++++++-
 .../codec/compression/JdkZlibDecoder.java     | 70 ++++++++++++++++---
 .../codec/compression/ZlibDecoder.java        | 65 +++++++++++++++++
 .../handler/codec/compression/JZlibTest.java  |  4 +-
 .../codec/compression/JdkZlibTest.java        |  4 +-
 .../codec/compression/ZlibCrossTest1.java     |  4 +-
 .../codec/compression/ZlibCrossTest2.java     |  4 +-
 .../handler/codec/compression/ZlibTest.java   | 57 ++++++++++++++-
 8 files changed, 245 insertions(+), 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
index f87344ab59..6c65cd5fe6 100644
--- 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 @@ package io.netty.handler.codec.compression;
 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 class JZlibDecoder extends ZlibDecoder {
      * @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 class JZlibDecoder extends ZlibDecoder {
      * @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 @@ public class JZlibDecoder extends ZlibDecoder {
             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 @@ public class JZlibDecoder extends ZlibDecoder {
             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
index 776d1e7816..7e694222aa 100644
--- 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 @@ public class JdkZlibDecoder extends ZlibDecoder {
      * 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 class JdkZlibDecoder extends ZlibDecoder {
      * 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 class JdkZlibDecoder extends ZlibDecoder {
      * 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 @@ public class JdkZlibDecoder extends ZlibDecoder {
             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 @@ public class JdkZlibDecoder extends ZlibDecoder {
                     }
                     break;
                 } else {
-                    decompressed.ensureWritable(inflater.getRemaining() << 1);
+                    decompressed = prepareDecompressBuffer(ctx, decompressed, inflater.getRemaining() << 1);
                 }
             }
 
@@ -238,6 +287,11 @@ public class JdkZlibDecoder extends ZlibDecoder {
         }
     }
 
+    @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
index d01bc6b4de..26fd3e74a7 100644
--- 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,6 +16,8 @@
 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;
 
 /**
@@ -23,9 +25,72 @@ import io.netty.handler.codec.ByteToMessageDecoder;
  */
 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
index 28f3919c60..015559ed3b 100644
--- 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 @@ public class JZlibTest extends ZlibTest {
     }
 
     @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
index 54a48a9cae..5ff19f1154 100644
--- 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 @@ public class JdkZlibTest extends ZlibTest {
     }
 
     @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
index 9e16e1a3dc..3c312742c8 100644
--- 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 @@ public class ZlibCrossTest1 extends ZlibTest {
     }
 
     @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
index 8717019ed9..00c6e18c42 100644
--- 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 @@ public class ZlibCrossTest2 extends ZlibTest {
     }
 
     @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/ZlibTest.java b/codec/src/test/java/io/netty/handler/codec/compression/ZlibTest.java
index 24f76b2a81..5e9d1288fe 100644
--- a/codec/src/test/java/io/netty/handler/codec/compression/ZlibTest.java
+++ b/codec/src/test/java/io/netty/handler/codec/compression/ZlibTest.java
@@ -15,7 +15,9 @@
  */
 package io.netty.handler.codec.compression;
 
+import io.netty.buffer.AbstractByteBufAllocator;
 import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufAllocator;
 import io.netty.buffer.ByteBufInputStream;
 import io.netty.buffer.Unpooled;
 import io.netty.channel.embedded.EmbeddedChannel;
@@ -88,8 +90,12 @@ public abstract class ZlibTest {
         rand.nextBytes(BYTES_LARGE);
     }
 
+    protected ZlibDecoder createDecoder(ZlibWrapper wrapper) {
+        return createDecoder(wrapper, 0);
+    }
+
     protected abstract ZlibEncoder createEncoder(ZlibWrapper wrapper);
-    protected abstract ZlibDecoder createDecoder(ZlibWrapper wrapper);
+    protected abstract ZlibDecoder createDecoder(ZlibWrapper wrapper, int maxAllocation);
 
     @Test
     public void testGZIP2() throws Exception {
@@ -345,6 +351,25 @@ public abstract class ZlibTest {
         testCompressLarge(ZlibWrapper.GZIP, ZlibWrapper.ZLIB_OR_NONE);
     }
 
+    @Test
+    public void testMaxAllocation() throws Exception {
+        int maxAllocation = 1024;
+        ZlibDecoder decoder = createDecoder(ZlibWrapper.ZLIB, maxAllocation);
+        EmbeddedChannel chDecoder = new EmbeddedChannel(decoder);
+        TestByteBufAllocator alloc = new TestByteBufAllocator(chDecoder.alloc());
+        chDecoder.config().setAllocator(alloc);
+
+        try {
+            chDecoder.writeInbound(Unpooled.wrappedBuffer(deflate(BYTES_LARGE)));
+            fail("decompressed size > maxAllocation, so should have thrown exception");
+        } catch (DecompressionException e) {
+            assertTrue(e.getMessage().startsWith("Decompression buffer has reached maximum size"));
+            assertEquals(maxAllocation, alloc.getMaxAllocation());
+            assertTrue(decoder.isClosed());
+            assertFalse(chDecoder.finish());
+        }
+    }
+
     private static byte[] gzip(byte[] bytes) throws IOException {
         ByteArrayOutputStream out = new ByteArrayOutputStream();
         GZIPOutputStream stream = new GZIPOutputStream(out);
@@ -360,4 +385,34 @@ public abstract class ZlibTest {
         stream.close();
         return out.toByteArray();
     }
+
+    private static final class TestByteBufAllocator extends AbstractByteBufAllocator {
+        private ByteBufAllocator wrapped;
+        private int maxAllocation;
+
+        TestByteBufAllocator(ByteBufAllocator wrapped) {
+            this.wrapped = wrapped;
+        }
+
+        public int getMaxAllocation() {
+            return maxAllocation;
+        }
+
+        @Override
+        public boolean isDirectBufferPooled() {
+            return wrapped.isDirectBufferPooled();
+        }
+
+        @Override
+        protected ByteBuf newHeapBuffer(int initialCapacity, int maxCapacity) {
+            maxAllocation = Math.max(maxAllocation, maxCapacity);
+            return wrapped.heapBuffer(initialCapacity, maxCapacity);
+        }
+
+        @Override
+        protected ByteBuf newDirectBuffer(int initialCapacity, int maxCapacity) {
+            maxAllocation = Math.max(maxAllocation, maxCapacity);
+            return wrapped.directBuffer(initialCapacity, maxCapacity);
+        }
+    }
 }
-- 
2.35.3

