validate rpc chunk responses (#36)

* Validate size of found chunks in cas/cache

5 files changed, 35 insertions, 20 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6f93f7d40..b3ad3d842 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,6 +14,7 @@
 - Improvement: Increase the retry count to 4 (5 attempts in total) when talking to Jupiter for oplog export/import
 - Improvement: Optimize `CompositeBuffer::ViewOrCopyRange` speeding up compressed buffer headers by not materializing entire source buffer
 - Improvement: Optimize `CompressedBuffer::GetRange()` with new `CompressedBuffer::ReadHeader()` that does one less read from source data resulting in a 30% perf increase.
+- Improvement: Validate lengths of chunks fetched from CAS/Cache store, if full chunk can not be retrieved, treat it as missing
 
 ## 5.4.2
 - Bugfix: Shared memory for zenserver state may hang around after all zenserver processes exit - make sure we find a valid entry in `zen up` before bailing
diff --git a/src/zenstore/blockstore.cpp b/src/zenstore/blockstore.cpp
index 69487f9dc..dfa852434 100644
--- a/src/zenstore/blockstore.cpp
+++ b/src/zenstore/blockstore.cpp
@@ -482,7 +482,11 @@ BlockStore::TryGetChunk(const BlockStoreLocation& Location) const
 	{
 		if (const Ref<BlockStoreFile>& Block = BlockIt->second; Block)
 		{
-			return Block->GetChunk(Location.Offset, Location.Size);
+			IoBuffer Chunk = Block->GetChunk(Location.Offset, Location.Size);
+			if (Chunk.GetSize() == Location.Size)
+			{
+				return Chunk;
+			}
 		}
 	}
 	return IoBuffer();
diff --git a/src/zenstore/cache/cachedisklayer.cpp b/src/zenstore/cache/cachedisklayer.cpp
index d897e26ce..3605f5582 100644
--- a/src/zenstore/cache/cachedisklayer.cpp
+++ b/src/zenstore/cache/cachedisklayer.cpp
@@ -1141,7 +1141,7 @@ ZenCacheDiskLayer::CacheBucket::GetInlineCacheValue(const DiskLocation& Loc) con
 }
 
 IoBuffer
-ZenCacheDiskLayer::CacheBucket::GetStandaloneCacheValue(ZenContentType ContentType, const IoHash& HashKey) const
+ZenCacheDiskLayer::CacheBucket::GetStandaloneCacheValue(const DiskLocation& Loc, const IoHash& HashKey) const
 {
 	ZEN_TRACE_CPU("Z$::Bucket::GetStandaloneCacheValue");
 
@@ -1152,9 +1152,11 @@ ZenCacheDiskLayer::CacheBucket::GetStandaloneCacheValue(ZenContentType ContentTy
 
 	if (IoBuffer Data = IoBufferBuilder::MakeFromFile(DataFilePath.ToPath()))
 	{
-		Data.SetContentType(ContentType);
-
-		return Data;
+		if (Data.GetSize() == Loc.Size())
+		{
+			Data.SetContentType(Loc.GetContentType());
+			return Data;
+		}
 	}
 
 	return {};
@@ -1211,7 +1213,7 @@ ZenCacheDiskLayer::CacheBucket::Get(const IoHash& HashKey, ZenCacheValue& OutVal
 		}
 		if (Location.IsFlagSet(DiskLocation::kStandaloneFile))
 		{
-			OutValue.Value = GetStandaloneCacheValue(Location.GetContentType(), HashKey);
+			OutValue.Value = GetStandaloneCacheValue(Location, HashKey);
 		}
 		else
 		{
@@ -1621,7 +1623,7 @@ ZenCacheDiskLayer::CacheBucket::ScrubStorage(ScrubContext& Ctx)
 				else
 				{
 					// Structured cache value
-					IoBuffer Buffer = GetStandaloneCacheValue(Loc.GetContentType(), HashKey);
+					IoBuffer Buffer = GetStandaloneCacheValue(Loc, HashKey);
 					if (!Buffer)
 					{
 						ReportBadKey(HashKey);
@@ -1880,7 +1882,7 @@ ZenCacheDiskLayer::CacheBucket::GatherReferences(GcContext& GcCtx)
 			IoBuffer Buffer;
 			if (Loc.IsFlagSet(DiskLocation::kStandaloneFile))
 			{
-				if (Buffer = GetStandaloneCacheValue(Loc.GetContentType(), Key); !Buffer)
+				if (Buffer = GetStandaloneCacheValue(Loc, Key); !Buffer)
 				{
 					continue;
 				}
@@ -2269,9 +2271,8 @@ ZenCacheDiskLayer::CacheBucket::GetValueDetails(RwLock::SharedLockScope& IndexLo
 	const BucketPayload& Payload = m_Payloads[Index];
 	if (Payload.Location.IsFlagSet(DiskLocation::kStructured))
 	{
-		IoBuffer	 Value = Payload.Location.IsFlagSet(DiskLocation::kStandaloneFile)
-								 ? GetStandaloneCacheValue(Payload.Location.GetContentType(), Key)
-								 : GetInlineCacheValue(Payload.Location);
+		IoBuffer	 Value = Payload.Location.IsFlagSet(DiskLocation::kStandaloneFile) ? GetStandaloneCacheValue(Payload.Location, Key)
+																					   : GetInlineCacheValue(Payload.Location);
 		CbObjectView Obj(Value.GetData());
 		Obj.IterateAttachments([&Attachments](CbFieldView Field) { Attachments.emplace_back(Field.AsAttachment()); });
 	}
@@ -3066,7 +3067,7 @@ public:
 		{
 			m_CacheBucket.m_IndexLock.WithExclusiveLock([&]() { m_CacheBucket.m_TrackedReferences = std::make_unique<HashSet>(); });
 
-			std::vector<IoHash> StandaloneKeys;
+			std::vector<std::pair<IoHash, DiskLocation>> StandaloneKeys;
 			{
 				std::vector<IoHash>						  InlineKeys;
 				std::unordered_map<uint32_t, std::size_t> BlockIndexToEntriesPerBlockIndex;
@@ -3099,7 +3100,7 @@ public:
 						const IoHash& Key = Entry.first;
 						if (Loc.IsFlagSet(DiskLocation::kStandaloneFile))
 						{
-							StandaloneKeys.push_back(Key);
+							StandaloneKeys.push_back(std::make_pair(Key, Loc));
 							continue;
 						}
 
@@ -3157,7 +3158,7 @@ public:
 					}
 				}
 			}
-			for (const IoHash& Key : StandaloneKeys)
+			for (const auto& It : StandaloneKeys)
 			{
 				if (Ctx.IsCancelledFlag.load())
 				{
@@ -3165,7 +3166,7 @@ public:
 					return;
 				}
 
-				IoBuffer Buffer = m_CacheBucket.GetStandaloneCacheValue(ZenContentType::kCbObject, Key);
+				IoBuffer Buffer = m_CacheBucket.GetStandaloneCacheValue(It.second, It.first);
 				if (!Buffer)
 				{
 					continue;
diff --git a/src/zenstore/filecas.cpp b/src/zenstore/filecas.cpp
index 428183827..6b4185ff4 100644
--- a/src/zenstore/filecas.cpp
+++ b/src/zenstore/filecas.cpp
@@ -729,19 +729,28 @@ FileCasStrategy::FindChunk(const IoHash& ChunkHash)
 
 	ZEN_ASSERT(m_IsInitialized);
 
+	uint64_t ExpectedSize = 0;
 	{
 		RwLock::SharedLockScope _(m_Lock);
-		if (!m_Index.contains(ChunkHash))
+		if (auto It = m_Index.find(ChunkHash); It != m_Index.end())
+		{
+			ExpectedSize = It->second.Size;
+		}
+		else
 		{
 			return {};
 		}
 	}
 
-	ShardingHelper Name(m_RootDirectory.c_str(), ChunkHash);
-
+	ShardingHelper			Name(m_RootDirectory.c_str(), ChunkHash);
 	RwLock::SharedLockScope _(LockForHash(ChunkHash));
 
-	return IoBufferBuilder::MakeFromFile(Name.ShardedPath.c_str());
+	if (IoBuffer Chunk = IoBufferBuilder::MakeFromFile(Name.ShardedPath.c_str()); Chunk.GetSize() == ExpectedSize)
+	{
+		return Chunk;
+	}
+
+	return {};
 }
 
 bool
diff --git a/src/zenstore/include/zenstore/cache/cachedisklayer.h b/src/zenstore/include/zenstore/cache/cachedisklayer.h
index 7aced67ad..471cc5dcd 100644
--- a/src/zenstore/include/zenstore/cache/cachedisklayer.h
+++ b/src/zenstore/include/zenstore/cache/cachedisklayer.h
@@ -326,7 +326,7 @@ public:
 
 		void	 BuildPath(PathBuilderBase& Path, const IoHash& HashKey) const;
 		void	 PutStandaloneCacheValue(const IoHash& HashKey, const ZenCacheValue& Value, std::span<IoHash> References);
-		IoBuffer GetStandaloneCacheValue(ZenContentType ContentType, const IoHash& HashKey) const;
+		IoBuffer GetStandaloneCacheValue(const DiskLocation& Loc, const IoHash& HashKey) const;
 		void	 PutInlineCacheValue(const IoHash& HashKey, const ZenCacheValue& Value, std::span<IoHash> References);
 		IoBuffer GetInlineCacheValue(const DiskLocation& Loc) const;
 		CacheValueDetails::ValueDetails GetValueDetails(RwLock::SharedLockScope&, const IoHash& Key, PayloadIndex Index) const;