From e1b1364a0d9bbd120c48d7df0badcda5861a26d2 Mon Sep 17 00:00:00 2001
From: Aditya Thebe <contact@adityathebe.com>
Date: Fri, 13 Sep 2024 18:30:31 +0545
Subject: [PATCH] feat: allow certain env vars to be passed on to the cmd

---
 shell/shell.go      | 22 ++++++++++++++++++++++
 shell/shell_test.go | 16 ++++++++++++----
 2 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/shell/shell.go b/shell/shell.go
index c618edec..7b5ae3af 100644
--- a/shell/shell.go
+++ b/shell/shell.go
@@ -22,6 +22,20 @@ import (
 	"github.com/samber/oops"
 )
 
+// List of env var keys that we pass on to the exec command
+var allowedEnvVars = map[string]struct{}{
+	"CLOUDSDK_PYTHON":                       {},
+	"DEBIAN_FRONTEND":                       {},
+	"DOTNET_SYSTEM_GLOBALIZATION_INVARIANT": {},
+	"HOME":                                  {},
+	"LC_CTYPE":                              {},
+	"PATH":                                  {},
+	"PS_INSTALL_FOLDER":                     {},
+	"PS_VERSION":                            {},
+	"PSModuleAnalysisCachePath":             {},
+	"USER":                                  {},
+}
+
 var checkoutLocks = utils.NamedLock{}
 
 type Exec struct {
@@ -76,6 +90,14 @@ func Run(ctx context.Context, exec Exec) (*ExecDetails, error) {
 
 	// Set to a non-nil empty slice to prevent access to current environment variables
 	cmd.Env = []string{}
+
+	for _, e := range os.Environ() {
+		key, _, ok := strings.Cut(e, "=")
+		if _, exists := allowedEnvVars[key]; exists && ok {
+			cmd.Env = append(cmd.Env, e)
+		}
+	}
+
 	if len(envParams.envs) != 0 {
 		ctx.Logger.V(6).Infof("using environment %s", logger.Pretty(envParams.envs))
 		cmd.Env = append(cmd.Env, envParams.envs...)
diff --git a/shell/shell_test.go b/shell/shell_test.go
index f9efc192..d6448cc1 100644
--- a/shell/shell_test.go
+++ b/shell/shell_test.go
@@ -4,6 +4,7 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/flanksource/commons/collections"
 	"github.com/flanksource/duty/context"
 	"github.com/flanksource/duty/types"
 	"github.com/samber/lo"
@@ -23,7 +24,7 @@ func TestEnv(t *testing.T) {
 					{Name: "mc_test_secret", ValueStatic: "abcdef"},
 				},
 			},
-			expectedVars: []string{"mc_test_secret=abcdef"},
+			expectedVars: []string{"mc_test_secret"},
 		},
 		{
 			name: "access multiple custom env vars",
@@ -34,7 +35,7 @@ func TestEnv(t *testing.T) {
 					{Name: "mc_test_secret_id", ValueStatic: "xyz"},
 				},
 			},
-			expectedVars: []string{"mc_test_secret_key=abc", "mc_test_secret_id=xyz"},
+			expectedVars: []string{"mc_test_secret_key", "mc_test_secret_id"},
 		},
 		{
 			name: "no access to process env",
@@ -69,8 +70,15 @@ func TestEnv(t *testing.T) {
 				return key != "PWD" && key != "SHLVL" && key != "_"
 			})
 
-			if !lo.Every(envVars, td.expectedVars) || !lo.Every(td.expectedVars, envVars) {
-				t.Errorf("expected %s, got %s", td.expectedVars, envVars)
+			envVarKeys := lo.Map(envVars, func(v string, _ int) string {
+				key, _, _ := strings.Cut(v, "=")
+				return key
+			})
+
+			expected := collections.MapKeys(allowedEnvVars)
+			expected = append(expected, td.expectedVars...)
+			if !lo.Every(expected, envVarKeys) {
+				t.Errorf("expected %s, got %s", td.expectedVars, envVarKeys)
 			}
 		})
 	}