From 609c27463683c5537cac5e518a386f9817c4a906 Mon Sep 17 00:00:00 2001
From: Mauro Morales <mauro.morales@spectrocloud.com>
Date: Thu, 29 Jun 2023 15:40:24 +0200
Subject: [PATCH 1/2] Remove validation of os-release attribute when source has
 been given for upgrades

Signed-off-by: Mauro Morales <mauro.morales@spectrocloud.com>
---
 internal/agent/upgrade.go | 29 ++++++++++++++---------------
 1 file changed, 14 insertions(+), 15 deletions(-)

diff --git a/internal/agent/upgrade.go b/internal/agent/upgrade.go
index b1e6ad5e..b6996438 100644
--- a/internal/agent/upgrade.go
+++ b/internal/agent/upgrade.go
@@ -79,10 +79,13 @@ func Upgrade(
 		}
 	}
 
-	discoveredImage := ""
-	bus.Manager.Response(events.EventVersionImage, func(p *pluggable.Plugin, r *pluggable.EventResponse) {
-		discoveredImage = r.Data
-	})
+	img := source
+
+	if img == "" {
+		bus.Manager.Response(events.EventVersionImage, func(p *pluggable.Plugin, r *pluggable.EventResponse) {
+			img = r.Data
+		})
+	}
 
 	_, err := bus.Manager.Publish(events.EventVersionImage, &events.VersionImagePayload{
 		Version: version,
@@ -91,18 +94,14 @@ func Upgrade(
 		return err
 	}
 
-	registry, err := utils.OSRelease("IMAGE_REPO")
-	if err != nil {
-		fmt.Printf("Cant find IMAGE_REPO key under /etc/os-release\n")
-		return err
-	}
+	if img == "" {
+		registry, err := utils.OSRelease("IMAGE_REPO")
+		if err != nil {
+			fmt.Printf("Can't find IMAGE_REPO key under /etc/os-release\n")
+			return err
+		}
 
-	img := fmt.Sprintf("%s:%s", registry, version)
-	if discoveredImage != "" {
-		img = discoveredImage
-	}
-	if source != "" {
-		img = source
+		img = fmt.Sprintf("%s:%s", registry, version)
 	}
 
 	if debug {

From 62eee9ffbfea95d69df71708f93fa792076c60ab Mon Sep 17 00:00:00 2001
From: Mauro Morales <mauro.morales@spectrocloud.com>
Date: Fri, 30 Jun 2023 11:37:50 +0300
Subject: [PATCH 2/2] Refactor

Signed-off-by: Dimitris Karakasilis <dimitris@spectrocloud.com>
---
 internal/agent/upgrade.go | 49 +++++++++++++++++++++++++--------------
 1 file changed, 31 insertions(+), 18 deletions(-)

diff --git a/internal/agent/upgrade.go b/internal/agent/upgrade.go
index b6996438..458a7887 100644
--- a/internal/agent/upgrade.go
+++ b/internal/agent/upgrade.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+
 	"github.com/Masterminds/semver/v3"
 	events "github.com/kairos-io/kairos-sdk/bus"
 	"github.com/kairos-io/kairos-sdk/collector"
@@ -80,28 +81,13 @@ func Upgrade(
 	}
 
 	img := source
-
-	if img == "" {
-		bus.Manager.Response(events.EventVersionImage, func(p *pluggable.Plugin, r *pluggable.EventResponse) {
-			img = r.Data
-		})
-	}
-
-	_, err := bus.Manager.Publish(events.EventVersionImage, &events.VersionImagePayload{
-		Version: version,
-	})
-	if err != nil {
-		return err
-	}
-
+	var err error
 	if img == "" {
-		registry, err := utils.OSRelease("IMAGE_REPO")
+		img, err = determineUpgradeImage(version)
 		if err != nil {
-			fmt.Printf("Can't find IMAGE_REPO key under /etc/os-release\n")
+			fmt.Println(err.Error())
 			return err
 		}
-
-		img = fmt.Sprintf("%s:%s", registry, version)
 	}
 
 	if debug {
@@ -145,3 +131,30 @@ func Upgrade(
 
 	return upgradeAction.Run()
 }
+
+// determineUpgradeImage asks the provider plugin for an image or constructs
+// it using version and data from /etc/os-release
+func determineUpgradeImage(version string) (string, error) {
+	var img string
+	bus.Manager.Response(events.EventVersionImage, func(p *pluggable.Plugin, r *pluggable.EventResponse) {
+		img = r.Data
+	})
+
+	_, err := bus.Manager.Publish(events.EventVersionImage, &events.VersionImagePayload{
+		Version: version,
+	})
+	if err != nil {
+		return "", err
+	}
+
+	if img != "" {
+		return img, nil
+	}
+
+	registry, err := utils.OSRelease("IMAGE_REPO")
+	if err != nil {
+		return "", fmt.Errorf("can't find IMAGE_REPO key under /etc/os-release %w", err)
+	}
+
+	return fmt.Sprintf("%s:%s", registry, version), nil
+}