Bring back the size calculation in Upgrade

it was removed by mistake

Signed-off-by: Dimitris Karakasilis <dimitris@karakasilis.me>

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
+}