Rework kboot/chainload flow to shut down before calling the next stage

Next stage boots now exit back to main() after replying to the proxy
command, allowing shutdown functions to be called. Introduces a new
P_VECTOR proxy op, distinct from P_CALL. The Python side is reworked
to remain compatible with older versions that do not support this.

Signed-off-by: Hector Martin <marcan@marcan.st>

proxyclient/chainload.py

@@ -127,11 +127,6 @@ if args.el1:
 
 print(f"Jumping to stub at 0x{stub.addr:x}")
 
-try:
-    p.mmu_shutdown()
-except:
-    pass
-
 p.reboot(stub.addr, new_base + bootargs_off, image_addr, new_base, image_size)
 
 iface.nop()

proxyclient/proxy.py

@@ -206,6 +206,9 @@ class UartInterface:
                     raise UartRemoteError("Reply error: Unknown error (%d)"%status)
             return data
 
+    def wait_boot(self):
+        self.reply(self.REQ_BOOT)
+
     def nop(self):
         self.cmd(self.REQ_NOP)
         self.reply(self.REQ_NOP)
@@ -259,12 +262,15 @@ class UartInterface:
 class ProxyError(RuntimeError):
     pass
 
-class ProxyCMDError(ProxyError):
+class ProxyReplyError(ProxyError):
     pass
 
 class ProxyRemoteError(ProxyError):
     pass
 
+class ProxyCommandError(ProxyRemoteError):
+    pass
+
 class AlignmentError(Exception):
     pass
 
@@ -283,6 +289,7 @@ class M1N1Proxy:
     P_GET_EXC_COUNT = 0x008
     P_EL0_CALL = 0x009
     P_EL1_CALL = 0x00a
+    P_VECTOR = 0x00b
 
     GUARD_OFF = 0
     GUARD_SKIP = 1
@@ -401,10 +408,10 @@ class M1N1Proxy:
         if reboot:
             return
         if rop != opcode:
-            raise ProxyCMDError("Reply opcode mismatch: Expected 0x%08x, got 0x%08x"%(opcode,rop))
+            raise ProxyReplyError("Reply opcode mismatch: Expected 0x%08x, got 0x%08x"%(opcode,rop))
         if status != self.S_OK:
             if status == self.S_BADCMD:
-                raise ProxyRemoteError("Reply error: Bad Command")
+                raise ProxyCommandError("Reply error: Bad Command")
             else:
                 raise ProxyRemoteError("Reply error: Unknown error (%d)"%status)
         return retval
@@ -441,11 +448,18 @@ class M1N1Proxy:
     def reboot(self, addr, *args, el1=False):
         if len(args) > 4:
             raise ValueError("Too many arguments")
-        self.request(self.P_EL1_CALL if el1 else self.P_CALL, addr, *args, reboot=True)
+        if el1:
-    def vector(self, addr, *args, el1=False):
+            self.request(self.P_EL1_CALL, addr, *args, no_reply=True)
-        if len(args) > 4:
+        else:
-            raise ValueError("Too many arguments")
+            try:
-        self.request(self.P_EL1_CALL if el1 else self.P_CALL, addr, *args, no_reply=True)
+                self.request(self.P_VECTOR, addr, *args)
+                self.iface.wait_boot()
+            except ProxyCommandError: # old m1n1 does not support P_VECTOR
+                try:
+                    self.mmu_shutdown()
+                except ProxyCommandError: # older m1n1 does not support MMU
+                    pass
+                self.request(self.P_CALL, addr, *args, reboot=True)
     def get_bootargs(self):
         return self.request(self.P_GET_BOOTARGS)
     def get_base(self):
@@ -643,7 +657,7 @@ class M1N1Proxy:
         self.request(self.P_FREE, ptr)
 
     def kboot_boot(self, kernel):
-        self.request(self.P_KBOOT_BOOT, kernel, no_reply=True)
+        self.request(self.P_KBOOT_BOOT, kernel)
     def kboot_set_bootargs(self, bootargs):
         self.request(self.P_KBOOT_SET_BOOTARGS, bootargs)
     def kboot_set_initrd(self, base, size):

src/kboot.c

@@ -27,8 +27,6 @@ static size_t initrd_size = 0;
         return -1;                                                                                 \
     } while (0)
 
-typedef void (*kernel_entry)(void *devtree, u64 rsv1, u64 rsv2, u64 rsv3);
-
 static int dt_set_chosen(void)
 {
 
@@ -291,12 +289,13 @@ int kboot_prepare_dt(void *fdt)
 
 int kboot_boot(void *kernel)
 {
-    mmu_shutdown();
+    printf("Preparing to boot kernel at %p with fdt at %p\n", kernel, dt);
-    exception_shutdown();
 
-    printf("Booting kernel at %p with fdt at %p\n", kernel, dt);
+    next_stage.entry = kernel;
+    next_stage.args[0] = (u64)dt;
+    next_stage.args[1] = 0;
+    next_stage.args[2] = 0;
+    next_stage.args[3] = 0;
 
-    ((kernel_entry)kernel)(dt, 0, 0, 0);
+    return 0;
-
-    panic("kernel returned\n");
 }

src/main.c

@@ -3,6 +3,7 @@
 #include "../config.h"
 
 #include "adt.h"
+#include "exception.h"
 #include "fb.h"
 #include "heapblock.h"
 #include "memory.h"
@@ -19,6 +20,8 @@
 
 #include "../build/build_tag.h"
 
+struct vector_args next_stage;
+
 void print_info(void)
 {
     printf("Device info:\n");
@@ -42,6 +45,24 @@ void print_info(void)
     printf("\n");
 }
 
+void run_actions(void)
+{
+    printf("Checking for payloads...\n");
+
+    if (payload_run() == 0) {
+        printf("Valid payload found\n");
+        return;
+    }
+
+    printf("No valid payload found\n");
+
+    usb_init();
+
+    printf("Running proxy...\n");
+
+    uartproxy_run();
+}
+
 void m1n1_main(void)
 {
     printf("\n\nm1n1 v%s\n", BUILD_TAG);
@@ -62,17 +83,24 @@ void m1n1_main(void)
     wdt_disable();
     pmgr_init();
 
-    printf("Checking for payloads...\n");
+    printf("Initialization complete.\n");
 
-    payload_run();
+    run_actions();
 
-    printf("No valid payload found\n");
+    if (!next_stage.entry) {
+        panic("Nothing to do!\n");
+    }
 
-    usb_init();
+    printf("Preparing to run next stage at %p...\n", next_stage.entry);
 
-    printf("Running proxy...\n");
+    mmu_shutdown();
-    uartproxy_run();
+    exception_shutdown();
+    usb_shutdown();
 
-    while (1)
+    printf("Vectoring to next stage...\n");
-        ;
+
+    next_stage.entry(next_stage.args[0], next_stage.args[1], next_stage.args[2],
+                     next_stage.args[3]);
+
+    panic("Next stage returned!\n");
 }

src/payload.c

@@ -166,7 +166,7 @@ static void *load_one_payload(void *start, size_t size)
     }
 }
 
-void payload_run(void)
+int payload_run(void)
 {
     void *p = _payload_start;
 
@@ -178,10 +178,11 @@ void payload_run(void)
 
         if (kboot_prepare_dt(fdt)) {
             printf("Failed to prepare FDT!");
-            return;
+            return -1;
         }
 
-        kboot_boot(kernel);
+        return kboot_boot(kernel);
-        printf("Failed to boot kernel!");
     }
+
+    return -1;
 }

src/payload.h

@@ -3,6 +3,6 @@
 #ifndef __PAYLOAD_H__
 #define __PAYLOAD_H__
 
-void payload_run(void);
+int payload_run(void);
 
 #endif

src/proxy.c

@@ -10,6 +10,7 @@
 #include "memory.h"
 #include "pmgr.h"
 #include "smp.h"
+#include "string.h"
 #include "tunables.h"
 #include "types.h"
 #include "uart.h"
@@ -23,8 +24,6 @@ int proxy_process(ProxyRequest *request, ProxyReply *reply)
 {
     enum exc_guard_t guard_save = exc_guard;
 
-    callfunc *f;
-
     reply->opcode = request->opcode;
     reply->status = S_OK;
     reply->retval = 0;
@@ -33,11 +32,12 @@ int proxy_process(ProxyRequest *request, ProxyReply *reply)
             break;
         case P_EXIT:
             return 0;
-        case P_CALL:
+        case P_CALL: {
-            f = (callfunc *)request->args[0];
+            generic_func *f = (generic_func *)request->args[0];
             reply->retval =
                 f(request->args[1], request->args[2], request->args[3], request->args[4]);
             break;
+        }
         case P_GET_BOOTARGS:
             reply->retval = boot_args_addr;
             break;
@@ -68,15 +68,17 @@ int proxy_process(ProxyRequest *request, ProxyReply *reply)
             exc_count = 0;
             break;
         case P_EL0_CALL:
-            f = (callfunc *)request->args[0];
             reply->retval = el0_call((void *)request->args[0], request->args[1], request->args[2],
                                      request->args[3], request->args[4]);
             break;
         case P_EL1_CALL:
-            f = (callfunc *)request->args[0];
             reply->retval = el1_call((void *)request->args[0], request->args[1], request->args[2],
                                      request->args[3], request->args[4]);
             break;
+        case P_VECTOR:
+            next_stage.entry = (generic_func *)request->args[0];
+            memcpy(next_stage.args, &request->args[1], 4 * sizeof(u64));
+            return 0;
 
         case P_WRITE64:
             exc_guard = GUARD_SKIP;
@@ -301,7 +303,8 @@ int proxy_process(ProxyRequest *request, ProxyReply *reply)
             break;
 
         case P_KBOOT_BOOT:
-            kboot_boot((void *)request->args[0]);
+            if (kboot_boot((void *)request->args[0]) == 0)
+                return 0;
             break;
         case P_KBOOT_SET_BOOTARGS:
             kboot_set_bootargs((void *)request->args[0]);

src/proxy.h

@@ -17,6 +17,7 @@ typedef enum {
     P_GET_EXC_COUNT,
     P_EL0_CALL,
     P_EL1_CALL,
+    P_VECTOR,
 
     P_WRITE64 = 0x100, // Generic register functions
     P_WRITE32,
@@ -107,8 +108,6 @@ typedef enum {
 #define S_OK     0
 #define S_BADCMD -1
 
-typedef u64(callfunc)(u64, u64, u64, u64);
-
 typedef struct {
     u64 opcode;
     u64 args[6];

src/usb.c

@@ -178,3 +178,14 @@ void usb_init(void)
         printf("USB%d: initialized at %p\n", i, iodev_usb[i].opaque);
     }
 }
+
+void usb_shutdown(void)
+{
+    for (int i = 0; i < USB_INSTANCES; i++) {
+        if (!iodev_usb[i].opaque)
+            continue;
+
+        printf("USB%d: shutdown\n", i);
+        usb_dwc3_shutdown(iodev_usb[i].opaque);
+    }
+}

src/usb.h

@@ -9,5 +9,6 @@
 dwc3_dev_t *usb_bringup(u32 idx);
 
 void usb_init(void);
+void usb_shutdown(void);
 
 #endif

src/utils.h

@@ -335,4 +335,13 @@ static inline int poll32(u64 addr, u32 mask, u32 target, u32 timeout)
     return -1;
 }
 
+typedef u64(generic_func)(u64, u64, u64, u64);
+
+struct vector_args {
+    generic_func *entry;
+    u64 args[4];
+};
+
+extern struct vector_args next_stage;
+
 #endif