1 files changed, 148 insertions, 0 deletions
diff --git a/tools/src/chainload.c b/tools/src/chainload.c
new file mode 100644
index 0000000..1dd7c68
--- /dev/null
+++ b/tools/src/chainload.c
@@ -0,0 +1,148 @@
+/* SPDX-License-Identifier: MIT */
+
+#include "../build/build_cfg.h"
+
+#include "chainload.h"
+#include "adt.h"
+#include "malloc.h"
+#include "memory.h"
+#include "nvme.h"
+#include "string.h"
+#include "types.h"
+#include "utils.h"
+#include "xnuboot.h"
+
+#ifdef CHAINLOADING
+int rust_load_image(const char *spec, void **image, size_t *size);
+#endif
+
+extern u8 _chainload_stub_start[];
+extern u8 _chainload_stub_end[];
+
+int chainload_image(void *image, size_t size, char **vars, size_t var_cnt)
+{
+    u64 new_base = (u64)_base;
+    size_t image_size = size;
+
+    printf("chainload: Preparing image...\n");
+
+    // m1n1 variables
+    for (size_t i = 0; i < var_cnt; i++)
+        image_size += strlen(vars[i]) + 1;
+
+    // pad to end payload
+    image_size += 4;
+    image_size = ALIGN_UP(image_size, SZ_16K);
+
+    // SEPFW
+    size_t sepfw_off = image_size;
+
+    int anode = adt_path_offset(adt, "/chosen/memory-map");
+    if (anode < 0) {
+        printf("chainload: /chosen/memory-map not found\n");
+        return -1;
+    }
+    u64 sepfw[2];
+    if (ADT_GETPROP_ARRAY(adt, anode, "SEPFW", sepfw) < 0) {
+        printf("chainload: Failed to find SEPFW\n");
+        return -1;
+    }
+
+    image_size += sepfw[1];
+    image_size = ALIGN_UP(image_size, SZ_16K);
+
+    // Bootargs
+    size_t bootargs_off = image_size;
+    const size_t bootargs_size = SZ_16K;
+    image_size += bootargs_size;
+
+    printf("chainload: Total image size: 0x%lx\n", image_size);
+
+    size_t stub_size = _chainload_stub_end - _chainload_stub_start;
+
+    void *new_image = malloc(image_size + stub_size);
+
+    // Copy m1n1
+    memcpy(new_image, image, size);
+
+    // Add vars
+    u8 *p = new_image + size;
+    for (size_t i = 0; i < var_cnt; i++) {
+        size_t len = strlen(vars[i]);
+
+        memcpy(p, vars[i], len);
+        p[len] = '\n';
+        p += len + 1;
+    }
+
+    // Add end padding
+    memset(p, 0, 4);
+
+    // Copy SEPFW
+    memcpy(new_image + sepfw_off, (void *)sepfw[0], sepfw[1]);
+
+    // Adjust ADT SEPFW address
+    sepfw[0] = new_base + sepfw_off;
+    if (adt_setprop(adt, anode, "SEPFW", &sepfw, sizeof(sepfw)) < 0) {
+        printf("chainload: Failed to set SEPFW prop\n");
+        free(new_image);
+        return -1;
+    }
+
+    // Copy bootargs
+    struct boot_args *new_boot_args = new_image + bootargs_off;
+    *new_boot_args = cur_boot_args;
+    new_boot_args->top_of_kernel_data = new_base + image_size;
+
+    // Copy chainload stub
+    void *stub = new_image + image_size;
+    memcpy(stub, _chainload_stub_start, stub_size);
+    dc_cvau_range(stub, stub_size);
+    ic_ivau_range(stub, stub_size);
+
+    // Set up next stage
+    next_stage.entry = stub;
+    next_stage.args[0] = new_base + bootargs_off;
+    next_stage.args[1] = (u64)new_image;
+    next_stage.args[2] = new_base;
+    next_stage.args[3] = image_size;
+    next_stage.args[4] = new_base + 0x800; // m1n1 entrypoint
+    next_stage.restore_logo = false;
+
+    return 0;
+}
+
+#ifdef CHAINLOADING
+
+int chainload_load(const char *spec, char **vars, size_t var_cnt)
+{
+    void *image;
+    size_t size;
+    int ret;
+
+    if (!nvme_init()) {
+        printf("chainload: NVME init failed\n");
+        return -1;
+    }
+
+    ret = rust_load_image(spec, &image, &size);
+    nvme_shutdown();
+    if (ret < 0)
+        return ret;
+
+    return chainload_image(image, size, vars, var_cnt);
+}
+
+#else
+
+int chainload_load(const char *spec, char **vars, size_t var_cnt)
+{
+    UNUSED(spec);
+    UNUSED(vars);
+    UNUSED(var_cnt);
+
+    printf("Chainloading files not supported in this build!\n");
+    return -1;
+}
+
+#endif