From 096d0a89ee4d6fc6163d5db103d15a27fa12689d Mon Sep 17 00:00:00 2001
From: wheremyfoodat <44909372+wheremyfoodat@users.noreply.github.com>
Date: Thu, 11 Jul 2024 22:22:33 +0300
Subject: [PATCH 1/2] Fix AES-CTR decryption for non-NCCHKey0 games

---
 include/loader/ncch.hpp  |  2 ++
 src/core/loader/ncch.cpp | 22 ++++++++++++++++------
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/include/loader/ncch.hpp b/include/loader/ncch.hpp
index 42ce1590..8e35643b 100644
--- a/include/loader/ncch.hpp
+++ b/include/loader/ncch.hpp
@@ -60,6 +60,8 @@ struct NCCH {
 	CodeSetInfo text, data, rodata;
 	FSInfo partitionInfo;
 
+	std::optional<Crypto::AESKey> primaryKey, secondaryKey;
+
 	// Contents of the .code file in the ExeFS
 	std::vector<u8> codeFile;
 	// Contains of the cart's save data
diff --git a/src/core/loader/ncch.cpp b/src/core/loader/ncch.cpp
index 3bf73e5d..a8e50101 100644
--- a/src/core/loader/ncch.cpp
+++ b/src/core/loader/ncch.cpp
@@ -29,6 +29,9 @@ bool NCCH::loadFromHeader(Crypto::AESEngine &aesEngine, IOFile& file, const FSIn
 	smdh.clear();
 	partitionInfo = info;
 
+	primaryKey = {};
+	secondaryKey = {};
+
 	size = u64(*(u32*)&header[0x104]) * mediaUnit; // TODO: Maybe don't type pun because big endian will break
 	exheaderSize = *(u32*)&header[0x180];
 
@@ -78,11 +81,11 @@ bool NCCH::loadFromHeader(Crypto::AESEngine &aesEngine, IOFile& file, const FSIn
 		if (!primaryResult.first || !secondaryResult.first) {
 			gotCryptoKeys = false;
 		} else {
-			Crypto::AESKey primaryKey = primaryResult.second;
-			Crypto::AESKey secondaryKey = secondaryResult.second;
+			primaryKey = primaryResult.second;
+			secondaryKey = secondaryResult.second;
 
 			EncryptionInfo encryptionInfoTmp;
-			encryptionInfoTmp.normalKey = primaryKey;
+			encryptionInfoTmp.normalKey = *primaryKey;
 			encryptionInfoTmp.initialCounter.fill(0);
 
 			for (std::size_t i = 1; i <= sizeof(std::uint64_t) - 1; i++) {
@@ -94,7 +97,7 @@ bool NCCH::loadFromHeader(Crypto::AESEngine &aesEngine, IOFile& file, const FSIn
 			encryptionInfoTmp.initialCounter[8] = 2;
 			exeFS.encryptionInfo = encryptionInfoTmp;
 
-			encryptionInfoTmp.normalKey = secondaryKey;
+			encryptionInfoTmp.normalKey = *secondaryKey;
 			encryptionInfoTmp.initialCounter[8] = 3;
 			romFS.encryptionInfo = encryptionInfoTmp;
 		}
@@ -201,13 +204,20 @@ bool NCCH::loadFromHeader(Crypto::AESEngine &aesEngine, IOFile& file, const FSIn
 					Helpers::panic("Second code file in a single NCCH partition. What should this do?\n");
 				}
 
+				// All files in ExeFS use the same IV, though .code uses the secondary key for decryption
+				// whereas .icon/.banner use the primary key.
+				FSInfo info = exeFS;
+				if (secondaryKey.has_value() && info.encryptionInfo.has_value()) {
+					info.encryptionInfo->normalKey = secondaryKey.value();
+				}
+
 				if (compressCode) {
 					std::vector<u8> tmp;
 					tmp.resize(fileSize);
 
 					// A file offset of 0 means our file is located right after the ExeFS header
 					// So in the ROM, files are located at (file offset + exeFS offset + exeFS header size)
-					readFromFile(file, exeFS, tmp.data(), fileOffset + exeFSHeaderSize, fileSize);
+					readFromFile(file, info, tmp.data(), fileOffset + exeFSHeaderSize, fileSize);
 					
 					// Decompress .code file from the tmp vector to the "code" vector
 					if (!CartLZ77::decompress(codeFile, tmp)) {
@@ -216,7 +226,7 @@ bool NCCH::loadFromHeader(Crypto::AESEngine &aesEngine, IOFile& file, const FSIn
 					}
 				} else {
 					codeFile.resize(fileSize);
-					readFromFile(file, exeFS, codeFile.data(), fileOffset + exeFSHeaderSize, fileSize);
+					readFromFile(file, info, codeFile.data(), fileOffset + exeFSHeaderSize, fileSize);
 				}
 			} else if (std::strcmp(name, "icon") == 0) {
 				// Parse icon file to extract region info and more in the future (logo, etc)

From e6084363152edff610951ec81fda0add720a47e1 Mon Sep 17 00:00:00 2001
From: wheremyfoodat <44909372+wheremyfoodat@users.noreply.github.com>
Date: Thu, 11 Jul 2024 22:27:05 +0300
Subject: [PATCH 2/2] Sanity check: Assert .code is encrypted before setting
 normal key

---
 src/core/loader/ncch.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/core/loader/ncch.cpp b/src/core/loader/ncch.cpp
index a8e50101..47d5a4c2 100644
--- a/src/core/loader/ncch.cpp
+++ b/src/core/loader/ncch.cpp
@@ -207,8 +207,8 @@ bool NCCH::loadFromHeader(Crypto::AESEngine &aesEngine, IOFile& file, const FSIn
 				// All files in ExeFS use the same IV, though .code uses the secondary key for decryption
 				// whereas .icon/.banner use the primary key.
 				FSInfo info = exeFS;
-				if (secondaryKey.has_value() && info.encryptionInfo.has_value()) {
-					info.encryptionInfo->normalKey = secondaryKey.value();
+				if (encrypted && secondaryKey.has_value() && info.encryptionInfo.has_value()) {
+					info.encryptionInfo->normalKey = *secondaryKey;
 				}
 
 				if (compressCode) {