feat(apply): magic rollback

.gitignore

@@ -2,8 +2,9 @@
 /result*
 /.direnv/
 
-# Binary destination from `make`
+# Binary destinations from `make`
 /nixos
+/activation-supervisor
 
 # For bubbletea log files
 *.log

Makefile

@@ -1,4 +1,5 @@
 APP_NAME := nixos
+SUPERVISOR_NAME := activation-supervisor
 BUILD_VAR_PKG := github.com/nix-community/nixos-cli/internal/build/vars
 
 VERSION ?= $(shell git describe --tags --always)
@@ -19,15 +20,23 @@ CGO_ENABLED ?= 0
 all: build
 
 .PHONY: build
-build:
+build: main supervisor
+
+.PHONY: main
+main:
 	@echo "building $(APP_NAME)..."
 	CGO_ENABLED=$(CGO_ENABLED) go build -o ./$(APP_NAME) -ldflags="$(LDFLAGS)" .
 
+.PHONY: supervisor
+supervisor:
+	@echo "building $(SUPERVISOR_NAME)..."
+	CGO_ENABLED=$(CGO_ENABLED) go build -o ./$(SUPERVISOR_NAME) -ldflags="$(LDFLAGS)" ./supervisor/
+
 .PHONY: clean
 clean:
 	@echo "cleaning up..."
 	go clean
-	rm -rf site/ man/
+	rm -rf ./$(APP_NAME), ./$(SUPERVISOR_NAME) site/ man/
 
 .PHONY: check
 check:

cmd/activate/run.go

@@ -365,9 +365,8 @@ func activateMain(cmd *cobra.Command, opts *cmdOpts.ActivateOpts) error {
 		return err
 	}
 
-	err = os.MkdirAll("/run/nixos", 0o755)
-	if err != nil {
-		log.Errorf("failed to create /run/nixos: %s", err)
+	if err := activation.EnsureActivationDirectoryExists(); err != nil {
+		log.Error(err)
 		return err
 	}
 

cmd/apply/apply.go

@@ -615,16 +615,38 @@ func applyMain(cmd *cobra.Command, opts *cmdOpts.ApplyOpts) error {
 		panic("unknown switch to configuration action to take, this is a bug")
 	}
 
-	err = activation.SwitchToConfiguration(targetHost, resultLocation, stcAction, &activation.SwitchToConfigurationOptions{
-		InstallBootloader: opts.InstallBootloader,
-		Specialisation:    specialisation,
-		UseRootCommand:    opts.RemoteRoot && targetHost.IsRemote(),
-		RootCommand:       cfg.RootCommand,
-	})
-	if err != nil {
-		rollbackProfile = true
-		log.Errorf("failed to switch to configuration: %v", err)
-		return err
+	useActivationSupervisor := shouldUseActivationSupervisor(cfg, targetHost, stcAction, resultLocation)
+	useRemoteRootForActivation := opts.RemoteRoot && targetHost.IsRemote()
+
+	if useActivationSupervisor {
+		// Let the supervisor handle the rollback if it exists.
+		rollbackProfile = false
+		err = activation.RunActivationSupervisor(targetHost, resultLocation, stcAction, &activation.RunActivationSupervisorOptions{
+			ProfileName:       opts.ProfileName,
+			InstallBootloader: opts.InstallBootloader,
+			Specialisation:    specialisation,
+			UseRootCommand:    useRemoteRootForActivation,
+			RootCommand:       cfg.RootCommand,
+		})
+		if err != nil {
+			log.Error("", err)
+			log.Warn("the target system should roll back soon")
+			return err
+		}
+	} else {
+		// Otherwise, just use the switch-to-configuration script directly
+		// and handle profile rollback ourselves.
+		err = activation.SwitchToConfiguration(targetHost, resultLocation, stcAction, &activation.SwitchToConfigurationOptions{
+			InstallBootloader: opts.InstallBootloader,
+			Specialisation:    specialisation,
+			UseRootCommand:    useRemoteRootForActivation,
+			RootCommand:       cfg.RootCommand,
+		})
+		if err != nil {
+			rollbackProfile = true
+			log.Errorf("failed to switch to configuration: %v", err)
+			return err
+		}
 	}
 
 	return nil
@@ -793,3 +815,19 @@ func getImageName(
 
 	return strings.TrimSpace(stdout.String()), nil
 }
+
+func shouldUseActivationSupervisor(cfg *settings.Settings, host system.System, action activation.SwitchToConfigurationAction, resultLocation string) bool {
+	if !cfg.AutoRollback || !host.IsRemote() {
+		return false
+	}
+
+	validAction := action == activation.SwitchToConfigurationActionBoot ||
+		action == activation.SwitchToConfigurationActionSwitch ||
+		action == activation.SwitchToConfigurationActionTest
+
+	if !validAction {
+		return false
+	}
+
+	return host.HasCommand(filepath.Join(resultLocation, "bin", "activation-supervisor"))
+}

internal/activation/activation.go

@@ -6,6 +6,7 @@ import (
 	"path/filepath"
 	"regexp"
 	"strconv"
+	"strings"
 
 	"github.com/nix-community/nixos-cli/internal/constants"
 	"github.com/nix-community/nixos-cli/internal/generation"
@@ -14,6 +15,10 @@ import (
 	"github.com/nix-community/nixos-cli/internal/system"
 )
 
+const (
+	RUNNING_ACTIVATION_SUPERVISOR = "NIXOS_CLI_RUNNING_ACTIVATION_SUPERVISOR"
+)
+
 // Parse the generation's `nixos-cli` configuration to find the default specialisation
 // for that generation.
 func FindDefaultSpecialisationFromConfig(s system.System, generationDirname string) (string, error) {
@@ -227,3 +232,115 @@ func SwitchToConfiguration(s system.CommandRunner, generationLocation string, ac
 	_, err := s.Run(cmd)
 	return err
 }
+
+// Create an activation trigger path name from a NixOS system
+// closure's location.
+//
+// Used for remote activation.
+func MakeActivationTriggerPath(systemLocation string) string {
+	// Obtain the cryptographic hash + nixos system closure name
+	basename := filepath.Base(systemLocation)
+
+	hash, _, found := strings.Cut(basename, "-")
+	if !found {
+		hash = strconv.Itoa(len(basename))
+	}
+
+	return filepath.Join(constants.NixOSActivationDirectory, "trigger", hash)
+}
+
+// Create the activation runtime directories with the required
+// structure.
+//
+// This creates the trigger directory as sticky, in case non-root users
+// need to activate things.
+func EnsureActivationDirectoryExists() error {
+	err := os.MkdirAll(constants.NixOSActivationDirectory, 0o755)
+	if err != nil {
+		return fmt.Errorf("failed to create %s: %s", constants.NixOSActivationDirectory, err)
+	}
+
+	triggerDirectory := filepath.Join(constants.NixOSActivationDirectory, "trigger")
+
+	err = os.MkdirAll(triggerDirectory, 0o1755)
+	if err != nil {
+		return fmt.Errorf("failed to create %s: %s", triggerDirectory, err)
+	}
+
+	return nil
+}
+
+type RunActivationSupervisorOptions struct {
+	ProfileName       string
+	InstallBootloader bool
+	Specialisation    string
+	UseRootCommand    bool
+	RootCommand       string
+}
+
+func RunActivationSupervisor(
+	s system.System,
+	generationLocation string,
+	action SwitchToConfigurationAction,
+	opts *RunActivationSupervisorOptions,
+) error {
+	log := s.Logger()
+
+	exePath := filepath.Join(generationLocation, "bin", "activation-supervisor")
+
+	currentGeneration, err := s.FS().ReadLink(constants.CurrentSystem)
+	if err != nil {
+		return err
+	}
+
+	argv := []string{
+		"systemd-run",
+		"--collect",
+		"--no-ask-password",
+		"--pipe",
+		"--quiet",
+		"--service-type=exec",
+		"--unit=nixos-cli-activation-supervisor-run",
+		"-E", "LOCALE_ARCHIVE",
+		"-E", fmt.Sprintf("%s=1", RUNNING_ACTIVATION_SUPERVISOR),
+	}
+
+	if opts.InstallBootloader {
+		argv = append(argv, "-E", "NIXOS_INSTALL_BOOTLOADER=1")
+	}
+
+	argv = append(argv, exePath, "run", action.String(), "--previous-gen", currentGeneration)
+
+	if opts.ProfileName != "" && opts.ProfileName != "system" {
+		argv = append(argv, "-p", opts.ProfileName)
+	}
+
+	if opts.Specialisation != "" {
+		argv = append(argv, "-s", opts.Specialisation)
+	}
+
+	if log.GetLogLevel() == logger.LogLevelDebug {
+		argv = append(argv, "-v")
+	}
+
+	log.CmdArray(argv)
+
+	cmd := system.NewCommand(argv[0], argv[1:]...)
+	if opts.UseRootCommand {
+		cmd.RunAsRoot(opts.RootCommand)
+	}
+
+	_, err = s.Run(cmd)
+	if err != nil {
+		return fmt.Errorf("activation supervisor exited abnormally: %v", err)
+	}
+
+	triggerPath := MakeActivationTriggerPath(generationLocation)
+
+	err = s.FS().CreateFile(triggerPath)
+	if err != nil {
+		log.Errorf("failed to create acknowledgement file on remote system: %v", err)
+	}
+
+	return nil
+}

internal/constants/constants.go

@@ -7,4 +7,5 @@ const (
 	CurrentSystem             = "/run/current-system"
 	NixOSMarker               = "/etc/NIXOS"
 	NixChannelDirectory       = NixProfileDirectory + "/per-user/root/channels"
+	NixOSActivationDirectory  = "/run/nixos"
 )

internal/system/fs.go

@@ -7,4 +7,5 @@ type Filesystem interface {
 	ReadLink(path string) (string, error)
 	MkdirAll(path string, perm os.FileMode) error
 	ReadFile(path string) ([]byte, error)
+	CreateFile(path string) error
 }

internal/system/local_fs.go

@@ -21,3 +21,13 @@ func (LocalFilesystem) ReadFile(path string) ([]byte, error) {
 func (LocalFilesystem) MkdirAll(path string, perm os.FileMode) error {
 	return os.MkdirAll(path, perm)
 }
+
+func (LocalFilesystem) CreateFile(path string) error {
+	f, err := os.Create(path)
+	if err != nil {
+		return err
+	}
+
+	_ = f.Close()
+	return nil
+}

internal/system/sftp_fs.go

@@ -37,3 +37,13 @@ func (f *SFTPFilesystem) ReadFile(path string) ([]byte, error) {
 func (f *SFTPFilesystem) MkdirAll(path string, perm os.FileMode) error {
 	return f.client.MkdirAll(path)
 }
+
+func (f *SFTPFilesystem) CreateFile(path string) error {
+	file, err := f.client.Create(path)
+	if err != nil {
+		return err
+	}
+
+	_ = file.Close()
+	return nil
+}

supervisor/main.go

@@ -0,0 +1,76 @@
+//go:build linux
+
+package main
+
+import (
+	"context"
+	"fmt"
+	"os"
+
+	"github.com/nix-community/nixos-cli/internal/activation"
+	"github.com/nix-community/nixos-cli/internal/build"
+	"github.com/nix-community/nixos-cli/internal/logger"
+	"github.com/nix-community/nixos-cli/internal/system"
+	"github.com/spf13/cobra"
+)
+
+const (
+	RUNNING_ACTIVATION_SUPERVISOR = "NIXOS_CLI_RUNNING_ACTIVATION_SUPERVISOR"
+)
+
+func mainCommand() *cobra.Command {
+	log := logger.NewConsoleLogger()
+
+	cmdCtx := logger.WithLogger(context.Background(), log)
+
+	cmd := &cobra.Command{
+		Use:          "activation-supervisor",
+		Short:        "activation-supervisor",
+		Long:         "nixos-cli activation supervisor for activating remote systems.",
+		Version:      build.Version(),
+		SilenceUsage: true,
+		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
+			s := system.NewLocalSystem(log)
+
+			if !s.IsNixOS() {
+				return fmt.Errorf("the activation supervisor is not supported on non-NixOS systems")
+			}
+
+			if os.Geteuid() != 0 {
+				return fmt.Errorf("this command must be ran as root")
+			}
+
+			if os.Getenv(RUNNING_ACTIVATION_SUPERVISOR) != "1" {
+				return fmt.Errorf("the activation supervisor is not meant to be ran directly by users")
+			}
+
+			if err := activation.EnsureActivationDirectoryExists(); err != nil {
+				return err
+			}
+
+			return nil
+		},
+		CompletionOptions: cobra.CompletionOptions{
+			DisableDefaultCmd: true,
+		},
+	}
+
+	cmd.AddCommand(RunCommand())
+	cmd.AddCommand(WatchdogCommand())
+
+	cmd.SetHelpCommand(&cobra.Command{Hidden: true})
+	cmd.SetHelpTemplate(cmd.HelpTemplate() + `
+This command is not meant to be ran directly. Consult nixos-cli-apply(1) for
+more information on how this is executed.
+`)
+
+	cmd.SetContext(cmdCtx)
+
+	return cmd
+}
+
+func main() {
+	if err := mainCommand().Execute(); err != nil {
+		os.Exit(1)
+	}
+}