Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Allow a limit to be set on the decompressed buffer size for ZlibDecoders #9924

Merged
merged 8 commits into from Jan 31, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
Expand Up @@ -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;

Expand All @@ -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.
rdicroce marked this conversation as resolved.
Show resolved Hide resolved
*
* @param maxAllocation
* Maximum size of the decompression buffer. Must be >= 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);
}

/**
Expand All @@ -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 >= 0.
* If zero, maximum size is decided by the {@link ByteBufAllocator}.
*
* @throws DecompressionException if failed to initialize zlib
*/
public JZlibDecoder(ZlibWrapper wrapper, int maxAllocation) {
rdicroce marked this conversation as resolved.
Show resolved Hide resolved
super(maxAllocation);

ObjectUtil.checkNotNull(wrapper, "wrapper");

int resultCode = z.init(ZlibUtil.convertWrapperType(wrapper));
Expand All @@ -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 >= 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);
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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;
}
}
Expand Up @@ -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;

Expand Down Expand Up @@ -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);
}

/**
Expand All @@ -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);
}

/**
Expand All @@ -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;
Expand Down Expand Up @@ -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()) {
Expand Down Expand Up @@ -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);
}
}

Expand Down Expand Up @@ -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);
Expand Down
Expand Up @@ -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() {
rdicroce marked this conversation as resolved.
Show resolved Hide resolved
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) {
rdicroce marked this conversation as resolved.
Show resolved Hide resolved
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) {
rdicroce marked this conversation as resolved.
Show resolved Hide resolved
// 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) {
}

}
Expand Up @@ -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);
}
}
Expand Up @@ -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)
Expand Down
Expand Up @@ -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);
}
}
Expand Up @@ -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)
Expand Down