Skip to content

Refactor BinaryTree and ExplodingInputStream #831

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.

Sign up for GitHub

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

Open
fkjellberg wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Refactor BinaryTree and ExplodingInputStream #831

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
86 changes: 35 additions & 51 deletions src/main/java/org/apache/commons/compress/archivers/zip/BinaryTree.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,9 +19,7 @@

package org.apache.commons.compress.archivers.zip;

import java.io.EOFException;
import java.io.IOException;
import java.io.InputStream;

import org.apache.commons.compress.archivers.ArchiveException;
import org.apache.commons.lang3.ArrayFill;
Expand All @@ -40,59 +38,47 @@ final class BinaryTree {
private static final int NODE = -2;

/**
* Decodes the packed binary tree from the specified stream.
* Creates a binary tree from the specified array of code lengths.
* <p>
* Each code length is the number of bits of the code for the symbol at that index, and must be between 1 and 16 (inclusive).
* </p>
*
* @param codeLengths code length per symbol, each value must be between 1 and 16 (inclusive).
* @return the binary tree.
* @throws ArchiveException if the code lengths array is empty, or if a code length is less than 1 or greater than 16.
*/
static BinaryTree decode(final InputStream inputStream, final int totalNumberOfValues) throws IOException {
if (totalNumberOfValues < 0) {
throw new IllegalArgumentException("totalNumberOfValues must be bigger than 0, is " + totalNumberOfValues);
}
// the first byte contains the size of the structure minus one
final int size = inputStream.read() + 1;
if (size == 0) {
throw new ArchiveException("Cannot read the size of the encoded tree, unexpected end of stream");
}

final byte[] encodedTree = org.apache.commons.compress.utils.IOUtils.readRange(inputStream, size);
if (encodedTree.length != size) {
throw new EOFException();
static BinaryTree fromCodeLengths(final int[] codeLengths) throws ArchiveException {
if (codeLengths.length == 0) {
throw new ArchiveException("Empty code length list");
}

/* The maximum bit length for a value (16 or lower) */
int maxLength = 0;

final int[] originalBitLengths = new int[totalNumberOfValues];
int pos = 0;
for (final byte b : encodedTree) {
// each byte encodes the number of values (upper 4 bits) for a bit length (lower 4 bits)
final int numberOfValues = ((b & 0xF0) >> 4) + 1;
if (pos + numberOfValues > totalNumberOfValues) {
throw new ArchiveException("Number of values exceeds given total number of values");
}
final int bitLength = (b & 0x0F) + 1;

for (int j = 0; j < numberOfValues; j++) {
originalBitLengths[pos++] = bitLength;
// validate and find max length
int maxCodeLength = 0;
for (int i = 0; i < codeLengths.length; i++) {
final int len = codeLengths[i];
if (len < 1 || len > 16) {
throw new ArchiveException(
String.format("Invalid code length at symbol %d: %d (expected in [1, 16])", i, len));
}

maxLength = Math.max(maxLength, bitLength);
maxCodeLength = Math.max(maxCodeLength, len);
}

final int oBitLengths = originalBitLengths.length;
// sort the array of bit lengths and memorize the permutation used to restore the order of the codes
final int[] permutation = new int[oBitLengths];
// sort the array of code lengths and memorize the permutation used to restore the order of the codes
final int[] permutation = new int[codeLengths.length];
for (int k = 0; k < permutation.length; k++) {
permutation[k] = k;
}

int c = 0;
final int[] sortedBitLengths = new int[oBitLengths];
for (int k = 0; k < oBitLengths; k++) {
final int[] sortedCodeLengths = new int[codeLengths.length];
for (int k = 1; k <= maxCodeLength; k++) {
// iterate over the values
for (int l = 0; l < oBitLengths; l++) {
for (int l = 0; l < codeLengths.length; l++) {
// look for the value in the original array
if (originalBitLengths[l] == k) {
if (codeLengths[l] == k) {
// put the value at the current position in the sorted array...
sortedBitLengths[c] = k;
sortedCodeLengths[c] = k;

// ...and memorize the permutation
permutation[c] = l;
Expand All @@ -105,27 +91,25 @@ static BinaryTree decode(final InputStream inputStream, final int totalNumberOfV
// decode the values of the tree
int code = 0;
int codeIncrement = 0;
int lastBitLength = 0;
int lastCodeLength = 0;

final int[] codes = new int[totalNumberOfValues];
final int[] codes = new int[codeLengths.length];

for (int i = totalNumberOfValues - 1; i >= 0; i--) {
for (int i = codeLengths.length - 1; i >= 0; i--) {
code += codeIncrement;
if (sortedBitLengths[i] != lastBitLength) {
lastBitLength = sortedBitLengths[i];
codeIncrement = 1 << 16 - lastBitLength;
if (sortedCodeLengths[i] != lastCodeLength) {
lastCodeLength = sortedCodeLengths[i];
codeIncrement = 1 << 16 - lastCodeLength;
}
codes[permutation[i]] = code;
}

// build the tree
final BinaryTree tree = new BinaryTree(maxLength);
final BinaryTree tree = new BinaryTree(maxCodeLength);

for (int k = 0; k < codes.length; k++) {
final int bitLength = originalBitLengths[k];
if (bitLength > 0) {
tree.addLeaf(0, Integer.reverse(codes[k] << 16), bitLength, k);
}
final int codeLength = codeLengths[k];
tree.addLeaf(0, Integer.reverse(codes[k] << 16), codeLength, k);
}

return tree;
Expand Down
71 changes: 68 additions & 3 deletions src/main/java/org/apache/commons/compress/archivers/zip/ExplodingInputStream.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@

import org.apache.commons.compress.archivers.ArchiveException;
import org.apache.commons.compress.utils.InputStreamStatistics;
import org.apache.commons.io.IOUtils;
import org.apache.commons.io.input.BoundedInputStream;
import org.apache.commons.io.input.CloseShieldInputStream;

Expand Down Expand Up @@ -176,11 +177,11 @@ private void init() throws IOException {
// we do not want to close in
try (BoundedInputStream cis = BoundedInputStream.builder().setInputStream(CloseShieldInputStream.wrap(in)).get()) {
if (numberOfTrees == 3) {
literalTree = BinaryTree.decode(cis, 256);
literalTree = readBinaryTree(cis, 256);
}

lengthTree = BinaryTree.decode(cis, 64);
distanceTree = BinaryTree.decode(cis, 64);
lengthTree = readBinaryTree(cis, 64);
distanceTree = readBinaryTree(cis, 64);
treeSizes += cis.getCount();
}

Expand All @@ -205,4 +206,68 @@ public int read() throws IOException {
return ret;
}

/**
* Reads the specified number of code lengths from the input stream and constructs a binary tree from them.
*
* @param inputStream the input stream to read the code lengths from.
* @param numberOfCodeLengths the number of code lengths to read.
* @return the binary tree.
* @throws IOException if an I/O error occurs.
*/
static BinaryTree readBinaryTree(final InputStream inputStream, final int numberOfCodeLengths) throws IOException {
return BinaryTree.fromCodeLengths(readCodeLengths(inputStream, numberOfCodeLengths));
}

/**
* Reads the code lengths from the specified stream.
* <p>
* The code lengths are stored in a run length encoded format, where each byte encodes the number of values (upper 4 bits) for a code length (lower 4
* bits). The first byte contains the size of the encoded code lengths minus one.
* </p>
* <p>
* Both the number of values and the code length are stored minus one, so each byte encodes 1 to 16 values with a code length between 1 and 16
* (inclusive).
* </p>
*
* @param inputStream the stream to read the code lengths from.
* @param numberOfCodeLengths the number of code lengths to read.
* @return array of code lengths, each value between 1 and 16 (inclusive).
* @throws IOException if an I/O error occurs.
*/
static int[] readCodeLengths(final InputStream inputStream, final int numberOfCodeLengths) throws IOException {
if (numberOfCodeLengths <= 0) {
throw new IllegalArgumentException("numberOfCodeLengths must be bigger than 0, is " + numberOfCodeLengths);
}

// The first byte contains the size of the encoded code lengths minus one
final int size = inputStream.read() + 1;
if (size == 0) {
throw new ArchiveException("Cannot read the size of the encoded tree, unexpected end of stream");
}

// Read the encoded code lengths from the stream
final byte[] encodedCodeLengths = IOUtils.toByteArray(inputStream, size);

// Decode the code lengths
final int[] codeLengths = new int[numberOfCodeLengths];
int pos = 0;
for (final byte b : encodedCodeLengths) {
// Each byte encodes the number of values (upper 4 bits) for a code length (lower 4 bits)
final int numberOfValues = ((b & 0xF0) >> 4) + 1;
if (pos + numberOfValues > numberOfCodeLengths) {
throw new ArchiveException("Number of values exceeds given total number of values");
}
final int codeLength = (b & 0x0F) + 1;
for (int j = 0; j < numberOfValues; j++) {
codeLengths[pos++] = codeLength;
}
}

if (pos != numberOfCodeLengths) {
throw new ArchiveException("Number of values does not match given total number of values");
}

return codeLengths;
}

}
93 changes: 71 additions & 22 deletions src/test/java/org/apache/commons/compress/archivers/zip/BinaryTreeTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,37 +19,69 @@

package org.apache.commons.compress.archivers.zip;

import static org.apache.commons.compress.AbstractTest.toByteArray;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertThrows;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;

import org.apache.commons.compress.archivers.ArchiveException;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;

class BinaryTreeTest {

@Test
void testDecode() throws IOException {
final InputStream in = new ByteArrayInputStream(new byte[] { 0x02, 0x42, 0x01, 0x13 });

final BinaryTree tree = BinaryTree.decode(in, 8);

assertNotNull(tree);

try (BitStream stream = new BitStream(new ByteArrayInputStream(new byte[] { (byte) 0x8D, (byte) 0xC5, (byte) 0x11, 0x00 }))) {
assertEquals(0, tree.read(stream));
assertEquals(1, tree.read(stream));
assertEquals(2, tree.read(stream));
assertEquals(3, tree.read(stream));
assertEquals(4, tree.read(stream));
assertEquals(5, tree.read(stream));
assertEquals(6, tree.read(stream));
assertEquals(7, tree.read(stream));
}
void testDecode1() throws IOException {
final int[] codeLengths = new int[] { 1, 1 };
// Symbol...........................: 0 1

final BinaryTree tree = BinaryTree.fromCodeLengths(codeLengths);

assertEquals(0, tree.read(createBitStream(0b0000_0001))); // xxxx xxx1
assertEquals(1, tree.read(createBitStream(0b0000_0000))); // xxxx xxx0
}

@Test
void testDecode2() throws IOException {
final int[] codeLengths = new int[] { 2, 2, 1 };
// Symbol...........................: 0 1 2

final BinaryTree tree = BinaryTree.fromCodeLengths(codeLengths);

assertEquals(2, tree.read(createBitStream(0b0000_0001))); // xxxx xxx1
assertEquals(0, tree.read(createBitStream(0b0000_0010))); // xxxx xx10
assertEquals(1, tree.read(createBitStream(0b0000_0000))); // xxxx xx00
}

@Test
void testDecode3() throws IOException {
final int[] codeLengths = new int[] { 3, 3, 3, 3, 3, 2, 4, 4 };
// Symbol...........................: 0 1 2 3 4 5 6 7

final BinaryTree tree = BinaryTree.fromCodeLengths(codeLengths);

assertEquals(0, tree.read(createBitStream(0b0000_0101))); // xxxx x101
assertEquals(1, tree.read(createBitStream(0b0000_0001))); // xxxx x001
assertEquals(2, tree.read(createBitStream(0b0000_0110))); // xxxx x110
assertEquals(3, tree.read(createBitStream(0b0000_0010))); // xxxx x010
assertEquals(4, tree.read(createBitStream(0b0000_0100))); // xxxx x100
assertEquals(5, tree.read(createBitStream(0b0000_0011))); // xxxx xx11
assertEquals(6, tree.read(createBitStream(0b0000_1000))); // xxxx 1000
assertEquals(7, tree.read(createBitStream(0b0000_0000))); // xxxx 0000
}

@Test
void testDecodeCodeLengthGreaterThanSymbolCount() throws IOException {
final int[] codeLengths = new int[] { 2, 2 };
// Symbol...........................: 0 1

final BinaryTree tree = BinaryTree.fromCodeLengths(codeLengths);

assertEquals(0, tree.read(createBitStream(0b0000_0010))); // xxxx xx10
assertEquals(1, tree.read(createBitStream(0b0000_0000))); // xxxx xx00
}

@Test
Expand All @@ -58,8 +90,25 @@ void testExceptions() {
binaryFinary.addLeaf(0, 0, 0, 1);
assertThrows(IllegalArgumentException.class, () -> binaryFinary.addLeaf(0, 0, 0, 1));

final InputStream is = new ByteArrayInputStream(new byte[] {});
assertThrows(ArchiveException.class, () -> BinaryTree.decode(is, 0));
assertThrows(ArchiveException.class, () -> new BinaryTree(4).read(new BitStream(new ByteArrayInputStream(new byte[] { 0 }))));
assertThrows(ArchiveException.class, () -> new BinaryTree(4).read(createBitStream(0)));
}

@Test
void testNoCodeLengths() {
final ArchiveException e = assertThrows(ArchiveException.class, () -> BinaryTree.fromCodeLengths(new int[0]),
"Expected ArchiveException for empty code length list");
assertEquals("Empty code length list", e.getMessage());
}

@ParameterizedTest
@ValueSource(ints = { Integer.MIN_VALUE, -1, 0, 17, Integer.MAX_VALUE })
void testCodeLengthOutOfRange(final int invalidLength) {
final ArchiveException e = assertThrows(ArchiveException.class,
() -> BinaryTree.fromCodeLengths(new int[] { 1, 2, invalidLength, 2 }));
assertEquals("Invalid code length at symbol 2: " + invalidLength + " (expected in [1, 16])", e.getMessage());
}

private BitStream createBitStream(final int... data) throws IOException {
return new BitStream(new ByteArrayInputStream(toByteArray(data)));
}
}
Loading