diff --git a/pkgs/os-specific/linux/nixos-rebuild/nixos-rebuild.sh b/pkgs/os-specific/linux/nixos-rebuild/nixos-rebuild.sh
index 695f173da8fc..795bedb99d4d 100755
--- a/pkgs/os-specific/linux/nixos-rebuild/nixos-rebuild.sh
+++ b/pkgs/os-specific/linux/nixos-rebuild/nixos-rebuild.sh
@@ -352,16 +352,6 @@ if [[ -z $flake && -e /etc/nixos/flake.nix && -z $noFlake ]]; then
     flake="$(dirname "$(readlink -f /etc/nixos/flake.nix)")"
 fi
 
-# Re-execute nixos-rebuild from the Nixpkgs tree.
-# FIXME: get nixos-rebuild from $flake.
-if [[ -z $_NIXOS_REBUILD_REEXEC && -n $canRun && -z $fast && -z $flake ]]; then
-    if p=$(runCmd nix-build --no-out-link --expr 'with import <nixpkgs/nixos> {}; config.system.build.nixos-rebuild' "${extraBuildFlags[@]}"); then
-        export _NIXOS_REBUILD_REEXEC=1
-        runCmd exec "$p/bin/nixos-rebuild" "${origArgs[@]}"
-        exit 1
-    fi
-fi
-
 # For convenience, use the hostname as the default configuration to
 # build from the flake.
 if [[ -n $flake ]]; then
@@ -380,6 +370,40 @@ if [[ -n $flake ]]; then
     fi
 fi
 
+
+tmpDir=$(mktemp -t -d nixos-rebuild.XXXXXX)
+
+cleanup() {
+    for ctrl in "$tmpDir"/ssh-*; do
+        ssh -o ControlPath="$ctrl" -O exit dummyhost 2>/dev/null || true
+    done
+    rm -rf "$tmpDir"
+}
+trap cleanup EXIT
+
+
+# Re-execute nixos-rebuild from the Nixpkgs tree.
+if [[ -z $_NIXOS_REBUILD_REEXEC && -n $canRun && -z $fast ]]; then
+    if [[ -z $flake ]]; then
+        if p=$(runCmd nix-build --no-out-link --expr 'with import <nixpkgs/nixos> {}; config.system.build.nixos-rebuild' "${extraBuildFlags[@]}"); then
+            SHOULD_REEXEC=1
+        fi
+    else
+        runCmd nix "${flakeFlags[@]}" build --out-link "${tmpDir}/nixos-rebuild" "$flake#$flakeAttr.config.system.build.nixos-rebuild" "${extraBuildFlags[@]}" "${lockFlags[@]}"
+        if p=$(readlink -e "${tmpDir}/nixos-rebuild"); then
+            SHOULD_REEXEC=1
+        fi
+    fi
+
+    if [[ -n $SHOULD_REEXEC ]]; then
+        export _NIXOS_REBUILD_REEXEC=1
+        # Manually call cleanup as the EXIT trap is not triggered when using exec
+        cleanup
+        runCmd exec "$p/bin/nixos-rebuild" "${origArgs[@]}"
+        exit 1
+    fi
+fi
+
 # Find configuration.nix and open editor instead of building.
 if [ "$action" = edit ]; then
     if [[ -z $flake ]]; then
@@ -394,19 +418,8 @@ if [ "$action" = edit ]; then
     exit 1
 fi
 
-
-tmpDir=$(mktemp -t -d nixos-rebuild.XXXXXX)
 SSHOPTS="$NIX_SSHOPTS -o ControlMaster=auto -o ControlPath=$tmpDir/ssh-%n -o ControlPersist=60"
 
-cleanup() {
-    for ctrl in "$tmpDir"/ssh-*; do
-        ssh -o ControlPath="$ctrl" -O exit dummyhost 2>/dev/null || true
-    done
-    rm -rf "$tmpDir"
-}
-trap cleanup EXIT
-
-
 # First build Nix, since NixOS may require a newer version than the
 # current one.
 if [[ -n "$rollback" || "$action" = dry-build ]]; then