From 3716bf0b825ec5d6d158c5b8598c47bc889cd06f Mon Sep 17 00:00:00 2001
From: jakcron <jkrcarr@gmail.com>
Date: Wed, 26 Sep 2018 21:48:38 +0800
Subject: [PATCH] [nstool] Fix critical bug in HashTreeWrappedIFile (bumped to
 v1.0.5)

---
 programs/nstool/makefile                      |  2 +-
 .../nstool/source/HashTreeWrappedIFile.cpp    | 69 ++++++++++---------
 programs/nstool/source/version.h              |  2 +-
 3 files changed, 38 insertions(+), 35 deletions(-)

diff --git a/programs/nstool/makefile b/programs/nstool/makefile
index d73e3b4..91999dc 100644
--- a/programs/nstool/makefile
+++ b/programs/nstool/makefile
@@ -34,7 +34,7 @@ else
 endif
 
 BIN_DIR = bin
-OUTPUT = $(BIN_DIR)/$(shell basename $(CURDIR))
+OUTPUT = $(BIN_DIR)/$(shell basename "$(CURDIR)")
 
 all: build
 
diff --git a/programs/nstool/source/HashTreeWrappedIFile.cpp b/programs/nstool/source/HashTreeWrappedIFile.cpp
index a69acb2..54168de 100644
--- a/programs/nstool/source/HashTreeWrappedIFile.cpp
+++ b/programs/nstool/source/HashTreeWrappedIFile.cpp
@@ -23,45 +23,49 @@ void HashTreeWrappedIFile::seek(size_t offset)
 
 void HashTreeWrappedIFile::read(byte_t* out, size_t len)
 {
-	size_t offset_in_start_block = getOffsetInBlock(mDataOffset);
-	size_t offset_in_end_block = getOffsetInBlock(offset_in_start_block + len);
-
-	size_t start_block = getOffsetBlock(mDataOffset);
-	size_t block_num = align(offset_in_start_block + len, mDataBlockSize) / mDataBlockSize;
-
-	size_t partial_last_block_num = block_num % mCacheBlockNum;
-	bool has_partial_block_num = partial_last_block_num > 0;
-	size_t read_iterations = (block_num / mCacheBlockNum) + has_partial_block_num;
-
-	size_t block_read_len;
-	size_t block_export_offset;
-	size_t block_export_size;
-	size_t block_export_pos = 0;
-	for (size_t i = 0; i < read_iterations; i++)
+	struct sBlockPosition 
 	{
-		// how many blocks to read from source file
-		block_read_len = (i+1 == read_iterations && has_partial_block_num) ? partial_last_block_num : mCacheBlockNum;
+		size_t index;
+		size_t pos;
+	} start_blk, end_blk;
 
-		// offset in this current read to copy from
-		block_export_offset = (i == 0) ? offset_in_start_block : 0;
+	start_blk.index = getOffsetBlock(mDataOffset);
+	start_blk.pos = getOffsetInBlock(mDataOffset);
 
-		// size of current read to copy
-		block_export_size = (block_read_len * mDataBlockSize) - block_export_offset;
+	end_blk.index = getOffsetBlock(mDataOffset + len);
+	end_blk.pos = getOffsetInBlock(mDataOffset + len);
+	if (end_blk.pos == 0 && len != 0)
+	{
+		end_blk.index -= 1;
+		end_blk.pos = mDataBlockSize;
+	}
 
-		// if last read, reduce the export size by one block less offset_in_end_block
-		if (i+1 == read_iterations)
-		{
-			block_export_size -= (mDataBlockSize - offset_in_end_block);
-		}
+	size_t total_blk_num = (end_blk.index - start_blk.index) + 1;
 
-		// read the blocks
-		readData(start_block + (i * mCacheBlockNum), block_read_len);
+	size_t read_blk_num = 0;
+	size_t cache_export_start_pos, cache_export_end_pos, cache_export_size;
+	size_t export_pos = 0;
+	for (size_t i = 0; i < total_blk_num; i += read_blk_num, export_pos += cache_export_size)
+	{
+		read_blk_num = _MIN(mCacheBlockNum, (total_blk_num - i));
+		readData(start_blk.index + i, read_blk_num);
 
-		// export the section of data that is relevant
-		memcpy(out + block_export_pos, mCache.data() + block_export_offset, block_export_size);
+		// if this is the first read, adjust offset
+		if (i == 0)
+			cache_export_start_pos = start_blk.pos;
+		else
+			cache_export_start_pos = 0;
+		
+		// if this is the last block, adjust end offset
+		if ((i + read_blk_num) == total_blk_num)
+			cache_export_end_pos = ((read_blk_num - 1) * mDataBlockSize) + end_blk.pos;
+		else
+			cache_export_end_pos = read_blk_num * mDataBlockSize;
 
-		// update export position
-		block_export_pos += block_export_size;
+		// determine cache export size
+		cache_export_size = cache_export_end_pos - cache_export_start_pos;
+
+		memcpy(out + export_pos, mCache.data() + cache_export_start_pos, cache_export_size);
 	}
 
 	// update offset
@@ -150,7 +154,6 @@ void HashTreeWrappedIFile::initialiseDataLayer(const HashTreeMeta& hdr)
 
 void HashTreeWrappedIFile::readData(size_t block_offset, size_t block_num)
 {
-	(*mData)->seek(block_offset * mDataBlockSize);
 	fnd::sha::sSha256Hash hash;
 
 	// determine read size
diff --git a/programs/nstool/source/version.h b/programs/nstool/source/version.h
index 24fccdf..4773fb4 100644
--- a/programs/nstool/source/version.h
+++ b/programs/nstool/source/version.h
@@ -3,5 +3,5 @@
 #define BIN_NAME	"nstool"
 #define VER_MAJOR	1
 #define VER_MINOR	0
-#define VER_PATCH   4
+#define VER_PATCH   5
 #define AUTHORS		"jakcron"
\ No newline at end of file