COMPRESS-382 and COMPRESS-386 -- add memoryLimit to CompressorStreamFactory and ZCompressorInputStream and LZMACompressorInputStream
Project: http://git-wip-us.apache.org/repos/asf/commons-compress/repo Commit: http://git-wip-us.apache.org/repos/asf/commons-compress/commit/7d73baf1 Tree: http://git-wip-us.apache.org/repos/asf/commons-compress/tree/7d73baf1 Diff: http://git-wip-us.apache.org/repos/asf/commons-compress/diff/7d73baf1 Branch: refs/heads/master Commit: 7d73baf10e435fcaa4927495afc185fb473c416b Parents: 3ecbdd7 Author: tballison <talli...@mitre.org> Authored: Fri Apr 21 19:21:42 2017 -0400 Committer: tballison <talli...@mitre.org> Committed: Fri Apr 21 19:21:42 2017 -0400 ---------------------------------------------------------------------- .../CompressorMemoryLimitException.java | 36 ++++++++++++++++ .../compressors/CompressorStreamFactory.java | 43 +++++++++++++++++--- .../lzma/LZMACompressorInputStream.java | 6 +-- .../compressors/lzw/LZWInputStream.java | 17 ++++++++ .../compressors/z/ZCompressorInputStream.java | 24 +++++++++-- .../compressors/DetectCompressorTestCase.java | 24 +++++++++++ src/test/resources/COMPRESS-386 | 1 + 7 files changed, 140 insertions(+), 11 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/commons-compress/blob/7d73baf1/src/main/java/org/apache/commons/compress/compressors/CompressorMemoryLimitException.java ---------------------------------------------------------------------- diff --git a/src/main/java/org/apache/commons/compress/compressors/CompressorMemoryLimitException.java b/src/main/java/org/apache/commons/compress/compressors/CompressorMemoryLimitException.java new file mode 100644 index 0000000..4c87d1e --- /dev/null +++ b/src/main/java/org/apache/commons/compress/compressors/CompressorMemoryLimitException.java @@ -0,0 +1,36 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.commons.compress.compressors; + +/** + * If a stream checks for estimated memory allocation, and the estimate + * goes above the memory limit, this is thrown. + * + * @since 1.14 + */ +public class CompressorMemoryLimitException extends CompressorException { + + public CompressorMemoryLimitException(String message) { + super(message); + } + + public CompressorMemoryLimitException(String message, Exception e) { + super(message, e); + } +} http://git-wip-us.apache.org/repos/asf/commons-compress/blob/7d73baf1/src/main/java/org/apache/commons/compress/compressors/CompressorStreamFactory.java ---------------------------------------------------------------------- diff --git a/src/main/java/org/apache/commons/compress/compressors/CompressorStreamFactory.java b/src/main/java/org/apache/commons/compress/compressors/CompressorStreamFactory.java index a29179c..a0a816f 100644 --- a/src/main/java/org/apache/commons/compress/compressors/CompressorStreamFactory.java +++ b/src/main/java/org/apache/commons/compress/compressors/CompressorStreamFactory.java @@ -57,6 +57,7 @@ import org.apache.commons.compress.utils.IOUtils; import org.apache.commons.compress.utils.Lists; import org.apache.commons.compress.utils.ServiceLoaderIterator; import org.apache.commons.compress.utils.Sets; +import org.tukaani.xz.MemoryLimitException; /** * <p> @@ -349,28 +350,52 @@ public class CompressorStreamFactory implements CompressorStreamProvider { */ private volatile boolean decompressConcatenated = false; + private final int memoryLimitInKb; /** * Create an instance with the decompress Concatenated option set to false. */ public CompressorStreamFactory() { this.decompressUntilEOF = null; + this.memoryLimitInKb = -1; } /** * Create an instance with the provided decompress Concatenated option. - * + * * @param decompressUntilEOF * if true, decompress until the end of the input; if false, stop * after the first stream and leave the input position to point * to the next byte after the stream. This setting applies to the * gzip, bzip2 and xz formats only. - * @since 1.10 + * + * @param memoryLimitInKb + * Some streams require allocation of potentially significant + * byte arrays/tables, and they can offer checks to prevent OOMs + * on corrupt files. Set the maximum allowed memory allocation in KBs. + * + * @since 1.14 */ - public CompressorStreamFactory(final boolean decompressUntilEOF) { + public CompressorStreamFactory(final boolean decompressUntilEOF, final int memoryLimitInKb) { this.decompressUntilEOF = Boolean.valueOf(decompressUntilEOF); // Also copy to existing variable so can continue to use that as the // current value this.decompressConcatenated = decompressUntilEOF; + this.memoryLimitInKb = memoryLimitInKb; + } + + + /** + * Create an instance with the provided decompress Concatenated option. + * + * @param decompressUntilEOF + * if true, decompress until the end of the input; if false, stop + * after the first stream and leave the input position to point + * to the next byte after the stream. This setting applies to the + * gzip, bzip2 and xz formats only. + * @since 1.10 + */ + public CompressorStreamFactory(final boolean decompressUntilEOF) { + this(decompressUntilEOF, -1); } /** @@ -510,7 +535,11 @@ public class CompressorStreamFactory implements CompressorStreamProvider { if (!LZMAUtils.isLZMACompressionAvailable()) { throw new CompressorException("LZMA compression is not available"); } - return new LZMACompressorInputStream(in); + try { + return new LZMACompressorInputStream(in, memoryLimitInKb); + } catch (MemoryLimitException e) { + throw new CompressorMemoryLimitException("exceeded calculated memory limit", e); + } } if (PACK200.equalsIgnoreCase(name)) { @@ -526,7 +555,11 @@ public class CompressorStreamFactory implements CompressorStreamProvider { } if (Z.equalsIgnoreCase(name)) { - return new ZCompressorInputStream(in); + try { + return new ZCompressorInputStream(in, memoryLimitInKb); + } catch (ZCompressorInputStream.IOExceptionWrappingMemoryLimitException e) { + throw new CompressorMemoryLimitException(e.getMessage()); + } } if (DEFLATE.equalsIgnoreCase(name)) { http://git-wip-us.apache.org/repos/asf/commons-compress/blob/7d73baf1/src/main/java/org/apache/commons/compress/compressors/lzma/LZMACompressorInputStream.java ---------------------------------------------------------------------- diff --git a/src/main/java/org/apache/commons/compress/compressors/lzma/LZMACompressorInputStream.java b/src/main/java/org/apache/commons/compress/compressors/lzma/LZMACompressorInputStream.java index ea8a498..0520e62 100644 --- a/src/main/java/org/apache/commons/compress/compressors/lzma/LZMACompressorInputStream.java +++ b/src/main/java/org/apache/commons/compress/compressors/lzma/LZMACompressorInputStream.java @@ -42,7 +42,7 @@ public class LZMACompressorInputStream extends CompressorInputStream { * * @param inputStream where to read the compressed data * - * @param memoryLimitKb calculated memory use threshold. Throws MemoryLimitException + * @param memoryLimitInKb calculated memory use threshold. Throws MemoryLimitException * if calculate memory use is above this threshold * * @throws IOException if the input is not in the .lzma format, @@ -53,9 +53,9 @@ public class LZMACompressorInputStream extends CompressorInputStream { * * @since 1.14 */ - public LZMACompressorInputStream(final InputStream inputStream, int memoryLimitKb) + public LZMACompressorInputStream(final InputStream inputStream, int memoryLimitInKb) throws IOException { - in = new LZMAInputStream(inputStream, memoryLimitKb); + in = new LZMAInputStream(inputStream, memoryLimitInKb); } /** {@inheritDoc} */ http://git-wip-us.apache.org/repos/asf/commons-compress/blob/7d73baf1/src/main/java/org/apache/commons/compress/compressors/lzw/LZWInputStream.java ---------------------------------------------------------------------- diff --git a/src/main/java/org/apache/commons/compress/compressors/lzw/LZWInputStream.java b/src/main/java/org/apache/commons/compress/compressors/lzw/LZWInputStream.java index 5967f66..b61d9a1 100644 --- a/src/main/java/org/apache/commons/compress/compressors/lzw/LZWInputStream.java +++ b/src/main/java/org/apache/commons/compress/compressors/lzw/LZWInputStream.java @@ -23,6 +23,7 @@ import java.io.InputStream; import java.nio.ByteOrder; import org.apache.commons.compress.compressors.CompressorInputStream; +import org.apache.commons.compress.compressors.CompressorMemoryLimitException; import org.apache.commons.compress.utils.BitInputStream; /** @@ -114,6 +115,22 @@ public abstract class LZWInputStream extends CompressorInputStream { /** * Initializes the arrays based on the maximum code size. * @param maxCodeSize maximum code size + * @param memoryLimitInKb maximum allowed table size in Kb + * @throws CompressorMemoryLimitException if maxTableSize is > memoryLimitInKb + */ + protected void initializeTables(final int maxCodeSize, final int memoryLimitInKb) + throws CompressorMemoryLimitException { + final int maxTableSize = 1 << maxCodeSize; + if (memoryLimitInKb > -1 && maxTableSize > memoryLimitInKb*1024) { + throw new CompressorMemoryLimitException("Tried to allocate "+maxTableSize + + " but memoryLimitInKb only allows "+(memoryLimitInKb*1024)); + } + initializeTables(maxCodeSize); + } + + /** + * Initializes the arrays based on the maximum code size. + * @param maxCodeSize maximum code size */ protected void initializeTables(final int maxCodeSize) { final int maxTableSize = 1 << maxCodeSize; http://git-wip-us.apache.org/repos/asf/commons-compress/blob/7d73baf1/src/main/java/org/apache/commons/compress/compressors/z/ZCompressorInputStream.java ---------------------------------------------------------------------- diff --git a/src/main/java/org/apache/commons/compress/compressors/z/ZCompressorInputStream.java b/src/main/java/org/apache/commons/compress/compressors/z/ZCompressorInputStream.java index 799e58d..22d7c0c 100644 --- a/src/main/java/org/apache/commons/compress/compressors/z/ZCompressorInputStream.java +++ b/src/main/java/org/apache/commons/compress/compressors/z/ZCompressorInputStream.java @@ -22,6 +22,7 @@ import java.io.IOException; import java.io.InputStream; import java.nio.ByteOrder; +import org.apache.commons.compress.compressors.CompressorMemoryLimitException; import org.apache.commons.compress.compressors.lzw.LZWInputStream; /** @@ -37,8 +38,9 @@ public class ZCompressorInputStream extends LZWInputStream { private final boolean blockMode; private final int maxCodeSize; private long totalCodesRead = 0; - - public ZCompressorInputStream(final InputStream inputStream) throws IOException { + + public ZCompressorInputStream(final InputStream inputStream, int memoryLimitInKb) + throws IOException { super(inputStream, ByteOrder.LITTLE_ENDIAN); final int firstByte = (int) in.readBits(8); final int secondByte = (int) in.readBits(8); @@ -51,9 +53,17 @@ public class ZCompressorInputStream extends LZWInputStream { if (blockMode) { setClearCode(DEFAULT_CODE_SIZE); } - initializeTables(maxCodeSize); + try { + initializeTables(maxCodeSize, memoryLimitInKb); + } catch (CompressorMemoryLimitException e) { + throw new IOExceptionWrappingMemoryLimitException(e.getMessage()); + } clearEntries(); } + + public ZCompressorInputStream(final InputStream inputStream) throws IOException { + this(inputStream, -1); + } private void clearEntries() { setTableSize((1 << 8) + (blockMode ? 1 : 0)); @@ -163,4 +173,12 @@ public class ZCompressorInputStream extends LZWInputStream { return length > 3 && signature[0] == MAGIC_1 && signature[1] == (byte) MAGIC_2; } + /** + * Wrapper that subclasses IOException to wrap a MemoryLimitException + */ + public static class IOExceptionWrappingMemoryLimitException extends IOException { + public IOExceptionWrappingMemoryLimitException(String message) { + super(message); + } + } } http://git-wip-us.apache.org/repos/asf/commons-compress/blob/7d73baf1/src/test/java/org/apache/commons/compress/compressors/DetectCompressorTestCase.java ---------------------------------------------------------------------- diff --git a/src/test/java/org/apache/commons/compress/compressors/DetectCompressorTestCase.java b/src/test/java/org/apache/commons/compress/compressors/DetectCompressorTestCase.java index 67abcd3..7cec5df 100644 --- a/src/test/java/org/apache/commons/compress/compressors/DetectCompressorTestCase.java +++ b/src/test/java/org/apache/commons/compress/compressors/DetectCompressorTestCase.java @@ -133,6 +133,10 @@ public final class DetectCompressorTestCase { assertEquals(CompressorStreamFactory.SNAPPY_FRAMED, detect("bla.tar.sz")); assertEquals(CompressorStreamFactory.Z, detect("bla.tar.Z")); + //make sure we don't oom on detect + assertEquals(CompressorStreamFactory.Z, detect("COMPRESS-386")); + assertEquals(CompressorStreamFactory.LZMA, detect("COMPRESS-382")); + try { CompressorStreamFactory.detect(new BufferedInputStream(new ByteArrayInputStream(new byte[0]))); fail("shouldn't be able to detect empty stream"); @@ -167,6 +171,26 @@ public final class DetectCompressorTestCase { } @Test + public void testMemoryLimit() throws Exception { + testMemoryLimit("COMPRESS-382"); + testMemoryLimit("COMPRESS-386"); + } + + private void testMemoryLimit(String fileName) throws IOException, CompressorException { + CompressorStreamFactory fac = new CompressorStreamFactory(true, + 100); + try (InputStream is = new BufferedInputStream( + new FileInputStream(getFile(fileName)))) { + InputStream compressorInputStream = fac.createCompressorInputStream(is); + fail("Should have thrown CompressorMemoryLimitException"); + } catch (CompressorMemoryLimitException e) { + + } + + } + + + @Test public void testOverride() { CompressorStreamFactory fac = new CompressorStreamFactory(); assertFalse(fac.getDecompressConcatenated()); http://git-wip-us.apache.org/repos/asf/commons-compress/blob/7d73baf1/src/test/resources/COMPRESS-386 ---------------------------------------------------------------------- diff --git a/src/test/resources/COMPRESS-386 b/src/test/resources/COMPRESS-386 new file mode 100644 index 0000000..36d7f52 --- /dev/null +++ b/src/test/resources/COMPRESS-386 @@ -0,0 +1 @@ + � B \ No newline at end of file
