chainload.py: Improve chainloading

- Support chainload-in-place
- Do Mach-O loading host-side
- Handle copying SEPFW (WIP, ADT not updated yet)

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

proxyclient/chainload.py

@@ -2,13 +2,16 @@
 
 import argparse, pathlib
 
-parser = argparse.ArgumentParser(description='(Linux) kernel loader for m1n1')
+parser = argparse.ArgumentParser(description='Mach-O loader for m1n1')
 parser.add_argument('payload', type=pathlib.Path)
 parser.add_argument('-1', '--el1', action="store_true")
+parser.add_argument('-s', '--sepfw', action="store_true")
 args = parser.parse_args()
 
 from setup import *
 from tgtypes import *
+import adt
+import asm
 
 payload = args.payload.read_bytes()
 
@@ -27,37 +30,92 @@ for cmd in obj.cmds:
 
 memory_size = vmax - vmin
 
-# FIXME: this will currently still waste the whole m1n1 size including payload area (64+MB) on each
-# chainload. The best way to fix this is to support in-place chainloading, which has other
-# advantages.
+image = bytearray(memory_size)
 
-try:
-    # Try to use the m1n1 heap to avoid wasting 128MB RAM on every load
-    new_base = p.memalign(0x10000, memory_size)
-except:
-    # Fall back to proxy heap, which will be at the right place in old versions
-    new_base = u.memalign(0x10000, memory_size)
+new_base = u.base
 
-for cmd in obj.cmds:
+def align(v, a=16384):
+    return (v + a - 1) & ~(a - 1)
+
+for cmdi, cmd in enumerate(obj.cmds):
+    is_m1n1 = None
     if cmd.cmd == LoadCmdType.SEGMENT_64:
-        if cmd.args.fileoff == 0:
-            continue # do not load mach-o headers, not needed for now
-        dest = cmd.args.vmaddr - vmin + new_base
+        if is_m1n1 is None:
+            is_m1n1 = cmd.args.segname == "_HDR"
+        dest = cmd.args.vmaddr - vmin
         end = min(len(payload), cmd.args.fileoff + cmd.args.filesize)
         size = end - cmd.args.fileoff
-        print("Loading %d bytes from 0x%x to 0x%x" % (size, cmd.args.fileoff, dest))
-        u.compressed_writemem(dest, payload[cmd.args.fileoff:end], True)
+        print(f"LOAD: {cmd.args.segname} {size} bytes from {cmd.args.fileoff:x} to {dest + new_base:x}")
+        image[dest:dest + size] = payload[cmd.args.fileoff:end]
         if cmd.args.vmsize > size:
             clearsize = cmd.args.vmsize - size
-            print("Zeroing %d bytes from 0x%x to 0x%x" % (clearsize, dest + size, dest + size + clearsize))
-            p.memset8(dest + size, 0, clearsize)
-
-p.dc_cvau(new_base, memory_size)
-p.ic_ivau(new_base, memory_size)
+            if cmd.args.segname == "PYLD":
+                print("SKIP: %d bytes from 0x%x to 0x%x" % (clearsize, dest + new_base + size, dest + new_base + size + clearsize))
+                memory_size -= clearsize - 4 # leave a payload end marker
+                image = image[:memory_size]
+            else:
+                print("ZERO: %d bytes from 0x%x to 0x%x" % (clearsize, dest + new_base + size, dest + new_base + size + clearsize))
+                image[dest + size:dest + cmd.args.vmsize] = bytes(clearsize)
 
 entry -= vmin
 entry += new_base
 
+if args.sepfw:
+    adt_base = u.ba.devtree - u.ba.virt_base + u.ba.phys_base
+    adt_size = u.ba.devtree_size
+    print(f"Fetching ADT ({adt_size} bytes)...")
+    adt = adt.load_adt(iface.readmem(adt_base, u.ba.devtree_size))
+
+    sepfw_start, sepfw_length = adt["chosen"]["memory-map"].SEPFW
+
+else:
+    sepfw_start, sepfw_length = 0, 0
+
+image_size = align(len(image))
+sepfw_off = image_size
+image_size += align(sepfw_length)
+bootargs_off = image_size
+image_size += 0x4000
+
+print(f"Total region size: 0x{image_size:x} bytes")
+image_addr = u.malloc(image_size)
+
+print(f"Loading kernel image (0x{len(image):x} bytes)...")
+u.compressed_writemem(image_addr, image, True)
+p.dc_cvau(image_addr, len(image))
+
+if args.sepfw:
+    print(f"Copying SEPFW (0x{sepfw_length:x} bytes)...")
+    p.memcpy8(image_addr + sepfw_off, sepfw_start, sepfw_length)
+
+print(f"Setting up bootargs...")
+tba = u.ba.copy()
+#tba.phys_base = new_base
+#tba.virt_base = 0xfffffe0010000000 + new_base & (32 * 1024 * 1024 - 1)
+tba.top_of_kdata = new_base + image_size
+
+iface.writemem(image_addr + bootargs_off, BootArgs.build(tba))
+
+print(f"Copying stub...")
+
+stub = asm.ARMAsm(f"""
+1:
+        ldp x4, x5, [x1], #8
+        stp x4, x5, [x2]
+        dc cvau, x2
+        ic ivau, x2
+        add x2, x2, #8
+        sub x3, x3, #8
+        cbnz x3, 1b
+
+        ldr x1, ={entry}
+        br x1
+""", image_addr + image_size)
+
+iface.writemem(stub.addr, stub.data)
+p.dc_cvau(stub.addr, stub.len)
+p.ic_ivau(stub.addr, stub.len)
+
 if args.el1:
     print("Setting up EL1 config")
 
@@ -67,14 +125,14 @@ if args.el1:
     # Unredirect IRQs/FIQs
     u.msr(HCR_EL2, u.mrs(HCR_EL2) & ~(3 << 3)) # ~(IMO | FMO)
 
-print("Jumping to 0x%x" % entry)
+print(f"Jumping to stub at 0x{stub.addr:x}")
 
 try:
     p.mmu_shutdown()
 except:
     pass
 
-p.reboot(entry, u.ba_addr, el1=args.el1)
+p.reboot(stub.addr, new_base + bootargs_off, image_addr, new_base, image_size)
 
 iface.nop()
 print("Proxy is alive again")