From 377d032e79259aaa8dc3d51c15e22e78e44551e0 Mon Sep 17 00:00:00 2001
From: Luis Davim <luis.davim@sendoso.com>
Date: Sat, 14 May 2022 14:13:11 +0100
Subject: [PATCH] feat: recursively expad env variables

Signed-off-by: Luis Davim <luis.davim@sendoso.com>
---
 internal/app/cli.go         |  9 ++----
 internal/app/state_files.go |  2 +-
 internal/app/utils.go       | 39 +++++++++++++++++++++++---
 internal/app/utils_test.go  | 56 +++++++++++++++++++++++++++++++++++++
 4 files changed, 94 insertions(+), 12 deletions(-)

diff --git a/internal/app/cli.go b/internal/app/cli.go
index 93cce664..f016c386 100644
--- a/internal/app/cli.go
+++ b/internal/app/cli.go
@@ -6,8 +6,6 @@ import (
 	"os"
 	"sort"
 	"strings"
-
-	"github.com/joho/godotenv"
 )
 
 const (
@@ -251,11 +249,8 @@ func (c *cli) readState(s *state) error {
 		}
 	}
 
-	if len(c.envFiles) != 0 {
-		err := godotenv.Overload(c.envFiles...)
-		if err != nil {
-			return fmt.Errorf("error loading env file: %w", err)
-		}
+	if err := prepareEnv(c.envFiles); err != nil {
+		return err
 	}
 
 	// wipe & create a temporary directory
diff --git a/internal/app/state_files.go b/internal/app/state_files.go
index bb06d1cc..d501248e 100644
--- a/internal/app/state_files.go
+++ b/internal/app/state_files.go
@@ -189,7 +189,7 @@ func (s *state) expand(relativeToFile string) {
 		if r.Chart != "" {
 			var download bool
 			// support env vars in path
-			r.Chart = os.ExpandEnv(r.Chart)
+			r.Chart = os.Expand(r.Chart, getEnv)
 			repoName := strings.Split(r.Chart, "/")[0]
 			_, isRepo := s.HelmRepos[repoName]
 			isRepo = isRepo || stringInSlice(repoName, s.PreconfiguredHelmRepos)
diff --git a/internal/app/utils.go b/internal/app/utils.go
index a5573600..87bccc4b 100644
--- a/internal/app/utils.go
+++ b/internal/app/utils.go
@@ -18,6 +18,8 @@ import (
 	"unicode/utf8"
 
 	"github.com/Masterminds/semver"
+	"github.com/joho/godotenv"
+
 	"github.com/Praqma/helmsman/internal/aws"
 	"github.com/Praqma/helmsman/internal/azure"
 	"github.com/Praqma/helmsman/internal/gcs"
@@ -140,16 +142,45 @@ func readFile(filepath string) string {
 	return string(data)
 }
 
+// getEnv fetches the value for an environment variable
+// recusively expanding the variable's value
+func getEnv(key string) string {
+	value := os.Getenv(key)
+	for envVar.MatchString(value) {
+		value = os.ExpandEnv(value)
+	}
+	return value
+}
+
+// prepareEnv loads dotenv files and recusively expands all environment variables
+func prepareEnv(envFiles []string) error {
+	if len(envFiles) != 0 {
+		err := godotenv.Overload(envFiles...)
+		if err != nil {
+			return fmt.Errorf("error loading env file: %w", err)
+		}
+	}
+	for _, e := range os.Environ() {
+		if !strings.Contains(e, "$") {
+			continue
+		}
+		e = os.Expand(e, getEnv)
+		pair := strings.SplitN(e, "=", 2)
+		os.Setenv(pair[0], pair[1])
+	}
+	return nil
+}
+
 // substituteEnv checks if a string has an env variable (contains '$'), then it returns its value
 // if the env variable is empty or unset, an empty string is returned
 // if the string does not contain '$', it is returned as is.
-func substituteEnv(name string) string {
-	if strings.Contains(name, "$") {
+func substituteEnv(str string) string {
+	if strings.Contains(str, "$") {
 		// add $$ escaping for $ strings
 		os.Setenv("HELMSMAN_DOLLAR", "$")
-		return os.ExpandEnv(strings.ReplaceAll(name, "$$", "${HELMSMAN_DOLLAR}"))
+		return os.ExpandEnv(strings.ReplaceAll(str, "$$", "${HELMSMAN_DOLLAR}"))
 	}
-	return name
+	return str
 }
 
 // validateEnvVars parses a string line-by-line and detect env variables in
diff --git a/internal/app/utils_test.go b/internal/app/utils_test.go
index 1e282016..57a0228f 100644
--- a/internal/app/utils_test.go
+++ b/internal/app/utils_test.go
@@ -5,6 +5,62 @@ import (
 	"testing"
 )
 
+func TestGetEnv(t *testing.T) {
+	tests := []struct {
+		name     string
+		expected func(key string) string
+		key      string
+	}{
+		{
+			name: "direct",
+			key:  "BAR",
+			expected: func(key string) string {
+				expected := "myValue"
+				os.Setenv(key, expected)
+				return expected
+			},
+		},
+		{
+			name: "string_with_var",
+			key:  "BAR",
+			expected: func(key string) string {
+				expected := "contains myValue"
+				os.Setenv("FOO", "myValue")
+				os.Setenv(key, "contains ${FOO}")
+				return expected
+			},
+		},
+		{
+			name: "nested_one_level",
+			key:  "BAR",
+			expected: func(key string) string {
+				expected := "myValue"
+				os.Setenv("FOO", expected)
+				os.Setenv(key, "${FOO}")
+				return expected
+			},
+		},
+		{
+			name: "nested_two_levels",
+			key:  "BAR",
+			expected: func(key string) string {
+				expected := "myValue"
+				os.Setenv("FOZ", expected)
+				os.Setenv("FOO", "$FOZ")
+				os.Setenv(key, "${FOO}")
+				return expected
+			},
+		},
+	}
+	for _, tt := range tests {
+		expected := tt.expected(tt.key)
+		value := getEnv(tt.key)
+		if value != expected {
+			t.Errorf("getEnv() - unexpected value: wanted: %s got: %s", expected, value)
+		}
+	}
+}
+
 func TestOciRefToFilename(t *testing.T) {
 	tests := []struct {
 		name string