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
---
 src/zenutil/packageformat.cpp | 101 +++++++++++++++++++++++++++++-------------
 1 file changed, 70 insertions(+), 31 deletions(-)

(limited to 'src/zenutil/packageformat.cpp')

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