From 06708ea1c044215eee37a2703b6764cf2e89d53b Mon Sep 17 00:00:00 2001
From: Dan Engelbrecht <dan.engelbrecht@epicgames.com>
Date: Sat, 13 Apr 2024 13:53:24 +0200
Subject: Validate input buffer size when trying to parse package message (#47)

* add validation of input buffer size when trying to parse package message
* avoid doing memcopy when parsing package message
---
 src/zenutil/packageformat.cpp | 46 +++++++++++++++++++++++--------------------
 1 file changed, 25 insertions(+), 21 deletions(-)

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

diff --git a/src/zenutil/packageformat.cpp b/src/zenutil/packageformat.cpp
index 2e0f2dc7c..3fa602a96 100644
--- a/src/zenutil/packageformat.cpp
+++ b/src/zenutil/packageformat.cpp
@@ -327,17 +327,14 @@ FormatPackageMessage(const CbPackage& Data, FormatFlags Flags, void* TargetProce
 bool
 IsPackageMessage(IoBuffer Payload)
 {
-	if (!Payload)
+	if (Payload.GetSize() < sizeof(CbPackageHeader))
 	{
 		return false;
 	}
 
-	BinaryReader Reader(Payload);
-
-	CbPackageHeader Hdr;
-	Reader.Read(&Hdr, sizeof Hdr);
-
-	if (Hdr.HeaderMagic != kCbPkgMagic)
+	BinaryReader		   Reader(Payload);
+	const CbPackageHeader* Hdr = reinterpret_cast<const CbPackageHeader*>(Reader.GetView(sizeof(CbPackageHeader)).GetData());
+	if (Hdr->HeaderMagic != kCbPkgMagic)
 	{
 		return false;
 	}
@@ -350,31 +347,32 @@ ParsePackageMessage(IoBuffer Payload, std::function<IoBuffer(const IoHash&, uint
 {
 	ZEN_TRACE_CPU("ParsePackageMessage");
 
-	if (!Payload)
+	if (Payload.GetSize() < sizeof(CbPackageHeader))
 	{
-		return {};
+		throw std::invalid_argument(fmt::format("invalid CbPackage, missing complete header (size {})", Payload.GetSize()));
 	}
 
 	BinaryReader Reader(Payload);
 
-	if (Payload.GetSize() < sizeof(CbPackageHeader))
+	const CbPackageHeader* Hdr = reinterpret_cast<const CbPackageHeader*>(Reader.GetView(sizeof(CbPackageHeader)).GetData());
+	if (Hdr->HeaderMagic != kCbPkgMagic)
 	{
-		throw std::invalid_argument(fmt::format("invalid CbPackage, missing complete header (size {})", Payload.GetSize()));
+		throw std::invalid_argument(
+			fmt::format("invalid CbPackage header magic, expected {0:x}, got {0:x}", static_cast<uint32_t>(kCbPkgMagic), Hdr->HeaderMagic));
 	}
+	Reader.Skip(sizeof(CbPackageHeader));
 
-	CbPackageHeader Hdr;
-	Reader.Read(&Hdr, sizeof Hdr);
+	const uint32_t ChunkCount = Hdr->AttachmentCount + 1;
 
-	if (Hdr.HeaderMagic != kCbPkgMagic)
+	if (Reader.Remaining() < sizeof(CbAttachmentEntry) * ChunkCount)
 	{
-		throw std::invalid_argument("invalid CbPackage header magic");
+		throw std::invalid_argument(fmt::format("invalid CbPackage, missing attachment entry data (need {} bytes, have {} bytes)",
+												sizeof(CbAttachmentEntry) * ChunkCount,
+												Reader.Remaining()));
 	}
-
-	const uint32_t ChunkCount = Hdr.AttachmentCount + 1;
-
-	std::unique_ptr<CbAttachmentEntry[]> AttachmentEntries{new CbAttachmentEntry[ChunkCount]};
-
-	Reader.Read(AttachmentEntries.get(), sizeof(CbAttachmentEntry) * ChunkCount);
+	const CbAttachmentEntry* AttachmentEntries =
+		reinterpret_cast<const CbAttachmentEntry*>(Reader.GetView(sizeof(CbAttachmentEntry) * ChunkCount).GetData());
+	Reader.Skip(sizeof(CbAttachmentEntry) * ChunkCount);
 
 	CbPackage Package;
 
@@ -390,6 +388,12 @@ ParsePackageMessage(IoBuffer Payload, std::function<IoBuffer(const IoHash&, uint
 		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);
 
-- 
cgit v1.2.3