systemd-boot-builder.py: add support for profiles (#26318)

* systemd-boot-builder.py: add support for profiles

This will also list the generations of other profiles than `system` in
the boot menu. See the documentation of the `--profile-name` option of
nixos-rebuild for more information on profiles.

* Fix errors introduced by previous commit

1 files changed, 55 insertions, 30 deletions

diff --git a/nixos/modules/system/boot/loader/systemd-boot/systemd-boot-builder.py b/nixos/modules/system/boot/loader/systemd-boot/systemd-boot-builder.py
index cb2a17f18f48..779005c0df52 100644
--- a/nixos/modules/system/boot/loader/systemd-boot/systemd-boot-builder.py
+++ b/nixos/modules/system/boot/loader/systemd-boot/systemd-boot-builder.py
@@ -11,35 +11,42 @@ import errno
 import warnings
 import ctypes
 libc = ctypes.CDLL("libc.so.6")
+import re
 
 def copy_if_not_exists(source, dest):
     if not os.path.exists(dest):
         shutil.copyfile(source, dest)
 
-def system_dir(generation):
-    return "/nix/var/nix/profiles/system-%d-link" % (generation)
+def system_dir(profile, generation):
+    if profile:
+        return "/nix/var/nix/profiles/system-profiles/%s-%d-link" % (profile, generation)
+    else:
+        return "/nix/var/nix/profiles/system-%d-link" % (generation)
 
-BOOT_ENTRY = """title NixOS
+BOOT_ENTRY = """title NixOS{profile}
 version Generation {generation}
 linux {kernel}
 initrd {initrd}
 options {kernel_params}
 """
 
-def write_loader_conf(generation):
+def write_loader_conf(profile, generation):
     with open("@efiSysMountPoint@/loader/loader.conf.tmp", 'w') as f:
         if "@timeout@" != "":
             f.write("timeout @timeout@\n")
-        f.write("default nixos-generation-%d\n" % generation)
+        if profile:
+            f.write("default nixos-%s-generation-%d\n" % (profile, generation))
+        else:
+            f.write("default nixos-generation-%d\n" % (generation))
         if not @editor@:
             f.write("editor 0");
     os.rename("@efiSysMountPoint@/loader/loader.conf.tmp", "@efiSysMountPoint@/loader/loader.conf")
 
-def profile_path(generation, name):
-    return os.readlink("%s/%s" % (system_dir(generation), name))
+def profile_path(profile, generation, name):
+    return os.readlink("%s/%s" % (system_dir(profile, generation), name))
 
-def copy_from_profile(generation, name, dry_run=False):
-    store_file_path = profile_path(generation, name)
+def copy_from_profile(profile, generation, name, dry_run=False):
+    store_file_path = profile_path(profile, generation, name)
     suffix = os.path.basename(store_file_path)
     store_dir = os.path.basename(os.path.dirname(store_file_path))
     efi_file_path = "/efi/nixos/%s-%s.efi" % (store_dir, suffix)
@@ -47,22 +54,26 @@ def copy_from_profile(generation, name, dry_run=False):
         copy_if_not_exists(store_file_path, "@efiSysMountPoint@%s" % (efi_file_path))
     return efi_file_path
 
-def write_entry(generation, machine_id):
-    kernel = copy_from_profile(generation, "kernel")
-    initrd = copy_from_profile(generation, "initrd")
+def write_entry(profile, generation, machine_id):
+    kernel = copy_from_profile(profile, generation, "kernel")
+    initrd = copy_from_profile(profile, generation, "initrd")
     try:
-        append_initrd_secrets = profile_path(generation, "append-initrd-secrets")
+        append_initrd_secrets = profile_path(profile, generation, "append-initrd-secrets")
         subprocess.check_call([append_initrd_secrets, "@efiSysMountPoint@%s" % (initrd)])
     except FileNotFoundError:
         pass
-    entry_file = "@efiSysMountPoint@/loader/entries/nixos-generation-%d.conf" % (generation)
-    generation_dir = os.readlink(system_dir(generation))
+    if profile:
+        entry_file = "@efiSysMountPoint@/loader/entries/nixos-%s-generation-%d.conf" % (profile, generation)
+    else:
+        entry_file = "@efiSysMountPoint@/loader/entries/nixos-generation-%d.conf" % (generation)
+    generation_dir = os.readlink(system_dir(profile, generation))
     tmp_path = "%s.tmp" % (entry_file)
     kernel_params = "systemConfig=%s init=%s/init " % (generation_dir, generation_dir)
     with open("%s/kernel-params" % (generation_dir)) as params_file:
         kernel_params = kernel_params + params_file.read()
     with open(tmp_path, 'w') as f:
-        f.write(BOOT_ENTRY.format(generation=generation,
+        f.write(BOOT_ENTRY.format(profile=" [" + profile + "]" if profile else "",
+                    generation=generation,
                     kernel=kernel,
                     initrd=initrd,
                     kernel_params=kernel_params))
@@ -77,29 +88,33 @@ def mkdir_p(path):
         if e.errno != errno.EEXIST or not os.path.isdir(path):
             raise
 
-def get_generations(profile):
+def get_generations(profile=None):
     gen_list = subprocess.check_output([
         "@nix@/bin/nix-env",
         "--list-generations",
         "-p",
-        "/nix/var/nix/profiles/%s" % (profile),
+        "/nix/var/nix/profiles/%s" % ("system-profiles/" + profile if profile else "system"),
         "--option", "build-users-group", ""],
         universal_newlines=True)
     gen_lines = gen_list.split('\n')
     gen_lines.pop()
-    return [ int(line.split()[0]) for line in gen_lines ]
+    return [ (profile, int(line.split()[0])) for line in gen_lines ]
 
 def remove_old_entries(gens):
-    slice_start = len("@efiSysMountPoint@/loader/entries/nixos-generation-")
-    slice_end = -1 * len(".conf")
+    rex_profile = re.compile("^@efiSysMountPoint@/loader/entries/nixos-(.*)-generation-.*\.conf$")
+    rex_generation = re.compile("^@efiSysMountPoint@/loader/entries/nixos.*-generation-(.*)\.conf$")
     known_paths = []
     for gen in gens:
-        known_paths.append(copy_from_profile(gen, "kernel", True))
-        known_paths.append(copy_from_profile(gen, "initrd", True))
-    for path in glob.iglob("@efiSysMountPoint@/loader/entries/nixos-generation-[1-9]*.conf"):
+        known_paths.append(copy_from_profile(*gen, "kernel", True))
+        known_paths.append(copy_from_profile(*gen, "initrd", True))
+    for path in glob.iglob("@efiSysMountPoint@/loader/entries/nixos*-generation-[1-9]*.conf"):
         try:
-            gen = int(path[slice_start:slice_end])
-            if not gen in gens:
+            if rex_profile.match(path):
+                prof = rex_profile.sub(r"\1", path)
+            else:
+                prof = "system"
+            gen = int(rex_generation.sub(r"\1", path))
+            if not (prof, gen) in gens:
                 os.unlink(path)
         except ValueError:
             pass
@@ -107,6 +122,14 @@ def remove_old_entries(gens):
         if not path in known_paths:
             os.unlink(path)
 
+def get_profiles():
+    if os.path.isdir("/nix/var/nix/profiles/system-profiles/"):
+        return [x
+            for x in os.listdir("/nix/var/nix/profiles/system-profiles/")
+            if not x.endswith("-link")]
+    else:
+        return []
+
 def main():
     parser = argparse.ArgumentParser(description='Update NixOS-related systemd-boot files')
     parser.add_argument('default_config', metavar='DEFAULT-CONFIG', help='The default NixOS config to boot')
@@ -141,12 +164,14 @@ def main():
     mkdir_p("@efiSysMountPoint@/efi/nixos")
     mkdir_p("@efiSysMountPoint@/loader/entries")
 
-    gens = get_generations("system")
+    gens = get_generations()
+    for profile in get_profiles():
+        gens += get_generations(profile)
     remove_old_entries(gens)
     for gen in gens:
-        write_entry(gen, machine_id)
-        if os.readlink(system_dir(gen)) == args.default_config:
-            write_loader_conf(gen)
+        write_entry(*gen, machine_id)
+        if os.readlink(system_dir(*gen)) == args.default_config:
+            write_loader_conf(*gen)
 
     # Since fat32 provides little recovery facilities after a crash,
     # it can leave the system in an unbootable state, when a crash/outage