From f477776b579e5c25915d9c74c0579f5056ecc4f9 Mon Sep 17 00:00:00 2001
From: Dimitris Karakasilis <dimitris@karakasilis.me>
Date: Fri, 29 Sep 2023 10:12:12 +0300
Subject: [PATCH] Bring back the size calculation in Upgrade

it was removed by mistake

Signed-off-by: Dimitris Karakasilis <dimitris@karakasilis.me>
---
 pkg/config/spec.go | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/pkg/config/spec.go b/pkg/config/spec.go
index 7c3acfa..e7b6b26 100644
--- a/pkg/config/spec.go
+++ b/pkg/config/spec.go
@@ -265,6 +265,11 @@ func NewUpgradeSpec(cfg *Config) (*v1.UpgradeSpec, error) {
 		State:      installState,
 	}
 
+	err = setSourceSize(cfg, spec)
+	if err != nil {
+		cfg.Logger.Warnf("Failed to infer size for images: %s", err.Error())
+	}
+
 	return spec, nil
 }
 
@@ -693,3 +698,27 @@ func isMounted(config *Config, part *v1.Partition) (bool, error) {
 	}
 	return !notMnt, nil
 }
+
+func setSourceSize(cfg *Config, spec *v1.UpgradeSpec) error {
+	var size int64
+	var err error
+
+	if spec.RecoveryUpgrade {
+		size, err = GetSourceSize(cfg, spec.Recovery.Source)
+	} else {
+		size, err = GetSourceSize(cfg, spec.Active.Source)
+	}
+	if err != nil {
+		return err
+	}
+
+	cfg.Logger.Infof("Setting image size to %dMb", size)
+	// On upgrade only the active or recovery will be upgraded, so we dont need to override passive
+	if spec.RecoveryUpgrade {
+		spec.Recovery.Size = uint(size)
+	} else {
+		spec.Active.Size = uint(size)
+	}
+
+	return nil
+}