From 4e558812ac59a0cc3a72d4d1c36473b81a308eeb Mon Sep 17 00:00:00 2001
From: Robert Hensing <robert@roberthensing.nl>
Date: Wed, 25 Sep 2024 11:12:15 +0200
Subject: [PATCH] nixos-rebuild: Call toplevel/bin/apply when available

By executing the whole switching process using one script, we avoid
some ssh roundtrips that would slow down deployment swith --target-host.

For the other benefits, see https://github.com/NixOS/nixpkgs/issues/266290
---
 .../linux/nixos-rebuild/nixos-rebuild.sh      | 53 +++++++++++++++++--
 1 file changed, 48 insertions(+), 5 deletions(-)

diff --git a/pkgs/os-specific/linux/nixos-rebuild/nixos-rebuild.sh b/pkgs/os-specific/linux/nixos-rebuild/nixos-rebuild.sh
index dc5ed4b2d02c..6236c923c8a9 100755
--- a/pkgs/os-specific/linux/nixos-rebuild/nixos-rebuild.sh
+++ b/pkgs/os-specific/linux/nixos-rebuild/nixos-rebuild.sh
@@ -790,7 +790,6 @@ if [ -z "$rollback" ]; then
             pathToConfig="$(nixFlakeBuild "$flake#$flakeAttr.config.system.build.toplevel" "${extraBuildFlags[@]}" "${lockFlags[@]}")"
         fi
         copyToTarget "$pathToConfig"
-        targetHostSudoCmd nix-env -p "$profile" --set "$pathToConfig"
     elif [[ "$action" = test || "$action" = build || "$action" = dry-build || "$action" = dry-activate ]]; then
         if [[ -z $buildingAttribute ]]; then
             pathToConfig="$(nixBuild $buildFile -A "${attr:+$attr.}config.system.build.toplevel" "${extraBuildFlags[@]}")"
@@ -841,12 +840,56 @@ else # [ -n "$rollback" ]
 fi
 
 
+hasApplyScript=
+# If we're doing a deployment-like action, we need to know whether the config has
+# an apply script. NixOS versions >= 24.11 should be deployed with toplevel/bin/apply.
+if [[ "$action" = switch || "$action" = boot || "$action" = test || "$action" = dry-activate ]]; then
+    hasApplyScriptOut="$(targetHostCmd sh -c "if test -e $pathToConfig/bin/apply; then echo __has-apply-script__; elif test -e $pathToConfig/bin; then echo __has-no-apply-script__; else echo $pathToConfig is gone; fi
+    ")"
+    # SSH can be messy (e.g. when user has a shell rc file that prints to stdout)
+    # So we only check for the substring
+    case "$hasApplyScriptOut" in
+        *__has-apply-script__*)
+            hasApplyScript=1
+            ;;
+        *__has-no-apply-script__*)
+            hasApplyScript=
+            ;;
+        *)
+            # Unlikely
+            echo "$hasApplyScriptOut" 1>&2
+            log "error: $pathToConfig could not be read"
+            exit 1
+            ;;
+    esac
+fi
+
+# switch|boot|test|dry-activate
+#
 # If we're not just building, then make the new configuration the boot
 # default and/or activate it now.
-if [[ "$action" = switch || "$action" = boot || "$action" = test || "$action" = dry-activate ]]; then
-    # Using systemd-run here to protect against PTY failures/network
-    # disconnections during rebuild.
-    # See: https://github.com/NixOS/nixpkgs/issues/39118
+if [[ -n "$hasApplyScript" ]] \
+    && [[ "$action" = switch || "$action" = boot || "$action" = test || "$action" = dry-activate ]]; then
+    cmd=("$pathToConfig/bin/apply" "$action" "--profile" "$profile")
+    if [[ -n "$specialisation" ]]; then
+        cmd+=("--specialisation" "$specialisation")
+    fi
+    if [[ -n "$installBootloader" ]]; then
+        cmd+=("--install-bootloader")
+    fi
+    targetHostSudoCmd "${cmd[@]}"
+elif [[ "$action" = switch || "$action" = boot || "$action" = test || "$action" = dry-activate ]]; then
+    # Legacy, without apply script, NixOS < 24.11
+
+    if [[ "$action" = switch || "$action" = boot ]]; then
+        if [[ -z "$rollback" ]]; then
+            : # We've already switched it so that hasApplyScript would check the right $pathToConfig
+        else
+            targetHostSudoCmd nix-env -p "$profile" --set "$pathToConfig"
+        fi
+    fi
+
+    # Legacy logic to support deploying NixOS <24.11; see hasApplyScript
     cmd=(
         "systemd-run"
         "-E" "LOCALE_ARCHIVE" # Will be set to new value early in switch-to-configuration script, but interpreter starts out with old value