Respect user defined/default sizes on upgrade

pkg/action/upgrade_test.go

@@ -18,22 +18,17 @@ package action_test
 
 import (
 	"bytes"
-	"encoding/json"
 	"fmt"
-	"os"
-	"path/filepath"
-	"strings"
-
-	"github.com/kairos-io/kairos-agent/v2/internal/agent"
 	"github.com/kairos-io/kairos-agent/v2/pkg/action"
 	agentConfig "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"
 	fsutils "github.com/kairos-io/kairos-agent/v2/pkg/utils/fs"
 	v1mock "github.com/kairos-io/kairos-agent/v2/tests/mocks"
-	"github.com/kairos-io/kairos-sdk/collector"
 	ghwMock "github.com/kairos-io/kairos-sdk/ghw/mocks"
 	sdkTypes "github.com/kairos-io/kairos-sdk/types"
+	"os"
+	"path/filepath"
 
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
@@ -55,10 +50,8 @@ var _ = Describe("Upgrade Actions test", func() {
 	var ghwTest ghwMock.GhwMock
 	var extractor *v1mock.FakeImageExtractor
 	var dummySourceFile string
-	var dummySourceSizeMb int64
 
 	BeforeEach(func() {
-		dummySourceSizeMb = 20
 		runner = v1mock.NewFakeRunner()
 		syscall = &v1mock.FakeSyscall{}
 		mounter = v1mock.NewErrorMounter()
@@ -153,34 +146,6 @@ var _ = Describe("Upgrade Actions test", func() {
 		AfterEach(func() {
 			ghwTest.Clean()
 		})
-		It("calculates the recovery source size correctly", func() {
-			dummySourceFile = createDummyFile(fs, dummySourceSizeMb)
-			upgradeConfig := agent.ExtraConfigUpgrade{}
-			upgradeConfig.Upgrade.Entry = constants.BootEntryRecovery
-			upgradeConfig.Upgrade.RecoverySystem.URI = fmt.Sprintf("file:%s", dummySourceFile)
-			d, err := json.Marshal(upgradeConfig)
-			Expect(err).ToNot(HaveOccurred())
-			cliConfig := string(d)
-
-			config, err := agentConfig.Scan(collector.Readers(strings.NewReader(cliConfig)))
-			Expect(err).ToNot(HaveOccurred())
-
-			agentConfig.WithFs(fs)(config)
-			agentConfig.WithRunner(runner)(config)
-			agentConfig.WithLogger(logger)(config)
-			agentConfig.WithMounter(mounter)(config)
-			agentConfig.WithSyscall(syscall)(config)
-			agentConfig.WithClient(client)(config)
-			agentConfig.WithCloudInitRunner(cloudInit)(config)
-			agentConfig.WithImageExtractor(extractor)(config)
-			agentConfig.WithPlatform("linux/amd64")(config)
-			config.ImageExtractor = extractor
-
-			spec, err = agentConfig.NewUpgradeSpec(config)
-			Expect(err).ShouldNot(HaveOccurred())
-			Expect(spec.Entry).To(Equal(constants.BootEntryRecovery))
-			Expect(spec.Recovery.Size).To(Equal(uint(100 + dummySourceSizeMb))) // We adding 100Mb on top
-		})
 		Describe(fmt.Sprintf("Booting from %s", constants.ActiveLabel), Label("active_label"), func() {
 			var err error
 			BeforeEach(func() {
@@ -659,25 +624,3 @@ var _ = Describe("Upgrade Actions test", func() {
 		})
 	})
 })
-
-func createDummyFile(fs vfs.FS, sizeMb int64) string {
-	fileSize := int64(sizeMb * 1024 * 1024)
-
-	tmpFile, err := os.CreateTemp("", "dummyfile_*.tmp")
-	Expect(err).ToNot(HaveOccurred())
-	tmpName := tmpFile.Name()
-	tmpFile.Close()
-	os.RemoveAll(tmpName)
-
-	dir := filepath.Dir(tmpName)
-	err = fs.Mkdir(dir, os.ModePerm)
-	Expect(err).ToNot(HaveOccurred())
-
-	f, err := fs.Create(tmpName)
-	Expect(err).ShouldNot(HaveOccurred())
-	err = f.Truncate(fileSize)
-	Expect(err).ShouldNot(HaveOccurred())
-	f.Close()
-
-	return tmpName
-}

pkg/config/config.go

@@ -392,10 +392,11 @@ func FilterKeys(d []byte) ([]byte, error) {
 }
 
 // ScanNoLogs is a wrapper around Scan that sets the logger to null
+// Also sets the NoLogs option to true by default
 func ScanNoLogs(opts ...collector.Option) (c *Config, err error) {
 	log := sdkTypes.NewNullLogger()
 	result := NewConfig(WithLogger(log))
-	return scan(result, opts...)
+	return scan(result, append(opts, collector.NoLogs)...)
 }
 
 // Scan is a wrapper around collector.Scan that sets the logger to the default Kairos logger

pkg/config/config_test.go

@@ -228,7 +228,7 @@ var _ = Describe("Schema", func() {
 			cleanup()
 		})
 		It("Scan can override options", func() {
-			c, err := Scan(collector.Readers(strings.NewReader(`uki-max-entries: 34`)), collector.NoLogs)
+			c, err := ScanNoLogs(collector.Readers(strings.NewReader(`uki-max-entries: 34`)))
 			Expect(err).ShouldNot(HaveOccurred())
 			Expect(c.UkiMaxEntries).To(Equal(34))
 		})
@@ -279,12 +279,12 @@ stages:
           groups:
             - "admin"
 `
-			config, err := pkgConfig.Scan(collector.Readers(strings.NewReader(cc)), collector.NoLogs)
+			config, err := pkgConfig.ScanNoLogs(collector.Readers(strings.NewReader(cc)))
 			Expect(err).ToNot(HaveOccurred())
 			Expect(config.CheckForUsers()).ToNot(HaveOccurred())
 		})
 		It("Fails if there is no user", func() {
-			config, err := pkgConfig.Scan()
+			config, err := pkgConfig.ScanNoLogs(collector.NoLogs)
 			Expect(err).ToNot(HaveOccurred())
 			Expect(config.CheckForUsers()).To(HaveOccurred())
 		})
@@ -297,7 +297,7 @@ stages:
         kairos:
           passwd: "kairos"
 `
-			config, err := pkgConfig.Scan(collector.Readers(strings.NewReader(cc)), collector.NoLogs)
+			config, err := pkgConfig.ScanNoLogs(collector.Readers(strings.NewReader(cc)))
 			Expect(err).ToNot(HaveOccurred())
 			Expect(config.CheckForUsers()).To(HaveOccurred())
 		})

pkg/config/spec.go

@@ -403,6 +403,14 @@ func NewUpgradeSpec(cfg *Config) (*v1.UpgradeSpec, error) {
 func setUpgradeSourceSize(cfg *Config, spec *v1.UpgradeSpec) error {
 	var size int64
 	var err error
+	var originalSize uint
+
+	// Store the default given size in the spec. This includes the user specified values which have already been marshalled in the spec
+	if spec.RecoveryUpgrade() {
+		originalSize = spec.Recovery.Size
+	} else {
+		originalSize = spec.Active.Size
+	}
 
 	var targetSpec *v1.Image
 	if spec.RecoveryUpgrade() {
@@ -412,16 +420,25 @@ func setUpgradeSourceSize(cfg *Config, spec *v1.UpgradeSpec) error {
 	}
 
 	if targetSpec.Source != nil && targetSpec.Source.IsEmpty() {
+		cfg.Logger.Debugf("No source specified for image, skipping size calculation")
 		return nil
 	}
 
 	size, err = GetSourceSize(cfg, targetSpec.Source)
 	if err != nil {
+		cfg.Logger.Warnf("Failed to infer size for images: %s", err.Error())
 		return err
 	}
 
-	cfg.Logger.Infof("Setting image size to %dMb", size)
-	targetSpec.Size = uint(size)
+	if uint(size) < originalSize {
+		cfg.Logger.Debugf("Calculated size (%dMB) is less than specified/default size (%dMB)", size, originalSize)
+		targetSpec.Size = originalSize
+	} else {
+		cfg.Logger.Debugf("Calculated size (%dMB) is higher than specified/default size (%dMB)", size, originalSize)
+		targetSpec.Size = uint(size)
+	}
+
+	cfg.Logger.Infof("Setting image size to %dMB", targetSpec.Size)
 
 	return nil
 }
@@ -789,7 +806,7 @@ func ReadUkiUpgradeSpecFromConfig(c *Config) (*v1.UpgradeUkiSpec, error) {
 // getSize will calculate the size of a file or symlink and will do nothing with directories
 // fileList: keeps track of the files visited to avoid counting a file more than once if it's a symlink. It could also be used as a way to filter some files
 // size: will be the memory that adds up all the files sizes. Meaning it could be initialized with a value greater than 0 if needed.
-func getSize(size *int64, fileList map[string]bool, path string, d fs.DirEntry, err error) error {
+func getSize(vfs v1.FS, size *int64, fileList map[string]bool, path string, d fs.DirEntry, err error) error {
 	if err != nil {
 		return err
 	}
@@ -802,7 +819,7 @@ func getSize(size *int64, fileList map[string]bool, path string, d fs.DirEntry,
 	if d.Type()&fs.ModeSymlink != 0 {
 		// If it's a symlink, get its target and calculate its size.
 		var err error
-		actualFilePath, err = os.Readlink(path)
+		actualFilePath, err = vfs.Readlink(path)
 		if err != nil {
 			return err
 		}
@@ -813,7 +830,7 @@ func getSize(size *int64, fileList map[string]bool, path string, d fs.DirEntry,
 		}
 	}
 
-	fileInfo, err := os.Stat(actualFilePath)
+	fileInfo, err := vfs.Stat(actualFilePath)
 	if os.IsNotExist(err) || fileList[actualFilePath] {
 		return nil
 	}
@@ -865,7 +882,7 @@ func GetSourceSize(config *Config, source *v1.ImageSource) (int64, error) {
 				// During install or upgrade outside kubernetes, we dont care about those dirs as they are not expected to be in the source dir
 				config.Logger.Logger.Debug().Str("path", path).Str("hostDir", hostDir).Msg("Skipping dir as it is a runtime directory under kubernetes (/proc, /dev or /run)")
 			} else {
-				v := getSize(&size, filesVisited, path, d, err)
+				v := getSize(config.Fs, &size, filesVisited, path, d, err)
 				return v
 			}
 

pkg/config/spec_test.go

@@ -20,6 +20,7 @@ import (
 	"bytes"
 	"os"
 	"path/filepath"
+	"strings"
 
 	"github.com/kairos-io/kairos-agent/v2/pkg/config"
 	"github.com/kairos-io/kairos-agent/v2/pkg/constants"
@@ -449,6 +450,45 @@ var _ = Describe("Types", Label("types", "config"), func() {
 					Expect(spec.Recovery.Source.IsEmpty()).To(BeTrue())
 					Expect(spec.Recovery.FS).To(Equal(constants.SquashFs))
 				})
+
+				It("sets image size to default value if not set", func() {
+					spec, err := config.NewUpgradeSpec(c)
+					Expect(err).ShouldNot(HaveOccurred())
+					Expect(spec.Active.Size).To(Equal(constants.ImgSize))
+				})
+
+				It("sets image size to provided value if set in the config and image is smaller", func() {
+					cfg, err := config.ScanNoLogs(collector.Readers(strings.NewReader("#cloud-config\nupgrade:\n  system:\n    size: 666\n")))
+					// Set manually the config collector in the cfg file before unmarshalling the spec
+					c.Config = cfg.Config
+					spec, err := config.NewUpgradeSpec(c)
+					Expect(err).ShouldNot(HaveOccurred())
+					Expect(spec.Active.Size).To(Equal(uint(666)))
+				})
+				It("sets image size to default value if not set in the config and image is smaller", func() {
+					cfg, err := config.ScanNoLogs(collector.Readers(strings.NewReader("#cloud-config\nupgrade:\n  system:\n    uri: dir:/\n")))
+					// Set manually the config collector in the cfg file before unmarshalling the spec
+					c.Config = cfg.Config
+					spec, err := config.NewUpgradeSpec(c)
+					Expect(err).ShouldNot(HaveOccurred())
+					Expect(spec.Active.Size).To(Equal(constants.ImgSize))
+				})
+
+				It("sets image size to the source if default is smaller", func() {
+					cfg, err := config.ScanNoLogs(collector.Readers(strings.NewReader("#cloud-config\nupgrade:\n  system:\n    uri: file:/tmp/waka\n")))
+					// Set manually the config collector in the cfg file before unmarshalling the spec
+					c.Config = cfg.Config
+					Expect(c.Fs.Mkdir("/tmp", 0777)).ShouldNot(HaveOccurred())
+					Expect(c.Fs.WriteFile("/tmp/waka", []byte("waka"), 0777)).ShouldNot(HaveOccurred())
+					Expect(c.Fs.Truncate("/tmp/waka", 5120*1024*1024)).ShouldNot(HaveOccurred())
+					spec, err := config.NewUpgradeSpec(c)
+					Expect(err).ShouldNot(HaveOccurred())
+					f, err := c.Fs.Stat("/tmp/waka")
+					Expect(err).ShouldNot(HaveOccurred())
+					// Make the same calculation as the code
+					Expect(spec.Active.Size).To(Equal(uint(f.Size()/1000/1000) + 100))
+				})
+
 			})
 		})
 		Describe("Config from cloudconfig", Label("cloud-config"), func() {

pkg/types/v1/common.go

@@ -36,7 +36,7 @@ const (
 // ImageSource represents the source from where an image is created for easy identification
 type ImageSource struct {
 	source  string `yaml:"source"`
-	srcType string
+	srcType string `yaml:"type"`
 }
 
 func (i ImageSource) Value() string {

pkg/types/v1/fs.go

@@ -37,4 +37,5 @@ type FS interface {
 	OpenFile(name string, flag int, perm fs.FileMode) (*os.File, error)
 	WriteFile(filename string, data []byte, perm os.FileMode) error
 	Rename(oldpath, newpath string) error
+	Truncate(name string, size int64) error
 }