From f4d3501f8c6e3712ba13f686febbc9b316374b83 Mon Sep 17 00:00:00 2001
From: jakcron <jkrcarr@gmail.com>
Date: Sun, 27 May 2018 18:16:42 +0800
Subject: [PATCH] [nstool] Fix bug in AesCtrWrappedIFile.

---
 programs/nstool/source/AesCtrWrappedIFile.cpp | 72 ++++++++++++-------
 programs/nstool/source/AesCtrWrappedIFile.h   |  7 +-
 2 files changed, 49 insertions(+), 30 deletions(-)

diff --git a/programs/nstool/source/AesCtrWrappedIFile.cpp b/programs/nstool/source/AesCtrWrappedIFile.cpp
index 8ce436e..81fd145 100644
--- a/programs/nstool/source/AesCtrWrappedIFile.cpp
+++ b/programs/nstool/source/AesCtrWrappedIFile.cpp
@@ -12,7 +12,7 @@ AesCtrWrappedIFile::AesCtrWrappedIFile(fnd::IFile* file, bool ownIfile, const cr
 	mBaseCtr(ctr),
 	mFileOffset(0)
 {
-	mScratch.alloc(kAesCtrScratchAllocSize);
+	mCache.alloc(kCacheSizeAllocSize);
 }
 
 AesCtrWrappedIFile::~AesCtrWrappedIFile()
@@ -35,22 +35,27 @@ void AesCtrWrappedIFile::seek(size_t offset)
 
 void AesCtrWrappedIFile::read(byte_t* out, size_t len)
 {
-	internalSeek();
+	//printf("[%x] AesCtrWrappedIFile::read(offset=0x%" PRIx64 ", size=0x%" PRIx64 ")\n", this, mFileOffset, len);
 
-	for (size_t i = 0; i < (len / kAesCtrScratchSize); i++)
-	{
-		mFile->read(mScratch.getBytes() + mBlockOffset, kAesCtrScratchSize);
-		crypto::aes::AesCtr(mScratch.getBytes(), kAesCtrScratchAllocSize, mKey.key, mCurrentCtr.iv, mScratch.getBytes());
-		memcpy(out + (i * kAesCtrScratchSize), mScratch.getBytes() + mBlockOffset, kAesCtrScratchSize);
-	}
+	size_t read_len;
+	size_t read_pos; 
 
-	if (len % kAesCtrScratchSize)
+	size_t cache_reads = (len / kCacheSize) + ((len % kCacheSize) != 0);
+
+	for (size_t i = 0; i < cache_reads; i++)
 	{
-		size_t read_len = len % kAesCtrScratchSize;
-		size_t read_pos = ((len / kAesCtrScratchSize) * kAesCtrScratchSize); 
-		mFile->read(mScratch.getBytes() + mBlockOffset, read_len);
-		crypto::aes::AesCtr(mScratch.getBytes(), kAesCtrScratchAllocSize, mKey.key, mCurrentCtr.iv, mScratch.getBytes());
-		memcpy(out + read_pos, mScratch.getBytes() + mBlockOffset, read_len);
+		read_len = MIN(len - (i * kCacheSize), kCacheSize);
+		read_pos = ((mFileOffset >> 4) << 4) + (i * kCacheSize);
+
+		//printf("[%x] AesCtrWrappedIFile::read() CACHE READ: readlen=%" PRIx64 "\n", this, read_len);
+		
+		mFile->seek(read_pos);
+		mFile->read(mCache.getBytes(), kCacheSizeAllocSize);
+
+		crypto::aes::AesIncrementCounter(mBaseCtr.iv, read_pos>>4, mCurrentCtr.iv);
+		crypto::aes::AesCtr(mCache.getBytes(), kCacheSizeAllocSize, mKey.key, mCurrentCtr.iv, mCache.getBytes());
+
+		memcpy(out + (i * kCacheSize), mCache.getBytes() + (mFileOffset & 0xf), read_len);
 	}
 
 	seek(mFileOffset + len);
@@ -62,10 +67,32 @@ void AesCtrWrappedIFile::read(byte_t* out, size_t offset, size_t len)
 	read(out, len);
 }
 
-void AesCtrWrappedIFile::write(const byte_t* out, size_t len)
+void AesCtrWrappedIFile::write(const byte_t* in, size_t len)
 {
-	internalSeek();
+	size_t write_len;
+	size_t write_pos; 
 
+	size_t cache_writes = (len / kCacheSize) + ((len % kCacheSize) != 0);
+
+	for (size_t i = 0; i < cache_writes; i++)
+	{
+		write_len = MIN(len - (i * kCacheSize), kCacheSize);
+		write_pos = ((mFileOffset >> 4) << 4) + (i * kCacheSize);
+
+		//printf("[%x] AesCtrWrappedIFile::read() CACHE READ: readlen=%" PRIx64 "\n", this, read_len);
+		
+		memcpy(mCache.getBytes() + (mFileOffset & 0xf), in + (i * kCacheSize), write_len);
+
+		crypto::aes::AesIncrementCounter(mBaseCtr.iv, write_pos>>4, mCurrentCtr.iv);
+		crypto::aes::AesCtr(mCache.getBytes(), kCacheSizeAllocSize, mKey.key, mCurrentCtr.iv, mCache.getBytes());
+
+		mFile->seek(write_pos);
+		mFile->write(mCache.getBytes(), kCacheSizeAllocSize);
+	}
+
+	seek(mFileOffset + len);
+
+	/*
 	for (size_t i = 0; i < (len / kAesCtrScratchSize); i++)
 	{
 		memcpy(mScratch.getBytes() + mBlockOffset, out + (i * kAesCtrScratchSize), kAesCtrScratchSize);
@@ -81,19 +108,12 @@ void AesCtrWrappedIFile::write(const byte_t* out, size_t len)
 		crypto::aes::AesCtr(mScratch.getBytes(), kAesCtrScratchAllocSize, mKey.key, mCurrentCtr.iv, mScratch.getBytes());
 		mFile->write(mScratch.getBytes() + mBlockOffset, write_len);
 	}
-
+	*/
 	seek(mFileOffset + len);
 }
 
-void AesCtrWrappedIFile::write(const byte_t* out, size_t offset, size_t len)
+void AesCtrWrappedIFile::write(const byte_t* in, size_t offset, size_t len)
 {
 	seek(offset);
-	write(out, len);
-}
-
-void AesCtrWrappedIFile::internalSeek()
-{
-	mFile->seek(mFileOffset);
-	crypto::aes::AesIncrementCounter(mBaseCtr.iv, mFileOffset>>4, mCurrentCtr.iv);
-	mBlockOffset = mFileOffset & 0xf;
+	write(in, len);
 }
\ No newline at end of file
diff --git a/programs/nstool/source/AesCtrWrappedIFile.h b/programs/nstool/source/AesCtrWrappedIFile.h
index 0d30754..2f83e1a 100644
--- a/programs/nstool/source/AesCtrWrappedIFile.h
+++ b/programs/nstool/source/AesCtrWrappedIFile.h
@@ -17,17 +17,16 @@ public:
 	void write(const byte_t* out, size_t offset, size_t len);
 private:
 	const std::string kModuleName = "AesCtrWrappedIFile";
-	static const size_t kAesCtrScratchSize = 0x1000000;
-	static const size_t kAesCtrScratchAllocSize = kAesCtrScratchSize + crypto::aes::kAesBlockSize;
+	static const size_t kCacheSize = 0x10000;
+	static const size_t kCacheSizeAllocSize = kCacheSize + crypto::aes::kAesBlockSize;
 
 	bool mOwnIFile;
 	fnd::IFile* mFile;
 	crypto::aes::sAes128Key mKey;
 	crypto::aes::sAesIvCtr mBaseCtr, mCurrentCtr;
-	size_t mBlockOffset;
 	size_t mFileOffset;
 
-	fnd::MemoryBlob mScratch;
+	fnd::MemoryBlob mCache;
 
 	void internalSeek();
 };
\ No newline at end of file