path: root/src/zenhttp/packageformat.cpp

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

#include <zenhttp/packageformat.h>

#include <zencore/compactbinarybuilder.h>
#include <zencore/compactbinarypackage.h>
#include <zencore/compactbinaryutil.h>
#include <zencore/compositebuffer.h>
#include <zencore/filesystem.h>
#include <zencore/fmtutils.h>
#include <zencore/iobuffer.h>
#include <zencore/iohash.h>
#include <zencore/logging.h>
#include <zencore/scopeguard.h>
#include <zencore/stream.h>
#include <zencore/testing.h>
#include <zencore/testutils.h>
#include <zencore/trace.h>

#include <span>
#include <vector>

#include <EASTL/fixed_vector.h>

#if ZEN_PLATFORM_WINDOWS
#	include <zencore/windows.h>
#endif

ZEN_THIRD_PARTY_INCLUDES_START
#include <tsl/robin_map.h>
ZEN_THIRD_PARTY_INCLUDES_END

namespace zen {

const std::string_view HandlePrefix(":?#:");

typedef eastl::fixed_vector<IoBuffer, 16> IoBufferVec_t;

IoBufferVec_t FormatPackageMessageInternal(const CbPackage& Data, FormatFlags Flags, void* TargetProcessHandle);

std::vector<IoBuffer>
FormatPackageMessage(const CbPackage& Data, void* TargetProcessHandle)
{
	return FormatPackageMessage(Data, FormatFlags::kDefault, TargetProcessHandle);
}
CompositeBuffer
FormatPackageMessageBuffer(const CbPackage& Data, void* TargetProcessHandle)
{
	return FormatPackageMessageBuffer(Data, FormatFlags::kDefault, TargetProcessHandle);
}

std::vector<IoBuffer>
FormatPackageMessage(const CbPackage& Data, FormatFlags Flags, void* TargetProcessHandle)
{
	auto Vec = FormatPackageMessageInternal(Data, Flags, TargetProcessHandle);
	return std::vector<IoBuffer>(begin(Vec), end(Vec));
}

CompositeBuffer
FormatPackageMessageBuffer(const CbPackage& Data, FormatFlags Flags, void* TargetProcessHandle)
{
	auto Vec = FormatPackageMessageInternal(Data, Flags, TargetProcessHandle);
	return CompositeBuffer(std::span{begin(Vec), end(Vec)});
}

static void
MarshalLocal(CbAttachmentEntry*&		  AttachmentInfo,
			 const std::string&			  Path8,
			 CbAttachmentReferenceHeader& LocalRef,
			 const IoHash&				  AttachmentHash,
			 bool						  IsCompressed,
			 IoBufferVec_t&				  ResponseBuffers)
{
	IoBuffer RefBuffer(sizeof(CbAttachmentReferenceHeader) + Path8.size());

	CbAttachmentReferenceHeader* RefHdr = RefBuffer.MutableData<CbAttachmentReferenceHeader>();
	*RefHdr++							= LocalRef;
	memcpy(RefHdr, Path8.data(), Path8.size());

	*AttachmentInfo++ = {.PayloadSize = RefBuffer.GetSize(),
						 .Flags		  = (IsCompressed ? uint32_t(CbAttachmentEntry::kIsCompressed) : 0u) | CbAttachmentEntry::kIsLocalRef,
						 .AttachmentHash = AttachmentHash};

	ResponseBuffers.emplace_back(std::move(RefBuffer));
};

static bool
IsLocalRef(tsl::robin_map<void*, std::string>& FileNameMap,
		   std::vector<void*>&				   DuplicatedHandles,
		   const CompositeBuffer&			   AttachmentBinary,
		   bool								   DenyPartialLocalReferences,
		   void*							   TargetProcessHandle,
		   CbAttachmentReferenceHeader&		   LocalRef,
		   std::string&						   Path8)
{
	const SharedBuffer&	  Segment = AttachmentBinary.GetSegments().front();
	IoBufferFileReference Ref;
	const IoBuffer&		  SegmentBuffer = Segment.AsIoBuffer();

	if (!SegmentBuffer.GetFileReference(Ref))
	{
		return false;
	}

	if (DenyPartialLocalReferences && !SegmentBuffer.IsWholeFile())
	{
		return false;
	}

	if (auto It = FileNameMap.find(Ref.FileHandle); It != FileNameMap.end())
	{
		Path8 = It->second;
	}
	else
	{
		bool UseFilePath = true;
#if ZEN_PLATFORM_WINDOWS
		if (TargetProcessHandle != nullptr)
		{
			HANDLE TargetHandle = INVALID_HANDLE_VALUE;
			BOOL   OK			= ::DuplicateHandle(GetCurrentProcess(),
										Ref.FileHandle,
										(HANDLE)TargetProcessHandle,
										&TargetHandle,
										FILE_GENERIC_READ,
										FALSE,
										0);
			if (OK)
			{
				DuplicatedHandles.push_back((void*)TargetHandle);
				Path8		= fmt::format("{}{}", HandlePrefix, reinterpret_cast<uint64_t>(TargetHandle));
				UseFilePath = false;
			}
		}
#else	// ZEN_PLATFORM_WINDOWS
		ZEN_UNUSED(TargetProcessHandle);
		ZEN_UNUSED(DuplicatedHandles);
		// Not supported on Linux/Mac. Could potentially use pidfd_getfd() but that requires a fairly new Linux kernel/includes and to
		// deal with access rights etc.
#endif	// ZEN_PLATFORM_WINDOWS
		if (UseFilePath)
		{
			ExtendablePathBuilder<256> LocalRefFile;
			std::error_code			   Ec;
			std::filesystem::path	   FilePath = PathFromHandle(Ref.FileHandle, Ec);
			if (Ec)
			{
				ZEN_WARN("Failed to get path for file handle {} in IsLocalRef check, reason '{}'", Ref.FileHandle, Ec.message());
				return false;
			}
			LocalRefFile.Append(std::filesystem::absolute(FilePath));
			Path8 = LocalRefFile.ToUtf8();
		}
		FileNameMap.insert_or_assign(Ref.FileHandle, Path8);
	}

	LocalRef.AbsolutePathLength = gsl::narrow<uint16_t>(Path8.size());
	LocalRef.PayloadByteOffset	= Ref.FileChunkOffset;
	LocalRef.PayloadByteSize	= Ref.FileChunkSize;

	return true;
};

IoBufferVec_t
FormatPackageMessageInternal(const CbPackage& Data, FormatFlags Flags, void* TargetProcessHandle)
{
	ZEN_TRACE_CPU("FormatPackageMessage");

	std::vector<void*> DuplicatedHandles;
#if ZEN_PLATFORM_WINDOWS
	auto _ = MakeGuard([&DuplicatedHandles, &TargetProcessHandle]() {
		if (TargetProcessHandle == nullptr)
		{
			return;
		}

		for (void* DuplicatedHandle : DuplicatedHandles)
		{
			HANDLE ClosingHandle;
			if (::DuplicateHandle((HANDLE)TargetProcessHandle,
								  (HANDLE)DuplicatedHandle,
								  GetCurrentProcess(),
								  &ClosingHandle,
								  0,
								  FALSE,
								  DUPLICATE_CLOSE_SOURCE | DUPLICATE_SAME_ACCESS) == TRUE)
			{
				::CloseHandle(ClosingHandle);
			}
		}
	});
#endif	// ZEN_PLATFORM_WINDOWS

	const std::span<const CbAttachment>& Attachments = Data.GetAttachments();
	IoBufferVec_t						 ResponseBuffers;

	ResponseBuffers.reserve(2 + Attachments.size());  // TODO: may want to use an additional fudge factor here to avoid growing since each
													  // attachment is likely to consist of several buffers

	IoBuffer AttachmentMetadataBuffer = IoBuffer{sizeof(CbPackageHeader) + sizeof(CbAttachmentEntry) * (Attachments.size() + /* root */ 1)};
	MutableMemoryView HeaderView	  = AttachmentMetadataBuffer.GetMutableView();
	// Fixed size header

	CbPackageHeader* Hdr = (CbPackageHeader*)HeaderView.GetData();
	*Hdr				 = {.HeaderMagic = kCbPkgMagic, .AttachmentCount = gsl::narrow<uint32_t>(Attachments.size())};
	HeaderView.MidInline(sizeof(CbPackageHeader));

	// Attachment metadata array
	CbAttachmentEntry* AttachmentInfo = reinterpret_cast<CbAttachmentEntry*>(HeaderView.GetData());
	ResponseBuffers.emplace_back(std::move(AttachmentMetadataBuffer));	// Attachment metadata

	// Root object

	IoBuffer RootIoBuffer = Data.GetObject().GetBuffer().AsIoBuffer();
	ZEN_ASSERT(RootIoBuffer.GetSize() > 0);
	*AttachmentInfo++ = {.PayloadSize = RootIoBuffer.Size(), .Flags = CbAttachmentEntry::kIsObject, .AttachmentHash = Data.GetObjectHash()};
	ResponseBuffers.emplace_back(std::move(RootIoBuffer));	// Root object

	// Attachment payloads
	tsl::robin_map<void*, std::string> FileNameMap;

	for (const CbAttachment& Attachment : Attachments)
	{
		if (Attachment.IsNull())
		{
			ZEN_NOT_IMPLEMENTED("Null attachments are not supported");
		}
		else if (const CompressedBuffer& AttachmentBuffer = Attachment.AsCompressedBinary())
		{
			const CompositeBuffer& Compressed	  = AttachmentBuffer.GetCompressed();
			IoHash				   AttachmentHash = Attachment.GetHash();

			// If the data is either not backed by a file, or there are multiple
			// fragments then we cannot marshal it by local reference. We might
			// want/need to extend this in the future to allow multiple chunk
			// segments to be marshaled at once

			bool MarshalByLocalRef = EnumHasAllFlags(Flags, FormatFlags::kAllowLocalReferences) && (Compressed.GetSegments().size() == 1);
			bool DenyPartialLocalReferences = EnumHasAllFlags(Flags, FormatFlags::kDenyPartialLocalReferences);
			CbAttachmentReferenceHeader LocalRef;
			std::string					Path8;

			if (MarshalByLocalRef)
			{
				MarshalByLocalRef = IsLocalRef(FileNameMap,
											   DuplicatedHandles,
											   Compressed,
											   DenyPartialLocalReferences,
											   TargetProcessHandle,
											   LocalRef,
											   Path8);
			}

			if (MarshalByLocalRef)
			{
				const bool IsCompressed = true;
				bool	   IsHandle		= false;
#if ZEN_PLATFORM_WINDOWS
				IsHandle = Path8.starts_with(HandlePrefix);
#endif
				MarshalLocal(AttachmentInfo, Path8, LocalRef, AttachmentHash, IsCompressed, ResponseBuffers);
				ZEN_DEBUG("Marshalled '{}' as file {} of {} bytes", Path8, IsHandle ? "handle" : "path", Compressed.GetSize());
			}
			else
			{
				*AttachmentInfo++ = {.PayloadSize	 = AttachmentBuffer.GetCompressedSize(),
									 .Flags			 = CbAttachmentEntry::kIsCompressed,
									 .AttachmentHash = AttachmentHash};

				std::span<const SharedBuffer> Segments = Compressed.GetSegments();
				ResponseBuffers.reserve(ResponseBuffers.size() + Segments.size() - 1);
				for (const SharedBuffer& Segment : Segments)
				{
					ZEN_ASSERT(Segment.GetSize() > 0);
					ResponseBuffers.emplace_back(Segment.AsIoBuffer());
				}
			}
		}
		else if (CbObject AttachmentObject = Attachment.AsObject())
		{
			IoBuffer ObjIoBuffer = AttachmentObject.GetBuffer().AsIoBuffer();
			ZEN_ASSERT(ObjIoBuffer.GetSize() > 0);
			*AttachmentInfo++ = {.PayloadSize	 = ObjIoBuffer.Size(),
								 .Flags			 = CbAttachmentEntry::kIsObject,
								 .AttachmentHash = Attachment.GetHash()};
			ResponseBuffers.emplace_back(std::move(ObjIoBuffer));
		}
		else if (const CompositeBuffer& AttachmentBinary = Attachment.AsCompositeBinary())
		{
			IoHash AttachmentHash = Attachment.GetHash();
			bool   MarshalByLocalRef =
				EnumHasAllFlags(Flags, FormatFlags::kAllowLocalReferences) && (AttachmentBinary.GetSegments().size() == 1);
			bool DenyPartialLocalReferences = EnumHasAllFlags(Flags, FormatFlags::kDenyPartialLocalReferences);

			CbAttachmentReferenceHeader LocalRef;
			std::string					Path8;

			if (MarshalByLocalRef)
			{
				MarshalByLocalRef = IsLocalRef(FileNameMap,
											   DuplicatedHandles,
											   AttachmentBinary,
											   DenyPartialLocalReferences,
											   TargetProcessHandle,
											   LocalRef,
											   Path8);
			}

			if (MarshalByLocalRef)
			{
				const bool IsCompressed = false;
				bool	   IsHandle		= false;
#if ZEN_PLATFORM_WINDOWS
				IsHandle = Path8.starts_with(HandlePrefix);
#endif
				MarshalLocal(AttachmentInfo, Path8, LocalRef, AttachmentHash, IsCompressed, ResponseBuffers);
				ZEN_DEBUG("Marshalled '{}' as file {} of {} bytes", Path8, IsHandle ? "handle" : "path", AttachmentBinary.GetSize());
			}
			else
			{
				*AttachmentInfo++ = {.PayloadSize = AttachmentBinary.GetSize(), .Flags = 0, .AttachmentHash = Attachment.GetHash()};

				std::span<const SharedBuffer> Segments = AttachmentBinary.GetSegments();
				ResponseBuffers.reserve(ResponseBuffers.size() + Segments.size() - 1);
				for (const SharedBuffer& Segment : Segments)
				{
					ZEN_ASSERT(Segment.GetSize() > 0);
					ResponseBuffers.emplace_back(Segment.AsIoBuffer());
				}
			}
		}
		else
		{
			ZEN_NOT_IMPLEMENTED("Unknown attachment kind");
		}
	}
	FileNameMap.clear();
#if ZEN_PLATFORM_WINDOWS
	DuplicatedHandles.clear();
#endif	// ZEN_PLATFORM_WINDOWS

	return ResponseBuffers;
}

bool
IsPackageMessage(IoBuffer Payload)
{
	if (Payload.GetSize() < sizeof(CbPackageHeader))
	{
		return false;
	}

	BinaryReader		   Reader(Payload);
	const CbPackageHeader* Hdr = reinterpret_cast<const CbPackageHeader*>(Reader.GetView(sizeof(CbPackageHeader)).GetData());
	if (Hdr->HeaderMagic != kCbPkgMagic)
	{
		return false;
	}

	return true;
}

CbPackage
ParsePackageMessage(IoBuffer Payload, std::function<IoBuffer(const IoHash&, uint64_t)> CreateBuffer)
{
	ZEN_TRACE_CPU("ParsePackageMessage");

	if (Payload.GetSize() < sizeof(CbPackageHeader))
	{
		throw std::invalid_argument(fmt::format("invalid CbPackage, missing complete header (size {})", Payload.GetSize()));
	}

	BinaryReader Reader(Payload);

	const CbPackageHeader* Hdr = reinterpret_cast<const CbPackageHeader*>(Reader.GetView(sizeof(CbPackageHeader)).GetData());
	if (Hdr->HeaderMagic != kCbPkgMagic)
	{
		throw std::invalid_argument(
			fmt::format("invalid CbPackage header magic, expected {0:x}, got {1:x}", static_cast<uint32_t>(kCbPkgMagic), Hdr->HeaderMagic));
	}
	Reader.Skip(sizeof(CbPackageHeader));

	const uint32_t ChunkCount = Hdr->AttachmentCount + 1;

	if (Reader.Remaining() < sizeof(CbAttachmentEntry) * ChunkCount)
	{
		throw std::invalid_argument(fmt::format("invalid CbPackage, missing attachment entry data (need {} bytes, have {} bytes)",
												sizeof(CbAttachmentEntry) * ChunkCount,
												Reader.Remaining()));
	}
	const CbAttachmentEntry* AttachmentEntries =
		reinterpret_cast<const CbAttachmentEntry*>(Reader.GetView(sizeof(CbAttachmentEntry) * ChunkCount).GetData());
	Reader.Skip(sizeof(CbAttachmentEntry) * ChunkCount);

	CbPackage Package;

	std::vector<CbAttachment> Attachments;
	Attachments.reserve(ChunkCount);  // Guessing here...

	tsl::robin_map<std::string, IoBuffer> PartialFileBuffers;

	std::vector<std::pair<uint32_t, std::string>> MalformedAttachments;

	for (uint32_t i = 0; i < ChunkCount; ++i)
	{
		const CbAttachmentEntry& Entry			= AttachmentEntries[i];
		const uint64_t			 AttachmentSize = Entry.PayloadSize;

		if (Reader.Remaining() < AttachmentSize)
		{
			throw std::invalid_argument(fmt::format("invalid CbPackage, missing attachment data (need {} bytes, have {} bytes)",
													AttachmentSize,
													Reader.Remaining()));
		}
		const IoBuffer AttachmentBuffer(Payload, Reader.CurrentOffset(), AttachmentSize);
		Reader.Skip(AttachmentSize);

		if (Entry.Flags & CbAttachmentEntry::kIsLocalRef)
		{
			// Marshal local reference - a "pointer" to the chunk backing file

			ZEN_ASSERT(AttachmentBuffer.Size() >= sizeof(CbAttachmentReferenceHeader));

			const CbAttachmentReferenceHeader* AttachRefHdr = AttachmentBuffer.Data<CbAttachmentReferenceHeader>();
			const char*						   PathPointer	= reinterpret_cast<const char*>(AttachRefHdr + 1);

			ZEN_ASSERT(AttachmentBuffer.Size() >= (sizeof(CbAttachmentReferenceHeader) + AttachRefHdr->AbsolutePathLength));
			std::string_view PathView(PathPointer, AttachRefHdr->AbsolutePathLength);

			IoBuffer FullFileBuffer;

			std::filesystem::path Path(Utf8ToWide(PathView));
			if (auto It = PartialFileBuffers.find(Path.string()); It != PartialFileBuffers.end())
			{
				FullFileBuffer = It->second;
			}
			else
			{
				if (PathView.starts_with(HandlePrefix))
				{
#if ZEN_PLATFORM_WINDOWS
					std::string_view		HandleString(PathView.substr(HandlePrefix.length()));
					std::optional<uint64_t> HandleNumber(ParseInt<uint64_t>(HandleString));
					if (HandleNumber.has_value())
					{
						HANDLE		   FileHandle = HANDLE(HandleNumber.value());
						ULARGE_INTEGER liFileSize;
						liFileSize.LowPart = ::GetFileSize(FileHandle, &liFileSize.HighPart);
						if (liFileSize.LowPart != INVALID_FILE_SIZE)
						{
							FullFileBuffer =
								IoBuffer(IoBuffer::File, (void*)FileHandle, 0, uint64_t(liFileSize.QuadPart), /*IsWholeFile*/ true);
							PartialFileBuffers.insert_or_assign(Path.string(), FullFileBuffer);
						}
					}
#else	// ZEN_PLATFORM_WINDOWS
		// Not supported on Linux/Mac. Could potentially use pidfd_getfd() but that requires a fairly new Linux kernel/includes
		// and to deal with acceess rights etc.
					ZEN_ASSERT(false);
#endif	// ZEN_PLATFORM_WINDOWS
				}
				else
				{
					FullFileBuffer = PartialFileBuffers.insert_or_assign(Path.string(), IoBufferBuilder::MakeFromFile(Path)).first->second;
				}
			}

			if (FullFileBuffer)
			{
				IoBuffer ChunkReference = AttachRefHdr->PayloadByteOffset == 0 && AttachRefHdr->PayloadByteSize == FullFileBuffer.GetSize()
											  ? FullFileBuffer
											  : IoBuffer(FullFileBuffer, AttachRefHdr->PayloadByteOffset, AttachRefHdr->PayloadByteSize);

				CompressedBuffer CompBuf(CompressedBuffer::FromCompressedNoValidate(std::move(ChunkReference)));
				if (CompBuf)
				{
					Attachments.emplace_back(CbAttachment(std::move(CompBuf), Entry.AttachmentHash));
				}
				else
				{
					MalformedAttachments.push_back(std::make_pair(i,
																  fmt::format("Invalid format in '{}' (offset {}, size {}) for {}",
																			  Path,
																			  AttachRefHdr->PayloadByteOffset,
																			  AttachRefHdr->PayloadByteSize,
																			  Entry.AttachmentHash)));
				}
			}
			else
			{
				MalformedAttachments.push_back(std::make_pair(i,
															  fmt::format("Unable to resolve chunk at '{}' (offset {}, size {}) for {}",
																		  Path,
																		  AttachRefHdr->PayloadByteOffset,
																		  AttachRefHdr->PayloadByteSize,
																		  Entry.AttachmentHash)));
			}
		}
		else if (Entry.Flags & CbAttachmentEntry::kIsCompressed)
		{
			if (Entry.Flags & CbAttachmentEntry::kIsObject)
			{
				CbObject AttachmentObject;

				CompressedBuffer CompBuf(CompressedBuffer::FromCompressedNoValidate(IoBuffer(AttachmentBuffer)));
				if (!CompBuf)
				{
					// First payload is always a compact binary object
					MalformedAttachments.push_back(
						std::make_pair(i,
									   fmt::format("Invalid format, expected compressed buffer for CbObject (size {}) for {}",
												   AttachmentBuffer.GetSize(),
												   Entry.AttachmentHash)));
				}
				else
				{
					CbValidateError ValidationError = CbValidateError::None;
					AttachmentObject				= ValidateAndReadCompactBinaryObject(std::move(CompBuf), ValidationError);
					if (ValidationError != CbValidateError::None)
					{
						MalformedAttachments.push_back(std::make_pair(
							i,
							fmt::format("Invalid format, CbObject for {}. Reason '{}'", Entry.AttachmentHash, ToString(ValidationError))));
					}
				}

				if (i == 0)
				{
					// First payload is always a compact binary object
					Package.SetObject(AttachmentObject);
				}
				else
				{
					Attachments.emplace_back(CbAttachment(AttachmentObject, Entry.AttachmentHash));
				}
			}
			else
			{
				CompressedBuffer CompBuf(CompressedBuffer::FromCompressedNoValidate(IoBuffer(AttachmentBuffer)));
				if (CompBuf)
				{
					Attachments.emplace_back(CbAttachment(std::move(CompBuf), Entry.AttachmentHash));
				}
				else
				{
					MalformedAttachments.push_back(
						std::make_pair(i,
									   fmt::format("Invalid format, expected compressed buffer for attachment (size {}) for {}",
												   AttachmentBuffer.GetSize(),
												   Entry.AttachmentHash)));
				}
			}
		}
		else /* not compressed */
		{
			if (Entry.Flags & CbAttachmentEntry::kIsObject)
			{
				CbValidateError ValidationError	 = CbValidateError::None;
				CbObject		AttachmentObject = ValidateAndReadCompactBinaryObject(std::move(AttachmentBuffer), ValidationError);
				if (ValidationError != CbValidateError::None)
				{
					MalformedAttachments.push_back(std::make_pair(
						i,
						fmt::format("Invalid format, CbObject for {}. Reason '{}'", Entry.AttachmentHash, ToString(ValidationError))));
				}

				if (i == 0)
				{
					Package.SetObject(AttachmentObject);
				}
				else
				{
					Attachments.emplace_back(CbAttachment(AttachmentObject, Entry.AttachmentHash));
				}
			}
			else if (AttachmentSize > 0)
			{
				// Make a copy of the buffer so the attachments don't reference the entire payload
				IoBuffer AttachmentBufferCopy = CreateBuffer(Entry.AttachmentHash, AttachmentSize);
				ZEN_ASSERT(AttachmentBufferCopy);
				ZEN_ASSERT(AttachmentBufferCopy.Size() == AttachmentSize);
				AttachmentBufferCopy.GetMutableView().CopyFrom(AttachmentBuffer.GetView());

				Attachments.emplace_back(CbAttachment(SharedBuffer{AttachmentBufferCopy}, Entry.AttachmentHash));
			}
			else
			{
				MalformedAttachments.push_back(
					std::make_pair(i, fmt::format("Invalid format, attachment of size zero detected for {}", Entry.AttachmentHash)));
			}
		}
	}
	PartialFileBuffers.clear();

	Package.AddAttachments(Attachments);

	using namespace std::literals;

	if (!MalformedAttachments.empty())
	{
		ExtendableStringBuilder<1024> SB;
		SB << (uint64_t)MalformedAttachments.size() << " malformed attachments in package message:\n";
		for (const auto& It : MalformedAttachments)
		{
			SB << "    #"sv << It.first << ": " << It.second << "\n";
		}
		ZEN_WARN("{}", SB.ToView());
		throw std::invalid_argument(SB.ToString());
	}

	return Package;
}

bool
ParsePackageMessageWithLegacyFallback(const IoBuffer& Response, CbPackage& OutPackage)
{
	if (IsPackageMessage(Response))
	{
		OutPackage = ParsePackageMessage(Response);
		return true;
	}
	return OutPackage.TryLoad(Response);
}

CbPackageReader::CbPackageReader() : m_CreateBuffer([](const IoHash&, uint64_t Size) -> IoBuffer { return IoBuffer{Size}; })
{
}

CbPackageReader::~CbPackageReader()
{
}

void
CbPackageReader::SetPayloadBufferCreator(std::function<IoBuffer(const IoHash& Cid, uint64_t Size)> CreateBuffer)
{
	m_CreateBuffer = CreateBuffer;
}

uint64_t
CbPackageReader::ProcessPackageHeaderData(const void* Data, uint64_t DataBytes)
{
	ZEN_ASSERT(m_CurrentState != State::kReadingBuffers);

	switch (m_CurrentState)
	{
		case State::kInitialState:
			ZEN_ASSERT(Data == nullptr);
			m_CurrentState = State::kReadingHeader;
			return sizeof m_PackageHeader;

		case State::kReadingHeader:
			ZEN_ASSERT(DataBytes == sizeof m_PackageHeader);
			memcpy(&m_PackageHeader, Data, sizeof m_PackageHeader);
			ZEN_ASSERT(m_PackageHeader.HeaderMagic == kCbPkgMagic);
			m_CurrentState = State::kReadingAttachmentEntries;
			m_AttachmentEntries.resize(m_PackageHeader.AttachmentCount + 1);
			return (m_PackageHeader.AttachmentCount + 1) * sizeof(CbAttachmentEntry);

		case State::kReadingAttachmentEntries:
			ZEN_ASSERT(DataBytes == ((m_PackageHeader.AttachmentCount + 1) * sizeof(CbAttachmentEntry)));
			memcpy(m_AttachmentEntries.data(), Data, DataBytes);

			for (CbAttachmentEntry& Entry : m_AttachmentEntries)
			{
				// This preallocates memory for payloads but note that for the local references
				// the caller will need to handle the payload differently (i.e it's a
				// CbAttachmentReferenceHeader not the actual payload)

				m_PayloadBuffers.emplace_back(IoBuffer{Entry.PayloadSize});
			}

			m_CurrentState = State::kReadingBuffers;
			return 0;

		default:
			ZEN_ASSERT(false);
			return 0;
	}
}

IoBuffer
CbPackageReader::MarshalLocalChunkReference(IoBuffer AttachmentBuffer)
{
	// Marshal local reference - a "pointer" to the chunk backing file

	ZEN_ASSERT(AttachmentBuffer.Size() >= sizeof(CbAttachmentReferenceHeader));

	const CbAttachmentReferenceHeader* AttachRefHdr = AttachmentBuffer.Data<CbAttachmentReferenceHeader>();
	const char8_t*					   PathPointer	= reinterpret_cast<const char8_t*>(AttachRefHdr + 1);

	ZEN_ASSERT(AttachmentBuffer.Size() >= (sizeof(CbAttachmentReferenceHeader) + AttachRefHdr->AbsolutePathLength));

	std::u8string_view PathView{PathPointer, AttachRefHdr->AbsolutePathLength};

	std::filesystem::path Path{PathView};

	IoBuffer ChunkReference = IoBufferBuilder::MakeFromFile(Path, AttachRefHdr->PayloadByteOffset, AttachRefHdr->PayloadByteSize);

	if (!ChunkReference)
	{
		// Unable to open chunk reference

		throw std::runtime_error(fmt::format("unable to resolve local reference to '{}' (offset {}, size {})",
											 PathToUtf8(Path),
											 AttachRefHdr->PayloadByteOffset,
											 AttachRefHdr->PayloadByteSize));
	}

	return ChunkReference;
};

void
CbPackageReader::Finalize()
{
	if (m_AttachmentEntries.empty())
	{
		return;
	}

	m_Attachments.reserve(m_AttachmentEntries.size() - 1);

	int CurrentAttachmentIndex = 0;
	for (CbAttachmentEntry& Entry : m_AttachmentEntries)
	{
		IoBuffer AttachmentBuffer = m_PayloadBuffers[CurrentAttachmentIndex];

		if (CurrentAttachmentIndex == 0)
		{
			// Root object
			if (Entry.Flags & CbAttachmentEntry::kIsObject)
			{
				if (Entry.Flags & CbAttachmentEntry::kIsLocalRef)
				{
					CbValidateError ValidateError = CbValidateError::None;
					m_RootObject = ValidateAndReadCompactBinaryObject(MarshalLocalChunkReference(AttachmentBuffer), ValidateError);
					if (ValidateError != CbValidateError::None)
					{
						throw std::runtime_error(fmt::format("Root object format is invalid, reason: '{}'", ToString(ValidateError)));
					}
				}
				else if (Entry.Flags & CbAttachmentEntry::kIsCompressed)
				{
					IoHash			 RawHash;
					uint64_t		 RawSize;
					CompressedBuffer Compressed = CompressedBuffer::FromCompressed(SharedBuffer(AttachmentBuffer), RawHash, RawSize);
					if (RawHash == Entry.AttachmentHash)
					{
						CbValidateError ValidateError = CbValidateError::None;
						m_RootObject				  = ValidateAndReadCompactBinaryObject(std::move(Compressed), ValidateError);
						if (ValidateError != CbValidateError::None)
						{
							throw std::runtime_error(fmt::format("Root object format is invalid, reason: '{}'", ToString(ValidateError)));
						}
					}
				}
				else
				{
					CbValidateError ValidateError = CbValidateError::None;
					m_RootObject				  = ValidateAndReadCompactBinaryObject(std::move(AttachmentBuffer), ValidateError);
					if (ValidateError != CbValidateError::None)
					{
						throw std::runtime_error(fmt::format("Root object format is invalid, reason: '{}'", ToString(ValidateError)));
					}
				}
			}
			else
			{
				throw std::runtime_error("missing or invalid root object");
			}
		}
		else if (Entry.Flags & CbAttachmentEntry::kIsLocalRef)
		{
			IoBuffer ChunkReference = MarshalLocalChunkReference(AttachmentBuffer);

			if (Entry.Flags & CbAttachmentEntry::kIsCompressed)
			{
				IoHash			 RawHash;
				uint64_t		 RawSize;
				CompressedBuffer Compressed = CompressedBuffer::FromCompressed(SharedBuffer(ChunkReference), RawHash, RawSize);
				if (RawHash == Entry.AttachmentHash)
				{
					m_Attachments.emplace_back(CbAttachment(Compressed, Entry.AttachmentHash));
				}
			}
			else
			{
				CompressedBuffer Compressed =
					CompressedBuffer::Compress(SharedBuffer(ChunkReference), OodleCompressor::NotSet, OodleCompressionLevel::None);
				m_Attachments.emplace_back(CbAttachment(std::move(Compressed), Compressed.DecodeRawHash()));
			}
		}

		++CurrentAttachmentIndex;
	}
}

/**
	______________________ _____________________________
	\__    ___/\_   _____//   _____/\__    ___/   _____/
	  |    |    |    __)_ \_____  \   |    |  \_____  \
	  |    |    |        \/        \  |    |  /        \
	  |____|   /_______  /_______  /  |____| /_______  /
					   \/        \/                  \/
 */

#if ZEN_WITH_TESTS

TEST_CASE("CbPackage.Serialization")
{
	// Make a test package

	CbAttachment Attach1{SharedBuffer::MakeView(MakeMemoryView("abcd"))};
	CbAttachment Attach2{SharedBuffer::MakeView(MakeMemoryView("efgh"))};

	CbObjectWriter Cbo;
	Cbo.AddAttachment("abcd", Attach1);
	Cbo.AddAttachment("efgh", Attach2);

	CbPackage Pkg;
	Pkg.AddAttachment(Attach1);
	Pkg.AddAttachment(Attach2);
	Pkg.SetObject(Cbo.Save());

	SharedBuffer   Buffer		  = FormatPackageMessageBuffer(Pkg).Flatten();
	const uint8_t* CursorPtr	  = reinterpret_cast<const uint8_t*>(Buffer.GetData());
	uint64_t	   RemainingBytes = Buffer.GetSize();

	auto ConsumeBytes = [&](uint64_t ByteCount) {
		ZEN_ASSERT(ByteCount <= RemainingBytes);
		void* ReturnPtr = (void*)CursorPtr;
		CursorPtr += ByteCount;
		RemainingBytes -= ByteCount;
		return ReturnPtr;
	};

	auto CopyBytes = [&](void* TargetBuffer, uint64_t ByteCount) {
		ZEN_ASSERT(ByteCount <= RemainingBytes);
		memcpy(TargetBuffer, CursorPtr, ByteCount);
		CursorPtr += ByteCount;
		RemainingBytes -= ByteCount;
	};

	CbPackageReader Reader;
	uint64_t		InitialRead = Reader.ProcessPackageHeaderData(nullptr, 0);
	uint64_t		NextBytes	= Reader.ProcessPackageHeaderData(ConsumeBytes(InitialRead), InitialRead);
	NextBytes					= Reader.ProcessPackageHeaderData(ConsumeBytes(NextBytes), NextBytes);
	auto Buffers				= Reader.GetPayloadBuffers();

	for (auto& PayloadBuffer : Buffers)
	{
		CopyBytes(PayloadBuffer.MutableData(), PayloadBuffer.GetSize());
	}

	Reader.Finalize();
}

TEST_CASE("CbPackage.EmptyObject")
{
	CbPackage Pkg;
	Pkg.SetObject({});
	std::vector<IoBuffer> Result = FormatPackageMessage(Pkg, nullptr);
}

TEST_CASE("CbPackage.LocalRef")
{
	ScopedTemporaryDirectory TempDir;

	auto Path1 = TempDir.Path() / "abcd";
	auto Path2 = TempDir.Path() / "efgh";

	{
		IoBuffer Buffer1 = IoBufferBuilder::MakeCloneFromMemory(MakeMemoryView("abcd"));
		IoBuffer Buffer2 = IoBufferBuilder::MakeCloneFromMemory(MakeMemoryView("efgh"));

		WriteFile(Path1, Buffer1);
		WriteFile(Path2, Buffer2);
	}

	// Make a test package

	IoBuffer FileBuffer1 = IoBufferBuilder::MakeFromFile(Path1);
	IoBuffer FileBuffer2 = IoBufferBuilder::MakeFromFile(Path2);

	CbAttachment Attach1{SharedBuffer(FileBuffer1)};
	CbAttachment Attach2{SharedBuffer(FileBuffer2)};

	CbObjectWriter Cbo;
	Cbo.AddAttachment("abcd", Attach1);
	Cbo.AddAttachment("efgh", Attach2);

	CbPackage Pkg;
	Pkg.AddAttachment(Attach1);
	Pkg.AddAttachment(Attach2);
	Pkg.SetObject(Cbo.Save());

	SharedBuffer   Buffer		  = FormatPackageMessageBuffer(Pkg, FormatFlags::kAllowLocalReferences).Flatten();
	const uint8_t* CursorPtr	  = reinterpret_cast<const uint8_t*>(Buffer.GetData());
	uint64_t	   RemainingBytes = Buffer.GetSize();

	auto ConsumeBytes = [&](uint64_t ByteCount) {
		ZEN_ASSERT(ByteCount <= RemainingBytes);
		void* ReturnPtr = (void*)CursorPtr;
		CursorPtr += ByteCount;
		RemainingBytes -= ByteCount;
		return ReturnPtr;
	};

	auto CopyBytes = [&](void* TargetBuffer, uint64_t ByteCount) {
		ZEN_ASSERT(ByteCount <= RemainingBytes);
		memcpy(TargetBuffer, CursorPtr, ByteCount);
		CursorPtr += ByteCount;
		RemainingBytes -= ByteCount;
	};

	CbPackageReader Reader;
	uint64_t		InitialRead = Reader.ProcessPackageHeaderData(nullptr, 0);
	uint64_t		NextBytes	= Reader.ProcessPackageHeaderData(ConsumeBytes(InitialRead), InitialRead);
	NextBytes					= Reader.ProcessPackageHeaderData(ConsumeBytes(NextBytes), NextBytes);
	auto Buffers				= Reader.GetPayloadBuffers();

	for (auto& PayloadBuffer : Buffers)
	{
		CopyBytes(PayloadBuffer.MutableData(), PayloadBuffer.GetSize());
	}

	Reader.Finalize();
}

void
forcelink_packageformat()
{
}

#endif

}  // namespace zen