From bcb83a79a78f36f872c7c6b1d96db97be1af9e59 Mon Sep 17 00:00:00 2001
From: vardhaman22 <vardhaman.surana@suse.com>
Date: Mon, 25 Sep 2023 21:02:14 +0530
Subject: [PATCH] changed etcd restoreImage to rke-tools for etcd >=3.5.7

---
 cluster/etcd.go      | 37 ++++++++++++++++++++++++++-
 cluster/etcd_test.go | 61 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 97 insertions(+), 1 deletion(-)
 create mode 100644 cluster/etcd_test.go

diff --git a/cluster/etcd.go b/cluster/etcd.go
index 897f440c..cd84a12a 100644
--- a/cluster/etcd.go
+++ b/cluster/etcd.go
@@ -16,6 +16,8 @@ import (
 	"golang.org/x/sync/errgroup"
 )
 
+const MinEtcdVersionWithDistrolessImage = "v3.5.7"
+
 func (c *Cluster) SnapshotEtcd(ctx context.Context, snapshotName string) error {
 	backupImage := c.getBackupImage()
 	containerTimeout := DefaultEtcdBackupConfigTimeout
@@ -225,13 +227,14 @@ func (c *Cluster) RestoreEtcdSnapshot(ctx context.Context, snapshotPath string)
 	// Start restore process on all etcd hosts
 	initCluster := services.GetEtcdInitialCluster(c.EtcdHosts)
 	backupImage := c.getBackupImage()
+	restoreImage := c.getRestoreImage()
 	for _, host := range c.EtcdHosts {
 		containerTimeout := DefaultEtcdBackupConfigTimeout
 		if c.Services.Etcd.BackupConfig != nil && c.Services.Etcd.BackupConfig.Timeout > 0 {
 			containerTimeout = c.Services.Etcd.BackupConfig.Timeout
 		}
 		newCtx := context.WithValue(ctx, docker.WaitTimeoutContextKey, containerTimeout)
-		if err := services.RestoreEtcdSnapshot(newCtx, host, c.PrivateRegistriesMap, c.SystemImages.Etcd, backupImage,
+		if err := services.RestoreEtcdSnapshot(newCtx, host, c.PrivateRegistriesMap, restoreImage, backupImage,
 			snapshotPath, initCluster, c.Services.Etcd, c.Version); err != nil {
 			return fmt.Errorf("[etcd] Failed to restore etcd snapshot: %v", err)
 		}
@@ -281,3 +284,35 @@ func (c *Cluster) getBackupImage() string {
 	logrus.Debugf("[etcd] Image used for etcd snapshot is: [%s]", rkeToolsImage)
 	return rkeToolsImage
 }
+
+func (c *Cluster) getRestoreImage() string {
+
+	// use etcd image for restore in case of custom system image
+	if !strings.Contains(c.SystemImages.Etcd, "rancher/mirrored-coreos-etcd") {
+		return c.SystemImages.Etcd
+	}
+
+	etcdImageTag, err := util.GetImageTagFromImage(c.SystemImages.Etcd)
+	if err != nil {
+		logrus.Errorf("[etcd] getRestoreImage: error extracting tag from etcd image: %v", err)
+		return ""
+	}
+
+	etcdVersion, err := util.StrToSemVer(etcdImageTag)
+	if err != nil {
+		logrus.Errorf("[etcd] getRestoreImage: error converting etcd image tag to semver: %v", err)
+		return ""
+	}
+
+	minEtcdVersionWithDistrolessImage, err := util.StrToSemVer(MinEtcdVersionWithDistrolessImage)
+	if err != nil {
+		logrus.Errorf("[etcd] getRestoreImage: error converting min distroless etcd image version to semver: %v", err)
+		return ""
+	}
+
+	if etcdVersion.LessThan(*minEtcdVersionWithDistrolessImage) {
+		return c.SystemImages.Etcd
+	}
+
+	return c.getBackupImage()
+}
diff --git a/cluster/etcd_test.go b/cluster/etcd_test.go
new file mode 100644
index 00000000..1a3d685c
--- /dev/null
+++ b/cluster/etcd_test.go
@@ -0,0 +1,61 @@
+package cluster
+
+import (
+	"context"
+	"testing"
+
+	"github.com/rancher/rke/metadata"
+	"github.com/rancher/rke/types"
+	v3 "github.com/rancher/rke/types"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetRestoreImage(t *testing.T) {
+	ctx := context.Background()
+
+	metadata.InitMetadata(ctx)
+
+	cluster := &Cluster{
+		RancherKubernetesEngineConfig: v3.RancherKubernetesEngineConfig{
+			SystemImages: types.RKESystemImages{
+				Etcd:   "rancher/mirrored-coreos-etcd:v3.5.7",
+				Alpine: "rancher/rke-tools:v0.1.90",
+			},
+		},
+	}
+
+	expectedRestoreImage := cluster.getBackupImage()
+	restoreImage := cluster.getRestoreImage()
+
+	assert.NotEmpty(t, restoreImage, "")
+	assert.Equal(t, expectedRestoreImage, restoreImage,
+		"expected restoreImage is different when etcd image version is v3.5.7")
+
+	cluster.SystemImages.Etcd = "rancher/mirrored-coreos-etcd:v3.5.8"
+
+	expectedRestoreImage = cluster.getBackupImage()
+	restoreImage = cluster.getRestoreImage()
+
+	assert.NotEmpty(t, restoreImage, "")
+	assert.Equal(t, expectedRestoreImage, restoreImage,
+		"expected restoreImage is different when etcd image version is greater than v3.5.7")
+
+	cluster.SystemImages.Etcd = "rancher/mirrored-coreos-etcd:v3.5.6"
+
+	expectedRestoreImage = cluster.SystemImages.Etcd
+	restoreImage = cluster.getRestoreImage()
+
+	assert.NotEmpty(t, restoreImage, "")
+	assert.Equal(t, expectedRestoreImage, restoreImage,
+		"expected restoreImage is different when etcd image version is less than v3.5.7")
+
+	// test for custom image
+	cluster.SystemImages.Etcd = "custom/mirrored-coreos-etcd:v3.5.7"
+
+	expectedRestoreImage = cluster.SystemImages.Etcd
+	restoreImage = cluster.getRestoreImage()
+
+	assert.NotEmpty(t, restoreImage, "")
+	assert.Equal(t, expectedRestoreImage, restoreImage,
+		"expected restoreImage is different when custom etcd image is used")
+}