main & co: Support upgrading m1n1 stage 1

This is finally useful now that we've run into and fixed NVMe timeout
related issues in stage 1.

Signed-off-by: Hector Martin <[email protected]>

Author: marcan

src/m1n1.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: MIT
+
+def build(src, dest, vars):
+    if isinstance(vars, (list, tuple)):
+        vars = b"".join(i.encode("ascii") + b"\n" for i in vars) + b"\0\0\0\0"
+
+    with open(src, "rb") as fd:
+        m1n1_data = fd.read()
+
+    with open(dest, "wb") as fd:
+        fd.write(m1n1_data + vars)
+
+def extract_vars(src):
+    with open(src, "rb") as fd:
+        m1n1_data = fd.read()
+
+    try:
+        vars = m1n1_data.split(b"STACKBOT")[1].split(b"\0")[0].decode("ascii")
+    except Exception:
+        return None
+
+    return [i for i in vars.split("\n") if i]
+
+def get_version(path):
+    data = open(path, "rb").read()
+    if b"##m1n1_ver##" in data:
+        return data.split(b"##m1n1_ver##")[1].split(b"\0")[0].decode("ascii")
+    else:
+        return None

src/main.py

@@ -3,7 +3,7 @@
 import os, os.path, shlex, subprocess, sys, time, termios, json, getpass
 from dataclasses import dataclass
 
-import system, osenum, stub, diskutil, osinstall, asahi_firmware
+import system, osenum, stub, diskutil, osinstall, asahi_firmware, m1n1
 from util import *
 
 PART_ALIGN = psize("1MiB")
@@ -89,6 +89,8 @@ def __init__(self):
         self.ipsw = None
         self.osins = None
         self.osi = None
+        self.m1n1 = "boot/m1n1.bin"
+        self.m1n1_ver = m1n1.get_version(self.m1n1)
 
     def input(self):
         self.flush_input()
@@ -307,30 +309,58 @@ def action_install_into_free(self, avail_free):
 
         self.do_install(os_size)
 
-    def action_resume(self, oses):
+    def action_resume_or_upgrade(self, oses, upgrade):
         choices = {str(i): f"{p.desc}\n      {str(o)}" for i, (p, o) in enumerate(oses)}
 
         if len(choices) > 1:
             print()
-            p_question("Choose an incomplete install to resume:")
+            if upgrade:
+                p_question("Choose an existing install to upgrade:")
+            else:
+                p_question("Choose an incomplete install to resume:")
             idx = self.choice("Installed OS", choices)
         else:
             idx = list(choices.keys())[0]
 
         self.part, osi = oses[int(idx)]
 
-        p_progress(f"Resuming installation into {self.part.name} ({self.part.label})")
+        if upgrade:
+            p_progress(f"Upgrading installation {self.part.name} ({self.part.label})")
+            p_info(f"  Old m1n1 stage 1 version: {osi.m1n1_ver}")
+            p_info(f"  New m1n1 stage 1 version: {self.m1n1_ver}")
+            print()
+        else:
+            p_progress(f"Resuming installation into {self.part.name} ({self.part.label})")
 
         self.ins = stub.StubInstaller(self.sysinfo, self.dutil, self.osinfo)
         if not self.ins.check_existing_install(osi):
-            p_error("The existing installation is missing files.")
-            p_message("This tool can only resume installations that completed the first")
-            p_message("stage of the installation process. If it was interrupted, please")
-            p_message("delete the partitions manually and reinstall from scratch.")
-            return True
+            op = "upgrade" if upgrade else "resume"
+            p_error(   "The existing installation is missing files.")
+            p_message(f"This tool can only {op} installations that completed the first")
+            p_message( "stage of the installation process. If it was interrupted, please")
+            p_message( "delete the partitions manually and reinstall from scratch.")
+            return False
 
         self.dutil.remount_rw(self.ins.osi.system)
 
+        if upgrade:
+            # Note: we get the vars out of the boot.bin in the system volume instead of the
+            # actual installed fuOS. This is arguably the better option, since it allows
+            # users to fix their install using this functionality if they messed up the boot
+            # object.
+            vars = m1n1.extract_vars(self.ins.boot_obj_path)
+            if vars is None:
+                p_error("Could not get variables from the installed m1n1")
+                p_message(f"Path: {self.ins.boot_obj_path}")
+                return False
+
+            p_progress(f"Transferring m1n1 variables:")
+            for v in vars:
+                p_info(f"  {v}")
+
+            print()
+            m1n1.build(self.m1n1, self.ins.boot_obj_path, vars)
+
         # Unhide the SystemVersion, if hidden
         self.ins.prepare_for_bless()
 
@@ -792,6 +822,7 @@ def main_loop(self):
         parts_empty_apfs = []
         parts_resizable = []
         oses_incomplete = []
+        oses_upgradable = []
 
         for i, p in enumerate(self.parts):
             if p.type in ("Apple_APFS_ISC",):
@@ -853,7 +884,9 @@ def main_loop(self):
                 else:
                     state += " "
                 p_plain(f"    OS: [{state}] {os}")
-                if os.stub and not (os.bp and os.bp.get("coih", None)):
+                if os.stub and os.m1n1_ver and os.m1n1_ver != self.m1n1_ver:
+                    oses_upgradable.append((p, os))
+                elif os.stub and not (os.bp and os.bp.get("coih", None)):
                     oses_incomplete.append((p, os))
 
         print()
@@ -879,8 +912,9 @@ def main_loop(self):
         if parts_resizable:
             actions["r"] = "Resize an existing partition to make space for a new OS"
             default = default or "r"
-        if self.sysinfo.boot_mode == "one true recoveryOS" and False:
-            actions["m"] = "Upgrade bootloader of an existing OS"
+        if oses_upgradable:
+            actions["m"] = "Upgrade m1n1 on an existing OS"
+            default = default or "m"
 
         if not actions:
             p_error("No actions available on this system.")
@@ -900,10 +934,9 @@ def main_loop(self):
         elif act == "r":
             return self.action_resize(parts_resizable)
         elif act == "m":
-            p_error("Unimplemented")
-            sys.exit(1)
+            return self.action_resume_or_upgrade(oses_upgradable, upgrade=True)
         elif act == "p":
-            return self.action_resume(oses_incomplete)
+            return self.action_resume_or_upgrade(oses_incomplete, upgrade=False)
         elif act == "q":
             return False
 

src/osenum.py

@@ -2,6 +2,7 @@
 import os, os.path, plistlib, subprocess, logging
 from dataclasses import dataclass
 
+import m1n1
 from util import *
 
 UUID_SROS = "3D3287DE-280D-4619-AAAB-D97469CA9C71"
@@ -204,9 +205,8 @@ def collect_os(self, part, volumes, vgid):
                                      osi.bp["nsih"],
                                      "System/Library/Caches/com.apple.kernelcaches",
                                      "kernelcache.custom." + coih)
-            fuos = open(fuos_path, "rb").read()
-            if b"##m1n1_ver##" in fuos:
-                osi.m1n1_ver = fuos.split(b"##m1n1_ver##")[1].split(b"\0")[0].decode("ascii")
+            osi.m1n1_ver = m1n1.get_version(fuos_path)
+            if osi.m1n1_ver:
                 logging.info(f"  m1n1 version found: {osi.m1n1_ver}")
 
         if b": Paired" in bps:

src/osinstall.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: MIT
 import os, shutil, sys, stat, subprocess, urlcache, zipfile, logging
 
+import m1n1
 from util import *
 
 class OSInstaller(PackageInstaller):
@@ -138,8 +139,6 @@ def install(self, boot_obj_path):
         next_object = self.template.get("next_object", None)
         logging.info(f"  Boot object: {boot_object}")
         logging.info(f"  Next object: {next_object}")
-        with open(os.path.join("boot", boot_object), "rb") as fd:
-            m1n1_data = fd.read()
 
         m1n1_vars = []
         if self.efi_part:
@@ -152,9 +151,6 @@ def install(self, boot_obj_path):
         for i in m1n1_vars:
             logging.info(f"  {i}")
 
-        m1n1_data += b"".join(i.encode("ascii") + b"\n" for i in m1n1_vars) + b"\0\0\0\0"
-
-        with open(boot_obj_path, "wb") as fd:
-            fd.write(m1n1_data)
+        m1n1.build(os.path.join("boot", boot_object), boot_obj_path, m1n1_vars)
 
         logging.info(f"Built boot object at {boot_obj_path}")