From 67068ade9085dedc73a45726cb6fa7aa6cae23c1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Robert=20M=C3=BCller?= <robytemueller@gmail.com>
Date: Wed, 18 Sep 2024 23:58:42 +0200
Subject: [PATCH] Fix snapshot builder creating too large snapshots

The snapshot builder was only considering the size of the data and of the new item being created in the `NewItem` function but not the size of the `CSnapshot` class and the size of the offsets (one `int` for each item including the new one). This could lead to snapshots being too large, which could cause the server to crash when the snapshots were copied into buffers of size `CSnapshot::MAX_SIZE`. However, it should be unlikely for this to happen normally, as the maximum number of snap items (`CSnapshot::MAX_ITEMS`, which is `1024`) is usually reached before the maximum snapshot size.

Also check in the `CSnapshot::IsValid` function that the snapshot is not too large (`CSnapshot::MAX_SIZE`) and check that the number of items is not too large (`CSnapshot::MAX_ITEMS`).
---
 src/engine/shared/snapshot.cpp | 30 ++++++++++++++++++++++--------
 1 file changed, 22 insertions(+), 8 deletions(-)

diff --git a/src/engine/shared/snapshot.cpp b/src/engine/shared/snapshot.cpp
index 76e920d6a..0e2b777f8 100644
--- a/src/engine/shared/snapshot.cpp
+++ b/src/engine/shared/snapshot.cpp
@@ -135,8 +135,15 @@ void CSnapshot::DebugDump() const
 bool CSnapshot::IsValid(size_t ActualSize) const
 {
 	// validate total size
-	if(ActualSize < sizeof(CSnapshot) || m_NumItems < 0 || m_DataSize < 0 || ActualSize != TotalSize())
+	if(ActualSize < sizeof(CSnapshot) ||
+		ActualSize > MAX_SIZE ||
+		m_NumItems < 0 ||
+		m_NumItems > MAX_ITEMS ||
+		m_DataSize < 0 ||
+		ActualSize != TotalSize())
+	{
 		return false;
+	}
 
 	// validate item offsets
 	const int *pOffsets = Offsets();
@@ -744,12 +751,15 @@ int *CSnapshotBuilder::GetItemData(int Key)
 int CSnapshotBuilder::Finish(void *pSnapData)
 {
 	// flatten and make the snapshot
+	dbg_assert(m_NumItems <= CSnapshot::MAX_ITEMS, "Too many snap items");
 	CSnapshot *pSnap = (CSnapshot *)pSnapData;
 	pSnap->m_DataSize = m_DataSize;
 	pSnap->m_NumItems = m_NumItems;
+	const size_t TotalSize = pSnap->TotalSize();
+	dbg_assert(TotalSize <= (size_t)CSnapshot::MAX_SIZE, "Snapshot too large");
 	mem_copy(pSnap->Offsets(), m_aOffsets, pSnap->OffsetSize());
 	mem_copy(pSnap->DataStart(), m_aData, m_DataSize);
-	return pSnap->TotalSize();
+	return TotalSize;
 }
 
 int CSnapshotBuilder::GetTypeFromIndex(int Index) const
@@ -794,11 +804,15 @@ void *CSnapshotBuilder::NewItem(int Type, int Id, int Size)
 		return nullptr;
 	}
 
-	if(m_DataSize + sizeof(CSnapshotItem) + Size >= CSnapshot::MAX_SIZE ||
-		m_NumItems + 1 >= CSnapshot::MAX_ITEMS)
+	if(m_NumItems >= CSnapshot::MAX_ITEMS)
+	{
+		return nullptr;
+	}
+
+	const size_t OffsetSize = (m_NumItems + 1) * sizeof(int);
+	const size_t ItemSize = sizeof(CSnapshotItem) + Size;
+	if(sizeof(CSnapshot) + OffsetSize + m_DataSize + ItemSize > CSnapshot::MAX_SIZE)
 	{
-		dbg_assert(m_DataSize < CSnapshot::MAX_SIZE, "too much data");
-		dbg_assert(m_NumItems < CSnapshot::MAX_ITEMS, "too many items");
 		return nullptr;
 	}
 
@@ -824,11 +838,11 @@ void *CSnapshotBuilder::NewItem(int Type, int Id, int Size)
 	else if(Type < 0)
 		return nullptr;
 
-	mem_zero(pObj, sizeof(CSnapshotItem) + Size);
 	pObj->m_TypeAndId = (Type << 16) | Id;
 	m_aOffsets[m_NumItems] = m_DataSize;
-	m_DataSize += sizeof(CSnapshotItem) + Size;
+	m_DataSize += ItemSize;
 	m_NumItems++;
 
+	mem_zero(pObj->Data(), Size);
 	return pObj->Data();
 }