1 files changed, 37 insertions, 14 deletions

diff --git a/src/zenutil/chunkblock.cpp b/src/zenutil/chunkblock.cpp
index a19cf5c1b..f3c14edc4 100644
--- a/src/zenutil/chunkblock.cpp
+++ b/src/zenutil/chunkblock.cpp
@@ -187,31 +187,54 @@ GenerateChunkBlock(std::vector<std::pair<IoHash, FetchChunkFunc>>&& FetchChunks,
 	return CompressedBlock;
 }
 
-bool
-IterateChunkBlock(const SharedBuffer&														  BlockPayload,
-				  std::function<void(CompressedBuffer&& Chunk, const IoHash& AttachmentHash)> Visitor,
-				  uint64_t&																	  OutHeaderSize)
+std::vector<uint32_t>
+ReadChunkBlockHeader(const MemoryView BlockView, uint64_t& OutHeaderSize)
 {
-	ZEN_ASSERT(BlockPayload);
-	if (BlockPayload.GetSize() < 1)
-	{
-		return false;
-	}
-
-	MemoryView	   BlockView = BlockPayload.GetView();
-	const uint8_t* ReadPtr	 = reinterpret_cast<const uint8_t*>(BlockView.GetData());
+	const uint8_t* ReadPtr = reinterpret_cast<const uint8_t*>(BlockView.GetData());
 	uint32_t	   NumberSize;
 	uint64_t	   ChunkCount = ReadVarUInt(ReadPtr, NumberSize);
 	ReadPtr += NumberSize;
-	std::vector<uint64_t> ChunkSizes;
+	std::vector<uint32_t> ChunkSizes;
 	ChunkSizes.reserve(ChunkCount);
 	while (ChunkCount--)
 	{
-		ChunkSizes.push_back(ReadVarUInt(ReadPtr, NumberSize));
+		if (ReadPtr >= BlockView.GetDataEnd())
+		{
+			throw std::runtime_error("Invalid block header, block data ended unexpectedly");
+		}
+		uint64_t ChunkSize = ReadVarUInt(ReadPtr, NumberSize);
+		if (ChunkSize > std::numeric_limits<uint32_t>::max())
+		{
+			throw std::runtime_error("Invalid block header, header data is corrupt");
+		}
+		if (ChunkSize < 1)
+		{
+			throw std::runtime_error("Invalid block header, header data is corrupt");
+		}
+		ChunkSizes.push_back(gsl::narrow<uint32_t>(ChunkSize));
 		ReadPtr += NumberSize;
 	}
 	uint64_t Offset = std::distance((const uint8_t*)BlockView.GetData(), ReadPtr);
 	OutHeaderSize	= Offset;
+	return ChunkSizes;
+}
+
+bool
+IterateChunkBlock(const SharedBuffer&														  BlockPayload,
+				  std::function<void(CompressedBuffer&& Chunk, const IoHash& AttachmentHash)> Visitor,
+				  uint64_t&																	  OutHeaderSize)
+{
+	ZEN_ASSERT(BlockPayload);
+	if (BlockPayload.GetSize() < 1)
+	{
+		return false;
+	}
+
+	MemoryView BlockView = BlockPayload.GetView();
+
+	std::vector<uint32_t> ChunkSizes = ReadChunkBlockHeader(BlockView, OutHeaderSize);
+	uint64_t			  Offset	 = OutHeaderSize;
+	OutHeaderSize					 = Offset;
 	for (uint64_t ChunkSize : ChunkSizes)
 	{
 		IoBuffer		 Chunk(BlockPayload.AsIoBuffer(), Offset, ChunkSize);