Allow nixos-rebuild boot to work with automated firmware updates

modules/flash-script.nix

@@ -9,6 +9,49 @@ 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. We check the exact
+      # path that we populate because it is possible for multiple capsule
+      # updates to be applied at once, so we don't want other files in this
+      # directory to influence our behavior.
+      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
+
+      # Jetpack 5.0 didn't expose this DMI variable,
+      if [[ ! -f /sys/devices/virtual/dmi/id/bios_version ]]; then
+        echo "Unable to determine current Jetson firmware version."
+        echo "You should reflash the firmware with the new version to ensure compatibility"
+        exit 1
+      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,19 +505,19 @@ 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
@@ -496,23 +539,7 @@ in
             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
+          ${lib.getExe updateFirmware}
         '';
     };
 

overlay-with-config.nix

@@ -34,6 +34,10 @@ final: prev: (
         else if lib.hasPrefix "xavier-" cfg.som then "0x19"
         else throw "Unknown SoC type";
 
+      otaUtils = prevJetpack.otaUtils.override {
+        inherit (config.boot.loader.efi) efiSysMountPoint;
+      };
+
       uefi-firmware = prevJetpack.uefi-firmware.override ({
         bootLogo = cfg.firmware.uefi.logo;
         debugMode = cfg.firmware.uefi.debugMode;

pkgs/ota-utils/default.nix

@@ -1,4 +1,4 @@
-{ lib, stdenvNoCC, bash, util-linux, e2fsprogs, tegra-eeprom-tool, l4tVersion }:
+{ lib, stdenvNoCC, bash, util-linux, e2fsprogs, tegra-eeprom-tool, l4tVersion, efiSysMountPoint ? "/boot" }:
 
 stdenvNoCC.mkDerivation {
   name = "ota-utils";
@@ -9,6 +9,10 @@ stdenvNoCC.mkDerivation {
   dontConfigure = true;
   dontBuild = true;
 
+  env = { inherit efiSysMountPoint; };
+
+  buildInputs = [ bash ];
+
   installPhase = ''
     runHook preInstall
 
@@ -19,6 +23,10 @@ stdenvNoCC.mkDerivation {
     cp ${./ota_helpers.func} $out/share/ota_helpers.func
     chmod +x $out/bin/ota-setup-efivars $out/bin/ota-apply-capsule-update $out/bin/ota-check-firmware
 
+    for path in $out/bin/ota-apply-capsule-update $out/share/ota_helpers.func; do
+      substituteInPlace "$path" --subst-var efiSysMountPoint
+    done
+
     for fname in ota-setup-efivars ota-apply-capsule-update; do
       substituteInPlace $out/bin/$fname \
         --replace "@ota_helpers@" "$out/share/ota_helpers.func"
@@ -28,8 +36,7 @@ stdenvNoCC.mkDerivation {
     substituteInPlace $out/bin/ota-check-firmware \
       --replace "@l4tVersion@" "${l4tVersion}"
 
-    # patchShebangs does not seem to work here for some reason
-    substituteInPlace $out/bin/* --replace '#!/usr/bin/env bash' '#!${bash}/bin/bash'
+    patchShebangs --host $out/bin
 
     runHook postInstall
   '';

pkgs/ota-utils/ota-apply-capsule-update.sh

@@ -8,17 +8,18 @@ capsuleFile=$1
 boardspec=$(tegra-boardspec)
 detect_can_write_runtime_uefi_vars "$boardspec"
 
-# Check for /boot being an ESP. On Xavier AGX, even though the efi vars need to
-# be written to an ESP on the emmc, capsule updates can still be written to an
-# ESP partition at /boot on other devices (e.g. nvme)
-if ! mountpoint -q /boot; then
-    echo "/boot is not mounted"
-    exit 1
+# Check for @efiSysMountPoint@ (defaults to /boot) being an ESP. On Xavier AGX,
+# even though the efi vars need to be written to an ESP on the emmc, capsule
+# updates can still be written to an ESP partition at @efiSysMountPoint@ on
+# other devices (e.g. nvme)
+if ! mountpoint -q @efiSysMountPoint@; then
+	echo "@efiSysMountPoint@ is not mounted"
+	exit 1
 fi
 
-mkdir -p /boot/EFI/UpdateCapsule
-cp "$capsuleFile" /boot/EFI/UpdateCapsule/TEGRA_BL.Cap
-sync /boot/EFI/UpdateCapsule/TEGRA_BL.Cap
+mkdir -p @efiSysMountPoint@/EFI/UpdateCapsule
+cp "$capsuleFile" @efiSysMountPoint@/EFI/UpdateCapsule/TEGRA_BL.Cap
+sync @efiSysMountPoint@/EFI/UpdateCapsule/TEGRA_BL.Cap
 
 set_efi_var OsIndications-8be4df61-93ca-11d2-aa0d-00e098032b8c "\x07\x00\x00\x00\x04\x00\x00\x00\x00\x00\x00\x00"
 

pkgs/ota-utils/ota_helpers.func

@@ -89,7 +89,7 @@ detect_can_write_runtime_uefi_vars() {
         noRuntimeUefiWrites=true
         espDir=/opt/nvidia/esp
     else
-        espDir=/boot
+        espDir=@efiSysMountPoint@
     fi
 }