EVEngine/src/modules/data/Compressor.cpp

#include "Compressor.h"
#include "common/config.h"
#include "common/Exception.h"

#include "lz4/lz4.h"
#include "lz4/lz4hc.h"

#include <zlib.h>
#include <vector>
#include <cstring>

namespace eve
{
namespace data
{

class LZ4Compressor : public Compressor
{
public:

	char *compress(std::string format, const char *data, size_t dataSize, int level, size_t &compressedSize) override
	{
		if (format != "lz4")
			throw eve::Exception("Invalid format (expecting LZ4)");

		if (dataSize > LZ4_MAX_INPUT_SIZE)
			throw eve::Exception("Data is too large for LZ4 compressor.");

		// We use a custom header to store some info with the compressed data.
		const size_t headersize = sizeof(uint32_t);

		int maxdestsize = LZ4_compressBound((int) dataSize);
		size_t maxsize = headersize + (size_t) maxdestsize;
		char *compressedbytes = nullptr;

		try
		{
			compressedbytes = new char[maxsize];
		}
		catch (std::bad_alloc &)
		{
			throw eve::Exception("Out of memory.");
		}

		// Store the size of the uncompressed data as a header.
#ifdef LOVE_BIG_ENDIAN
		// Make sure it's little-endian for storage.
		*(uint32_t *) compressedbytes = swapuint32((uint32_t) dataSize);
#else
		*(uint32_t *) compressedbytes = (uint32_t) dataSize;
#endif

		// Use LZ4-HC for compression level 9 and higher.
		int csize = 0;
		if (level > 8)
			csize = LZ4_compress_HC(data, compressedbytes + headersize, (int) dataSize, maxdestsize, LZ4HC_CLEVEL_DEFAULT);
		else
			csize = LZ4_compress_default(data, compressedbytes + headersize, (int) dataSize, maxdestsize);

		if (csize <= 0)
		{
			delete[] compressedbytes;
			throw eve::Exception("Could not LZ4-compress data.");
		}

		// We allocated space for the maximum possible amount of data, but the
		// actual compressed size might be much smaller, so we should shrink the
		// data buffer if so.
		if ((double) maxsize / (double) (csize + headersize) >= 1.2)
		{
			char *cbytes = new (std::nothrow) char[csize + headersize];
			if (cbytes)
			{
				memcpy(cbytes, compressedbytes, csize + headersize);
				delete[] compressedbytes;
				compressedbytes = cbytes;
			}
		}

		compressedSize = (size_t) csize + headersize;
		return compressedbytes;
	}

	char *decompress(std::string format, const char *data, size_t dataSize, size_t &decompressedSize) override
	{
		if (format != "lz4")
			throw eve::Exception("Invalid format (expecting LZ4)");

		const size_t headersize = sizeof(uint32_t);
		char *rawbytes = nullptr;

		if (dataSize < headersize)
			throw eve::Exception("Invalid LZ4-compressed data size.");

		// Extract the original uncompressed size (stored in our custom header.)
#ifdef LOVE_BIG_ENDIAN
		// Convert from stored little-endian to big-endian.
		uint32_t rawsize = swapuint32(*(uint32_t *) data);
#else
		uint32_t rawsize = *(uint32_t *) data;
#endif

		try
		{
			rawbytes = new char[rawsize];
		}
		catch (std::bad_alloc &)
		{
			throw eve::Exception("Out of memory.");
		}

		// If the uncompressed size is passed in as an argument (non-zero) and
		// it matches the header's stored size, then we assume it's 100% accurate
		// and we use a more efficient decompression function.
		if (decompressedSize > 0 && decompressedSize == (size_t) rawsize)
		{
			// We don't use the header here, but we need to account for its size.
			if (LZ4_decompress_fast(data + headersize, rawbytes, (int) decompressedSize) < 0)
			{
				delete[] rawbytes;
				throw eve::Exception("Could not decompress LZ4-compressed data.");
			}
		}
		else
		{
			// Account for our custom header's size in the decompress arguments.
			int result = LZ4_decompress_safe(data + headersize, rawbytes,
			                                 (int) (dataSize - headersize), rawsize);

			if (result < 0)
			{
				delete[] rawbytes;
				throw eve::Exception("Could not decompress LZ4-compressed data.");
			}

			decompressedSize = (size_t) result;
		}

		return rawbytes;
	}

	bool isSupported(std::string format) const override
	{
		return format == "lz4";
	}

}; // LZ4Compressor


class zlibCompressor : public Compressor
{
private:

	// The following three functions are mostly copied from the zlib source
	// (compressBound, compress2, and uncompress), but modified to support both
	// zlib and gzip.

	uLong zlibCompressBound(std::string format, uLong sourceLen)
	{
		uLong size = sourceLen + (sourceLen >> 12) + (sourceLen >> 14) + (sourceLen >> 25) + 13;

		// The gzip header is slightly larger than the zlib header.
		if (format == "gzip")
			size += 18 - 6;

		return size;
	}

	int zlibCompress(std::string format, Bytef *dest, uLongf *destLen, const Bytef *source, uLong sourceLen, int level)
	{
		z_stream stream = {};

		stream.next_in = (Bytef *) source;
		stream.avail_in = (uInt) sourceLen;

		stream.next_out = dest;
		stream.avail_out = (uInt) (*destLen);

		int windowbits = 15;
		if (format == "gzip")
			windowbits += 16; // This tells zlib to use a gzip header.
		else if (format == "deflate")
			windowbits = -windowbits;

		int err = deflateInit2(&stream, level, Z_DEFLATED, windowbits, 8, Z_DEFAULT_STRATEGY);

		if (err != Z_OK)
			return err;

		err = deflate(&stream, Z_FINISH);

		if (err != Z_STREAM_END)
		{
			deflateEnd(&stream);
			return err == Z_OK ? Z_BUF_ERROR : err;
		}

		*destLen = stream.total_out;

		return deflateEnd(&stream);
	}

	int zlibDecompress(std::string format, Bytef *dest, uLongf *destLen, const Bytef *source, uLong sourceLen)
	{
		z_stream stream = {};

		stream.next_in = (Bytef *) source;
		stream.avail_in = (uInt) sourceLen;

		stream.next_out = dest;
		stream.avail_out = (uInt) (*destLen);

		// 15 is the default. Adding 32 makes zlib auto-detect the header type.
		int windowbits = 15 + 32;

		if (format == "deflate")
			windowbits = -15;

		int err = inflateInit2(&stream, windowbits);

		if (err != Z_OK)
			return err;

		err = inflate(&stream, Z_FINISH);

		if (err != Z_STREAM_END)
		{
			inflateEnd(&stream);
			if (err == Z_NEED_DICT || (err == Z_BUF_ERROR && stream.avail_in == 0))
				return Z_DATA_ERROR;
			return err;
		}

		*destLen = stream.total_out;

		return inflateEnd(&stream);
	}

public:

	char *compress(std::string format, const char *data, size_t dataSize, int level, size_t &compressedSize) override
	{
		if (!isSupported(format))
			throw eve::Exception("Invalid format (expecting zlib or gzip)");

		if (level < 0)
			level = Z_DEFAULT_COMPRESSION;
		else if (level > 9)
			level = 9;

		uLong maxsize = zlibCompressBound(format, (uLong) dataSize);
		char *compressedbytes = nullptr;

		try
		{
			compressedbytes = new char[maxsize];
		}
		catch (std::bad_alloc &)
		{
			throw eve::Exception("Out of memory.");
		}

		uLongf destlen = maxsize;
		int status = zlibCompress(format, (Bytef *) compressedbytes, &destlen, (const Bytef *) data, (uLong) dataSize, level);

		if (status != Z_OK)
		{
			delete[] compressedbytes;
			throw eve::Exception("Could not zlib/gzip-compress data.");
		}

		// We allocated space for the maximum possible amount of data, but the
		// actual compressed size might be much smaller, so we should shrink the
		// data buffer if so.
		if ((double) maxsize / (double) destlen >= 1.3)
		{
			char *cbytes = new (std::nothrow) char[destlen];
			if (cbytes)
			{
				memcpy(cbytes, compressedbytes, destlen);
				delete[] compressedbytes;
				compressedbytes = cbytes;
			}
		}

		compressedSize = (size_t) destlen;
		return compressedbytes;
	}

	char *decompress(std::string format, const char *data, size_t dataSize, size_t &decompressedSize) override
	{
		if (!isSupported(format))
			throw eve::Exception("Invalid format (expecting zlib or gzip)");

		char *rawbytes = nullptr;

		// We might know the output size before decompression. If not, we guess.
		size_t rawsize = decompressedSize > 0 ? decompressedSize : dataSize * 2;

		// Repeatedly try to decompress with an increasingly large output buffer.
		while (true)
		{
			try
			{
				rawbytes = new char[rawsize];
			}
			catch (std::bad_alloc &)
			{
				throw eve::Exception("Out of memory.");
			}

			uLongf destLen = (uLongf) rawsize;
			int status = zlibDecompress(format, (Bytef *) rawbytes, &destLen, (const Bytef *) data, (uLong) dataSize);

			if (status == Z_OK)
			{
				decompressedSize = (size_t) destLen;
				break;
			}
			else if (status != Z_BUF_ERROR)
			{
				// For any error other than "not enough room", throw an exception.
				delete[] rawbytes;
				throw eve::Exception("Could not decompress zlib/gzip-compressed data.");
			}

			// Not enough room in the output buffer: try again with a larger size.
			delete[] rawbytes;
			rawsize *= 2;
		}

		return rawbytes;
	}

	bool isSupported(std::string format) const override
	{
		return format == "zlib" || format == "gzip" || format == "deflate";
	}

}; // zlibCompressor

Compressor *Compressor::getCompressor(std::string format)
{
	static LZ4Compressor lz4compressor;
	static zlibCompressor zlibcompressor;

	Compressor *compressors[] = {&lz4compressor, &zlibcompressor};

	for (Compressor *c : compressors)
	{
		if (c->isSupported(format))
			return c;
	}

	return nullptr;
}

} // data
} // eve