// Copyright Epic Games, Inc. All Rights Reserved.

#include "zenhttp/zipfs.h"

#include <zencore/logging.h>

ZEN_THIRD_PARTY_INCLUDES_START
#include <zlib.h>
ZEN_THIRD_PARTY_INCLUDES_END

namespace zen {

//////////////////////////////////////////////////////////////////////////
namespace {

#if ZEN_COMPILER_MSC
#	pragma warning(push)
#	pragma warning(disable : 4200)
#endif

	using ZipInt16 = uint16_t;

	struct ZipInt32
	{
				 operator uint32_t() const { return *(uint32_t*)Parts; }
		uint16_t Parts[2];
	};

	struct EocdRecord
	{
		enum : uint32_t
		{
			Magic = 0x0605'4b50,
		};
		ZipInt32 Signature;
		ZipInt16 ThisDiskIndex;
		ZipInt16 CdStartDiskIndex;
		ZipInt16 CdRecordThisDiskCount;
		ZipInt16 CdRecordCount;
		ZipInt32 CdSize;
		ZipInt32 CdOffset;
		ZipInt16 CommentSize;
		char	 Comment[];
	};

	struct CentralDirectoryRecord
	{
		enum : uint32_t
		{
			Magic = 0x0201'4b50,
		};

		ZipInt32 Signature;
		ZipInt16 VersionMadeBy;
		ZipInt16 VersionRequired;
		ZipInt16 Flags;
		ZipInt16 CompressionMethod;
		ZipInt16 LastModTime;
		ZipInt16 LastModDate;
		ZipInt32 Crc32;
		ZipInt32 CompressedSize;
		ZipInt32 OriginalSize;
		ZipInt16 FileNameLength;
		ZipInt16 ExtraFieldLength;
		ZipInt16 CommentLength;
		ZipInt16 DiskIndex;
		ZipInt16 InternalFileAttr;
		ZipInt32 ExternalFileAttr;
		ZipInt32 Offset;
		char	 FileName[];
	};

	struct LocalFileHeader
	{
		enum : uint32_t
		{
			Magic = 0x0403'4b50,
		};

		ZipInt32 Signature;
		ZipInt16 VersionRequired;
		ZipInt16 Flags;
		ZipInt16 CompressionMethod;
		ZipInt16 LastModTime;
		ZipInt16 LastModDate;
		ZipInt32 Crc32;
		ZipInt32 CompressedSize;
		ZipInt32 OriginalSize;
		ZipInt16 FileNameLength;
		ZipInt16 ExtraFieldLength;
		char	 FileName[];
	};

#if ZEN_COMPILER_MSC
#	pragma warning(pop)
#endif

}  // namespace

//////////////////////////////////////////////////////////////////////////
ZipFs::ZipFs(IoBuffer&& Buffer)
{
	// Treat the input buffer as attacker-controlled. Every offset, size, and
	// trailer length is validated against View.GetSize() before it is used to
	// form a pointer. All additions are performed in uint64_t to prevent 32-bit
	// wrap.
	const MemoryView View = Buffer.GetView();
	const uint8_t*	 Base = static_cast<const uint8_t*>(View.GetData());
	const size_t	 Size = View.GetSize();

	if (Size < sizeof(EocdRecord))
	{
		return;
	}

	const size_t	  EocdOffset = Size - sizeof(EocdRecord);
	const EocdRecord* Eocd		 = reinterpret_cast<const EocdRecord*>(Base + EocdOffset);

	// We only support a zip whose EOCD sits at the very end of the buffer — no
	// trailing comment, no Zip64.
	if (Eocd->Signature != EocdRecord::Magic)
	{
		return;
	}
	if (Eocd->ThisDiskIndex == 0xffff)
	{
		return;
	}

	const uint32_t CdOffsetRel	 = Eocd->CdOffset;
	const uint32_t CdSize		 = Eocd->CdSize;
	const uint16_t CdRecordCount = Eocd->CdRecordCount;

	// Central directory must fit strictly before the EOCD. Derive the archive
	// origin from the EOCD's declared layout so any pre-zip padding in the
	// buffer is accounted for; LFH offsets are relative to this origin.
	if (uint64_t(CdOffsetRel) + uint64_t(CdSize) > EocdOffset)
	{
		return;
	}

	const uint8_t* ArchiveStart = Base + (EocdOffset - CdOffsetRel - CdSize);
	const uint8_t* CdCursor		= ArchiveStart + CdOffsetRel;
	const uint8_t* CdEnd		= CdCursor + CdSize;

	for (uint32_t Record = 0; Record < CdRecordCount; ++Record)
	{
		if (size_t(CdEnd - CdCursor) < sizeof(CentralDirectoryRecord))
		{
			return;
		}
		const CentralDirectoryRecord& Cd = *reinterpret_cast<const CentralDirectoryRecord*>(CdCursor);
		if (Cd.Signature != CentralDirectoryRecord::Magic)
		{
			return;
		}

		const uint16_t NameLen	  = Cd.FileNameLength;
		const uint16_t ExtraLen	  = Cd.ExtraFieldLength;
		const uint16_t CommentLen = Cd.CommentLength;
		const uint32_t Trailer	  = uint32_t(NameLen) + uint32_t(ExtraLen) + uint32_t(CommentLen);
		if (size_t(CdEnd - CdCursor) - sizeof(CentralDirectoryRecord) < Trailer)
		{
			return;
		}

		const uint16_t Compression = Cd.CompressionMethod;
		const uint32_t Compressed  = Cd.CompressedSize;
		const uint32_t Original	   = Cd.OriginalSize;
		const uint32_t LfhOffset   = Cd.Offset;

		const bool AcceptableCompression = (Compression == 0) || (Compression == 8);
		const bool HasContent			 = Original > 0;
		if (AcceptableCompression && HasContent)
		{
			// LFH header must fit inside the [ArchiveStart, CdOffsetRel) region
			// (i.e. the pre-CD body). The LFH's own name + extra + compressed
			// payload must also fit in that region.
			const uint64_t LfhEndRel = uint64_t(LfhOffset) + sizeof(LocalFileHeader);
			if (LfhEndRel <= CdOffsetRel)
			{
				const LocalFileHeader* Lfh = reinterpret_cast<const LocalFileHeader*>(ArchiveStart + LfhOffset);
				const uint64_t		   DataStartRel =
					LfhEndRel + uint64_t(uint16_t(Lfh->FileNameLength)) + uint64_t(uint16_t(Lfh->ExtraFieldLength));
				const uint64_t DataEndRel = DataStartRel + uint64_t(Compressed);
				if (DataEndRel <= CdOffsetRel)
				{
					const uint8_t*	 FileData = ArchiveStart + DataStartRel;
					std::string_view FileName(Cd.FileName, NameLen);

					FileItem Item;
					Item.View			   = MemoryView{FileData, size_t(0)};
					Item.CompressionMethod = Compression;
					Item.CompressedSize	   = Compressed;
					Item.UncompressedSize  = Original;
					m_Files.insert({FileName, std::move(Item)});
				}
			}
		}

		CdCursor += sizeof(CentralDirectoryRecord) + Trailer;
	}

	m_Buffer = std::move(Buffer);
}

//////////////////////////////////////////////////////////////////////////
IoBuffer
ZipFs::GetFile(const std::string_view& FileName) const
{
	{
		RwLock::SharedLockScope _(m_FilesLock);

		FileMap::const_iterator Iter = m_Files.find(FileName);
		if (Iter == m_Files.end())
		{
			return {};
		}

		const FileItem& Item = Iter->second;
		if (Item.View.GetSize() > 0)
		{
			return IoBuffer(IoBuffer::Wrap, Item.View.GetData(), Item.View.GetSize());
		}
	}

	RwLock::ExclusiveLockScope _(m_FilesLock);

	FileItem& Item = m_Files.find(FileName)->second;
	if (Item.View.GetSize() > 0)
	{
		return IoBuffer(IoBuffer::Wrap, Item.View.GetData(), Item.View.GetSize());
	}

	const uint8_t* FileData = static_cast<const uint8_t*>(Item.View.GetData());

	if (Item.CompressionMethod == 0)
	{
		// Stored - point directly into the buffer
		Item.View = MemoryView(FileData, Item.UncompressedSize);
	}
	else
	{
		// Deflate - decompress using zlib
		Item.DecompressedData = IoBuffer(Item.UncompressedSize);

		z_stream Stream	 = {};
		Stream.next_in	 = const_cast<Bytef*>(FileData);
		Stream.avail_in	 = Item.CompressedSize;
		Stream.next_out	 = (Bytef*)Item.DecompressedData.GetMutableView().GetData();
		Stream.avail_out = Item.UncompressedSize;

		// Use raw inflate (-MAX_WBITS) since zip stores raw deflate streams
		if (inflateInit2(&Stream, -MAX_WBITS) != Z_OK)
		{
			ZEN_WARN("failed to initialize inflate for '{}'", FileName);
			return {};
		}

		int Result = inflate(&Stream, Z_FINISH);
		inflateEnd(&Stream);

		if (Result != Z_STREAM_END)
		{
			ZEN_WARN("failed to decompress '{}' (zlib error {})", FileName, Result);
			return {};
		}

		Item.View = Item.DecompressedData.GetView();
	}

	return IoBuffer(IoBuffer::Wrap, Item.View.GetData(), Item.View.GetSize());
}

}  // namespace zen