From 97418ca9a109c324baeb05b607021d82ded76eee Mon Sep 17 00:00:00 2001
From: TuxSH <1922548+TuxSH@users.noreply.github.com>
Date: Tue, 11 Jul 2023 00:33:32 +0200
Subject: [PATCH] Properly support arbitrarily-sized and uncompressed
 TwlBg/AgbBg

As we forgot to patch LGY k11, we were limited by the current size of the KIP.
That is not the case anymore and LGY k11 does support uncompressed KIPs.
---
 arm9/source/firm.c    | 21 ++++++++++++++++-----
 arm9/source/patches.c | 36 ++++++++++++++++++++++++++++++------
 arm9/source/patches.h |  3 ++-
 3 files changed, 48 insertions(+), 12 deletions(-)

diff --git a/arm9/source/firm.c b/arm9/source/firm.c
index 2829a8b..a2281c8 100755
--- a/arm9/source/firm.c
+++ b/arm9/source/firm.c
@@ -269,6 +269,8 @@ static inline void mergeSection0(FirmwareType firmType, u32 firmVersion, bool lo
     u32 srcModuleSize,
         nbModules = 0;
 
+    bool isLgyFirm = firmType == TWL_FIRM || firmType == AGB_FIRM;
+
     struct
     {
         char name[8];
@@ -311,7 +313,7 @@ static inline void mergeSection0(FirmwareType firmType, u32 firmVersion, bool lo
     u8 *dst = firm->section[0].address;
     const char *extModuleSizeError = "The external FIRM modules are too large.";
     // SAFE_FIRM only for N3DS and only if ENABLESAFEFIRMROSALINA is on
-    u32 maxModuleSize = (firmType == NATIVE_FIRM || firmType == SAFE_FIRM) ? 0x80000 : 0x600000;
+    u32 maxModuleSize = !isLgyFirm ? 0x80000 : 0x600000;
     for(u32 i = 0, dstModuleSize; i < nbModules; i++, dst += dstModuleSize, maxModuleSize -= dstModuleSize)
     {
         if(loadFromStorage)
@@ -344,11 +346,20 @@ static inline void mergeSection0(FirmwareType firmType, u32 firmVersion, bool lo
         memcpy(dst, moduleList[i].src, dstModuleSize);
     }
 
-    //4) Patch NATIVE_FIRM/SAFE_FIRM (N3DS) if necessary
-    if(nbModules == 6)
+    //4) Patch kernel to take module size into account
+    u32 newKipSectionSize = dst - firm->section[0].address;
+    u32 oldKipSectionSize = firm->section[0].size;
+    u8 *kernel11Addr = (u8 *)firm + firm->section[1].offset;
+    u32 kernel11Size = firm->section[1].size;
+    if (isLgyFirm)
     {
-        if(patchK11ModuleLoading(firm->section[0].size, dst - firm->section[0].address, (u8 *)firm + firm->section[1].offset, firm->section[1].size) != 0)
-            error("Failed to inject custom sysmodule");
+        if (patchK11ModuleLoadingLgy(newKipSectionSize, kernel11Addr, kernel11Size) != 0)
+            error("Failed to load sysmodules");
+    }
+    else
+    {
+        if (patchK11ModuleLoading(oldKipSectionSize, newKipSectionSize, nbModules, kernel11Addr, kernel11Size) != 0)
+            error("Failed to load sysmodules");
     }
 }
 
diff --git a/arm9/source/patches.c b/arm9/source/patches.c
index e01a631..5a179ff 100644
--- a/arm9/source/patches.c
+++ b/arm9/source/patches.c
@@ -463,7 +463,7 @@ u32 patchCheckForDevCommonKey(u8 *pos, u32 size)
     return 0;
 }
 
-u32 patchK11ModuleLoading(u32 section0size, u32 modulesSize, u8 *pos, u32 size)
+u32 patchK11ModuleLoading(u32 oldKipSectionSize, u32 newKipSectionSize, u32 numKips, u8 *pos, u32 size)
 {
     static const u8 moduleLoadingPattern[]  = {0xE2, 0x05, 0x00, 0x57},
                     modulePidPattern[] = {0x06, 0xA0, 0xE1, 0xF2}; //GetSystemInfo
@@ -472,20 +472,44 @@ u32 patchK11ModuleLoading(u32 section0size, u32 modulesSize, u8 *pos, u32 size)
 
     if(off == NULL) return 1;
 
-    off[1]++;
+    off[1] = (u8)numKips;
 
     u32 *off32;
     for(off32 = (u32 *)(off - 3); *off32 != 0xE59F0000; off32++);
     off32 += 2;
-    off32[1] = off32[0] + modulesSize;
-    for(; *off32 != section0size; off32++);
-    *off32 = ((modulesSize + 0x1FF) >> 9) << 9;
+    off32[1] = off32[0] + newKipSectionSize;
+    for(; *off32 != oldKipSectionSize; off32++);
+    *off32 = ((newKipSectionSize + 0x1FF) >> 9) << 9;
 
     off = memsearch(pos, modulePidPattern, size, 4);
 
     if(off == NULL) return 1;
 
-    off[0xB] = 6;
+    off[0xB] = (u8)numKips;
+
+    return 0;
+}
+
+u32 patchK11ModuleLoadingLgy(u32 newKipSectionSize, u8 *pos, u32 size)
+{
+    // Patch the function where TwlBg/AgbBg is copied from 18000000 (VRAM) to 21000000 (FCRAM).
+    // This is where we can also automatically obtain the section size
+
+    u16 *off = (u16 *)pos;
+    for (; (u8 *)off < pos + size && (off[0] != 0x06C9 || off[1] != 0x0600); off++);
+    if ((u8 *)off >= pos + size)
+        return 3;
+
+    off += 7;
+    u32 oldKipSectionSize = *(u32 *)off;
+    *(u32 *)off = newKipSectionSize;
+    off += 2;
+
+    u32 *off2 = (u32 *)off;
+    for (; (u8 *)off2 < pos + size && *off2 != oldKipSectionSize; off2++);
+    if ((u8 *)off2 >= pos + size)
+        return 4;
+    *off2 = newKipSectionSize;
 
     return 0;
 }
diff --git a/arm9/source/patches.h b/arm9/source/patches.h
index dae3598..a612165 100644
--- a/arm9/source/patches.h
+++ b/arm9/source/patches.h
@@ -50,7 +50,8 @@ u32 patchTitleInstallMinVersionChecks(u8 *pos, u32 size, u32 firmVersion);
 u32 patchZeroKeyNcchEncryptionCheck(u8 *pos, u32 size);
 u32 patchNandNcchEncryptionCheck(u8 *pos, u32 size);
 u32 patchCheckForDevCommonKey(u8 *pos, u32 size);
-u32 patchK11ModuleLoading(u32 section0size, u32 modulesSize, u8 *startPos, u32 size);
+u32 patchK11ModuleLoading(u32 oldKipSectionSize, u32 newKipSectionSize, u32 numKips, u8 *pos, u32 size);
+u32 patchK11ModuleLoadingLgy(u32 newKipSectionSize, u8 *pos, u32 size);
 u32 patchArm9ExceptionHandlersInstall(u8 *pos, u32 size);
 u32 patchSvcBreak9(u8 *pos, u32 size, u32 kernel9Address);
 u32 patchKernel9Panic(u8 *pos, u32 size);