From 33313f208db77c663a9d5a91e06a47a88709d900 Mon Sep 17 00:00:00 2001
From: jakcron <jkrcarr@gmail.com>
Date: Tue, 22 May 2018 19:45:40 +0800
Subject: [PATCH] [nstool] Where not guarenteed to control file position, IFile
 wrappers assert their position.

---
 programs/nstool/source/AesCtrWrappedIFile.cpp | 28 +++++++++++++------
 programs/nstool/source/AesCtrWrappedIFile.h   |  3 ++
 .../nstool/source/OffsetAdjustedIFile.cpp     |  4 +--
 3 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/programs/nstool/source/AesCtrWrappedIFile.cpp b/programs/nstool/source/AesCtrWrappedIFile.cpp
index 134ded9..8ce436e 100644
--- a/programs/nstool/source/AesCtrWrappedIFile.cpp
+++ b/programs/nstool/source/AesCtrWrappedIFile.cpp
@@ -1,19 +1,16 @@
 #include "AesCtrWrappedIFile.h"
 
 AesCtrWrappedIFile::AesCtrWrappedIFile(fnd::IFile* file, const crypto::aes::sAes128Key& key, const crypto::aes::sAesIvCtr& ctr) :
-	mOwnIFile(false),
-	mFile(file),
-	mKey(key),
-	mBaseCtr(ctr)
+	AesCtrWrappedIFile(file, false, key, ctr)
 {
-	mScratch.alloc(kAesCtrScratchAllocSize);
 }
 
 AesCtrWrappedIFile::AesCtrWrappedIFile(fnd::IFile* file, bool ownIfile, const crypto::aes::sAes128Key& key, const crypto::aes::sAesIvCtr& ctr) :
 	mOwnIFile(ownIfile),
 	mFile(file),
 	mKey(key),
-	mBaseCtr(ctr)
+	mBaseCtr(ctr),
+	mFileOffset(0)
 {
 	mScratch.alloc(kAesCtrScratchAllocSize);
 }
@@ -33,13 +30,13 @@ size_t AesCtrWrappedIFile::size()
 
 void AesCtrWrappedIFile::seek(size_t offset)
 {
-	mFile->seek(offset);
-	crypto::aes::AesIncrementCounter(mBaseCtr.iv, offset>>4, mCurrentCtr.iv);
-	mBlockOffset = offset & 0xf;
+	mFileOffset = offset;
 }
 
 void AesCtrWrappedIFile::read(byte_t* out, size_t len)
 {
+	internalSeek();
+
 	for (size_t i = 0; i < (len / kAesCtrScratchSize); i++)
 	{
 		mFile->read(mScratch.getBytes() + mBlockOffset, kAesCtrScratchSize);
@@ -55,6 +52,8 @@ void AesCtrWrappedIFile::read(byte_t* out, size_t len)
 		crypto::aes::AesCtr(mScratch.getBytes(), kAesCtrScratchAllocSize, mKey.key, mCurrentCtr.iv, mScratch.getBytes());
 		memcpy(out + read_pos, mScratch.getBytes() + mBlockOffset, read_len);
 	}
+
+	seek(mFileOffset + len);
 }
 
 void AesCtrWrappedIFile::read(byte_t* out, size_t offset, size_t len)
@@ -65,6 +64,8 @@ void AesCtrWrappedIFile::read(byte_t* out, size_t offset, size_t len)
 
 void AesCtrWrappedIFile::write(const byte_t* out, size_t len)
 {
+	internalSeek();
+
 	for (size_t i = 0; i < (len / kAesCtrScratchSize); i++)
 	{
 		memcpy(mScratch.getBytes() + mBlockOffset, out + (i * kAesCtrScratchSize), kAesCtrScratchSize);
@@ -80,10 +81,19 @@ 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)
 {
 	seek(offset);
 	write(out, len);
+}
+
+void AesCtrWrappedIFile::internalSeek()
+{
+	mFile->seek(mFileOffset);
+	crypto::aes::AesIncrementCounter(mBaseCtr.iv, mFileOffset>>4, mCurrentCtr.iv);
+	mBlockOffset = mFileOffset & 0xf;
 }
\ No newline at end of file
diff --git a/programs/nstool/source/AesCtrWrappedIFile.h b/programs/nstool/source/AesCtrWrappedIFile.h
index 3eaaaad..0d30754 100644
--- a/programs/nstool/source/AesCtrWrappedIFile.h
+++ b/programs/nstool/source/AesCtrWrappedIFile.h
@@ -25,6 +25,9 @@ private:
 	crypto::aes::sAes128Key mKey;
 	crypto::aes::sAesIvCtr mBaseCtr, mCurrentCtr;
 	size_t mBlockOffset;
+	size_t mFileOffset;
 
 	fnd::MemoryBlob mScratch;
+
+	void internalSeek();
 };
\ No newline at end of file
diff --git a/programs/nstool/source/OffsetAdjustedIFile.cpp b/programs/nstool/source/OffsetAdjustedIFile.cpp
index 4690217..e4dcd49 100644
--- a/programs/nstool/source/OffsetAdjustedIFile.cpp
+++ b/programs/nstool/source/OffsetAdjustedIFile.cpp
@@ -43,7 +43,7 @@ void OffsetAdjustedIFile::read(byte_t* out, size_t len)
 	// assert proper position in file
 	mFile->seek(mCurrentOffset + mBaseOffset);
 	mFile->read(out, len);
-	mCurrentOffset += len;
+	seek(mCurrentOffset + len);
 }
 
 void OffsetAdjustedIFile::read(byte_t* out, size_t offset, size_t len)
@@ -57,7 +57,7 @@ void OffsetAdjustedIFile::write(const byte_t* out, size_t len)
 	// assert proper position in file
 	mFile->seek(mCurrentOffset + mBaseOffset);
 	mFile->write(out, len);
-	mCurrentOffset += len;
+	seek(mCurrentOffset + len);
 }
 
 void OffsetAdjustedIFile::write(const byte_t* out, size_t offset, size_t len)