Allow nixos-rebuild boot to work with automated firmware updates

In 9cea389, we broke the ability for
using `nixos-rebuild boot` to apply capsule updates. This change adds
that ability back. The systemd service is kept as we still want capsule
updates to be applied in cases where a different bootloader is being
used (or perhaps some other niche setup).

modules/flash-script.nix

@@ -9,6 +9,39 @@ let
     types;
 
   cfg = config.hardware.nvidia-jetpack;
+
+  canUpdateFirmware = cfg.firmware.autoUpdate && cfg.som != null && cfg.flashScriptOverrides.targetBoard != null;
+
+  updateFirmware = pkgs.writeShellApplication {
+    name = "update-jetson-firmware";
+    runtimeInputs = [ pkgs.nvidia-jetpack.otaUtils ];
+    text = ''
+      # This directory is populated by ota-apply-capsule-update, don't run
+      # if we already have a capsule update present on the ESP.
+      if [[ -e ${config.boot.loader.efi.efiSysMountPoint}/EFI/UpdateCapsule/TEGRA_BL.Cap ]]; then
+        echo "Existing capsule update for platform firmware exists, exiting"
+        exit 0
+      fi
+
+      CUR_VER=$(cat /sys/devices/virtual/dmi/id/bios_version)
+      NEW_VER=${pkgs.nvidia-jetpack.l4tVersion}
+
+      if [[ "$CUR_VER" != "$NEW_VER" ]]; then
+        echo "Current Jetson firmware version is: $CUR_VER"
+        echo "New Jetson firmware version is: $NEW_VER"
+        echo
+
+        # Set efi vars here as well as in systemd service, in case we're
+        # upgrading from an older nixos generation that doesn't have the
+        # systemd service. Plus, this ota-setup-efivars will be from the
+        # generation we're switching to, which can contain additional
+        # fixes/improvements.
+        ota-setup-efivars ${cfg.flashScriptOverrides.targetBoard}
+
+        ota-apply-capsule-update ${pkgs.nvidia-jetpack.uefiCapsuleUpdate}
+      fi
+    '';
+  };
 in
 {
   imports = with lib; [
@@ -462,58 +495,20 @@ in
       serviceConfig.ExecStart = "${pkgs.nvidia-jetpack.otaUtils}/bin/ota-setup-efivars ${cfg.flashScriptOverrides.targetBoard}";
     };
 
-    systemd.services.firmware-update = lib.mkIf (cfg.firmware.autoUpdate && cfg.som != null && cfg.flashScriptOverrides.targetBoard != null) {
+    # Include the capsule-on-disk firmware update method with the bootloader
+    # installation process so that firmware updates work with "nixos-rebuild boot".
+    boot.loader = lib.mkIf canUpdateFirmware {
+      systemd-boot.extraInstallCommands = lib.getExe updateFirmware;
+      grub.extraInstallCommands = lib.getExe updateFirmware;
+    };
+
+    systemd.services.firmware-update = lib.mkIf canUpdateFirmware {
       wantedBy = [ "multi-user.target" ];
-      path = [ pkgs.nvidia-jetpack.otaUtils ];
       after = [
         "${utils.escapeSystemdPath config.boot.loader.efi.efiSysMountPoint}.mount"
         "opt-nvidia-esp.mount"
       ];
-      unitConfig = {
-        ConditionPathExists = "/sys/devices/virtual/dmi/id/bios_version";
-        # This directory is populated by ota-apply-capsule-update, don't run
-        # if we already have a capsule update present on the ESP.
-        ConditionDirectoryNotEmpty = "!${config.boot.loader.efi.efiSysMountPoint}/EFI/UpdateCapsule";
-      };
-      script =
-        # NOTE: Our intention is to not apply any capsule update if the
-        # user's intention is to "test" a new nixos config without having it
-        # persist across reboots. "nixos-rebuild test" does not append a new
-        # generation to /nix/var/nix/profiles for the system profile, so we
-        # can compare that symlink to /run/current-system to see if our
-        # current active config has been persisted as a generation. Note that
-        # this check _may_ break down if not using nixos-rebuild and using
-        # switch-to-configuration directly, however it is well-documented
-        # that a user would need to self-manage their system profile's
-        # generations if switching a system in that manner.
-        lib.optionalString config.system.switch.enable ''
-          if [[ -L /nix/var/nix/profiles/system ]]; then
-            latest_generation=$(readlink -f /nix/var/nix/profiles/system)
-            current_system=$(readlink -f /run/current-system)
-            if [[ $latest_generation == /nix/store* ]] && [[ $latest_generation != "$current_system" ]]; then
-              echo "Skipping capsule update, current active system not persisted to /nix/var/nix/profiles/system"
-              exit 0
-            fi
-          fi
-        '' + ''
-          CUR_VER=$(cat /sys/devices/virtual/dmi/id/bios_version)
-          NEW_VER=${pkgs.nvidia-jetpack.l4tVersion}
-
-          if [[ "$CUR_VER" != "$NEW_VER" ]]; then
-            echo "Current Jetson firmware version is: $CUR_VER"
-            echo "New Jetson firmware version is: $NEW_VER"
-            echo
-
-            # Set efi vars here as well as in systemd service, in case we're
-            # upgrading from an older nixos generation that doesn't have the
-            # systemd service. Plus, this ota-setup-efivars will be from the
-            # generation we're switching to, which can contain additional
-            # fixes/improvements.
-            ota-setup-efivars ${cfg.flashScriptOverrides.targetBoard}
-
-            ota-apply-capsule-update ${pkgs.nvidia-jetpack.uefiCapsuleUpdate}
-          fi
-        '';
+      serviceConfig.ExecStart = lib.getExe updateFirmware;
     };
 
     environment.systemPackages = lib.mkIf (cfg.firmware.autoUpdate && cfg.som != null && cfg.flashScriptOverrides.targetBoard != null) [