From b251436a619c7389fdfaff42cef78a374257d935 Mon Sep 17 00:00:00 2001
From: Nick Chan <towinchenmi@gmail.com>
Date: Mon, 4 Nov 2024 15:28:58 +0800
Subject: [PATCH] Put m1n1 under gBootArgs->topOfKernelData

m1n1's heap allocation goes upward from gBootArgs->topOfKernelData, it is
better to put m1n1 under gTopOfKernelData so m1n1 cannot overwrite itself
with allocations.

Signed-off-by: Nick Chan <towinchenmi@gmail.com>
---
 src/kernel/entry.c | 32 +++++++++++++++++++-------------
 src/shell/main.c   | 14 ++++++++++++++
 2 files changed, 33 insertions(+), 13 deletions(-)

diff --git a/src/kernel/entry.c b/src/kernel/entry.c
index 51c81c3c..c9953036 100644
--- a/src/kernel/entry.c
+++ b/src/kernel/entry.c
@@ -290,6 +290,7 @@ __attribute__((noinline)) void pongo_entry_cached()
     Description: entry point in llktrw
 
 */
+extern void* m1n1_boot(void);
 extern void set_exception_stack_core0();
 extern void lowlevel_set_identity(void);
 extern _Noreturn void jump_to_image_extended(uint64_t image, uint64_t args, void* tramp, uint64_t original_image);
@@ -311,26 +312,31 @@ _Noreturn void pongo_entry(uint64_t *kernel_args, void *entryp, void (*exit_to_e
     lowlevel_cleanup();
 
     uint64_t pfr0, entryOff = 0;
-    void* tramp = NULL;
-    if (gBootFlag != BOOT_FLAG_RAW)
-        tramp = (void*)((gTopOfKernelData + 0x3fffULL) & ~0x3fffULL);
+
+    if (gBootFlag > BOOT_FLAG_HOOK) {
+        // We're in EL1 here, but we might need to go back to EL3
+        __asm__ volatile("mrs %0, id_aa64pfr0_el1" : "=r"(pfr0));
+        if((pfr0 & 0xf000) != 0)
+        {
+            __asm__ volatile("smc 0"); // elevate to EL3
+        }
+    }
 
     switch (gBootFlag) {
-        case BOOT_FLAG_M1N1:
-            entryOff = 0x800;
+        case BOOT_FLAG_M1N1: {
+            /* IMPORTANT: argument parsing order is undefined in C ! */
+            void* entry = m1n1_boot();
+            /* m1n1_boot() makes gTopOfKernelData usable as-is */
+            jump_to_image_extended((uint64_t)entry, (uint64_t)gBootArgs, (void*)gTopOfKernelData, (uint64_t)gEntryPoint);
+            break;
+        }
         case BOOT_FLAG_RAW:
-            // We're in EL1 here, but we might need to go back to EL3
-            __asm__ volatile("mrs %0, id_aa64pfr0_el1" : "=r"(pfr0));
-            if((pfr0 & 0xf000) != 0)
-            {
-                __asm__ volatile("smc 0"); // elevate to EL3
-            }
             // XXX: We should really replace loader_xfer_recv_data with something dedicated here.
-            jump_to_image_extended(((uint64_t)loader_xfer_recv_data) - kCacheableView + 0x800000000 + entryOff, (uint64_t)gBootArgs, tramp, (uint64_t)gEntryPoint);
+            jump_to_image_extended(((uint64_t)loader_xfer_recv_data) - kCacheableView + 0x800000000 + entryOff, (uint64_t)gBootArgs, 0, (uint64_t)gEntryPoint);
             break;
         default:
             xnu_boot();
-            exit_to_el1_image(gBootArgs, gEntryPoint, tramp);
+            exit_to_el1_image(gBootArgs, gEntryPoint, (void*)((gTopOfKernelData + 0x3fffULL) & ~0x3fffULL));
             break;
     }
     screen_puts("didn't boot?!");
diff --git a/src/shell/main.c b/src/shell/main.c
index b6baf221..3ca40c4b 100644
--- a/src/shell/main.c
+++ b/src/shell/main.c
@@ -47,11 +47,15 @@ void pongo_boot_raw() {
 }
 
 extern char gFWVersion[256];
+int gM1N1Size = 0;
+void* gM1N1Stage = NULL;
 void pongo_boot_m1n1() {
     if (!loader_xfer_recv_count) {
         iprintf("please upload a raw m1n1.bin before issuing this command\n");
         return;
     }
+    gM1N1Size = loader_xfer_recv_count;
+    gM1N1Stage = loader_xfer_recv_data - kCacheableView + 0x800000000;
     loader_xfer_recv_count = 0;
     char *fwversion = dt_get_prop("/chosen", "firmware-version", NULL);
     strlcpy(fwversion, gFWVersion, 256);
@@ -60,6 +64,16 @@ void pongo_boot_m1n1() {
     task_yield();
 }
 
+void* m1n1_boot(void) {
+    /* Make sure we are aligned */
+    gTopOfKernelData = (gTopOfKernelData + 0x3fffULL) & ~0x3fffULL;
+    memcpy((void*)gTopOfKernelData, gM1N1Stage, gM1N1Size);
+    void* M1N1Entry = (void*)(gTopOfKernelData + 0x800);
+    gTopOfKernelData = (gTopOfKernelData + gM1N1Size + 0x3fffULL) & ~0x3fffULL;
+    gBootArgs->topOfKernelData = gTopOfKernelData;
+    return M1N1Entry;
+}
+
 void* ramdisk_buf;
 uint32_t ramdisk_size;