From 0bbd1fb43bbd9f878a2aa326ef06f2dc503a3b3f Mon Sep 17 00:00:00 2001
From: zousar <2936246+zousar@users.noreply.github.com>
Date: Wed, 26 Feb 2025 11:49:05 -0700
Subject: Enforce Overwrite Prevention According To Cache Policy

Overwrite with differing value should be denied if QueryLocal is not present and StoreLocal is present.  Overwrite with equal value should succeed regardless of policy flags.
---
 src/zenstore/cache/cachedisklayer.cpp | 93 +++++++++++++++++++++++++++++++++--
 1 file changed, 89 insertions(+), 4 deletions(-)

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

diff --git a/src/zenstore/cache/cachedisklayer.cpp b/src/zenstore/cache/cachedisklayer.cpp
index 25f68330a..eaed1f64e 100644
--- a/src/zenstore/cache/cachedisklayer.cpp
+++ b/src/zenstore/cache/cachedisklayer.cpp
@@ -1793,16 +1793,98 @@ ZenCacheDiskLayer::CacheBucket::Get(const IoHash& HashKey, ZenCacheValue& OutVal
 	}
 }
 
-void
+static bool
+ValueMatchesRawSizeAndHash(const ZenCacheValue& Value, uint64_t RawSize, const std::function<IoHash()>& RawHashProvider)
+{
+	if ((Value.RawSize != 0) || (Value.RawHash != IoHash::Zero))
+	{
+		return ((RawSize == Value.RawSize) && (RawHashProvider() == Value.RawHash));
+	}
+	else if (Value.Value.GetContentType() == ZenContentType::kCompressedBinary)
+	{
+		uint64_t ValueRawSize = 0;
+		IoHash	 ValueRawHash = IoHash::Zero;
+		return CompressedBuffer::ValidateCompressedHeader(Value.Value, ValueRawHash, ValueRawSize) && (RawSize == ValueRawSize) &&
+			   (RawHashProvider() == ValueRawHash);
+	}
+
+	return (RawSize == Value.Value.GetSize()) && (RawHashProvider() == IoHash::HashBuffer(Value.Value));
+}
+
+static bool
+ValueMatchesValue(const ZenCacheValue& Value1, const ZenCacheValue& Value2)
+{
+	if ((Value1.RawSize != 0) || (Value1.RawHash != IoHash::Zero))
+	{
+		return ValueMatchesRawSizeAndHash(Value2, Value1.RawSize, [&Value1]() { return Value1.RawHash; });
+	}
+	else if (Value1.Value.GetContentType() == ZenContentType::kCompressedBinary)
+	{
+		uint64_t Value1RawSize = 0;
+		IoHash	 Value1RawHash = IoHash::Zero;
+		return CompressedBuffer::ValidateCompressedHeader(Value1.Value, Value1RawHash, Value1RawSize) &&
+			   ValueMatchesRawSizeAndHash(Value2, Value1RawSize, [Value1RawHash]() { return Value1RawHash; });
+	}
+
+	return ValueMatchesRawSizeAndHash(Value2, Value1.Value.GetSize(), [&Value1]() { return IoHash::HashBuffer(Value1.Value); });
+}
+
+bool
 ZenCacheDiskLayer::CacheBucket::Put(const IoHash&		 HashKey,
 									const ZenCacheValue& Value,
 									std::span<IoHash>	 References,
+									bool				 Overwrite,
 									PutBatchHandle*		 OptionalBatchHandle)
 {
 	ZEN_TRACE_CPU("Z$::Bucket::Put");
 
 	metrics::RequestStats::Scope $(m_PutOps, Value.Value.Size());
 
+	if (!Overwrite)
+	{
+		RwLock::SharedLockScope IndexLock(m_IndexLock);
+		auto					It = m_Index.find(HashKey);
+		if (It != m_Index.end())
+		{
+			PayloadIndex EntryIndex	  = It.value();
+			m_AccessTimes[EntryIndex] = GcClock::TickCount();
+			DiskLocation Location	  = m_Payloads[EntryIndex].Location;
+
+			bool				 ComparisonComplete = false;
+			const BucketPayload* Payload			= &m_Payloads[EntryIndex];
+			if (Payload->MetaData)
+			{
+				const BucketMetaData& MetaData = m_MetaDatas[Payload->MetaData];
+				if (MetaData)
+				{
+					if (!ValueMatchesRawSizeAndHash(Value, MetaData.RawSize, [&MetaData]() { return MetaData.RawHash; }))
+					{
+						if (OptionalBatchHandle)
+						{
+							OptionalBatchHandle->OutResults.push_back(false);
+						}
+						return false;
+					}
+					ComparisonComplete = true;
+				}
+			}
+
+			if (!ComparisonComplete)
+			{
+				IndexLock.ReleaseNow();
+				ZenCacheValue ExistingValue;
+				if (Get(HashKey, ExistingValue) && !ValueMatchesValue(Value, ExistingValue))
+				{
+					if (OptionalBatchHandle)
+					{
+						OptionalBatchHandle->OutResults.push_back(false);
+					}
+					return false;
+				}
+			}
+		}
+	}
+
 	if (Value.Value.Size() >= m_Configuration.LargeObjectThreshold)
 	{
 		PutStandaloneCacheValue(HashKey, Value, References);
@@ -1817,6 +1899,7 @@ ZenCacheDiskLayer::CacheBucket::Put(const IoHash&		 HashKey,
 	}
 
 	m_DiskWriteCount++;
+	return true;
 }
 
 uint64_t
@@ -3835,21 +3918,23 @@ ZenCacheDiskLayer::Get(std::string_view InBucket, const IoHash& HashKey, GetBatc
 	}
 }
 
-void
+bool
 ZenCacheDiskLayer::Put(std::string_view		InBucket,
 					   const IoHash&		HashKey,
 					   const ZenCacheValue& Value,
 					   std::span<IoHash>	References,
+					   bool					Overwrite,
 					   PutBatchHandle*		OptionalBatchHandle)
 {
 	ZEN_TRACE_CPU("Z$::Put");
-
+	bool RetVal = false;
 	if (CacheBucket* Bucket = GetOrCreateBucket(InBucket); Bucket != nullptr)
 	{
 		CacheBucket::PutBatchHandle* BucketBatchHandle = OptionalBatchHandle == nullptr ? nullptr : OptionalBatchHandle->GetHandle(Bucket);
-		Bucket->Put(HashKey, Value, References, BucketBatchHandle);
+		RetVal										   = Bucket->Put(HashKey, Value, References, Overwrite, BucketBatchHandle);
 		TryMemCacheTrim();
 	}
+	return RetVal;
 }
 
 void
-- 
cgit v1.2.3