haiku.git

/*
 * Copyright 2009-2014, Ingo Weinhold, ingo_weinhold@gmx.de.
 * Distributed under the terms of the MIT License.
 */


#include <package/hpkg/v1/PackageDataReader.h>

#include <string.h>

#include <algorithm>
#include <new>

#include <DataIO.h>
#include <package/hpkg/BufferPool.h>
#include <package/hpkg/PoolBuffer.h>
#include <package/hpkg/v1/HPKGDefsPrivate.h>
#include <package/hpkg/v1/PackageData.h>
#include <ZlibCompressionAlgorithm.h>


namespace BPackageKit {

namespace BHPKG {

namespace V1 {


using BHPKG::BPrivate::PoolBufferPutter;


// minimum/maximum zlib chunk size we consider sane
static const size_t kMinSaneZlibChunkSize = 1024;
static const size_t kMaxSaneZlibChunkSize = 10 * 1024 * 1024;

// maximum number of entries in the zlib offset table buffer
static const uint32 kMaxZlibOffsetTableBufferSize = 512;

static const size_t kUncompressedReaderBufferSize
	= B_HPKG_DEFAULT_DATA_CHUNK_SIZE_ZLIB;


// #pragma mark - PackageDataReader


class PackageDataReader : public BAbstractBufferedDataReader {
public:
	PackageDataReader(BDataReader* dataReader)
		:
		fDataReader(dataReader)
	{
	}

	virtual ~PackageDataReader()
	{
	}

	virtual status_t Init(const BPackageData& data) = 0;

protected:
			BDataReader*			fDataReader;
};


// #pragma mark - UncompressedPackageDataReader


class UncompressedPackageDataReader : public PackageDataReader {
public:
	UncompressedPackageDataReader(BDataReader* dataReader,
		BBufferPool* bufferPool)
		:
		PackageDataReader(dataReader),
		fBufferPool(bufferPool)
	{
	}

	status_t Init(const BPackageData& data)
	{
		fOffset = data.Offset();
		fSize = data.UncompressedSize();
		return B_OK;
	}

	virtual status_t ReadData(off_t offset, void* buffer, size_t size)
	{
		if (size == 0)
			return B_OK;

		if (offset < 0)
			return B_BAD_VALUE;

		if ((uint64)offset > fSize || size > fSize - offset)
			return B_BAD_VALUE;

		return fDataReader->ReadData(fOffset + offset, buffer, size);
	}

	virtual status_t ReadDataToOutput(off_t offset, size_t size,
		BDataIO* output)
	{
		if (size == 0)
			return B_OK;

		if (offset < 0)
			return B_BAD_VALUE;

		if ((uint64)offset > fSize || size > fSize - offset)
			return B_BAD_VALUE;

		// get a temporary buffer
		PoolBuffer* buffer = fBufferPool->GetBuffer(
			kUncompressedReaderBufferSize);
		if (buffer == NULL)
			return B_NO_MEMORY;
		PoolBufferPutter bufferPutter(fBufferPool, &buffer);

		while (size > 0) {
			// read into the buffer
			size_t toRead = std::min(size, buffer->Size());
			status_t error = fDataReader->ReadData(fOffset + offset,
				buffer->Buffer(), toRead);
			if (error != B_OK)
				return error;

			// write to the output
			error = output->WriteExactly(buffer->Buffer(), toRead);
			if (error != B_OK)
				return error;

			offset += toRead;
			size -= toRead;
		}

		return B_OK;
	}

private:
	BBufferPool*	fBufferPool;
	uint64			fOffset;
	uint64			fSize;
};


// #pragma mark - ZlibPackageDataReader


class ZlibPackageDataReader : public PackageDataReader {
public:
	ZlibPackageDataReader(BDataReader* dataReader, BBufferPool* bufferPool)
		:
		PackageDataReader(dataReader),
		fBufferPool(bufferPool),
		fUncompressBuffer(NULL),
		fOffsetTable(NULL)
	{
	}

	~ZlibPackageDataReader()
	{
		delete[] fOffsetTable;

		fBufferPool->PutBuffer(&fUncompressBuffer);
	}

	status_t Init(const BPackageData& data)
	{
		fOffset = data.Offset();
		fCompressedSize = data.CompressedSize();
		fUncompressedSize = data.UncompressedSize();
		fChunkSize = data.ChunkSize();

		// validate chunk size
		if (fChunkSize == 0)
			fChunkSize = B_HPKG_DEFAULT_DATA_CHUNK_SIZE_ZLIB;
		if (fChunkSize < kMinSaneZlibChunkSize
			|| fChunkSize > kMaxSaneZlibChunkSize) {
			return B_BAD_DATA;
		}

		fChunkCount = (fUncompressedSize + (fChunkSize - 1)) / fChunkSize;
		fOffsetTableSize = (fChunkCount - 1) * sizeof(uint64);
		if (fOffsetTableSize >= fCompressedSize)
			return B_BAD_DATA;

		// allocate a buffer for the offset table
		if (fChunkCount > 1) {
			fOffsetTableBufferEntryCount = std::min(fChunkCount - 1,
				(uint64)kMaxZlibOffsetTableBufferSize);
			fOffsetTable = new(std::nothrow) uint64[
				fOffsetTableBufferEntryCount];
			if (fOffsetTable == NULL)
				return B_NO_MEMORY;

			fOffsetTableIndex = -1;
				// mark the table content invalid
		} else
			fChunkSize = fUncompressedSize;

		// mark uncompressed content invalid
		fUncompressedChunk = -1;

		return B_OK;
	}

	virtual status_t ReadDataToOutput(off_t offset, size_t size,
		BDataIO* output)
	{
		// check offset and size
		if (size == 0)
			return B_OK;

		if (offset < 0)
			return B_BAD_VALUE;

		if ((uint64)offset > fUncompressedSize
			|| size > fUncompressedSize - offset) {
			return B_BAD_VALUE;
		}

		// get our uncompressed chunk buffer back, if possible
		bool newBuffer;
		if (fBufferPool->GetBuffer(fChunkSize, &fUncompressBuffer, &newBuffer)
				== NULL) {
			return B_NO_MEMORY;
		}
		PoolBufferPutter uncompressBufferPutter(fBufferPool,
			&fUncompressBuffer);

		if (newBuffer)
			fUncompressedChunk = -1;

		// uncompress
		int64 chunkIndex = offset / fChunkSize;
		off_t chunkOffset = chunkIndex * fChunkSize;
		size_t inChunkOffset = offset - chunkOffset;

		while (size > 0) {
			// read and uncompress the chunk
			status_t error = _ReadChunk(chunkIndex);
			if (error != B_OK)
				return error;

			// write data to output
			size_t toCopy = std::min(size, (size_t)fChunkSize - inChunkOffset);
			error = output->WriteExactly(
				(uint8*)fUncompressBuffer->Buffer() + inChunkOffset, toCopy);
			if (error != B_OK)
				return error;

			size -= toCopy;

			chunkIndex++;
			chunkOffset += fChunkSize;
			inChunkOffset = 0;
		}

		return B_OK;
	}

private:
	status_t _ReadChunk(int64 chunkIndex)
	{
		if (chunkIndex == fUncompressedChunk)
			return B_OK;

		// get the chunk offset and size
		uint64 offset = 0;
		uint32 compressedSize = 0;
		status_t error = _GetCompressedChunkOffsetAndSize(chunkIndex, offset,
			compressedSize);
		if (error != B_OK)
			return error;

		uint32 uncompressedSize = (uint64)chunkIndex + 1 < fChunkCount
			? fChunkSize : fUncompressedSize - chunkIndex * fChunkSize;

		// read the chunk
		if (compressedSize == uncompressedSize) {
			// the chunk is not compressed -- read it directly into the
			// uncompressed buffer
			error = fDataReader->ReadData(offset, fUncompressBuffer->Buffer(),
				compressedSize);
		} else {
			// read to a read buffer and uncompress
			PoolBuffer* readBuffer = fBufferPool->GetBuffer(fChunkSize);
			if (readBuffer == NULL)
				return B_NO_MEMORY;
			PoolBufferPutter readBufferPutter(fBufferPool, readBuffer);

			error = fDataReader->ReadData(offset, readBuffer->Buffer(),
				compressedSize);
			if (error != B_OK)
				return error;

			iovec compressed = { readBuffer->Buffer(), compressedSize },
				uncompressed = { fUncompressBuffer->Buffer(), uncompressedSize };
			error = BZlibCompressionAlgorithm().DecompressBuffer(
				compressed, uncompressed);
			if (error == B_OK && uncompressed.iov_len != uncompressedSize)
				error = B_BAD_DATA;
		}

		if (error != B_OK) {
			// error reading/decompressing data -- mark the cached data invalid
			fUncompressedChunk = -1;
			return error;
		}

		fUncompressedChunk = chunkIndex;
		return B_OK;
	}

	status_t _GetCompressedChunkOffsetAndSize(int64 chunkIndex, uint64& _offset,
		uint32& _size)
	{
		// get the offset
		uint64 offset;
		if (chunkIndex == 0) {
			// first chunk is at 0
			offset = 0;
		} else {
			status_t error = _GetCompressedChunkRelativeOffset(chunkIndex,
				offset);
			if (error != B_OK)
				return error;
		}

		// get the end offset
		uint64 endOffset;
		if ((uint64)chunkIndex + 1 == fChunkCount) {
			// last chunk end with the end of the data
			endOffset = fCompressedSize - fOffsetTableSize;
		} else {
			status_t error = _GetCompressedChunkRelativeOffset(chunkIndex + 1,
				endOffset);
			if (error != B_OK)
				return error;
		}

		// sanity check
		if (endOffset < offset)
			return B_BAD_DATA;

		_offset = fOffset + fOffsetTableSize + offset;
		_size = endOffset - offset;
		return B_OK;
	}

	status_t _GetCompressedChunkRelativeOffset(int64 chunkIndex,
		uint64& _offset)
	{
		if (fOffsetTableIndex < 0 || fOffsetTableIndex > chunkIndex
			|| fOffsetTableIndex + fOffsetTableBufferEntryCount <= chunkIndex) {
			// read the table at the given index, or, if we can, the whole table
			int64 readAtIndex = fChunkCount - 1 > fOffsetTableBufferEntryCount
				? chunkIndex : 1;
			uint32 entriesToRead = std::min(
				(uint64)fOffsetTableBufferEntryCount,
				fChunkCount - readAtIndex);

			status_t error = fDataReader->ReadData(
				fOffset + (readAtIndex - 1) * sizeof(uint64),
				fOffsetTable, entriesToRead * sizeof(uint64));
			if (error != B_OK) {
				fOffsetTableIndex = -1;
				return error;
			}

			fOffsetTableIndex = readAtIndex;
		}

		// get and check the offset
		_offset = fOffsetTable[chunkIndex - fOffsetTableIndex];
		if (_offset > fCompressedSize - fOffsetTableSize)
			return B_BAD_DATA;

		return B_OK;
	}

private:
	BBufferPool*	fBufferPool;
	PoolBuffer*		fUncompressBuffer;
	int64			fUncompressedChunk;

	uint64			fOffset;
	uint64			fUncompressedSize;
	uint64			fCompressedSize;
	uint64			fOffsetTableSize;
	uint64			fChunkCount;
	uint32			fChunkSize;
	uint32			fOffsetTableBufferEntryCount;
	uint64*			fOffsetTable;
	int32			fOffsetTableIndex;
};


// #pragma mark - PackageDataHeapReader


class PackageDataInlineReader : public BBufferDataReader {
public:
	PackageDataInlineReader(const BPackageData& data)
		:
		BBufferDataReader(data.InlineData(), data.UncompressedSize()),
		fData(data)
	{
	}

private:
	BPackageData	fData;
};


// #pragma mark - BPackageDataReaderFactory


BPackageDataReaderFactory::BPackageDataReaderFactory(BBufferPool* bufferPool)
	:
	fBufferPool(bufferPool)
{
}


status_t
BPackageDataReaderFactory::CreatePackageDataReader(BDataReader* dataReader,
	const BPackageData& data, BAbstractBufferedDataReader*& _reader)
{
	if (data.IsEncodedInline()) {
		BAbstractBufferedDataReader* reader
			= new(std::nothrow) PackageDataInlineReader(data);
		if (reader == NULL)
			return B_NO_MEMORY;

		_reader = reader;
		return B_OK;
	}

	PackageDataReader* reader;

	switch (data.Compression()) {
		case B_HPKG_COMPRESSION_NONE:
			reader = new(std::nothrow) UncompressedPackageDataReader(
				dataReader, fBufferPool);
			break;
		case B_HPKG_COMPRESSION_ZLIB:
			reader = new(std::nothrow) ZlibPackageDataReader(dataReader,
				fBufferPool);
			break;
		default:
			return B_BAD_VALUE;
	}

	if (reader == NULL)
		return B_NO_MEMORY;

	status_t error = reader->Init(data);
	if (error != B_OK) {
		delete reader;
		return error;
	}

	_reader = reader;
	return B_OK;
}


}	// namespace V1

}	// namespace BHPKG

}	// namespace BPackageKit