fix: allow the context on env load to be canceled [IDE-1791]

pkg/envvars/environment.go

@@ -4,11 +4,12 @@ import (
 	"context"
 	"os"
 	"os/exec"
-	"path/filepath"
 	"runtime"
 	"strings"
+	"sync"
 	"time"
 
+	"github.com/rs/zerolog"
 	"github.com/subosito/gotenv"
 
 	"github.com/snyk/go-application-framework/pkg/utils"
@@ -20,14 +21,93 @@ const (
 	ShellEnvVarName = "SHELL"
 )
 
+var mu sync.Mutex
+
+//nolint:containedctx // This options struct is local to a single call and not stored beyond function execution.
+type loadConfiguredEnvironmentOptions struct {
+	ctx               context.Context
+	customConfigFiles []string
+	logger            *zerolog.Logger
+}
+
+type loadConfiguredEnvironmentOption func(opts *loadConfiguredEnvironmentOptions)
+
+func WithContext(ctx context.Context) loadConfiguredEnvironmentOption {
+	return func(opts *loadConfiguredEnvironmentOptions) {
+		opts.ctx = ctx
+	}
+}
+
+// WithCustomConfigFiles sets the list of custom config files to load.
+// All paths must be absolute; use [utils.MakeRelativePathsAbsolute] to resolve relative paths before calling.
+func WithCustomConfigFiles(absoluteFilePaths []string) loadConfiguredEnvironmentOption {
+	return func(opts *loadConfiguredEnvironmentOptions) {
+		opts.customConfigFiles = absoluteFilePaths
+	}
+}
+
+func WithLogger(logger *zerolog.Logger) loadConfiguredEnvironmentOption {
+	return func(opts *loadConfiguredEnvironmentOptions) {
+		if logger == nil {
+			return
+		}
+		opts.logger = logger
+	}
+}
+
 // LoadConfiguredEnvironment updates the environment with user and local configuration.
 // First Bash's env is read (as a fallback), then the user's preferred SHELL's env is read, then the configuration files.
 // The Bash env PATH is appended to the existing PATH (as a fallback), any other new PATH read is prepended (preferential).
+//
+// Deprecated: Use LoadConfiguredEnvironmentWithOptions instead.
 func LoadConfiguredEnvironment(customConfigFiles []string, workingDirectory string) {
-	bashOutput := getEnvFromShell("bash")
+	absoluteFilePaths := utils.MakeRelativePathsAbsolute(workingDirectory, customConfigFiles)
+	LoadConfiguredEnvironmentWithOptions(WithCustomConfigFiles(absoluteFilePaths))
+}
+
+// LoadConfiguredEnvironmentWithOptions updates the environment with user and local configuration.
+// First Bash's env is read (as a fallback), then the user's preferred SHELL's env is read, then the configuration files.
+// The Bash env PATH is appended to the existing PATH (as a fallback), any other new PATH read is prepended (preferential).
+func LoadConfiguredEnvironmentWithOptions(opts ...loadConfiguredEnvironmentOption) {
+	options := loadConfiguredEnvironmentOptions{
+		ctx:    context.Background(),
+		logger: utils.Ptr(zerolog.Nop()),
+	}
+	for _, opt := range opts {
+		opt(&options)
+	}
+	logger := options.logger.With().Str("method", "LoadConfiguredEnvironment").Logger()
+
+	mu.Lock()
+	defer func() {
+		mu.Unlock()
+		if options.ctx.Err() == nil {
+			logger.Trace().Msg("Loaded configured environment")
+		} else {
+			logger.Trace().Err(options.ctx.Err()).Msg("Loading configured environment was canceled")
+		}
+	}()
+	logger.Debug().Msg("Loading configured environment")
+
+	// Check the context hasn't been canceled before loading the environment.
+	if ctxErr := options.ctx.Err(); ctxErr != nil {
+		return
+	}
+
+	bashOutput := getEnvFromShell(options.ctx, options.logger, "bash")
+
+	// Check the context hadn't been canceled while loading the Bash environment.
+	if ctxErr := options.ctx.Err(); ctxErr != nil {
+		return
+	}
 
 	// this is applied at the end always, as it does not overwrite existing variables
-	defer func() { _ = gotenv.Apply(strings.NewReader(bashOutput)) }() //nolint:errcheck // we can't do anything with the error
+	defer func() {
+		applyErr := gotenv.Apply(strings.NewReader(bashOutput))
+		if applyErr != nil {
+			logger.Trace().Err(applyErr).Msg("Failed to apply environment variables from Bash")
+		}
+	}()
 
 	bashEnv := gotenv.Parse(strings.NewReader(bashOutput))
 
@@ -37,8 +117,17 @@ func LoadConfiguredEnvironment(customConfigFiles []string, workingDirectory stri
 
 	specificShell, ok := bashEnv[ShellEnvVarName]
 	if ok {
-		fromSpecificShell := getEnvFromShell(specificShell)
-		_ = gotenv.Apply(strings.NewReader(fromSpecificShell)) //nolint:errcheck // we can't do anything with the error
+		fromSpecificShell := getEnvFromShell(options.ctx, options.logger, specificShell)
+
+		// Check the context hadn't been canceled while loading the user's preferred shell environment.
+		if ctxErr := options.ctx.Err(); ctxErr != nil {
+			return
+		}
+
+		applyErr := gotenv.Apply(strings.NewReader(fromSpecificShell))
+		if applyErr != nil {
+			logger.Trace().Err(applyErr).Str("shell", specificShell).Msg("Failed to apply environment variables from the user's preferred shell")
+		}
 
 		specificShellEnv := gotenv.Parse(strings.NewReader(fromSpecificShell))
 		if specificShellPATH, ok := specificShellEnv[PathEnvVarName]; ok {
@@ -47,9 +136,10 @@ func LoadConfiguredEnvironment(customConfigFiles []string, workingDirectory stri
 	}
 
 	// process config files
-	for _, file := range customConfigFiles {
-		if !filepath.IsAbs(file) {
-			file = filepath.Join(workingDirectory, file)
+	for _, file := range options.customConfigFiles {
+		// Check the context hadn't been canceled while loading config files.
+		if options.ctx.Err() != nil {
+			return
 		}
 		loadFile(file)
 	}
@@ -86,7 +176,8 @@ var shellWhiteList = map[string]bool{
 	"/usr/bin/bash": true,
 }
 
-func getEnvFromShell(shell string) string {
+func getEnvFromShell(ctx context.Context, logger *zerolog.Logger, shell string) string {
+	funcLogger := logger.With().Str("method", "getEnvFromShell").Str("shell", shell).Logger()
 	// under windows, the shell environment is irrelevant
 	if runtime.GOOS == "windows" {
 		return ""
@@ -96,14 +187,17 @@ func getEnvFromShell(shell string) string {
 		return ""
 	}
 
-	ctx, cancelFunc := context.WithTimeout(context.Background(), 5*time.Second)
+	ctx, cancelFunc := context.WithTimeout(ctx, 5*time.Second)
 	defer cancelFunc()
 
+	funcLogger.Trace().Msg("get env from shell")
 	// deepcode ignore CommandInjection: false positive
 	env, err := exec.CommandContext(ctx, shell, "--login", "-i", "-c", "printenv && exit").Output()
 	if err != nil {
+		funcLogger.Trace().Err(err).Msg("failed to get env from shell")
 		return ""
 	}
+	funcLogger.Trace().Msg("got env from shell")
 
 	return string(env)
 }

pkg/envvars/environment_test.go

@@ -1,13 +1,16 @@
 package envvars
 
 import (
+	"bytes"
+	"context"
 	"fmt"
 	"math/rand"
 	"os"
 	"path/filepath"
 	"strconv"
 	"testing"
 
+	"github.com/rs/zerolog"
 	"github.com/stretchr/testify/require"
 )
 
@@ -129,6 +132,9 @@ func TestLoadFile(t *testing.T) {
 
 func TestLoadConfiguredEnvironment(t *testing.T) {
 	t.Run("should load default config files", func(t *testing.T) {
+		// Ensure we clean up the PATH back to how it was after the test.
+		t.Setenv("PATH", os.Getenv("PATH"))
+
 		dir := t.TempDir()
 		uniqueEnvVarConfigFile, absEnvVarConfigFile := setupTestFile(t, "1", dir)
 		uniqueEnvVarDotSnykEnv, absEnvVarDotSnykEnvFile := setupTestFile(t, ".snyk.env", dir)
@@ -140,6 +146,9 @@ func TestLoadConfiguredEnvironment(t *testing.T) {
 		err = os.Chdir(dir)
 		require.NoError(t, err)
 
+		// Don't allow Bash, so we skip the shell loading, it isn't relevant for this test.
+		setShellAllowedForTest(t, "bash", false)
+
 		LoadConfiguredEnvironment(files, dir)
 
 		require.Equal(t, uniqueEnvVarConfigFile, os.Getenv(uniqueEnvVarConfigFile))
@@ -151,6 +160,90 @@ func TestLoadConfiguredEnvironment(t *testing.T) {
 	})
 }
 
+func TestLoadConfiguredEnvironmentWithOptions(t *testing.T) {
+	t.Run("should load custom config files", func(t *testing.T) {
+		// Ensure we clean up the PATH back to how it was after the test.
+		t.Setenv("PATH", os.Getenv("PATH"))
+
+		dir := t.TempDir()
+		uniqueEnvVar, absEnvVarConfigFile := setupTestFile(t, ".custom.env", dir)
+
+		// Don't allow Bash, so we skip the shell loading, it isn't relevant for this test.
+		setShellAllowedForTest(t, "bash", false)
+
+		LoadConfiguredEnvironmentWithOptions(WithCustomConfigFiles([]string{absEnvVarConfigFile}))
+
+		require.Equal(t, uniqueEnvVar, os.Getenv(uniqueEnvVar))
+	})
+
+	t.Run("should return early when context is canceled", func(t *testing.T) {
+		// Ensure we clean up the PATH back to how it was after the test.
+		t.Setenv("PATH", os.Getenv("PATH"))
+
+		dir := t.TempDir()
+		uniqueEnvVar, absEnvVarConfigFile := setupTestFile(t, ".custom.env", dir)
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+
+		// Don't allow Bash, so we skip the shell loading (although in theory we shouldn't hit it anyway), it isn't relevant for this test.
+		setShellAllowedForTest(t, "bash", false)
+
+		LoadConfiguredEnvironmentWithOptions(
+			WithContext(ctx),
+			WithCustomConfigFiles([]string{absEnvVarConfigFile}),
+		)
+
+		require.Empty(t, os.Getenv(uniqueEnvVar))
+	})
+
+	t.Run("should use the passed logger", func(t *testing.T) {
+		// Ensure we clean up the PATH back to how it was after the test.
+		t.Setenv("PATH", os.Getenv("PATH"))
+
+		var logsBuffer bytes.Buffer
+		logger := zerolog.New(&logsBuffer).Level(zerolog.DebugLevel)
+
+		// Use a canceled context so we skip the actual loading, it isn't relevant for this test.
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		// Don't allow Bash, so we skip the shell loading (although in theory we shouldn't hit it anyway), it isn't relevant for this test.
+		setShellAllowedForTest(t, "bash", false)
+
+		LoadConfiguredEnvironmentWithOptions(WithContext(ctx), WithLogger(&logger))
+
+		require.NotEmpty(t, logsBuffer.String(), "Expected something to be logged.")
+	})
+}
+
+func TestGetEnvFromShell(t *testing.T) {
+	t.Run("should return empty for non-allowlisted shell", func(t *testing.T) {
+		// Ensure we clean up the PATH back to how it was after the test.
+		t.Setenv("PATH", os.Getenv("PATH"))
+
+		logger := zerolog.Nop()
+
+		env := getEnvFromShell(context.Background(), &logger, "not-allowlisted-shell")
+
+		require.Empty(t, env)
+	})
+}
+
+//nolint:unparam // Parameterized shell support keeps helper reusable when future tests need other allowlisted shells.
+func setShellAllowedForTest(t *testing.T, shell string, allowed bool) {
+	t.Helper()
+	originalValue, hasOriginalValue := shellWhiteList[shell]
+	shellWhiteList[shell] = allowed
+
+	t.Cleanup(func() {
+		if hasOriginalValue {
+			shellWhiteList[shell] = originalValue
+			return
+		}
+
+		delete(shellWhiteList, shell)
+	})
+}
+
 func setupTestFile(t *testing.T, fileName string, dir string) (string, string) {
 	t.Helper()
 	uniqueEnvVar := strconv.Itoa(rand.Int())

pkg/utils/paths.go

@@ -0,0 +1,17 @@
+package utils
+
+import "path/filepath"
+
+// MakeRelativePathsAbsolute resolves any relative paths in the given slice
+// by joining them with baseDirPath. Already-absolute paths are left unchanged.
+func MakeRelativePathsAbsolute(baseDirPath string, paths []string) []string {
+	result := make([]string, len(paths))
+	for i, p := range paths {
+		if filepath.IsAbs(p) {
+			result[i] = p
+		} else {
+			result[i] = filepath.Join(baseDirPath, p)
+		}
+	}
+	return result
+}

pkg/utils/paths_test.go

@@ -0,0 +1,62 @@
+package utils
+
+import (
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_MakeRelativePathsAbsolute(t *testing.T) {
+	baseDir := t.TempDir()
+	differentAbsDir := t.TempDir()
+
+	t.Run("resolves relative paths", func(t *testing.T) {
+		input := []string{".snyk.env", ".envrc"}
+		result := MakeRelativePathsAbsolute(baseDir, input)
+
+		assert.Equal(t, []string{
+			filepath.Join(baseDir, ".snyk.env"),
+			filepath.Join(baseDir, ".envrc"),
+		}, result)
+	})
+
+	t.Run("leaves absolute paths unchanged", func(t *testing.T) {
+		absPath := filepath.Join(differentAbsDir, "config.env")
+		input := []string{absPath}
+		result := MakeRelativePathsAbsolute(baseDir, input)
+
+		assert.Equal(t, []string{absPath}, result)
+	})
+
+	t.Run("handles mix of relative and absolute", func(t *testing.T) {
+		absPath := filepath.Join(differentAbsDir, "config.env")
+		input := []string{absPath, ".snyk.env"}
+		result := MakeRelativePathsAbsolute(baseDir, input)
+
+		assert.Equal(t, []string{
+			absPath,
+			filepath.Join(baseDir, ".snyk.env"),
+		}, result)
+	})
+
+	t.Run("returns empty slice for empty input", func(t *testing.T) {
+		result := MakeRelativePathsAbsolute(baseDir, []string{})
+		assert.Empty(t, result)
+	})
+
+	t.Run("returns empty slice for nil input", func(t *testing.T) {
+		result := MakeRelativePathsAbsolute(baseDir, nil)
+		assert.Empty(t, result)
+	})
+
+	t.Run("does not modify original slice", func(t *testing.T) {
+		input := []string{".snyk.env", ".envrc"}
+		original := make([]string, len(input))
+		copy(original, input)
+
+		MakeRelativePathsAbsolute(baseDir, input)
+
+		assert.Equal(t, original, input)
+	})
+}