Properly format windows env-vars, redeploy windows plans if they fail…

… up to 5 times
HarrisonWAffel · Oct 10, 2024 · 7f259ef · 7f259ef
1 parent bf39f57
commit 7f259ef
Show file tree

Hide file tree

Showing 14 changed files with 285 additions and 77 deletions.
diff --git a/pkg/apis/rke.cattle.io/v1/plan/plan.go b/pkg/apis/rke.cattle.io/v1/plan/plan.go
@@ -47,12 +47,13 @@ type Secret struct {
 }
 
 type OneTimeInstruction struct {
-	Name       string   `json:"name,omitempty"`
-	Image      string   `json:"image,omitempty"`
-	Env        []string `json:"env,omitempty"`
-	Args       []string `json:"args,omitempty"`
-	Command    string   `json:"command,omitempty"`
-	SaveOutput bool     `json:"saveOutput,omitempty"`
+	Name                              string   `json:"name,omitempty"`
+	Image                             string   `json:"image,omitempty"`
+	Env                               []string `json:"env,omitempty"`
+	Args                              []string `json:"args,omitempty"`
+	Command                           string   `json:"command,omitempty"`
+	SaveOutput                        bool     `json:"saveOutput,omitempty"`
+	ResetFailureCountOnServiceRestart bool     `json:"resetFailuresOnServiceRestart,omitempty"`
 }
 
 type PeriodicInstruction struct {

diff --git a/pkg/capr/common.go b/pkg/capr/common.go
@@ -674,3 +674,54 @@ func PreBootstrap(mgmtCluster *v3.Cluster) bool {
 
 	return !v3.ClusterConditionPreBootstrapped.IsTrue(mgmtCluster)
 }
+
+// FormatWindowsEnvVar accepts a corev1.EnvVar and returns a string to be used in either
+// a Powershell script or the Rancher planner, indicated by the isPlanVariable parameter.
+// This function automatically configures the '$env:' prefix for a given environment variable,
+// automatically prefixes boolean values with '$', and surrounds string variables with double quotes as
+// needed. If the provided variable name incorrectly uses either '$env:' or '$' for the given isPlanVariable
+// value, it will be removed.
+func FormatWindowsEnvVar(envVar corev1.EnvVar, isPlanVariable bool) string {
+	lowerValue := strings.ToLower(envVar.Value)
+	isBool := lowerValue == "$true" || lowerValue == "$false" ||
+		lowerValue == "true" || lowerValue == "false"
+
+	// remove any user provided prefixes and suffixes
+	if strings.HasPrefix(envVar.Name, "$env:") {
+		envVar.Name = strings.TrimPrefix(envVar.Name, "$env:")
+	}
+
+	if strings.HasPrefix(envVar.Value, "\"") {
+		envVar.Value = strings.TrimPrefix(envVar.Value, "\"")
+	}
+
+	if strings.HasSuffix(envVar.Value, "\"") {
+		envVar.Value = strings.TrimSuffix(envVar.Value, "\"")
+	}
+
+	if !isBool {
+		format := ""
+		if isPlanVariable {
+			format = "%s=%s"
+		} else {
+			// Non-boolean variables are always treated as strings,
+			// even numbers
+			format = "$env:%s=\"%s\""
+		}
+		return fmt.Sprintf(format, envVar.Name, envVar.Value)
+	}
+
+	if !strings.HasPrefix(envVar.Value, "$") && !isPlanVariable {
+		envVar.Value = "$" + envVar.Value
+	}
+
+	if strings.HasPrefix(envVar.Value, "$") && isPlanVariable {
+		envVar.Value = strings.TrimPrefix(envVar.Value, "$")
+	}
+
+	if isPlanVariable {
+		return fmt.Sprintf("%s=%s", envVar.Name, envVar.Value)
+	}
+
+	return fmt.Sprintf("$env:%s=%s", envVar.Name, envVar.Value)
+}
diff --git a/pkg/capr/common_test.go b/pkg/capr/common_test.go
@@ -1,6 +1,7 @@
 package capr
 
 import (
+	corev1 "k8s.io/api/core/v1"
 	"reflect"
 	"testing"
 
@@ -313,3 +314,76 @@ func TestCompressInterface(t *testing.T) {
 		})
 	}
 }
+
+func TestFormatWindowsEnvVar(t *testing.T) {
+	tests := []struct {
+		Name           string
+		EnvVar         corev1.EnvVar
+		IsPlanVar      bool
+		ExpectedString string
+	}{
+		{
+			Name: "Basic String",
+			EnvVar: corev1.EnvVar{
+				Name:  "BASIC_STRING",
+				Value: "ABC123",
+			},
+			IsPlanVar:      false,
+			ExpectedString: "$env:BASIC_STRING=\"ABC123\"",
+		},
+		{
+			Name: "Basic Bool",
+			EnvVar: corev1.EnvVar{
+				Name:  "BASIC_BOOL",
+				Value: "true",
+			},
+			IsPlanVar:      false,
+			ExpectedString: "$env:BASIC_BOOL=$true",
+		},
+		{
+			Name: "Basic Plan String",
+			EnvVar: corev1.EnvVar{
+				Name:  "PLAN_STRING",
+				Value: "VALUE",
+			},
+			IsPlanVar:      true,
+			ExpectedString: "PLAN_STRING=VALUE",
+		},
+		{
+			Name: "Basic Plan Bool",
+			EnvVar: corev1.EnvVar{
+				Name:  "PLAN_BOOL",
+				Value: "true",
+			},
+			IsPlanVar:      true,
+			ExpectedString: "PLAN_BOOL=true",
+		},
+		{
+			Name: "Plan name Mistakenly includes $env:",
+			EnvVar: corev1.EnvVar{
+				Name:  "$env:PLAN_BOOL",
+				Value: "true",
+			},
+			IsPlanVar:      true,
+			ExpectedString: "PLAN_BOOL=true",
+		},
+		{
+			Name: "Plan value Mistakenly Includes $",
+			EnvVar: corev1.EnvVar{
+				Name:  "PLAN_BOOL",
+				Value: "$true",
+			},
+			IsPlanVar:      true,
+			ExpectedString: "PLAN_BOOL=true",
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.Name, func(t *testing.T) {
+			out := FormatWindowsEnvVar(tc.EnvVar, tc.IsPlanVar)
+			if out != tc.ExpectedString {
+				t.Fatalf("Expected %s, got %s", tc.ExpectedString, out)
+			}
+		})
+	}
+}
diff --git a/pkg/capr/installer/installer.go b/pkg/capr/installer/installer.go
@@ -8,6 +8,7 @@ import (
 	"os"
 	"strings"
 
+	"github.com/rancher/rancher/pkg/capr"
 	"github.com/rancher/rancher/pkg/settings"
 	"github.com/rancher/rancher/pkg/systemtemplate"
 	"github.com/rancher/rancher/pkg/tls"
@@ -135,9 +136,15 @@ func WindowsInstallScript(ctx context.Context, token string, envVars []corev1.En
 	binaryURL := ""
 	if settings.WinsAgentVersion.Get() != "" {
 		if settings.ServerURL.Get() != "" {
-			binaryURL = fmt.Sprintf("$env:CATTLE_AGENT_BINARY_BASE_URL=\"%s/assets\"", settings.ServerURL.Get())
+			binaryURL = capr.FormatWindowsEnvVar(corev1.EnvVar{
+				Name:  "CATTLE_AGENT_BINARY_BASE_URL",
+				Value: fmt.Sprintf("%s/assets", settings.ServerURL.Get()),
+			}, false)
 		} else if defaultHost != "" {
-			binaryURL = fmt.Sprintf("$env:CATTLE_AGENT_BINARY_BASE_URL=\"https://%s/assets\"", defaultHost)
+			binaryURL = capr.FormatWindowsEnvVar(corev1.EnvVar{
+				Name:  "CATTLE_AGENT_BINARY_BASE_URL",
+				Value: fmt.Sprintf("https://%s/assets", defaultHost),
+			}, false)
 		}
 	}
 
@@ -156,22 +163,38 @@ func WindowsInstallScript(ctx context.Context, token string, envVars []corev1.En
 	if v, ok := ctx.Value(tls.InternalAPI).(bool); ok && v {
 		ca = systemtemplate.InternalCAChecksum()
 	}
+
 	if ca != "" {
-		ca = "$env:CATTLE_CA_CHECKSUM=\"" + ca + "\""
+		ca = capr.FormatWindowsEnvVar(corev1.EnvVar{
+			Name:  "CATTLE_CA_CHECKSUM",
+			Value: ca,
+		}, false)
 	}
+
+	var tokenEnvVar string
 	if token != "" {
-		token = "$env:CATTLE_ROLE_NONE=\"true\"\n$env:CATTLE_TOKEN=\"" + token + "\""
+		tokenEnvVar = capr.FormatWindowsEnvVar(corev1.EnvVar{
+			Name:  "CATTLE_TOKEN",
+			Value: token,
+		}, false)
 	}
+
 	envVarBuf := &strings.Builder{}
 	for _, envVar := range envVars {
 		if envVar.Value == "" {
 			continue
 		}
-		envVarBuf.WriteString(fmt.Sprintf("$env:%s=\"%s\"\n", envVar.Name, envVar.Value))
+		envVarBuf.WriteString(capr.FormatWindowsEnvVar(corev1.EnvVar{
+			Name:  envVar.Name,
+			Value: envVar.Value,
+		}, false))
 	}
 	server := ""
 	if settings.ServerURL.Get() != "" {
-		server = fmt.Sprintf("$env:CATTLE_SERVER=\"%s\"", settings.ServerURL.Get())
+		server = capr.FormatWindowsEnvVar(corev1.EnvVar{
+			Name:  "CATTLE_SERVER",
+			Value: settings.ServerURL.Get(),
+		}, false)
 	}
 
 	strictVerify := "false"
@@ -192,8 +215,9 @@ $env:CSI_PROXY_URL = "%s"
 $env:CSI_PROXY_VERSION = "%s"
 $env:CSI_PROXY_KUBELET_PATH = "C:%s/bin/kubelet.exe"
 $env:STRICT_VERIFY = "%s"
+$env:CATTLE_ROLE_NONE="true"
 
 Invoke-WinsInstaller @PSBoundParameters
 exit 0
-`, data, envVarBuf.String(), binaryURL, server, ca, token, csiProxyURL, csiProxyVersion, dataDir, strictVerify)), nil
+`, data, envVarBuf.String(), binaryURL, server, ca, tokenEnvVar, csiProxyURL, csiProxyVersion, dataDir, strictVerify)), nil
 }
diff --git a/pkg/capr/planner/etcdcreate.go b/pkg/capr/planner/etcdcreate.go
@@ -65,9 +65,7 @@ func (p *Planner) runEtcdSnapshotCreate(controlPlane *rkev1.RKEControlPlane, tok
 // Notably, this function will blatantly ignore drain and concurrency options, as during an etcd snapshot operation, there is no necessity to drain nodes.
 func (p *Planner) runEtcdSnapshotManagementServiceStart(controlPlane *rkev1.RKEControlPlane, tokensSecret plan.Secret, clusterPlan *plan.Plan, include roleFilter, operation string) error {
 	// Generate and deliver desired plan for the bootstrap/init node first.
-	if err := p.reconcile(controlPlane, tokensSecret, clusterPlan, true, bootstrapTier, isEtcd, isNotInitNodeOrIsDeleting,
-		"1", "",
-		controlPlane.Spec.UpgradeStrategy.ControlPlaneDrainOptions); err != nil {
+	if err := p.reconcile(controlPlane, tokensSecret, clusterPlan, true, bootstrapTier, isEtcd, isNotInitNodeOrIsDeleting, "1", "", controlPlane.Spec.UpgradeStrategy.ControlPlaneDrainOptions, -1, 1, false); err != nil {
 		return err
 	}
 
@@ -84,7 +82,7 @@ func (p *Planner) runEtcdSnapshotManagementServiceStart(controlPlane *rkev1.RKEC
 		if isInitNodeOrDeleting(entry) {
 			continue
 		}
-		plan, joinedServer, err := p.desiredPlan(controlPlane, tokensSecret, entry, joinServer)
+		plan, joinedServer, err := p.desiredPlan(controlPlane, tokensSecret, entry, joinServer, false)
 		if err != nil {
 			return err
 		}

diff --git a/pkg/capr/planner/etcdrestore.go b/pkg/capr/planner/etcdrestore.go
@@ -230,7 +230,7 @@ func (p *Planner) runEtcdSnapshotPostRestorePodCleanupPlan(controlPlane *rkev1.R
 	}
 	initNode := initNodes[0]
 
-	initNodePlan, _, err := p.desiredPlan(controlPlane, tokensSecret, initNode, "")
+	initNodePlan, _, err := p.desiredPlan(controlPlane, tokensSecret, initNode, "", false)
 	if err != nil {
 		return err
 	}
@@ -262,7 +262,7 @@ func (p *Planner) runEtcdSnapshotPostRestorePodCleanupPlan(controlPlane *rkev1.R
 
 	controlPlaneEntry := controlPlaneEntries[0]
 
-	firstControlPlanePlan, joinedServer, err := p.desiredPlan(controlPlane, tokensSecret, controlPlaneEntry, joinServer)
+	firstControlPlanePlan, joinedServer, err := p.desiredPlan(controlPlane, tokensSecret, controlPlaneEntry, joinServer, false)
 	if err != nil {
 		return err
 	}
@@ -280,7 +280,7 @@ func (p *Planner) runEtcdSnapshotPostRestoreNodeCleanupPlan(controlPlane *rkev1.
 	}
 	initNode := initNodes[0]
 
-	initNodePlan, _, err := p.desiredPlan(controlPlane, tokensSecret, initNode, "")
+	initNodePlan, _, err := p.desiredPlan(controlPlane, tokensSecret, initNode, "", false)
 	if err != nil {
 		return err
 	}