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

#include "zenhttp/zipfs.h"

#include <zencore/iobuffer.h>

#if ZEN_WITH_TESTS

ZEN_THIRD_PARTY_INCLUDES_START
#	include <doctest/doctest.h>
#	include <zlib.h>
ZEN_THIRD_PARTY_INCLUDES_END

#	include <cstring>
#	include <vector>

namespace zen {
void
zipfs_test_forcelink()
{
}
}  // namespace zen

TEST_SUITE_BEGIN("http.zipfs");

namespace {

// Helpers to build a minimal zip file in memory
struct ZipBuilder
{
	std::vector<uint8_t> Data;

	struct Entry
	{
		std::string Name;
		uint32_t	LocalHeaderOffset;
		uint16_t	CompressionMethod;
		uint32_t	CompressedSize;
		uint32_t	UncompressedSize;
	};

	std::vector<Entry> Entries;

	void Append(const void* Src, size_t Size)
	{
		const uint8_t* Bytes = (const uint8_t*)Src;
		Data.insert(Data.end(), Bytes, Bytes + Size);
	}

	void AppendU16(uint16_t V) { Append(&V, 2); }
	void AppendU32(uint32_t V) { Append(&V, 4); }

	void AddFile(const std::string& Name, const void* Content, size_t ContentSize, bool Deflate)
	{
		std::vector<uint8_t> FileData;
		uint16_t			 Method = 0;

		if (Deflate)
		{
			// Compress with raw deflate (no zlib/gzip header)
			uLongf				 BoundSize = compressBound((uLong)ContentSize);
			std::vector<uint8_t> TempBuf(BoundSize);

			z_stream Stream	 = {};
			Stream.next_in	 = (Bytef*)Content;
			Stream.avail_in	 = (uInt)ContentSize;
			Stream.next_out	 = TempBuf.data();
			Stream.avail_out = (uInt)TempBuf.size();

			deflateInit2(&Stream, Z_DEFAULT_COMPRESSION, Z_DEFLATED, -MAX_WBITS, 8, Z_DEFAULT_STRATEGY);
			deflate(&Stream, Z_FINISH);
			deflateEnd(&Stream);

			TempBuf.resize(Stream.total_out);
			FileData = std::move(TempBuf);
			Method	 = 8;
		}
		else
		{
			FileData.assign((const uint8_t*)Content, (const uint8_t*)Content + ContentSize);
		}

		Entry E;
		E.Name				= Name;
		E.LocalHeaderOffset = (uint32_t)Data.size();
		E.CompressionMethod = Method;
		E.CompressedSize	= (uint32_t)FileData.size();
		E.UncompressedSize	= (uint32_t)ContentSize;
		Entries.push_back(E);

		// Local file header
		AppendU32(0x04034b50);			   // signature
		AppendU16(20);					   // version needed
		AppendU16(0);					   // flags
		AppendU16(Method);				   // compression method
		AppendU16(0);					   // last mod time
		AppendU16(0);					   // last mod date
		AppendU32(0);					   // crc32 (not validated by ZipFs)
		AppendU32(E.CompressedSize);	   // compressed size
		AppendU32(E.UncompressedSize);	   // uncompressed size
		AppendU16((uint16_t)Name.size());  // file name length
		AppendU16(0);					   // extra field length
		Append(Name.data(), Name.size());  // file name
		Append(FileData.data(), FileData.size());
	}

	zen::IoBuffer Build()
	{
		uint32_t CdOffset = (uint32_t)Data.size();

		for (const Entry& E : Entries)
		{
			// Central directory record
			AppendU32(0x02014b50);				 // signature
			AppendU16(20);						 // version made by
			AppendU16(20);						 // version needed
			AppendU16(0);						 // flags
			AppendU16(E.CompressionMethod);		 // compression method
			AppendU16(0);						 // last mod time
			AppendU16(0);						 // last mod date
			AppendU32(0);						 // crc32
			AppendU32(E.CompressedSize);		 // compressed size
			AppendU32(E.UncompressedSize);		 // uncompressed size
			AppendU16((uint16_t)E.Name.size());	 // file name length
			AppendU16(0);						 // extra field length
			AppendU16(0);						 // comment length
			AppendU16(0);						 // disk index
			AppendU16(0);						 // internal file attr
			AppendU32(0);						 // external file attr
			AppendU32(E.LocalHeaderOffset);		 // offset
			Append(E.Name.data(), E.Name.size());
		}

		uint32_t CdSize = (uint32_t)Data.size() - CdOffset;

		// End of central directory record
		AppendU32(0x06054b50);				  // signature
		AppendU16(0);						  // this disk
		AppendU16(0);						  // cd start disk
		AppendU16((uint16_t)Entries.size());  // cd records this disk
		AppendU16((uint16_t)Entries.size());  // cd records total
		AppendU32(CdSize);					  // cd size
		AppendU32(CdOffset);				  // cd offset
		AppendU16(0);						  // comment length

		zen::IoBuffer Buffer(Data.size());
		std::memcpy(Buffer.GetMutableView().GetData(), Data.data(), Data.size());
		return Buffer;
	}
};

}  // namespace

TEST_CASE("zipfs.stored")
{
	const char* Content = "Hello, World!";

	ZipBuilder Zip;
	Zip.AddFile("test.txt", Content, std::strlen(Content), false);

	zen::ZipFs Fs(Zip.Build());

	zen::IoBuffer Result = Fs.GetFile("test.txt");
	REQUIRE(Result);
	CHECK(Result.GetView().GetSize() == std::strlen(Content));
	CHECK(std::memcmp(Result.GetView().GetData(), Content, std::strlen(Content)) == 0);
}

TEST_CASE("zipfs.deflate")
{
	const char* Content = "This is some content that will be deflate compressed in the zip file.";

	ZipBuilder Zip;
	Zip.AddFile("compressed.txt", Content, std::strlen(Content), true);

	zen::ZipFs Fs(Zip.Build());

	zen::IoBuffer Result = Fs.GetFile("compressed.txt");
	REQUIRE(Result);
	CHECK(Result.GetView().GetSize() == std::strlen(Content));
	CHECK(std::memcmp(Result.GetView().GetData(), Content, std::strlen(Content)) == 0);
}

TEST_CASE("zipfs.mixed")
{
	const char* StoredContent  = "stored content";
	const char* DeflateContent = "deflate content that is compressed";

	ZipBuilder Zip;
	Zip.AddFile("stored.txt", StoredContent, std::strlen(StoredContent), false);
	Zip.AddFile("deflated.txt", DeflateContent, std::strlen(DeflateContent), true);

	zen::ZipFs Fs(Zip.Build());

	zen::IoBuffer Stored = Fs.GetFile("stored.txt");
	REQUIRE(Stored);
	CHECK(Stored.GetView().GetSize() == std::strlen(StoredContent));
	CHECK(std::memcmp(Stored.GetView().GetData(), StoredContent, std::strlen(StoredContent)) == 0);

	zen::IoBuffer Deflated = Fs.GetFile("deflated.txt");
	REQUIRE(Deflated);
	CHECK(Deflated.GetView().GetSize() == std::strlen(DeflateContent));
	CHECK(std::memcmp(Deflated.GetView().GetData(), DeflateContent, std::strlen(DeflateContent)) == 0);
}

TEST_CASE("zipfs.not_found")
{
	const char* Content = "data";

	ZipBuilder Zip;
	Zip.AddFile("exists.txt", Content, std::strlen(Content), false);

	zen::ZipFs Fs(Zip.Build());

	zen::IoBuffer Result = Fs.GetFile("missing.txt");
	CHECK(!Result);
}

//////////////////////////////////////////////////////////////////////////
// Malformed / attacker-shaped inputs — the parser must refuse these rather
// than read out of bounds. Field offsets below mirror the in-memory layout:
//
//   EocdRecord:                         CentralDirectoryRecord:
//     +0  Signature        (4)             +0  Signature        (4)
//     +4  ThisDiskIndex    (2)             +20 CompressedSize   (4)
//     +6  CdStartDiskIndex (2)             +28 FileNameLength   (2)
//     +8  CdRecordsThis    (2)             +42 Offset           (4)
//     +10 CdRecords        (2)
//     +12 CdSize           (4)
//     +16 CdOffset         (4)
//     +20 CommentSize      (2)             = 46 bytes header
//     = 22 bytes

namespace {
constexpr size_t kEocdSize	   = 22;
constexpr size_t kEocdOffCdSz  = 12;
constexpr size_t kEocdOffCdOff = 16;
constexpr size_t kCdrOffCompSz = 20;
constexpr size_t kCdrOffNameLn = 28;
constexpr size_t kCdrOffLfhOff = 42;

template<typename T>
void
WriteAt(zen::IoBuffer& Buf, size_t Offset, T Value)
{
	std::memcpy(static_cast<uint8_t*>(Buf.GetMutableView().GetData()) + Offset, &Value, sizeof(Value));
}

template<typename T>
T
ReadAt(const zen::IoBuffer& Buf, size_t Offset)
{
	T Out;
	std::memcpy(&Out, static_cast<const uint8_t*>(Buf.GetView().GetData()) + Offset, sizeof(Out));
	return Out;
}

size_t
EocdOffset(const zen::IoBuffer& Buf)
{
	return Buf.GetView().GetSize() - kEocdSize;
}

size_t
CdOffset(const zen::IoBuffer& Buf)
{
	return ReadAt<uint32_t>(Buf, EocdOffset(Buf) + kEocdOffCdOff);
}
}  // namespace

TEST_CASE("zipfs.buffer_smaller_than_eocd")
{
	zen::IoBuffer Tiny(10);
	std::memset(Tiny.GetMutableView().GetData(), 0, Tiny.GetView().GetSize());
	zen::ZipFs Fs(std::move(Tiny));
	CHECK(!Fs.GetFile("anything"));
}

TEST_CASE("zipfs.bad_eocd_magic")
{
	ZipBuilder Zip;
	Zip.AddFile("test.txt", "x", 1, false);
	zen::IoBuffer Buf = Zip.Build();

	WriteAt<uint32_t>(Buf, EocdOffset(Buf), 0xdeadbeef);

	zen::ZipFs Fs(std::move(Buf));
	CHECK(!Fs.GetFile("test.txt"));
}

TEST_CASE("zipfs.cd_offset_past_eocd")
{
	ZipBuilder Zip;
	Zip.AddFile("test.txt", "x", 1, false);
	zen::IoBuffer Buf = Zip.Build();

	WriteAt<uint32_t>(Buf, EocdOffset(Buf) + kEocdOffCdOff, 0xFFFF0000u);

	zen::ZipFs Fs(std::move(Buf));
	CHECK(!Fs.GetFile("test.txt"));
}

TEST_CASE("zipfs.cd_size_past_eocd")
{
	ZipBuilder Zip;
	Zip.AddFile("test.txt", "x", 1, false);
	zen::IoBuffer Buf = Zip.Build();

	WriteAt<uint32_t>(Buf, EocdOffset(Buf) + kEocdOffCdSz, 0xFFFF0000u);

	zen::ZipFs Fs(std::move(Buf));
	CHECK(!Fs.GetFile("test.txt"));
}

TEST_CASE("zipfs.bad_cd_record_signature")
{
	ZipBuilder Zip;
	Zip.AddFile("test.txt", "x", 1, false);
	zen::IoBuffer Buf = Zip.Build();

	WriteAt<uint32_t>(Buf, CdOffset(Buf), 0xdeadbeef);

	zen::ZipFs Fs(std::move(Buf));
	CHECK(!Fs.GetFile("test.txt"));
}

TEST_CASE("zipfs.oversize_filename_length")
{
	ZipBuilder Zip;
	Zip.AddFile("test.txt", "x", 1, false);
	zen::IoBuffer Buf = Zip.Build();

	WriteAt<uint16_t>(Buf, CdOffset(Buf) + kCdrOffNameLn, 0xFFFF);

	zen::ZipFs Fs(std::move(Buf));
	CHECK(!Fs.GetFile("test.txt"));
}

TEST_CASE("zipfs.lfh_offset_past_buffer")
{
	ZipBuilder Zip;
	Zip.AddFile("test.txt", "x", 1, false);
	zen::IoBuffer Buf = Zip.Build();

	WriteAt<uint32_t>(Buf, CdOffset(Buf) + kCdrOffLfhOff, 0xFFFFFFFF);

	zen::ZipFs Fs(std::move(Buf));
	CHECK(!Fs.GetFile("test.txt"));
}

TEST_CASE("zipfs.compressed_size_past_buffer")
{
	ZipBuilder Zip;
	Zip.AddFile("test.txt", "x", 1, false);
	zen::IoBuffer Buf = Zip.Build();

	WriteAt<uint32_t>(Buf, CdOffset(Buf) + kCdrOffCompSz, 0xFFFF0000u);

	zen::ZipFs Fs(std::move(Buf));
	CHECK(!Fs.GetFile("test.txt"));
}

TEST_SUITE_END();

#endif	// ZEN_WITH_TESTS