From 80cd4ae05281d97965a785d44f1ef86370fb845f Mon Sep 17 00:00:00 2001
From: Paul Hollinsky <phollinsky@intrepidcs.com>
Date: Fri, 19 Nov 2021 14:14:55 -0500
Subject: [PATCH] Disk: Fix out of bounds accesses in base ReadDriver

---
 disk/diskreaddriver.cpp | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/disk/diskreaddriver.cpp b/disk/diskreaddriver.cpp
index 0111a2b..5948b5d 100644
--- a/disk/diskreaddriver.cpp
+++ b/disk/diskreaddriver.cpp
@@ -14,19 +14,21 @@ optional<uint64_t> DiskReadDriver::readLogicalDisk(Communication& com, device_ev
 
 	const uint32_t idealBlockSize = getBlockSizeBounds().second;
 	const uint64_t startBlock = pos / idealBlockSize;
-	const uint64_t posWithinFirstBlock = pos % idealBlockSize;
-	const uint64_t blocks = amount / idealBlockSize + (amount % idealBlockSize ? 1 : 0) + (posWithinFirstBlock ? 1 : 0);
+	const uint32_t posWithinFirstBlock = static_cast<uint32_t>(pos % idealBlockSize);
+	uint64_t blocks = amount / idealBlockSize + (amount % idealBlockSize ? 1 : 0);
+	if(blocks * idealBlockSize - posWithinFirstBlock < amount)
+		blocks++; // We need one more block to get the last partial block's worth
 	uint64_t blocksProcessed = 0;
 
 	while(blocksProcessed < blocks && timeout >= std::chrono::milliseconds::zero()) {
 		const uint64_t currentBlock = startBlock + blocksProcessed;
 
-		const int intoOffset = std::min<int>((blocksProcessed * idealBlockSize) - posWithinFirstBlock, 0);
-		const auto posWithinCurrentBlock = (blocksProcessed ? 0 : posWithinFirstBlock);
-		auto curAmt = idealBlockSize - posWithinCurrentBlock;
-		const auto amountLeft = amount - *ret;
+		const uint64_t intoOffset = std::max<uint64_t>((blocksProcessed * idealBlockSize) - posWithinFirstBlock, 0);
+		const uint32_t posWithinCurrentBlock = (blocksProcessed ? 0 : posWithinFirstBlock);
+		uint32_t curAmt = idealBlockSize - posWithinCurrentBlock;
+		const auto amountLeft = amount - ret.value_or(0);
 		if(curAmt > amountLeft)
-			curAmt = amountLeft;
+			curAmt = static_cast<uint32_t>(amountLeft);
 
 		const bool useAlignedReadBuffer = (posWithinCurrentBlock != 0 || curAmt != idealBlockSize);
 		if(useAlignedReadBuffer && alignedReadBuffer.size() < idealBlockSize)
@@ -34,7 +36,7 @@ optional<uint64_t> DiskReadDriver::readLogicalDisk(Communication& com, device_ev
 
 		auto start = std::chrono::high_resolution_clock::now();
 		auto readAmount = readLogicalDiskAligned(com, report, currentBlock * idealBlockSize,
-			useAlignedReadBuffer ? alignedReadBuffer.data() : into + intoOffset, idealBlockSize, timeout);
+			useAlignedReadBuffer ? alignedReadBuffer.data() : (into + intoOffset), idealBlockSize, timeout);
 		timeout -= std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::high_resolution_clock::now() - start);
 
 		if(!readAmount.has_value() || *readAmount == 0) {
@@ -50,7 +52,7 @@ optional<uint64_t> DiskReadDriver::readLogicalDisk(Communication& com, device_ev
 
 		if(!ret)
 			ret.emplace();
-		*ret += *readAmount;
+		*ret += std::min<uint64_t>(*readAmount, curAmt);
 		blocksProcessed++;
 	}