ObjFW  Diff

#ifndef INFLATE64
# import "OFInflateStream.h"
#else
# import "OFInflate64Stream.h"
# define OFInflateStream OFInflate64Stream
#endif



#import "OFInitializationFailedException.h"
#import "OFInvalidFormatException.h"
#import "OFNotOpenException.h"
#import "OFOutOfMemoryException.h"

#define BUFFER_SIZE		  OF_INFLATE_STREAM_BUFFER_SIZE

#define MAX_BITS 15

enum state {
	BLOCK_HEADER,
	UNCOMPRESSED_BLOCK_HEADER,
	UNCOMPRESSED_BLOCK,
	HUFFMAN_TREE,
	HUFFMAN_BLOCK
};

enum huffman_state {
	WRITE_VALUE,
	AWAIT_CODE,
	AWAIT_LENGTH_EXTRA_BITS,
	AWAIT_DISTANCE,
	AWAIT_DISTANCE_EXTRA_BITS,
	PROCESS_PAIR
};

struct huffman_tree {
	struct huffman_tree *leafs[2];
	uint16_t value;
};

#ifndef INFLATE64
static const uint8_t numDistanceCodes = 30;
static const uint8_t lengthCodes[29] = {
	/* indices are -257, values -3 */
	0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32, 40, 48, 56,
	64, 80, 96, 112, 128, 160, 192, 224, 255
};

	0, 0, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9, 10,
	10, 11, 11, 12, 12, 13, 13, 14, 14
};
#endif
static const uint8_t codeLengthsOrder[19] = {
	16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15
};
static struct huffman_tree *fixedLitLenTree, *fixedDistTree;

static bool
tryReadBits(OFInflateStream *stream, uint16_t *bits, uint8_t count)
{
	uint16_t ret = stream->_savedBits;

	assert(stream->_savedBitsLength < count);

	for (uint8_t i = stream->_savedBitsLength; i < count; i++) {
		if OF_UNLIKELY (stream->_bitIndex == 8) {
			if (stream->_bufferIndex < stream->_bufferLength)
				stream->_byte =
				    stream->_buffer[stream->_bufferIndex++];
			else {
				size_t length = [stream->_stream
				    readIntoBuffer: stream->_buffer

	stream->_savedBits = 0;
	stream->_savedBitsLength = 0;
	*bits = ret;

	return true;
}

static struct huffman_tree *
newTree(void)
{
	struct huffman_tree *tree;

	if ((tree = malloc(sizeof(*tree))) == NULL)
		@throw [OFOutOfMemoryException
		    exceptionWithRequestedSize: sizeof(*tree)];

	tree->leafs[0] = tree->leafs[1] = NULL;
	tree->value = 0xFFFF;

	return tree;
}

static void
treeInsert(struct huffman_tree *tree, uint16_t code, uint8_t length,
    uint16_t value)
{
	while (length > 0) {
		uint8_t bit;

		length--;
		bit = (code & (1 << length)) >> length;

		if (tree->leafs[bit] == NULL)
			tree->leafs[bit] = newTree();

		tree = tree->leafs[bit];
	}

	tree->value = value;
}

static struct huffman_tree *
constructTree(uint8_t lengths[], uint16_t count)
{
	struct huffman_tree *tree;
	uint16_t lengthCount[MAX_BITS + 1] = { 0 };
	uint16_t code, maxCode = 0, nextCode[MAX_BITS + 1];

	for (uint16_t i = 0; i < count; i++) {
		uint8_t length = lengths[i];

		if OF_UNLIKELY (length > MAX_BITS)
			@throw [OFInvalidFormatException exception];

		if (length > 0) {
			lengthCount[length]++;
			maxCode = i;
		}
	}

	code = 0;
	for (size_t i = 1; i <= MAX_BITS; i++) {
		code = (code + lengthCount[i - 1]) << 1;
		nextCode[i] = code;
	}

	tree = newTree();

	for (uint16_t i = 0; i <= maxCode; i++) {
		uint8_t length = lengths[i];

		if (length > 0)
			treeInsert(tree, nextCode[length]++, length, i);
	}

	return tree;
}

static bool
walkTree(OFInflateStream *stream, struct huffman_tree **tree, uint16_t *value)
{
	struct huffman_tree *iter = *tree;
	uint16_t bits;

	while (iter->value == 0xFFFF) {
		if OF_UNLIKELY (!tryReadBits(stream, &bits, 1)) {
			*tree = iter;
			return false;
		}

		if OF_UNLIKELY (iter->leafs[bits] == NULL)
			@throw [OFInvalidFormatException exception];

		iter = iter->leafs[bits];
	}

	*value = iter->value;
	return true;
}

static void
releaseTree(struct huffman_tree *tree)
{
	for (uint8_t i = 0; i < 2; i++)
		if OF_LIKELY (tree->leafs[i] != NULL)
			releaseTree(tree->leafs[i]);

	free(tree);
}

@implementation OFInflateStream
+ (void)initialize
{
	uint8_t lengths[288];

	if (self != [OFInflateStream class])
		return;

	for (uint16_t i = 0; i <= 143; i++)
		lengths[i] = 8;
	for (uint16_t i = 144; i <= 255; i++)
		lengths[i] = 9;
	for (uint16_t i = 256; i <= 279; i++)
		lengths[i] = 7;
	for (uint16_t i = 280; i <= 287; i++)
		lengths[i] = 8;

	fixedLitLenTree = constructTree(lengths, 288);

	for (uint16_t i = 0; i <= 31; i++)
		lengths[i] = 5;

	fixedDistTree = constructTree(lengths, 32);
}

#ifndef INFLATE64
+ (instancetype)streamWithStream: (OFStream *)stream
{
	return [[[self alloc] initWithStream: stream] autorelease];
}

- (instancetype)init
{
	OF_INVALID_INIT_METHOD
}

- (instancetype)initWithStream: (OFStream *)stream
{
	self = [super init];


	_stream = [stream retain];

	/* 0-7 address the bit, 8 means fetch next byte */
	_bitIndex = 8;

#ifdef INFLATE64
	_slidingWindowMask = 0xFFFF;
#else
	_slidingWindowMask = 0x7FFF;
#endif









	return self;
}

- (void)dealloc
{
	[self close];

	if (_state == HUFFMAN_TREE)
		if (_context.huffmanTree.codeLenTree != NULL)

			releaseTree(_context.huffmanTree.codeLenTree);

	if (_state == HUFFMAN_TREE || _state == HUFFMAN_BLOCK) {
		if (_context.huffman.litLenTree != fixedLitLenTree)
			releaseTree(_context.huffman.litLenTree);
		if (_context.huffman.distTree != fixedDistTree)
			releaseTree(_context.huffman.distTree);
	}

	[super dealloc];
}
#endif

- (size_t)lowlevelReadIntoBuffer: (void *)buffer_
			  length: (size_t)length
{
	uint8_t *buffer = buffer_;
	uint16_t bits, tmp;
	uint16_t value;
	size_t bytesWritten = 0;
	uint8_t *slidingWindow;
	uint16_t slidingWindowIndex;

	if (_stream == nil)
		@throw [OFNotOpenException exceptionWithObject: self];

	if (_atEndOfStream)
		return 0;

		tmp = (length < (size_t)CTX.length - CTX.position
		    ? (uint16_t)length : CTX.length - CTX.position);

		tmp = (uint16_t)[_stream readIntoBuffer: buffer + bytesWritten
						 length: tmp];

		if OF_UNLIKELY (_slidingWindow == NULL) {
			_slidingWindow =
			    [self allocMemoryWithSize: _slidingWindowMask + 1];
			/* Avoid leaking data */
			memset(_slidingWindow, 0, _slidingWindowMask + 1);
		}

		slidingWindow = _slidingWindow;
		slidingWindowIndex = _slidingWindowIndex;
		for (uint16_t i = 0; i < tmp; i++) {
			slidingWindow[slidingWindowIndex] =
			    buffer[bytesWritten + i];
			slidingWindowIndex = (slidingWindowIndex + 1) &
			    _slidingWindowMask;
		}
		_slidingWindowIndex = slidingWindowIndex;

			}

			if OF_LIKELY (CTX.lengths == NULL) {
				CTX.lengths = [self allocMemoryWithSize: 19];
				memset(CTX.lengths, 0, 19);
			}

			for (uint16_t i = CTX.receivedCount;
			    i < CTX.codeLenCodesCount + 4; i++) {
				if OF_UNLIKELY (!tryReadBits(self, &bits, 3)) {
					CTX.receivedCount = i;
					return bytesWritten;
				}

				CTX.lengths[codeLengthsOrder[i]] = bits;
			}


			CTX.codeLenTree = constructTree(CTX.lengths, 19);
			CTX.treeIter = CTX.codeLenTree;

			[self freeMemory: CTX.lengths];
			CTX.lengths = NULL;
			CTX.receivedCount = 0;
			CTX.value = 0xFF;
		}

		if OF_LIKELY (CTX.lengths == NULL)
			CTX.lengths = [self allocMemoryWithSize:
			    CTX.litLenCodesCount + CTX.distCodesCount + 258];

		for (uint16_t i = CTX.receivedCount;
		    i < CTX.litLenCodesCount + CTX.distCodesCount + 258;) {
			uint8_t j, count;

			if OF_LIKELY (CTX.value == 0xFF) {
				if OF_UNLIKELY (!walkTree(self, &CTX.treeIter,
				    &value)) {
					CTX.receivedCount = i;
					return bytesWritten;
				}

				CTX.treeIter = CTX.codeLenTree;

				if (value < 16) {

			for (j = 0; j < count; j++)
				CTX.lengths[i++] = value;

			CTX.value = 0xFF;
		}

		releaseTree(CTX.codeLenTree);
		CTX.codeLenTree = NULL;

		CTX.litLenTree = constructTree(CTX.lengths,
		    CTX.litLenCodesCount + 257);
		CTX.distTree = constructTree(
		    CTX.lengths + CTX.litLenCodesCount + 257,
		    CTX.distCodesCount + 1);

		[self freeMemory: CTX.lengths];

		/*
		 * litLenTree and distTree are at the same location in

			if OF_UNLIKELY (CTX.state == WRITE_VALUE) {
				if OF_UNLIKELY (length == 0)
					return bytesWritten;

				buffer[bytesWritten++] = CTX.value;
				length--;

				if (_slidingWindow == NULL) {
					_slidingWindow = [self
					    allocMemoryWithSize:
					    _slidingWindowMask + 1];
					/* Avoid leaking data */
					memset(_slidingWindow, 0,
					    _slidingWindowMask + 1);
				}

				_slidingWindow[_slidingWindowIndex] = CTX.value;
				_slidingWindowIndex =
				    (_slidingWindowIndex + 1) &
				    _slidingWindowMask;

				CTX.state = AWAIT_CODE;
				CTX.treeIter = CTX.litLenTree;

				CTX.state = AWAIT_DISTANCE;
				CTX.treeIter = CTX.distTree;
			}

			/* Distance of length distance pair */
			if (CTX.state == AWAIT_DISTANCE) {
				if OF_UNLIKELY (!walkTree(self, &CTX.treeIter,
				    &value))
					return bytesWritten;

				if OF_UNLIKELY (value >= numDistanceCodes)
					@throw [OFInvalidFormatException
					    exception];

				CTX.distance = distanceCodes[value];

				CTX.distance += bits;

				CTX.state = PROCESS_PAIR;
			}

			/* Length distance pair */
			if (CTX.state == PROCESS_PAIR) {
				uint16_t j;

				if OF_UNLIKELY (_slidingWindow == NULL)
					@throw [OFInvalidFormatException
					    exception];

				for (j = 0; j < CTX.length; j++) {
					uint16_t idx;

					if OF_UNLIKELY (length == 0) {
						CTX.length -= j;
						return bytesWritten;
					}

					    _slidingWindowMask;
				}

				CTX.state = AWAIT_CODE;
				CTX.treeIter = CTX.litLenTree;
			}


			if OF_UNLIKELY (!walkTree(self, &CTX.treeIter, &value))
				return bytesWritten;

			/* End of block */
			if OF_UNLIKELY (value == 256) {
				if (CTX.litLenTree != fixedLitLenTree)
					releaseTree(CTX.litLenTree);
				if (CTX.distTree != fixedDistTree)
					releaseTree(CTX.distTree);

				_state = BLOCK_HEADER;
				goto start;
			}

			/* Literal byte */
			if (value < 256) {
				if OF_UNLIKELY (length == 0) {
					CTX.state = WRITE_VALUE;
					CTX.value = value;
					return bytesWritten;
				}

				buffer[bytesWritten++] = value;
				length--;

				if (_slidingWindow == NULL) {
					_slidingWindow = [self
					    allocMemoryWithSize:
					    _slidingWindowMask + 1];
					/* Avoid leaking data */
					memset(_slidingWindow, 0,
					    _slidingWindowMask + 1);
				}

				_slidingWindow[_slidingWindowIndex] = value;
				_slidingWindowIndex =
				    (_slidingWindowIndex + 1) &
				    _slidingWindowMask;

				CTX.treeIter = CTX.litLenTree;
				continue;