From f63b88e96674ea8f99e1008d28224ea21a415f3a Mon Sep 17 00:00:00 2001
From: Dan Engelbrecht <dan.engelbrecht@epicgames.com>
Date: Tue, 6 Aug 2024 14:32:51 +0200
Subject: validate cbobject before iterating for attachments to avoid crash on
 malformed data (#101)

---
 src/zenstore/cache/cachedisklayer.cpp | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

(limited to 'src/zenstore/cache/cachedisklayer.cpp')

diff --git a/src/zenstore/cache/cachedisklayer.cpp b/src/zenstore/cache/cachedisklayer.cpp
index f865e1c3c..159f54422 100644
--- a/src/zenstore/cache/cachedisklayer.cpp
+++ b/src/zenstore/cache/cachedisklayer.cpp
@@ -3386,9 +3386,12 @@ ZenCacheDiskLayer::CacheBucket::RemoveExpiredData(GcCtx& Ctx, GcStats& Stats)
 bool
 ZenCacheDiskLayer::CacheBucket::GetReferencesLocked(GcCtx& Ctx, std::vector<IoHash>& OutReferences)
 {
-	auto GetAttachments = [&](const void* CbObjectData) {
-		CbObjectView Obj(CbObjectData);
-		Obj.IterateAttachments([&](CbFieldView Field) { OutReferences.emplace_back(Field.AsAttachment()); });
+	auto GetAttachments = [&](MemoryView Data) {
+		if (ValidateCompactBinary(Data, CbValidateMode::Default) == CbValidateError::None)
+		{
+			CbObjectView Obj(Data.GetData());
+			Obj.IterateAttachments([&](CbFieldView Field) { OutReferences.emplace_back(Field.AsAttachment()); });
+		}
 	};
 
 	std::vector<std::pair<IoHash, DiskLocation>> StandaloneKeys;
@@ -3446,13 +3449,13 @@ ZenCacheDiskLayer::CacheBucket::GetReferencesLocked(GcCtx& Ctx, std::vector<IoHa
 				InlineLocations,
 				ChunkIndexes,
 				[&](size_t ChunkIndex, const void* Data, uint64_t Size) {
-					ZEN_UNUSED(ChunkIndex, Size);
-					GetAttachments(Data);
+					ZEN_UNUSED(ChunkIndex);
+					GetAttachments(MemoryView(Data, Size));
 					return !Ctx.IsCancelledFlag.load();
 				},
 				[&](size_t ChunkIndex, BlockStoreFile& File, uint64_t Offset, uint64_t Size) {
 					ZEN_UNUSED(ChunkIndex);
-					GetAttachments(File.GetChunk(Offset, Size).GetData());
+					GetAttachments(File.GetChunk(Offset, Size).GetView());
 					return !Ctx.IsCancelledFlag.load();
 				});
 
@@ -3472,7 +3475,7 @@ ZenCacheDiskLayer::CacheBucket::GetReferencesLocked(GcCtx& Ctx, std::vector<IoHa
 		IoBuffer Buffer = GetStandaloneCacheValue(It.second, It.first);
 		if (Buffer)
 		{
-			GetAttachments(Buffer.GetData());
+			GetAttachments(Buffer.GetView());
 		}
 	}
 	return true;
-- 
cgit v1.2.3