From f4e16025c1889ab2fffa27504762c0f3abdde54e Mon Sep 17 00:00:00 2001
From: Paul Hollinsky <phollinsky@intrepidcs.com>
Date: Thu, 10 Mar 2022 18:03:39 -0500
Subject: [PATCH] Disk: ExtendedExtractor: Implement retries

---
 disk/extextractordiskreaddriver.cpp           | 65 ++++++++++++++++++-
 .../icsneo/disk/extextractordiskreaddriver.h  |  3 +
 2 files changed, 65 insertions(+), 3 deletions(-)

diff --git a/disk/extextractordiskreaddriver.cpp b/disk/extextractordiskreaddriver.cpp
index 90ae73c..876629c 100644
--- a/disk/extextractordiskreaddriver.cpp
+++ b/disk/extextractordiskreaddriver.cpp
@@ -4,6 +4,7 @@
 #include "icsneo/api/lifetime.h"
 #include <cstring>
 
+//#define ICSNEO_EXTENDED_EXTRACTOR_DEBUG_PRINTS
 #ifdef ICSNEO_EXTENDED_EXTRACTOR_DEBUG_PRINTS
 #include <iostream>
 #endif
@@ -13,7 +14,6 @@ using namespace icsneo::Disk;
 
 optional<uint64_t> ExtExtractorDiskReadDriver::readLogicalDiskAligned(Communication& com, device_eventhandler_t report,
 	uint64_t pos, uint8_t* into, uint64_t amount, std::chrono::milliseconds timeout) {
-	static std::shared_ptr<MessageFilter> NeoMemorySDRead = std::make_shared<MessageFilter>(Network::NetID::NeoMemorySDRead);
 
 	if(amount > getBlockSizeBounds().second)
 		return nullopt;
@@ -24,6 +24,19 @@ optional<uint64_t> ExtExtractorDiskReadDriver::readLogicalDiskAligned(Communicat
 	if(pos % getBlockSizeBounds().first != 0)
 		return nullopt;
 
+	optional<uint64_t> ret;
+	while(timeout > std::chrono::milliseconds(0) && !ret.has_value()) {
+		auto start = std::chrono::steady_clock::now();
+		ret = attemptReadLogicalDiskAligned(com, report, pos, into, amount, std::chrono::milliseconds(100));
+		timeout -= std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::steady_clock::now() - start);
+	}
+	return ret;
+}
+
+optional<uint64_t> ExtExtractorDiskReadDriver::attemptReadLogicalDiskAligned(Communication& com, device_eventhandler_t report,
+	uint64_t pos, uint8_t* into, uint64_t amount, std::chrono::milliseconds timeout) {
+	static std::shared_ptr<MessageFilter> NeoMemorySDRead = std::make_shared<MessageFilter>(Network::NetID::NeoMemorySDRead);
+
 	if(cachePos != pos || std::chrono::steady_clock::now() > cachedAt + CacheTime) {
 		uint64_t sector = pos / SectorSize;
 
@@ -38,13 +51,28 @@ optional<uint64_t> ExtExtractorDiskReadDriver::readLogicalDiskAligned(Communicat
 		uint16_t receiving = 0; // How much are we about to get before another header or completion
 		uint64_t received = 0;
 		uint16_t receivedCurrent = 0;
+		size_t skipping = 0;
 		std::vector<uint8_t> header;
 		std::unique_lock<std::mutex> lk(m);
 		bool error = !com.redirectRead([&](std::vector<uint8_t>&& data) {
 			std::unique_lock<std::mutex> lk2(m);
-			size_t offset = 0;
+			if(error) {
+				lk2.unlock();
+				cv.notify_all();
+				return;
+			}
+
+			if(skipping > data.size()) {
+				skipping -= data.size();
+				return;
+			}
+			size_t offset = skipping;
+			skipping = 0;
 			while(offset < data.size()) {
 				size_t left = data.size() - offset;
+				#ifdef ICSNEO_EXTENDED_EXTRACTOR_DEBUG_PRINTS
+				std::cout << "Going to process " << left << " bytes" << std::endl;
+				#endif
 				if(header.size() != HeaderLength) {
 					if(header.empty() && left && data[offset] != 0xaa) {
 						#ifdef ICSNEO_EXTENDED_EXTRACTOR_DEBUG_PRINTS
@@ -75,6 +103,37 @@ optional<uint64_t> ExtExtractorDiskReadDriver::readLogicalDiskAligned(Communicat
 					#endif
 					offset += headerLeft;
 
+					if(header[1] == uint8_t(Network::NetID::RED)) {
+						#ifdef ICSNEO_EXTENDED_EXTRACTOR_DEBUG_PRINTS
+						std::cout << "Got extended response " << int(offset) << std::endl;
+						#endif
+						// This is the extended command response, not all devices send this
+						// If we got it, we need to figure out how much more data to ignore
+						uint16_t length = (header[2] + (header[3] << 8));
+						// Try for another header after this, regardless how much we choose
+						// to skip and how we skip it
+						header.clear();
+						if(length <= 6) {
+							#ifdef ICSNEO_EXTENDED_EXTRACTOR_DEBUG_PRINTS
+							std::cout << "Incorrect extended response length " << int(length) << ' ' << int(offset) << std::endl;
+							#endif
+							error = true;
+							lk2.unlock();
+							cv.notify_all();
+							return;
+						}
+						length -= 7;
+						#ifdef ICSNEO_EXTENDED_EXTRACTOR_DEBUG_PRINTS
+						std::cout << "Skipping " << int(length) << ' ' << int(left) << std::endl;
+						#endif
+						if(left < length) {
+							skipping = length - left;
+							return;
+						}
+						offset += length;
+						continue;
+					}
+
 					// The device tells us how much it's sending us before the next header
 					receiving = (header[5] | (header[6] << 8));
 					#ifdef ICSNEO_EXTENDED_EXTRACTOR_DEBUG_PRINTS
@@ -151,4 +210,4 @@ optional<uint64_t> ExtExtractorDiskReadDriver::readLogicalDiskAligned(Communicat
 
 	memcpy(into, cache.data(), size_t(amount));
 	return amount;
-}
\ No newline at end of file
+}
diff --git a/include/icsneo/disk/extextractordiskreaddriver.h b/include/icsneo/disk/extextractordiskreaddriver.h
index b52359a..3809a21 100644
--- a/include/icsneo/disk/extextractordiskreaddriver.h
+++ b/include/icsneo/disk/extextractordiskreaddriver.h
@@ -35,6 +35,9 @@ private:
 
 	optional<uint64_t> readLogicalDiskAligned(Communication& com, device_eventhandler_t report,
 		uint64_t pos, uint8_t* into, uint64_t amount, std::chrono::milliseconds timeout) override;
+
+	optional<uint64_t> attemptReadLogicalDiskAligned(Communication& com, device_eventhandler_t report,
+		uint64_t pos, uint8_t* into, uint64_t amount, std::chrono::milliseconds timeout);
 };
 
 } // namespace Disk