Refactor methods and implement uki upgrade

Signed-off-by: Dimitris Karakasilis <dimitris@spectrocloud.com>

pkg/uki/common.go

@@ -0,0 +1,121 @@
+package uki
+
+import (
+	"fmt"
+	"io"
+	"os"
+	"strings"
+
+	v1 "github.com/kairos-io/kairos-agent/v2/pkg/types/v1"
+	fsutils "github.com/kairos-io/kairos-agent/v2/pkg/utils/fs"
+	sdkutils "github.com/kairos-io/kairos-sdk/utils"
+	"github.com/sanity-io/litter"
+)
+
+const UnassignedArtifactRole = "norole"
+
+// overwriteArtifactSetRole first deletes all artifacts prefixed with newRole
+// (because they are going to be replaced, e.g. old "passive") and then installs
+// the artifacts prefixed with oldRole as newRole.
+// E.g. removes "passive" and moved "active" to "passive"
+// This is a step that should happen before a new passive is installed on upgrades.
+func overwriteArtifactSetRole(fs v1.FS, dir, oldRole, newRole string, logger v1.Logger) error {
+	if err := removeArtifactSetWithRole(fs, dir, newRole); err != nil {
+		return fmt.Errorf("deleting role %s: %w", newRole, err)
+	}
+
+	if err := copyArtifactSetRole(fs, dir, oldRole, newRole, logger); err != nil {
+		return fmt.Errorf("copying artifact set role from %s to %s: %w", oldRole, newRole, err)
+	}
+
+	return nil
+}
+
+// copy the source file but rename the base name to as
+func copyArtifact(source, oldRole, newRole string) (string, error) {
+	newName := strings.ReplaceAll(source, oldRole, newRole)
+	return newName, copyFile(source, newName)
+}
+
+func removeArtifactSetWithRole(fs v1.FS, artifactDir, role string) error {
+	return fsutils.WalkDirFs(fs, artifactDir, func(path string, info os.DirEntry, err error) error {
+		if !info.IsDir() && strings.HasPrefix(info.Name(), role) {
+			return os.Remove(path)
+		}
+
+		return nil
+	})
+}
+
+func copyArtifactSetRole(fs v1.FS, artifactDir, oldRole, newRole string, logger v1.Logger) error {
+	return fsutils.WalkDirFs(fs, artifactDir, func(path string, info os.DirEntry, err error) error {
+		if !strings.HasPrefix(info.Name(), oldRole) {
+			return nil
+		}
+
+		newPath, err := copyArtifact(path, oldRole, newRole)
+		if err != nil {
+			return fmt.Errorf("copying artifact from %s to %s: %w", path, newPath, err)
+		}
+		if strings.HasSuffix(path, ".conf") {
+			if err := replaceRoleInKey(newPath, "efi", oldRole, newRole, logger); err != nil {
+				return err
+			}
+		}
+
+		return nil
+	})
+}
+
+func replaceRoleInKey(path, key, oldRole, newRole string, logger v1.Logger) (err error) {
+	// Extract the values
+	conf, err := sdkutils.SystemdBootConfReader(path)
+	if err != nil {
+		logger.Errorf("Error reading conf file %s: %s", path, err)
+		return err
+	}
+	logger.Debugf("Conf file %s has values %v", path, litter.Sdump(conf))
+
+	_, hasKey := conf[key]
+	if !hasKey {
+		return fmt.Errorf("no %s entry in .conf file", key)
+	}
+
+	conf[key] = strings.ReplaceAll(conf[key], oldRole, newRole)
+	newContents := ""
+	for k, v := range conf {
+		newContents = fmt.Sprintf("%s%s %s\n", newContents, k, v)
+	}
+	logger.Debugf("Conf file %s new values %v", path, litter.Sdump(conf))
+
+	return os.WriteFile(path, []byte(newContents), os.ModePerm)
+}
+
+func copyFile(src, dst string) error {
+	sourceFile, err := os.Open(src)
+	if err != nil {
+		panic(err)
+	}
+	defer sourceFile.Close()
+
+	destinationFile, err := os.Create(dst)
+	if err != nil {
+		panic(err)
+	}
+	defer destinationFile.Close()
+
+	if _, err = io.Copy(destinationFile, sourceFile); err != nil {
+		return err
+	}
+
+	// Flushes any buffered data to the destination file
+	if err = destinationFile.Sync(); err != nil {
+		return err
+	}
+
+	if err = sourceFile.Close(); err != nil {
+		return err
+	}
+
+	return destinationFile.Close()
+}

pkg/uki/install.go

@@ -2,7 +2,6 @@ package uki
 
 import (
 	"fmt"
-	"io"
 	"os"
 	"path/filepath"
 	"strings"
@@ -16,7 +15,6 @@ import (
 	fsutils "github.com/kairos-io/kairos-agent/v2/pkg/utils/fs"
 	events "github.com/kairos-io/kairos-sdk/bus"
 	sdkutils "github.com/kairos-io/kairos-sdk/utils"
-	"github.com/sanity-io/litter"
 )
 
 type InstallAction struct {
@@ -76,6 +74,9 @@ func (i *InstallAction) Run() (err error) {
 		return err
 	}
 
+	// TODO: Check if the size of the files we are going to copy, will fit in the
+	// partition. If not stop here.
+
 	// Copy the efi file into the proper dir
 	_, err = e.DumpSource(i.spec.Partitions.EFI.MountPoint, i.spec.Active.Source)
 	if err != nil {
@@ -128,19 +129,19 @@ func (i *InstallAction) Run() (err error) {
 		return err
 	}
 
-	for _, entry := range []string{"active", "passive", "recovery"} {
-		if err = installArtifactSet(i.cfg.Fs, i.spec.Partitions.EFI.MountPoint, entry, i.cfg.Logger); err != nil {
-			return fmt.Errorf("installing artifact set %s: %w", entry, err)
+	for _, role := range []string{"active", "passive", "recovery"} {
+		if err = copyArtifactSetRole(i.cfg.Fs, i.spec.Partitions.EFI.MountPoint, UnassignedArtifactRole, role, i.cfg.Logger); err != nil {
+			return fmt.Errorf("installing the new artifact set as %s: %w", role, err)
 		}
 	}
 
 	loaderConfPath := filepath.Join(i.spec.Partitions.EFI.MountPoint, "loader", "loader.conf")
-	if err = replacePlaceholders(loaderConfPath, "default", "active", i.cfg.Logger); err != nil {
+	if err = replaceRoleInKey(loaderConfPath, "default", UnassignedArtifactRole, "active", i.cfg.Logger); err != nil {
 		return err
 	}
 
-	if err = removeArtifactSet(i.cfg.Fs, i.spec.Partitions.EFI.MountPoint); err != nil {
-		return fmt.Errorf("removing artifact set: %w", err)
+	if err = removeArtifactSetWithRole(i.cfg.Fs, i.spec.Partitions.EFI.MountPoint, UnassignedArtifactRole); err != nil {
+		return fmt.Errorf("removing artifact set with role %s: %w", UnassignedArtifactRole, err)
 	}
 
 	// after install hook happens after install (this is for compatibility with normal install, so users can reuse their configs)
@@ -157,70 +158,6 @@ func (i *InstallAction) Run() (err error) {
 
 	return hook.Run(*i.cfg, i.spec, hook.AfterUkiInstall...)
 }
-func copy(src, dst string) error {
-	sourceFile, err := os.Open(src)
-	if err != nil {
-		panic(err)
-	}
-	defer sourceFile.Close()
-
-	destinationFile, err := os.Create(dst)
-	if err != nil {
-		panic(err)
-	}
-	defer destinationFile.Close()
-
-	if _, err = io.Copy(destinationFile, sourceFile); err != nil {
-		return err
-	}
-
-	// Flushes any buffered data to the destination file
-	if err = destinationFile.Sync(); err != nil {
-		return err
-	}
-
-	if err = sourceFile.Close(); err != nil {
-		return err
-	}
-
-	return destinationFile.Close()
-}
-
-// copy the source file but ranme the base name to as
-func copyArtifact(source string, as string) (string, error) {
-	newName := strings.ReplaceAll(source, "artifact", as)
-	return newName, copy(source, newName)
-}
-
-func removeArtifactSet(fs v1.FS, artifactDir string) error {
-	return fsutils.WalkDirFs(fs, artifactDir, func(path string, info os.DirEntry, err error) error {
-		if !info.IsDir() && strings.HasPrefix(info.Name(), "artifact") {
-			return os.Remove(path)
-		}
-
-		return nil
-	})
-}
-
-func installArtifactSet(fs v1.FS, artifactDir, as string, logger v1.Logger) error {
-	return fsutils.WalkDirFs(fs, artifactDir, func(path string, info os.DirEntry, err error) error {
-		if !strings.HasPrefix(info.Name(), "artifact") {
-			return nil
-		}
-
-		newPath, err := copyArtifact(path, as)
-		if err != nil {
-			return fmt.Errorf("copying artifact from %s to %s: %w", path, newPath, err)
-		}
-		if strings.HasSuffix(path, ".conf") {
-			if err := replacePlaceholders(newPath, "efi", as, logger); err != nil {
-				return err
-			}
-		}
-
-		return nil
-	})
-}
 
 func (i *InstallAction) SkipEntry(path string, conf map[string]string) (err error) {
 	// If match, get the efi file and remove it
@@ -244,36 +181,6 @@ func (i *InstallAction) SkipEntry(path string, conf map[string]string) (err erro
 	return err
 }
 
-func (i *InstallAction) replaceFilenamePlaceholder(path, replaceString string) (err error) {
-	newName := strings.ReplaceAll(path, "artifact", replaceString)
-
-	return os.Rename(path, newName)
-}
-
-func replacePlaceholders(path, key, replaceString string, logger v1.Logger) (err error) {
-	// Extract the values
-	conf, err := sdkutils.SystemdBootConfReader(path)
-	if err != nil {
-		logger.Errorf("Error reading conf file %s: %s", path, err)
-		return err
-	}
-	logger.Debugf("Conf file %s has values %v", path, litter.Sdump(conf))
-
-	_, hasKey := conf[key]
-	if !hasKey {
-		return fmt.Errorf("no %s entry in .conf file", key)
-	}
-
-	conf[key] = strings.ReplaceAll(conf[key], "artifact", replaceString)
-	newContents := ""
-	for k, v := range conf {
-		newContents = fmt.Sprintf("%s%s %s\n", newContents, k, v)
-	}
-	logger.Debugf("Conf file %s new values %v", path, litter.Sdump(conf))
-
-	return os.WriteFile(path, []byte(newContents), os.ModePerm)
-}
-
 // Hook is RunStage wrapper that only adds logic to ignore errors
 // in case v1.Config.Strict is set to false
 func Hook(config *config.Config, hook string) error {

pkg/uki/upgrade.go

@@ -1,7 +1,9 @@
 package uki
 
 import (
-	"github.com/Masterminds/semver/v3"
+	"fmt"
+	"path/filepath"
+
 	"github.com/kairos-io/kairos-agent/v2/pkg/config"
 	"github.com/kairos-io/kairos-agent/v2/pkg/constants"
 	"github.com/kairos-io/kairos-agent/v2/pkg/elemental"
@@ -9,11 +11,6 @@ import (
 	elementalUtils "github.com/kairos-io/kairos-agent/v2/pkg/utils"
 	events "github.com/kairos-io/kairos-sdk/bus"
 	"github.com/kairos-io/kairos-sdk/utils"
-	"github.com/sanity-io/litter"
-	"os"
-	"path/filepath"
-	"sort"
-	"strings"
 )
 
 type UpgradeAction struct {
@@ -39,37 +36,11 @@ func (i *UpgradeAction) Run() (err error) {
 		return err
 	}
 	cleanup.Push(umount)
-	// TODO: Check number of existing UKI files
-	efiFiles, err := i.getEfiFiles()
-	if err != nil {
-		return err
-	}
-	i.cfg.Logger.Infof("Found %d UKI files", len(efiFiles))
-	if len(efiFiles) > i.cfg.UkiMaxEntries && i.cfg.UkiMaxEntries > 0 {
-		i.cfg.Logger.Infof("Found %d UKI files, which is over max entries allowed(%d) removing the oldest one", len(efiFiles), i.cfg.UkiMaxEntries)
-		versionList := semver.Collection{}
-		for _, f := range efiFiles {
-			versionList = append(versionList, semver.MustParse(f))
-		}
-		// Sort it so the oldest one is first
-		sort.Sort(versionList)
-		i.cfg.Logger.Debugf("All versions found: %s", litter.Sdump(versionList))
-		// Remove the oldest one
-		i.cfg.Logger.Infof("Removing: %s", filepath.Join(i.spec.EfiPartition.MountPoint, "EFI", "kairos", versionList[0].Original()))
-		err = i.cfg.Fs.Remove(filepath.Join(i.spec.EfiPartition.MountPoint, "EFI", "kairos", versionList[0].Original()))
-		if err != nil {
-			return err
-		}
-		// Remove the conf file as well
-		i.cfg.Logger.Infof("Removing: %s", filepath.Join(i.spec.EfiPartition.MountPoint, "loader", "entries", versionList[0].String()+".conf"))
-		// Don't care about errors here, systemd-boot will ignore any configs if it cant find the efi file mentioned in it
-		e := i.cfg.Fs.Remove(filepath.Join(i.spec.EfiPartition.MountPoint, "loader", "entries", versionList[0].String()+".conf"))
-		if e != nil {
-			i.cfg.Logger.Warnf("Failed to remove conf file: %s", e)
-		}
-	} else {
-		i.cfg.Logger.Infof("Found %d UKI files, which is under max entries allowed(%d) not removing any", len(efiFiles), i.cfg.UkiMaxEntries)
-	}
+
+	// TODO: Get the size of the efi partition and decide if the images can fit
+	// before trying the upgrade.
+	// If we decide to first copy and then rotate, we need ~4 times the size of
+	// the artifact set [TBD]
 
 	// Dump artifact to efi dir
 	_, err = e.DumpSource(constants.UkiEfiDir, i.spec.Active.Source)
@@ -77,23 +48,29 @@ func (i *UpgradeAction) Run() (err error) {
 		return err
 	}
 
+	// Rotate first
+	err = overwriteArtifactSetRole(i.cfg.Fs, constants.UkiEfiDir, "active", "passive", i.cfg.Logger)
+	if err != nil {
+		return fmt.Errorf("rotating active to passive: %w", err)
+	}
+
+	// Install the new artifacts as "active"
+	err = overwriteArtifactSetRole(i.cfg.Fs, constants.UkiEfiDir, UnassignedArtifactRole, "active", i.cfg.Logger)
+	if err != nil {
+		return fmt.Errorf("installing the new artifacts as active: %w", err)
+	}
+
+	loaderConfPath := filepath.Join(constants.UkiEfiDir, "loader", "loader.conf")
+	if err = replaceRoleInKey(loaderConfPath, "default", UnassignedArtifactRole, "active", i.cfg.Logger); err != nil {
+		return err
+	}
+
+	if err = removeArtifactSetWithRole(i.cfg.Fs, constants.UkiEfiDir, UnassignedArtifactRole); err != nil {
+		return fmt.Errorf("removing artifact set: %w", err)
+	}
+
 	_ = elementalUtils.RunStage(i.cfg, "kairos-uki-upgrade.after")
 	_ = events.RunHookScript("/usr/bin/kairos-agent.uki.upgrade.after.hook") //nolint:errcheck
 
 	return nil
 }
-
-func (i *UpgradeAction) getEfiFiles() ([]string, error) {
-	var efiFiles []string
-	files, err := os.ReadDir(filepath.Join(i.spec.EfiPartition.MountPoint, "EFI", "kairos"))
-	if err != nil {
-		return efiFiles, err
-	}
-
-	for _, file := range files {
-		if !file.IsDir() && strings.HasSuffix(file.Name(), ".efi") {
-			efiFiles = append(efiFiles, file.Name())
-		}
-	}
-	return efiFiles, nil
-}