1 files changed, 73 insertions, 8 deletions

diff --git a/src/zencore/compress.cpp b/src/zencore/compress.cpp
index 6aa0adce0..a0e91f908 100644
--- a/src/zencore/compress.cpp
+++ b/src/zencore/compress.cpp
@@ -78,12 +78,20 @@ struct BufferHeader
 		BufferHeader Header;
 		if (sizeof(BufferHeader) <= CompressedData.GetSize())
 		{
-			//			if (CompressedData.GetSegments()[0].AsIoBuffer().IsWholeFile())
-			//			{
-			//				ZEN_ASSERT(true);
-			//			}
-			CompositeBuffer::Iterator It;
-			CompressedData.CopyTo(MakeMutableMemoryView(&Header, &Header + 1), It);
+			// Fast path: the overwhelmingly common case is that the 64-byte header sits entirely
+			// within the first segment and that segment is plain memory. Skip the iterator and
+			// the sub-range IoBuffer wrapper (which would otherwise heap-allocate an IoBufferCore).
+			const std::span<const SharedBuffer> Segments = CompressedData.GetSegments();
+			const SharedBuffer&					First	 = Segments.front();
+			if (sizeof(BufferHeader) <= First.GetSize() && !First.IsExtended())
+			{
+				MakeMutableMemoryView(&Header, &Header + 1).CopyFrom(First.GetView().Left(sizeof(BufferHeader)));
+			}
+			else
+			{
+				CompositeBuffer::Iterator It;
+				CompressedData.CopyTo(MakeMutableMemoryView(&Header, &Header + 1), It);
+			}
 			Header.ByteSwap();
 		}
 		return Header;
@@ -837,6 +845,15 @@ BlockDecoder::DecompressToStream(
 	{
 		return false;
 	}
+	// RawOffset+RawSize-1 below underflows when RawSize is 0, and the
+	// BlockCount-0 / BlockSize-0 arithmetic is only defined when the header
+	// has already been validated (see IsHeaderValid). Guard both here as
+	// defence in depth.
+	if (RawSize == 0 || Header.BlockCount == 0 || Header.BlockSizeExponent >= 32 || RawOffset > Header.TotalRawSize ||
+		RawSize > Header.TotalRawSize - RawOffset)
+	{
+		return false;
+	}
 
 	const uint64_t BlockSize = uint64_t(1) << Header.BlockSizeExponent;
 
@@ -1386,6 +1403,50 @@ GetDecoder(CompressionMethod Method)
 	}
 }
 
+// Sanity-check a header that was just read from an untrusted buffer before
+// any of the decode arithmetic (1 << BlockSizeExponent, BlockCount*BlockSize,
+// divides by BlockSize, etc.) is performed. Must be called after the magic,
+// decoder and CRC checks pass.
+static bool
+IsHeaderValid(const BufferHeader& Header)
+{
+	// 1 << BlockSizeExponent is UB for Exponent >= 64 and wildly impractical
+	// below that. Real producers use <= 24 (16 MiB blocks); cap at 32 for
+	// headroom while staying well below the UB boundary.
+	if (Header.BlockSizeExponent >= 32)
+	{
+		return false;
+	}
+
+	// Only the block-based methods use BlockCount / BlockSizeExponent. The
+	// None method keeps them zero.
+	if (Header.Method != CompressionMethod::None)
+	{
+		// A non-empty buffer needs at least one block.
+		if (Header.BlockCount == 0)
+		{
+			return Header.TotalRawSize == 0;
+		}
+
+		const uint64_t BlockSize = uint64_t(1) << Header.BlockSizeExponent;
+
+		// BlockCount * BlockSize must not overflow and must fit TotalRawSize
+		// in the half-open range ((BlockCount - 1) * BlockSize, BlockCount * BlockSize].
+		if (Header.BlockCount > (UINT64_MAX / BlockSize))
+		{
+			return false;
+		}
+		const uint64_t MaxRawSize = uint64_t(Header.BlockCount) * BlockSize;
+		const uint64_t MinRawSize = MaxRawSize - BlockSize;
+		if (Header.TotalRawSize > MaxRawSize || Header.TotalRawSize <= MinRawSize)
+		{
+			return false;
+		}
+	}
+
+	return true;
+}
+
 //////////////////////////////////////////////////////////////////////////
 
 bool
@@ -1426,6 +1487,10 @@ ReadHeader(const CompositeBuffer& CompressedData, BufferHeader& OutHeader, Uniqu
 	{
 		return false;
 	}
+	if (!IsHeaderValid(OutHeader))
+	{
+		return false;
+	}
 	uint64_t FullHeaderSize = Decoder->GetHeaderSize(OutHeader);
 	if (FullHeaderSize > CompressedDataSize)
 	{
@@ -1520,7 +1585,7 @@ TryReadHeader(DecoderContext& Context, Archive& Ar, FHeader& OutHeader, MemoryVi
 
 		FHeader* const HeaderCopy = static_cast<FHeader*>(HeaderView.GetData());
 		HeaderCopy->ByteSwap();
-		if (Header.Crc32 == FHeader::CalculateCrc32(HeaderView))
+		if (Header.Crc32 == FHeader::CalculateCrc32(HeaderView) && IsHeaderValid(Header))
 		{
 			Context.HeaderOffset  = uint64_t(Offset);
 			Context.HeaderSize	  = HeaderSize;
@@ -1560,7 +1625,7 @@ TryReadHeader(DecoderContext& Context, const CompositeBuffer& Buffer, FHeader& O
 		const MemoryView HeaderView = Buffer.ViewOrCopyRange(0, HeaderSize, Context.Header, [](uint64_t Size) {
 			return UniqueBuffer::Alloc(zen::Max(NextPow2(Size), DefaultHeaderSize));
 		});
-		if (Header.Crc32 == FHeader::CalculateCrc32(HeaderView))
+		if (Header.Crc32 == FHeader::CalculateCrc32(HeaderView) && IsHeaderValid(Header))
 		{
 			Context.HeaderOffset = 0;
 			Context.HeaderSize	 = HeaderSize;