From 72c1b20d2be226d55551354f7f1c83a3ec73dcfc Mon Sep 17 00:00:00 2001
From: Alwin Joshy <joshyalwin@gmail.com>
Date: Tue, 4 Feb 2025 12:54:45 +1100
Subject: [PATCH] loader: add relocation logging and move magic

Signed-off-by: Alwin Joshy <joshyalwin@gmail.com>
---
 loader/src/aarch64/crt0.S   | 15 +++++++++++----
 loader/src/loader.c         | 11 ++++++++++-
 loader/src/riscv/crt0.S     | 11 +++++++++--
 tool/microkit/src/loader.rs |  4 ++--
 4 files changed, 32 insertions(+), 9 deletions(-)

diff --git a/loader/src/aarch64/crt0.S b/loader/src/aarch64/crt0.S
index 6c3c4d24..7470b1f8 100644
--- a/loader/src/aarch64/crt0.S
+++ b/loader/src/aarch64/crt0.S
@@ -43,6 +43,13 @@ fixup_image_base:
     cmp     x0, x1
     beq     image_ok
 
+    /* Log that we are being relocated */
+    bl log_relocation
+
+    /* Restore x0 and x1 */
+    ldr     x0, =LINK_ADDRESS
+    adr     x1, _start
+
     /* Sanity check: We don't want to overwrite ourselves! We assume that
      * everything between _start (src_start) and _bss_end (src_end) is important (i.e.
      * something that might be run while relocating) but allow overlap for
@@ -52,11 +59,11 @@ fixup_image_base:
     add     x2, x2, #:lo12:_bss_end
 
     /* The loader_data is directly after _bss_end, with the first
-     * value being the loader_data struct. The first field of this
-     * struct is the size of the loader_data region, so we add
-     * this to _bss_end to get the real end of the image
+     * value being the loader_data struct. The second field of
+     * this struct is the size of the loader_data region, so we
+     * add this to _bss_end to get the real end of the image
      */
-    ldr     x3, [x2]
+    ldr     x3, [x2, #+8]
     add     x2, x2, x3
     sub     x2, x2, x1
 
diff --git a/loader/src/loader.c b/loader/src/loader.c
index e18c9db9..f5c219c1 100644
--- a/loader/src/loader.c
+++ b/loader/src/loader.c
@@ -58,8 +58,8 @@ struct region {
 };
 
 struct loader_data {
-    uintptr_t size;
     uintptr_t magic;
+    uintptr_t size;
     uintptr_t flags;
     uintptr_t kernel_entry;
     uintptr_t ui_p_reg_start;
@@ -746,6 +746,15 @@ static inline void enable_mmu(void)
 }
 #endif
 
+void log_relocation(uint64_t reloc_addr, uint64_t curr_addr)
+{
+    puts("LDR|INFO: loader is being relocated. Currently at: ");
+    puthex64(curr_addr);
+    puts(". Moving to: ");
+    puthex64(reloc_addr);
+    puts(".\n");
+}
+
 int main(void)
 {
 #if defined(BOARD_zcu102)
diff --git a/loader/src/riscv/crt0.S b/loader/src/riscv/crt0.S
index 81ce4c91..cb108cf1 100644
--- a/loader/src/riscv/crt0.S
+++ b/loader/src/riscv/crt0.S
@@ -113,6 +113,13 @@ fixup_image_base:
   la a1, _start
   beq a0, a1, image_ok
 
+  /* Log that we are being relocated */
+  jal log_relocation
+
+  /* Restore a0 and a1 */
+  li a0, LINK_ADDRESS
+  la a1, _start
+
   /* Sanity check: We don't want to overwrite ourselves! We assume that
    * everything between _start (src_start) and _bss_end (src_end) is important (i.e.
    * something that might be run while relocating) but allow overlap for
@@ -121,11 +128,11 @@ fixup_image_base:
   la a2, _bss_end
 
    /* The loader_data is directly after _bss_end, with the first
-   * value being the loader_data struct. The first field of this
+   * value being the loader_data struct. The second field of this
    * struct is the size of the loader_data region, so we add
    * this to _bss_end to get the real end of the image
    */
-  lw a3, (a2)
+  lw a3, 8(a2)
   add a2, a2, a3
   sub a2, a2, a1
 
diff --git a/tool/microkit/src/loader.rs b/tool/microkit/src/loader.rs
index 559c034c..94be093b 100644
--- a/tool/microkit/src/loader.rs
+++ b/tool/microkit/src/loader.rs
@@ -108,8 +108,8 @@ struct LoaderRegion64 {
 
 #[repr(C)]
 struct LoaderHeader64 {
-    size: u64,
     magic: u64,
+    size: u64,
     flags: u64,
     kernel_entry: u64,
     ui_p_reg_start: u64,
@@ -298,8 +298,8 @@ impl<'a> Loader<'a> {
             });
 
         let header = LoaderHeader64 {
-            size,
             magic,
+            size,
             flags,
             kernel_entry,
             ui_p_reg_start,