Try to fix hooks (#718)

* fix hooks


---------

Signed-off-by: Itxaka <itxaka@kairos.io>

internal/agent/hooks/bundles.go

@@ -5,7 +5,6 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
-	"strings"
 	"syscall"
 
 	config "github.com/kairos-io/kairos-agent/v2/pkg/config"
@@ -14,7 +13,6 @@ import (
 	"github.com/kairos-io/kairos-sdk/bundles"
 	"github.com/kairos-io/kairos-sdk/machine"
 	"github.com/kairos-io/kairos-sdk/utils"
-	kcrypt "github.com/kairos-io/kcrypt/pkg/lib"
 )
 
 // BundlePostInstall install bundles just after installation
@@ -41,29 +39,6 @@ func (b BundlePostInstall) Run(c config.Config, _ v1.Spec) error {
 		_ = machine.Umount(constants.OEMPath)
 	}()
 
-	// Path if we have encrypted persistent and we are not on UKI
-	if len(c.Install.Encrypt) != 0 {
-		err := kcrypt.UnlockAll(false)
-		if err != nil {
-			return err
-		}
-		// Close all the unencrypted partitions at the end!
-		defer func() {
-			for _, p := range c.Install.Encrypt {
-				_, _ = utils.SH("udevadm trigger --type=all || udevadm trigger")
-				syscall.Sync()
-				c.Logger.Debugf("Closing unencrypted /dev/disk/by-label/%s", p)
-				out, err := utils.SH(fmt.Sprintf("cryptsetup close /dev/disk/by-label/%s", p))
-				// There is a known error with cryptsetup that it can't close the device because of a semaphore
-				// doesnt seem to affect anything as the device is closed as expected so we ignore it if it matches the
-				// output of the error
-				if err != nil && !strings.Contains(out, "incorrect semaphore state") {
-					c.Logger.Warnf("could not close /dev/disk/by-label/%s: %s", p, out)
-				}
-			}
-		}()
-	}
-
 	_, _ = utils.SH("udevadm trigger --type=all || udevadm trigger")
 	syscall.Sync()
 	err := c.Syscall.Mount(filepath.Join("/dev/disk/by-label", constants.PersistentLabel), constants.UsrLocalPath, "ext4", 0, "")

internal/agent/hooks/finish.go

@@ -2,12 +2,6 @@ package hook
 
 import (
 	"fmt"
-	"os"
-	"regexp"
-	"strconv"
-	"strings"
-	"time"
-
 	"github.com/kairos-io/kairos-agent/v2/pkg/config"
 	"github.com/kairos-io/kairos-agent/v2/pkg/constants"
 	v1 "github.com/kairos-io/kairos-agent/v2/pkg/types/v1"
@@ -16,11 +10,115 @@ import (
 	"github.com/kairos-io/kairos-sdk/machine"
 	"github.com/kairos-io/kairos-sdk/utils"
 	kcrypt "github.com/kairos-io/kcrypt/pkg/lib"
+	"os"
+	"regexp"
+	"strconv"
+	"strings"
+	"syscall"
+	"time"
 )
 
-type KcryptUKI struct{}
+// Finish is a hook that runs after the install process.
+// It is used to encrypt partitions and run the BundlePostInstall, CustomMounts and CopyLogs hooks
+type Finish struct{}
 
-func (k KcryptUKI) Run(c config.Config, spec v1.Spec) error {
+func (k Finish) Run(c config.Config, spec v1.Spec) error {
+	var err error
+	if len(c.Install.Encrypt) != 0 || internalutils.IsUki() {
+		c.Logger.Logger.Info().Msg("Running encrypt hook")
+
+		if internalutils.IsUki() {
+			err = EncryptUKI(c, spec)
+		} else {
+			err = Encrypt(c, spec)
+		}
+
+		// They return with partitions unlocked so make sure to lock them before we end
+		defer lockPartitions(c)
+
+		if err != nil {
+			c.Logger.Logger.Error().Err(err).Msg("could not encrypt partitions")
+			return err
+		}
+		c.Logger.Logger.Info().Msg("Finished encrypt hook")
+	}
+
+	// Now that we have everything encrypted and ready if needed
+	err = BundlePostInstall{}.Run(c, spec)
+	if err != nil {
+		c.Logger.Logger.Warn().Err(err).Msg("could not copy run bundles post install")
+		if c.FailOnBundleErrors {
+			return err
+		}
+	}
+	err = CustomMounts{}.Run(c, spec)
+	if err != nil {
+		c.Logger.Logger.Warn().Err(err).Msg("could not create custom mounts")
+	}
+	err = CopyLogs{}.Run(c, spec)
+	if err != nil {
+		c.Logger.Logger.Warn().Err(err).Msg("could not copy logs")
+	}
+	return nil
+}
+
+// Encrypt is a hook that encrypts partitions using kcrypt for non uki.
+// It will unmount OEM and PERSISTENT and return with them unmounted
+func Encrypt(c config.Config, _ v1.Spec) error {
+	// Start with unmounted partitions
+	_ = machine.Umount(constants.OEMDir)        //nolint:errcheck
+	_ = machine.Umount(constants.PersistentDir) //nolint:errcheck
+
+	// Config passed during install ends up here, so we need to read it, try to mount it
+	_ = machine.Mount("COS_OEM", "/oem")
+	defer func() {
+		err := syscall.Unmount(constants.OEMPath, 0)
+		if err != nil {
+			c.Logger.Warnf("could not unmount Oem partition: %s", err)
+		}
+	}()
+
+	for _, p := range c.Install.Encrypt {
+		_, err := kcrypt.Luksify(p, c.Logger)
+		if err != nil {
+			c.Logger.Errorf("could not encrypt partition: %s", err)
+			return err
+		}
+	}
+
+	_ = kcrypt.UnlockAllWithLogger(false, c.Logger)
+
+	for _, p := range c.Install.Encrypt {
+		for i := 0; i < 10; i++ {
+			c.Logger.Infof("Waiting for unlocked partition %s to appear", p)
+			_, _ = utils.SH("sync")
+			part, _ := utils.SH(fmt.Sprintf("blkid -L %s", p))
+			if part == "" {
+				c.Logger.Infof("Partition %s not found, waiting %d seconds before retrying", p, i)
+				time.Sleep(time.Duration(i) * time.Second)
+				// Retry the unlock as well, because maybe the partition was not refreshed on time for unlock to unlock it
+				// So no matter how many tries we do, it will still be locked and will never appear
+				err := kcrypt.UnlockAllWithLogger(false, c.Logger)
+				if err != nil {
+					c.Logger.Debugf("UnlockAll returned: %s", err)
+				}
+				if i == 9 {
+					c.Logger.Errorf("Partition %s not unlocked/found after 10 retries", p)
+					return fmt.Errorf("partition %s not unlocked/found after 10 retries", p)
+				}
+				continue
+			}
+			c.Logger.Infof("Partition found, continuing")
+			break
+		}
+	}
+	return nil
+
+}
+
+// EncryptUKI encrypts the partitions using kcrypt in uki mode
+// It will unmount OEM and PERSISTENT and return with them unmounted
+func EncryptUKI(c config.Config, spec v1.Spec) error {
 	// pre-check for systemd version, we need something higher or equal to 252
 	run, err := utils.SH("systemctl --version | head -1 | awk '{ print $2}'")
 	systemdVersion := strings.TrimSpace(string(run))
@@ -48,7 +146,7 @@ func (k KcryptUKI) Run(c config.Config, spec v1.Spec) error {
 	// If systemd version is less than 252 return
 	if systemdVersionInt < 252 {
 		c.Logger.Infof("systemd version is %s, we need 252 or higher for encrypting partitions", systemdVersion)
-		return nil
+		return fmt.Errorf("systemd version is %s, we need 252 or higher for encrypting partitions", systemdVersion)
 	}
 
 	// Check for a TPM 2.0 device as its needed to encrypt
@@ -57,11 +155,9 @@ func (k KcryptUKI) Run(c config.Config, spec v1.Spec) error {
 	_, err = os.Stat("/dev/tpmrm0")
 	if err != nil {
 		c.Logger.Warnf("Skipping partition encryption, could not find TPM 2.0 device at /dev/tpmrm0")
-		return nil
+		return fmt.Errorf("Skipping partition encryption, could not find TPM 2.0 device at /dev/tpmrm0")
 	}
 
-	c.Logger.Logger.Debug().Msg("Running KcryptUKI hook")
-
 	// We always encrypt OEM and PERSISTENT under UKI
 	// If mounted, unmount it
 	_ = machine.Umount(constants.OEMDir)        //nolint:errcheck
@@ -92,16 +188,11 @@ func (k KcryptUKI) Run(c config.Config, spec v1.Spec) error {
 	for _, p := range append([]string{constants.OEMLabel, constants.PersistentLabel}, c.Install.Encrypt...) {
 		c.Logger.Infof("Encrypting %s", p)
 		_ = os.Setenv("SYSTEMD_LOG_LEVEL", "debug")
-		err := kcrypt.LuksifyMeasurements(p, c.BindPublicPCRs, c.BindPCRs, c.Logger)
+		err = kcrypt.LuksifyMeasurements(p, c.BindPublicPCRs, c.BindPCRs, c.Logger)
 		_ = os.Unsetenv("SYSTEMD_LOG_LEVEL")
 		if err != nil {
 			c.Logger.Errorf("could not encrypt partition: %s", err)
-			if c.FailOnBundleErrors {
+			return err
-				return err
-			}
-			// Give time to show the error
-			time.Sleep(10 * time.Second)
-			return nil // do not error out
 		}
 		c.Logger.Infof("Done encrypting %s", p)
 	}
@@ -109,31 +200,20 @@ func (k KcryptUKI) Run(c config.Config, spec v1.Spec) error {
 	_, _ = utils.SH("sync")
 
 	_ = os.Setenv("SYSTEMD_LOG_LEVEL", "debug")
+
 	err = kcrypt.UnlockAllWithLogger(true, c.Logger)
 
 	_ = os.Unsetenv("SYSTEMD_LOG_LEVEL")
 	if err != nil {
+		lockPartitions(c)
 		c.Logger.Errorf("could not unlock partitions: %s", err)
 		return err
 	}
-	// Close the unlocked partitions after dealing with them, otherwise we leave them open and they can be mounted by anyone
-	defer func() {
-		for _, p := range append([]string{constants.OEMLabel, constants.PersistentLabel}, c.Install.Encrypt...) {
-			c.Logger.Debugf("Closing unencrypted /dev/disk/by-label/%s", p)
-			out, err := utils.SH(fmt.Sprintf("cryptsetup close /dev/disk/by-label/%s", p))
-			// There is a known error with cryptsetup that it can't close the device because of a semaphore
-			// doesnt seem to affect anything as the device is closed as expected so we ignore it if it matches the
-			// output of the error
-			if err != nil && !strings.Contains(out, "incorrect semaphore state") {
-				c.Logger.Errorf("could not close /dev/disk/by-label/%s: %s", p, out)
-			}
-		}
-	}()
 
 	// Here it can take the oem partition a bit of time to appear after unlocking so we need to retry a couple of time with some waiting
 	// retry + backoff
 	// Check all encrypted partitions are unlocked
-	for _, p := range append([]string{constants.OEMLabel, constants.PersistentLabel}) {
+	for _, p := range []string{constants.OEMLabel, constants.PersistentLabel} {
 		for i := 0; i < 10; i++ {
 			c.Logger.Infof("Waiting for unlocked partition %s to appear", p)
 			_, _ = utils.SH("sync")
@@ -158,29 +238,21 @@ func (k KcryptUKI) Run(c config.Config, spec v1.Spec) error {
 		}
 	}
 
+	// Mount the unlocked oem partition
 	err = machine.Mount(constants.OEMLabel, constants.OEMDir)
 	if err != nil {
 		return err
 	}
+	// Copy back the contents of the oem partition that we saved before encrypting
 	err = internalutils.SyncData(c.Logger, c.Runner, c.Fs, tmpDir, constants.OEMDir, []string{}...)
 	if err != nil {
 		return err
 	}
+	// Unmount the oem partition and leave everything unmounted
 	err = machine.Umount(constants.OEMDir)
 	if err != nil {
 		return err
 	}
 
-	c.Logger.Logger.Debug().Msg("Finish KcryptUKI hook")
-	// We now have the partitions unlocked and ready, lets call the other hooks here instead of closing and reopening them each time
-	err = BundlePostInstall{}.Run(c, spec)
-	if err != nil {
-		return err
-	}
-	_ = CustomMounts{}.Run(c, spec)
-	err = CopyLogs{}.Run(c, spec)
-	if err != nil {
-		return err
-	}
 	return nil
 }

internal/agent/hooks/hook.go

@@ -1,48 +1,54 @@
 package hook
 
 import (
+	"fmt"
 	config "github.com/kairos-io/kairos-agent/v2/pkg/config"
 	v1 "github.com/kairos-io/kairos-agent/v2/pkg/types/v1"
+	"github.com/kairos-io/kairos-sdk/utils"
+	"strings"
 )
 
 type Interface interface {
 	Run(c config.Config, spec v1.Spec) error
 }
 
-var AfterInstall = []Interface{
+// FinishInstall is a list of hooks that run when the install process is finished completely.
-	&GrubOptions{}, // Set custom GRUB options
+// Its mean for options that are not related to the install process itself
-	&BundlePostInstall{},
+var FinishInstall = []Interface{
-	&CustomMounts{},
+	&GrubOptions{}, // Set custom GRUB options in OEM partition
-	&CopyLogs{},
+	&Lifecycle{},   // Handles poweroff/reboot by config options
+}
+
+// FinishReset is a list of hooks that run when the reset process is finished completely.
+var FinishReset = []Interface{
+	&CopyLogs{},  // Try to copy the reset logs to the persistent partition
 	&Lifecycle{}, // Handles poweroff/reboot by config options
 }
 
-var AfterReset = []Interface{
+// FinishUpgrade is a list of hooks that run when the upgrade process is finished completely.
-	&CopyLogs{},
+var FinishUpgrade = []Interface{
-	&Lifecycle{},
+	&Lifecycle{}, // Handles poweroff/reboot by config options
-}
-
-var AfterUpgrade = []Interface{
-	&Lifecycle{},
 }
 
+// FirstBoot is a list of hooks that run on the first boot of the node.
 var FirstBoot = []Interface{
 	&BundleFirstBoot{},
 	&GrubPostInstallOptions{},
 }
 
-// AfterUkiInstall sets which Hooks to run after uki runs the install action
+// FinishUKIInstall is a list of hooks that run when the install process is finished completely.
-var AfterUkiInstall = []Interface{
+// Its mean for options that are not related to the install process itself
-	&SysExtPostInstall{},
+var FinishUKIInstall = []Interface{
-	&Lifecycle{},
+	&SysExtPostInstall{}, // Installs sysexts into the EFI partition
+	&Lifecycle{},         // Handles poweroff/reboot by config options
 }
 
-var UKIEncryptionHooks = []Interface{
+// PostInstall is a list of hooks that run after the install process has run.
-	&KcryptUKI{},
+// Runs things that need to be done before we run other post install stages like
-}
+// encrypting partitions, copying the install logs or installing bundles
-
+// Most of this options are optional so they are not run by default unless specified int he config
-var EncryptionHooks = []Interface{
+var PostInstall = []Interface{
-	&Kcrypt{},
+	&Finish{},
 }
 
 func Run(c config.Config, spec v1.Spec, hooks ...Interface) error {
@@ -53,3 +59,18 @@ func Run(c config.Config, spec v1.Spec, hooks ...Interface) error {
 	}
 	return nil
 }
+
+// lockPartitions will try to close all the partitions that are unencrypted.
+func lockPartitions(c config.Config) {
+	for _, p := range c.Install.Encrypt {
+		_, _ = utils.SH("udevadm trigger --type=all || udevadm trigger")
+		c.Logger.Debugf("Closing unencrypted /dev/disk/by-label/%s", p)
+		out, err := utils.SH(fmt.Sprintf("cryptsetup close /dev/disk/by-label/%s", p))
+		// There is a known error with cryptsetup that it can't close the device because of a semaphore
+		// doesnt seem to affect anything as the device is closed as expected so we ignore it if it matches the
+		// output of the error
+		if err != nil && !strings.Contains(out, "incorrect semaphore state") {
+			c.Logger.Debugf("could not close /dev/disk/by-label/%s: %s", p, out)
+		}
+	}
+}

internal/agent/hooks/kcrypt.go

@@ -1,49 +0,0 @@
-package hook
-
-import (
-	"github.com/kairos-io/kairos-agent/v2/pkg/config"
-	"github.com/kairos-io/kairos-agent/v2/pkg/constants"
-	v1 "github.com/kairos-io/kairos-agent/v2/pkg/types/v1"
-	"github.com/kairos-io/kairos-sdk/machine"
-	kcrypt "github.com/kairos-io/kcrypt/pkg/lib"
-	"path/filepath"
-	"syscall"
-)
-
-type Kcrypt struct{}
-
-func (k Kcrypt) Run(c config.Config, _ v1.Spec) error {
-	if len(c.Install.Encrypt) == 0 {
-		return nil
-	}
-	c.Logger.Logger.Info().Msg("Running encrypt hook")
-
-	// We need to unmount the persistent partition to encrypt it
-	// we dont know the state here so we better try
-	err := machine.Umount(filepath.Join("/dev/disk/by-label", constants.PersistentLabel)) //nolint:errcheck
-	if err != nil {
-		c.Logger.Errorf("could not unmount persistent partition: %s", err)
-		return err
-	}
-
-	// Config passed during install ends up here, so we need to read it
-	_ = machine.Mount("COS_OEM", "/oem")
-	defer func() {
-		err := syscall.Unmount(constants.OEMPath, 0)
-		if err != nil {
-			c.Logger.Errorf("could not unmount Oem partition: %s", err)
-		}
-	}()
-
-	for _, p := range c.Install.Encrypt {
-		_, err := kcrypt.Luksify(p, c.Logger)
-		if err != nil {
-			c.Logger.Errorf("could not encrypt partition: %s", err)
-			if c.FailOnBundleErrors {
-				return err
-			}
-		}
-	}
-	c.Logger.Logger.Info().Msg("Finished encrypt hook")
-	return nil
-}

internal/agent/hooks/logs.go

@@ -2,6 +2,9 @@ package hook
 
 import (
 	"fmt"
+	"path/filepath"
+	"syscall"
+
 	"github.com/kairos-io/kairos-agent/v2/pkg/config"
 	"github.com/kairos-io/kairos-agent/v2/pkg/constants"
 	v1 "github.com/kairos-io/kairos-agent/v2/pkg/types/v1"
@@ -9,20 +12,22 @@ import (
 	fsutils "github.com/kairos-io/kairos-agent/v2/pkg/utils/fs"
 	"github.com/kairos-io/kairos-sdk/machine"
 	"github.com/kairos-io/kairos-sdk/utils"
-	kcrypt "github.com/kairos-io/kcrypt/pkg/lib"
-	"path/filepath"
-	"strings"
-	"syscall"
 )
 
-// CopyLogs copies all current logs to the persistent partition.
-// useful during install to keep the livecd logs
-// best effort, no error handling
 type CopyLogs struct{}
 
+// Run for CopyLogs copies all current logs to the persistent partition.
+// useful during install to keep the livecd logs. Its also run during reset
+// best effort, no error handling
 func (k CopyLogs) Run(c config.Config, _ v1.Spec) error {
+	// TODO: If we have encryption under RESET we need to make sure to:
+	// - Unlock the partitions
+	// - Mount OEM so we can read the config for encryption (remote server)
+	// - Mount the persistent partition
 	c.Logger.Logger.Debug().Msg("Running CopyLogs hook")
 	_ = machine.Umount(constants.PersistentDir)
+	_ = machine.Umount(constants.OEMDir)
+	_ = machine.Umount(constants.OEMPath)
 
 	// Config passed during install ends up here, kcrypt challenger needs to read it if we are using a server for encryption
 	_ = machine.Mount(constants.OEMLabel, constants.OEMPath)
@@ -30,32 +35,12 @@ func (k CopyLogs) Run(c config.Config, _ v1.Spec) error {
 		_ = machine.Umount(constants.OEMPath)
 	}()
 
-	// Path if we have encrypted persistent
-	if len(c.Install.Encrypt) != 0 {
-		err := kcrypt.UnlockAll(false)
-		if err != nil {
-			return err
-		}
-		// Close all the unencrypted partitions at the end!
-		defer func() {
-			for _, p := range c.Install.Encrypt {
-				c.Logger.Debugf("Closing unencrypted /dev/disk/by-label/%s", p)
-				out, err := utils.SH(fmt.Sprintf("cryptsetup close /dev/disk/by-label/%s", p))
-				// There is a known error with cryptsetup that it can't close the device because of a semaphore
-				// doesnt seem to affect anything as the device is closed as expected so we ignore it if it matches the
-				// output of the error
-				if err != nil && !strings.Contains(out, "incorrect semaphore state") {
-					c.Logger.Errorf("could not close /dev/disk/by-label/%s: %s", p, out)
-				}
-			}
-		}()
-	}
-
 	_, _ = utils.SH("udevadm trigger --type=all || udevadm trigger")
+	_ = utils.MkdirAll(c.Fs, constants.PersistentDir, 0755)
 	err := c.Syscall.Mount(filepath.Join("/dev/disk/by-label", constants.PersistentLabel), constants.PersistentDir, "ext4", 0, "")
 	if err != nil {
 		fmt.Printf("could not mount persistent: %s\n", err)
-		return err
+		return nil
 	}
 
 	defer func() {

internal/agent/hooks/mounts.go

@@ -24,23 +24,21 @@ func saveCloudConfig(name config.Stage, yc yip.YipConfig) error {
 	return os.WriteFile(filepath.Join("/oem", fmt.Sprintf("10_%s.yaml", name)), yipYAML, 0400)
 }
 
-// Read the keys sections ephemeral_mounts and bind mounts from install key in the cloud config.
+// Run Read the keys sections ephemeral_mounts and bind mounts from install key in the cloud config.
 // If not empty write an environment file to /run/cos/custom-layout.env.
 // That env file is in turn read by /overlay/files/system/oem/11_persistency.yaml in fs.after stage.
 func (cm CustomMounts) Run(c config.Config, _ v1.Spec) error {
-
-	//fmt.Println("Custom mounts hook")
-	//fmt.Println(strings.Join(c.Install.BindMounts, " "))
-	//fmt.Println(strings.Join(c.Install.EphemeralMounts, " "))
-
 	if len(c.Install.BindMounts) == 0 && len(c.Install.EphemeralMounts) == 0 {
 		return nil
 	}
 	c.Logger.Logger.Debug().Msg("Running CustomMounts hook")
 
-	machine.Mount("COS_OEM", "/oem") //nolint:errcheck
+	err := machine.Mount("COS_OEM", "/oem")
+	if err != nil {
+		return err
+	}
 	defer func() {
-		machine.Umount("/oem") //nolint:errcheck
+		_ = machine.Umount("/oem")
 	}()
 
 	var mountsList = map[string]string{}
@@ -48,15 +46,22 @@ func (cm CustomMounts) Run(c config.Config, _ v1.Spec) error {
 	mountsList["CUSTOM_BIND_MOUNTS"] = strings.Join(c.Install.BindMounts, " ")
 	mountsList["CUSTOM_EPHEMERAL_MOUNTS"] = strings.Join(c.Install.EphemeralMounts, " ")
 
-	config := yip.YipConfig{Stages: map[string][]schema.Stage{
+	cfg := yip.YipConfig{
-		"rootfs": []yip.Stage{{
+		Stages: map[string][]schema.Stage{
-			Name:            "user_custom_mounts",
+			"rootfs": {
-			EnvironmentFile: "/run/cos/custom-layout.env",
+				{
-			Environment:     mountsList,
+					Name:            "user_custom_mounts",
-		}},
+					EnvironmentFile: "/run/cos/custom-layout.env",
-	}}
+					Environment:     mountsList,
+				},
+			},
+		},
+	}
 
-	saveCloudConfig("user_custom_mounts", config) //nolint:errcheck
+	err = saveCloudConfig("user_custom_mounts", cfg)
+	if err != nil {
+		return err
+	}
 	c.Logger.Logger.Debug().Msg("Finish CustomMounts hook")
 	return nil
 }

internal/agent/reset.go

@@ -61,7 +61,7 @@ func reset(reboot, unattended, resetOem bool, dir ...string) error {
 
 	bus.Manager.Publish(sdk.EventAfterReset, sdk.EventPayload{}) //nolint:errcheck
 
-	return hook.Run(*cfg, resetSpec, hook.AfterReset...)
+	return hook.Run(*cfg, resetSpec, hook.FinishReset...)
 }
 
 func resetUki(reboot, unattended, resetOem bool, dir ...string) error {
@@ -92,7 +92,7 @@ func resetUki(reboot, unattended, resetOem bool, dir ...string) error {
 
 	bus.Manager.Publish(sdk.EventAfterReset, sdk.EventPayload{}) //nolint:errcheck
 
-	return hook.Run(*cfg, resetSpec, hook.AfterReset...)
+	return hook.Run(*cfg, resetSpec, hook.FinishReset...)
 }
 
 // sharedReset is the common reset code for both uki and non-uki

internal/agent/upgrade.go

@@ -115,7 +115,7 @@ func upgrade(sourceImageURL string, dirs []string, upgradeEntry string, strictVa
 		return err
 	}
 
-	return hook.Run(*c, upgradeSpec, hook.AfterUpgrade...)
+	return hook.Run(*c, upgradeSpec, hook.FinishUpgrade...)
 }
 
 func upgradeUki(sourceImageURL string, dirs []string, upgradeEntry string, strictValidations bool) error {
@@ -148,7 +148,7 @@ func upgradeUki(sourceImageURL string, dirs []string, upgradeEntry string, stric
 		return err
 	}
 
-	return hook.Run(*c, upgradeSpec, hook.AfterUpgrade...)
+	return hook.Run(*c, upgradeSpec, hook.FinishUpgrade...)
 }
 
 func getConfig(sourceImageURL string, dirs []string, upgradeEntry string, strictValidations bool) (*config.Config, error) {

pkg/action/install.go

@@ -270,7 +270,7 @@ func (i InstallAction) Run() (err error) {
 		return err
 	}
 
-	err = hook.Run(*i.cfg, i.spec, hook.EncryptionHooks...)
+	err = hook.Run(*i.cfg, i.spec, hook.PostInstall...)
 	if err != nil {
 		return err
 	}
@@ -307,5 +307,5 @@ func (i InstallAction) Run() (err error) {
 	_ = utils.RunStage(i.cfg, "kairos-install.after")
 	_ = events.RunHookScript("/usr/bin/kairos-agent.install.after.hook") //nolint:errcheck
 
-	return hook.Run(*i.cfg, i.spec, hook.AfterInstall...)
+	return hook.Run(*i.cfg, i.spec, hook.FinishInstall...)
 }

pkg/config/spec.go

@@ -399,7 +399,7 @@ func NewUpgradeSpec(cfg *Config) (*v1.UpgradeSpec, error) {
 	}
 
 	if spec.Active.Source.IsDocker() {
-                cfg.Logger.Infof("Checking if OCI image %s exists", spec.Active.Source.Value())
+		cfg.Logger.Infof("Checking if OCI image %s exists", spec.Active.Source.Value())
 		_, err := crane.Manifest(spec.Active.Source.Value())
 		if err != nil {
 			if strings.Contains(err.Error(), "MANIFEST_UNKNOWN") {

pkg/uki/install.go

@@ -199,7 +199,7 @@ func (i *InstallAction) Run() (err error) {
 		i.cfg.Logger.Warnf("selecting active boot entry: %s", err.Error())
 	}
 
-	err = hook.Run(*i.cfg, i.spec, hook.UKIEncryptionHooks...)
+	err = hook.Run(*i.cfg, i.spec, hook.PostInstall...)
 	if err != nil {
 		i.cfg.Logger.Errorf("running uki encryption hooks: %s", err.Error())
 		return err
@@ -223,7 +223,7 @@ func (i *InstallAction) Run() (err error) {
 		i.cfg.Logger.Errorf("running kairos-uki-install.after hook script: %s", err.Error())
 	}
 
-	return hook.Run(*i.cfg, i.spec, hook.AfterUkiInstall...)
+	return hook.Run(*i.cfg, i.spec, hook.FinishUKIInstall...)
 }
 
 func (i *InstallAction) SkipEntry(path string, conf map[string]string) (err error) {

pkg/uki/reset.go

@@ -2,7 +2,6 @@ package uki
 
 import (
 	"fmt"
-
 	"github.com/kairos-io/kairos-agent/v2/pkg/action"
 	"github.com/kairos-io/kairos-agent/v2/pkg/config"
 	"github.com/kairos-io/kairos-agent/v2/pkg/constants"
@@ -35,14 +34,12 @@ func (r *ResetAction) Run() (err error) {
 	cleanup := utils.NewCleanStack()
 	defer func() { err = cleanup.Cleanup(err) }()
 
-	// Unmount partitions if any is already mounted before formatting
+	// At this point, both partitions are unlocked but they might not be mounted, like persistent
-	err = e.UnmountPartitions(r.spec.Partitions.PartitionsByMountPoint(true))
+	// And the /dev/disk/by-label are not pointing to the proper ones
-	if err != nil {
+	// We need to manually trigger udev to make sure the symlinks are correct
-		r.cfg.Logger.Errorf("unmounting partitions: %s", err.Error())
+	_, err = utils.SH("udevadm trigger --type=all || udevadm trigger")
-		return err
+	_, err = utils.SH("udevadm settle")
-	}
 
-	// Reformat persistent partition
 	if r.spec.FormatPersistent {
 		persistent := r.spec.Partitions.Persistent
 		if persistent != nil {
@@ -58,11 +55,20 @@ func (r *ResetAction) Run() (err error) {
 	if r.spec.FormatOEM {
 		oem := r.spec.Partitions.OEM
 		if oem != nil {
+			err = e.UnmountPartition(oem)
+			if err != nil {
+				return err
+			}
 			err = e.FormatPartition(oem)
 			if err != nil {
 				r.cfg.Logger.Errorf("formatting OEM partition: %s", err.Error())
 				return err
 			}
+			// Mount it back, as oem is mounted during recovery, keep everything as is
+			err = e.MountPartition(oem)
+			if err != nil {
+				return err
+			}
 		}
 	}
 
@@ -101,14 +107,6 @@ func (r *ResetAction) Run() (err error) {
 		return err
 	}
 
-	if mnt, err := elementalUtils.IsMounted(r.cfg, r.spec.Partitions.OEM); !mnt && err == nil {
-		err = e.MountPartition(r.spec.Partitions.OEM)
-		if err != nil {
-			r.cfg.Logger.Errorf("mounting oem partition: %s", err.Error())
-			return err
-		}
-	}
-
 	err = Hook(r.cfg, constants.AfterResetHook)
 	if err != nil {
 		r.cfg.Logger.Errorf("running after install hook: %s", err.Error())