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

3 files changed, 41 insertions, 22 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f598bee07..82a9fad26 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,4 +1,8 @@
 ##
+- Improvement: add validation of input buffer size when trying to parse package message
+- Improvement: avoid doing memcopy when parsing package message
+
+## 5.4.4
 - Bugfix: Get raw size for compressed chunks correctly for `/prj/{project}/oplog/{log}/chunkinfos`
 - Bugfix: Fix log of Success/Failure for oplog import
 - Bugfix: Use proper API when checking oplog export blob existence in Jupiter
diff --git a/src/zencore/include/zencore/stream.h b/src/zencore/include/zencore/stream.h
index a28290041..fae393666 100644
--- a/src/zencore/include/zencore/stream.h
+++ b/src/zencore/include/zencore/stream.h
@@ -70,14 +70,25 @@ public:
 
 	inline void Read(void* DataPtr, size_t ByteCount)
 	{
+		ZEN_ASSERT(m_Offset + ByteCount <= m_BufferSize);
 		memcpy(DataPtr, m_BufferBase + m_Offset, ByteCount);
 		m_Offset += ByteCount;
 	}
+	inline MemoryView GetView(size_t ByteCount) const
+	{
+		ZEN_ASSERT(m_Offset + ByteCount <= m_BufferSize);
+		return MemoryView((const void*)(m_BufferBase + m_Offset), (const void*)(m_BufferBase + m_Offset + ByteCount));
+	}
 
 	inline uint64_t Size() const { return m_BufferSize; }
 	inline uint64_t GetSize() const { return Size(); }
 	inline uint64_t CurrentOffset() const { return m_Offset; }
-	inline void		Skip(size_t ByteCount) { m_Offset += ByteCount; };
+	inline uint64_t Remaining() const { return m_BufferSize - m_Offset; }
+	inline void		Skip(size_t ByteCount)
+	{
+		ZEN_ASSERT(m_Offset + ByteCount <= m_BufferSize);
+		m_Offset += ByteCount;
+	};
 
 protected:
 	const uint8_t* m_BufferBase;
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);