CWE-ou-minkata/Sources/Plasma/PubUtilLib/plCompression/plZlibStream.cpp

/*==LICENSE==*

CyanWorlds.com Engine - MMOG client, server and tools
Copyright (C) 2011  Cyan Worlds, Inc.

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.

You can contact Cyan Worlds, Inc. by email legal@cyan.com
 or by snail mail at:
      Cyan Worlds, Inc.
      14617 N Newport Hwy
      Mead, WA   99021

*==LICENSE==*/
#include "plZlibStream.h"
#include <zlib.h>

voidpf ZlibAlloc(voidpf opaque, uInt items, uInt size)
{
	return ALLOC(items*size);
}

void ZlibFree(voidpf opaque, voidpf address)
{
	FREE(address);
}

plZlibStream::plZlibStream() : fOutput(nil), fZStream(nil), fHeader(kNeedMoreData), fDecompressedOk(false)
{
}

plZlibStream::~plZlibStream()
{
	hsAssert(!fOutput && !fZStream, "plZlibStream not closed");
}

hsBool plZlibStream::Open(const char* filename, const char* mode)
{
	wchar* wFilename = hsStringToWString(filename);
	wchar* wMode = hsStringToWString(mode);
	hsBool ret = Open(wFilename, wMode);
	delete [] wFilename;
	delete [] wMode;
	return ret;
}

hsBool plZlibStream::Open(const wchar* filename, const wchar* mode)
{
	fFilename = filename;
	fMode = mode;

	fOutput = NEW(hsUNIXStream);
	return fOutput->Open(filename, L"wb");
}

hsBool plZlibStream::Close()
{
	if (fOutput)
	{
		fOutput->Close();
		DEL(fOutput);
		fOutput = nil;
	}
	if (fZStream)
	{
		z_streamp zstream = (z_streamp)fZStream;
		inflateEnd(zstream);
		DEL(zstream);
		fZStream = nil;
	}

	return true;
}

UInt32 plZlibStream::Write(UInt32 byteCount, const void* buffer)
{
	UInt8* byteBuf = (UInt8*)buffer;

	// Check if we've read in the full gzip header yet
	if (fHeader == kNeedMoreData)
	{
		int cutAmt = IValidateGzHeader(byteCount, buffer);

		// Once we've read in the whole header, cut out whatever part of it is
		// in this buffer, leaving raw compressed data
		if (cutAmt != 0)
		{
			byteCount -= cutAmt;
			byteBuf += cutAmt;
		}
	}

	if (fHeader == kInvalidHeader)
		return 0;

	if (fHeader == kValidHeader)
	{
		ASSERT(fOutput);
		ASSERT(fZStream);
		z_streamp zstream = (z_streamp)fZStream;
		zstream->avail_in = byteCount;
		zstream->next_in = byteBuf;

		char outBuf[2048];
		while (zstream->avail_in != 0)
		{
			zstream->avail_out = sizeof(outBuf);
			zstream->next_out = (UInt8*)outBuf;

			UInt32 amtWritten = zstream->total_out;

			int ret = inflate(zstream, Z_NO_FLUSH);

			bool inflateErr = (ret == Z_NEED_DICT || ret == Z_DATA_ERROR ||
								ret == Z_STREAM_ERROR || ret == Z_MEM_ERROR || ret == Z_BUF_ERROR);
			// If we have a decompression error, just fail
			if (inflateErr)
			{
				hsAssert(!inflateErr, "Error in inflate");
				fHeader = kInvalidHeader;
				break;
			}

			amtWritten = zstream->total_out - amtWritten;
			fOutput->Write(amtWritten, outBuf);

			// If zlib says we hit the end of the stream, ignore avail_in
			if (ret == Z_STREAM_END)
			{
				fDecompressedOk = true;
				break;
			}
		}
	}

	return byteCount;
}

int plZlibStream::IValidateGzHeader(UInt32 byteCount, const void* buffer)
{
	// Ripped from gzio.cpp
	#define HEAD_CRC     0x02
	#define EXTRA_FIELD  0x04
	#define ORIG_NAME    0x08
	#define COMMENT      0x10
	#define RESERVED     0xE0
	static int gz_magic[2] = {0x1f, 0x8b}; // gzip magic header
	
	#define CheckForEnd() if (s.AtEnd()) { fHeader = kNeedMoreData; return 0; }

	int i;

	int initCacheSize = fHeaderCache.size();
	for (i = 0; i < byteCount; i++)
		fHeaderCache.push_back(((UInt8*)buffer)[i]);
	hsReadOnlyStream s(fHeaderCache.size(), &fHeaderCache[0]);
	
	// Check the gzip magic header
	for (i = 0; i < 2; i++)
	{
		UInt8 c = s.ReadByte();
		if (c != gz_magic[i])
		{
			CheckForEnd();
			fHeader = kInvalidHeader;
			return 0;
		}
	}
	
	int method = s.ReadByte();
	int flags = s.ReadByte();
	if (method != Z_DEFLATED || (flags & RESERVED) != 0)
	{
		CheckForEnd();
		fHeader = kInvalidHeader;
		return 0;
	}
	
	// Discard time, xflags and OS code:
	for (i = 0; i < 6; i++)
		s.ReadByte();
	CheckForEnd();
	
	if ((flags & EXTRA_FIELD) != 0)
	{
		// skip the extra field
		UInt16 len = s.ReadSwap16();
		while (len-- != 0 && s.ReadByte())
		{
			CheckForEnd();
		}
	}
	
	int c;

	// skip the original file name
	if ((flags & ORIG_NAME) != 0)
	{
		while ((c = s.ReadByte()) != 0)
		{
			CheckForEnd();
		}
	}
	
	// skip the .gz file comment
	if ((flags & COMMENT) != 0)
	{
		while ((c = s.ReadByte()) != 0)
		{
			CheckForEnd();
		}
	}

	// skip the header crc
	if ((flags & HEAD_CRC) != 0)
	{
		s.ReadSwap16();
		CheckForEnd();
	}
	
	CheckForEnd();
	
	UInt32 headerSize = s.GetPosition();
	UInt32 clipBuffer = headerSize - initCacheSize;
	
	// Initialize the zlib stream
	z_streamp zstream = NEW(z_stream_s);
	memset(zstream, 0, sizeof(z_stream_s));
	zstream->zalloc = ZlibAlloc;
	zstream->zfree = ZlibFree;
	zstream->opaque = nil;
	zstream->avail_in = byteCount - clipBuffer;
	zstream->next_in = (UInt8*)&fHeaderCache[clipBuffer];
	// Gotta use inflateInit2, because there's no header for zlib to look at.
	// The -15 tells it to not try and find a header
	bool initOk = (inflateInit2(zstream, -15) == Z_OK);
	fZStream = zstream;

	fHeaderCache.clear();

	if (!initOk)
	{
		hsAssert(0, "Zip init failed");
		fHeader = kInvalidHeader;
		return 0;
	}
	ASSERT(fZStream);

	fHeader = kValidHeader;
	return clipBuffer;
}

hsBool plZlibStream::AtEnd()
{
	hsAssert(0, "AtEnd not supported");
	return true;
}

UInt32 plZlibStream::Read(UInt32 byteCount, void* buffer)
{
	hsAssert(0, "Read not supported");
	return 0;
}

void plZlibStream::Skip(UInt32 deltaByteCount)
{
	hsAssert(0, "Skip not supported");
}

void plZlibStream::Rewind()
{
	// hack so rewind will work (someone thought it would be funny to not implement base class functions)
	Close();
	Open(fFilename.c_str(), fMode.c_str());
	fHeader = kNeedMoreData;
	fDecompressedOk = false;
}

void plZlibStream::FastFwd()
{
	hsAssert(0, "FastFwd not supported");
}

UInt32 plZlibStream::GetEOF()
{
	hsAssert(0, "GetEOF not supported");
	return 0;
}