From 7d2421dd0aaebf75a2cbe5d6a0827f1d064c0c3f Mon Sep 17 00:00:00 2001
From: Dan Engelbrecht <dan.engelbrecht@epicgames.com>
Date: Thu, 4 Apr 2024 13:37:17 +0200
Subject: hardening parsepackagemessage (#38)

* hardening of ParsePackageMessage and extended details when malformed attachments are detected
---
 CHANGELOG.md                          |   1 +
 src/zenserver-test/zenserver-test.cpp |  10 ++--
 src/zenutil/packageformat.cpp         | 101 +++++++++++++++++++++++-----------
 3 files changed, 75 insertions(+), 37 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b3ad3d842..f1e31e680 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -15,6 +15,7 @@
 - Improvement: Optimize `CompositeBuffer::ViewOrCopyRange` speeding up compressed buffer headers by not materializing entire source buffer
 - Improvement: Optimize `CompressedBuffer::GetRange()` with new `CompressedBuffer::ReadHeader()` that does one less read from source data resulting in a 30% perf increase.
 - Improvement: Validate lengths of chunks fetched from CAS/Cache store, if full chunk can not be retrieved, treat it as missing
+- Improvement: Hardening of ParsePackageMessage and added extended details of all malformed attachments detected
 
 ## 5.4.2
 - Bugfix: Shared memory for zenserver state may hang around after all zenserver processes exit - make sure we find a valid entry in `zen up` before bailing
diff --git a/src/zenserver-test/zenserver-test.cpp b/src/zenserver-test/zenserver-test.cpp
index 412a2d26a..325b15e3f 100644
--- a/src/zenserver-test/zenserver-test.cpp
+++ b/src/zenserver-test/zenserver-test.cpp
@@ -1228,12 +1228,10 @@ TEST_CASE("zcache.rpc")
 		if (Result.status_code == 200)
 		{
 			CbPackage Response = ParsePackageMessage(zen::IoBuffer(zen::IoBuffer::Wrap, Result.text.data(), Result.text.size()));
-			if (!Response.IsNull())
-			{
-				OutResult.Response = std::move(Response);
-				CHECK(OutResult.Result.Parse(OutResult.Response));
-				OutResult.Success = true;
-			}
+			CHECK(!Response.IsNull());
+			OutResult.Response = std::move(Response);
+			CHECK(OutResult.Result.Parse(OutResult.Response));
+			OutResult.Success = true;
 		}
 
 		return OutResult;
diff --git a/src/zenutil/packageformat.cpp b/src/zenutil/packageformat.cpp
index 7c284a4e6..2e0f2dc7c 100644
--- a/src/zenutil/packageformat.cpp
+++ b/src/zenutil/packageformat.cpp
@@ -357,6 +357,11 @@ ParsePackageMessage(IoBuffer Payload, std::function<IoBuffer(const IoHash&, uint
 
 	BinaryReader Reader(Payload);
 
+	if (Payload.GetSize() < sizeof(CbPackageHeader))
+	{
+		throw std::invalid_argument(fmt::format("invalid CbPackage, missing complete header (size {})", Payload.GetSize()));
+	}
+
 	CbPackageHeader Hdr;
 	Reader.Read(&Hdr, sizeof Hdr);
 
@@ -378,8 +383,8 @@ ParsePackageMessage(IoBuffer Payload, std::function<IoBuffer(const IoHash&, uint
 
 	tsl::robin_map<std::string, IoBuffer> PartialFileBuffers;
 
-	// TODO: Throwing before this loop completes could result in leaking handles as we might not have picked up all the handles in the
-	// message
+	std::vector<std::pair<uint32_t, std::string>> MalformedAttachments;
+
 	for (uint32_t i = 0; i < ChunkCount; ++i)
 	{
 		const CbAttachmentEntry& Entry			= AttachmentEntries[i];
@@ -438,30 +443,34 @@ ParsePackageMessage(IoBuffer Payload, std::function<IoBuffer(const IoHash&, uint
 				}
 			}
 
-			if (!FullFileBuffer)
+			if (FullFileBuffer)
 			{
-				// Unable to open chunk reference
-				throw std::runtime_error(fmt::format("unable to resolve chunk #{} at '{}' (offset {}, size {})",
-													 i,
-													 Path,
-													 AttachRefHdr->PayloadByteOffset,
-													 AttachRefHdr->PayloadByteSize));
-			}
-
-			IoBuffer ChunkReference = AttachRefHdr->PayloadByteOffset == 0 && AttachRefHdr->PayloadByteSize == FullFileBuffer.GetSize()
-										  ? FullFileBuffer
-										  : IoBuffer(FullFileBuffer, AttachRefHdr->PayloadByteOffset, AttachRefHdr->PayloadByteSize);
+				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)
+				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 {})",
+																			  Path,
+																			  AttachRefHdr->PayloadByteOffset,
+																			  AttachRefHdr->PayloadByteSize)));
+				}
+			}
+			else
 			{
-				throw std::invalid_argument(fmt::format("invalid format for chunk #{} at '{}' (offset {}, size {})",
-														i,
-														Path,
-														AttachRefHdr->PayloadByteOffset,
-														AttachRefHdr->PayloadByteSize));
+				MalformedAttachments.push_back(std::make_pair(i,
+															  fmt::format("Unable to resolve chunk at '{}' (offset {}, size {})",
+																		  Path,
+																		  AttachRefHdr->PayloadByteOffset,
+																		  AttachRefHdr->PayloadByteSize)));
 			}
-			Attachments.emplace_back(CbAttachment(std::move(CompBuf), Entry.AttachmentHash));
 		}
 		else if (Entry.Flags & CbAttachmentEntry::kIsCompressed)
 		{
@@ -470,26 +479,39 @@ ParsePackageMessage(IoBuffer Payload, std::function<IoBuffer(const IoHash&, uint
 				if (i == 0)
 				{
 					CompressedBuffer CompBuf(CompressedBuffer::FromCompressedNoValidate(IoBuffer(AttachmentBuffer)));
-					if (!CompBuf)
+					if (CompBuf)
+					{
+						Package.SetObject(LoadCompactBinaryObject(std::move(CompBuf)));
+					}
+					else
 					{
-						throw std::invalid_argument(fmt::format("invalid format for chunk #{} expected compressed buffer for CbObject", i));
+						// 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 {})", AttachmentBuffer.GetSize())));
 					}
-					// First payload is always a compact binary object
-					Package.SetObject(LoadCompactBinaryObject(std::move(CompBuf)));
 				}
 				else
 				{
-					ZEN_NOT_IMPLEMENTED("Object attachments are not currently supported");
+					MalformedAttachments.push_back(
+						std::make_pair(i,
+									   fmt::format("Invalid format, compressed object attachments are not currently supported (size {})",
+												   AttachmentBuffer.GetSize())));
 				}
 			}
 			else
 			{
 				CompressedBuffer CompBuf(CompressedBuffer::FromCompressedNoValidate(IoBuffer(AttachmentBuffer)));
-				if (!CompBuf)
+				if (CompBuf)
+				{
+					Attachments.emplace_back(CbAttachment(std::move(CompBuf), Entry.AttachmentHash));
+				}
+				else
 				{
-					throw std::invalid_argument(fmt::format("invalid format for chunk #{} expected compressed buffer for attachment", i));
+					MalformedAttachments.push_back(std::make_pair(
+						i,
+						fmt::format("Invalid format, expected compressed buffer for attachment (size {})", AttachmentBuffer.GetSize())));
 				}
-				Attachments.emplace_back(CbAttachment(std::move(CompBuf), Entry.AttachmentHash));
 			}
 		}
 		else /* not compressed */
@@ -502,7 +524,10 @@ ParsePackageMessage(IoBuffer Payload, std::function<IoBuffer(const IoHash&, uint
 				}
 				else
 				{
-					ZEN_NOT_IMPLEMENTED("Object attachments are not currently supported");
+					MalformedAttachments.push_back(
+						std::make_pair(i,
+									   fmt::format("Invalid format, object attachments are not currently supported (size {})",
+												   AttachmentBuffer.GetSize())));
 				}
 			}
 			else
@@ -521,6 +546,20 @@ ParsePackageMessage(IoBuffer Payload, std::function<IoBuffer(const IoHash&, uint
 
 	Package.AddAttachments(Attachments);
 
+	using namespace std::literals;
+
+	if (!MalformedAttachments.empty())
+	{
+		StringBuilder<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;
 }
 
-- 
cgit v1.2.3