improve flexibility of limactl shell cmd

Signed-off-by: olalekan odukoya <[email protected]>

Author: olamilekan000

hack/bats/tests/preserve-env.bats

@@ -93,13 +93,13 @@ local_setup() {
     assert_line BAR=bar
 }
 
-@test 'wildcard does only work at the end of the pattern' {
+@test 'wildcard works at the start of the pattern' {
     export LIMA_SHELLENV_BLOCK="*FOO"
     export FOO=foo
     export BARFOO=barfoo
     run -0 limactl shell --preserve-env "$NAME" printenv
-    assert_line FOO=foo
-    assert_line BARFOO=barfoo
+    refute_line --regexp '^BARFOO='
+    refute_line --regexp '^FOO='
 }
 
 @test 'block list can use a , separated list with whitespace ignored' {
@@ -114,37 +114,74 @@ local_setup() {
     assert_line BARBAZ=barbaz
 }
 
-@test 'allow list overrides block list but blocks everything else' {
-    export LIMA_SHELLENV_ALLOW=SSH_FOO
+@test 'allow list can use a , separated list with whitespace ignored' {
+    export LIMA_SHELLENV_ALLOW="SSH_FOO, , BAR*, LD_UID"
     export SSH_FOO=ssh_foo
     export SSH_BAR=ssh_bar
+    export SSH_BLOCK=ssh_block
     export BAR=bar
+    export BARBAZ=barbaz
+    export LD_UID=randomuid
     run -0 limactl shell --preserve-env "$NAME" printenv
+    
     assert_line SSH_FOO=ssh_foo
+    assert_line BAR=bar
+    assert_line BARBAZ=barbaz
+    assert_line LD_UID=randomuid
+    
     refute_line --regexp '^SSH_BAR='
-    refute_line --regexp '^BAR='
+    refute_line --regexp '^SSH_BLOCK='
 }
 
-@test 'allow list can use a , separated list with whitespace ignored' {
-    export LIMA_SHELLENV_ALLOW="FOO*, , BAR"
+@test 'wildcard patterns work in all positions' {
+    export LIMA_SHELLENV_BLOCK="*FOO*BAR*"
     export FOO=foo
     export FOOBAR=foobar
+    export FOOXYZBAR=fooxyzbar
+    export FOOBAZ=foobaz
+    export BAZBAR=bazbar
     export BAR=bar
-    export BARBAZ=barbaz
+    export XFOOYBARZDOTCOM=xfooybarzdotcom
+    export NORMAL_VAR=normal_var
+    export UNRELATED=unrelated
     run -0 limactl shell --preserve-env "$NAME" printenv
-    assert_line FOO=foo
-    assert_line FOOBAR=foobar
+    
+    # Should block *FOO*BAR* pattern
+    refute_line --regexp '^FOOBAR='
+    refute_line --regexp '^FOOXYZBAR='    
+    refute_line --regexp '^XFOOYBARZDOTCOM='
+    
+    # Should allow variables that don't match any pattern
+    assert_line FOOBAZ=foobaz
+    assert_line NORMAL_VAR=normal_var
+    assert_line UNRELATED=unrelated
+    assert_line BAZBAR=bazbar
     assert_line BAR=bar
-    refute_line --regexp '^BARBAZ='
+    assert_line FOO=foo
 }
 
-@test 'setting both allow list and block list generates a warning' {
-    export LIMA_SHELLENV_ALLOW=FOO
-    export LIMA_SHELLENV_BLOCK=BAR
-    export FOO=foo
-    run -0 --separate-stderr limactl shell --preserve-env "$NAME" printenv FOO
-    assert_output foo
-    assert_stderr --regexp 'level=warning msg="Both LIMA_SHELLENV_BLOCK and LIMA_SHELLENV_ALLOW are set'
+@test 'allowlist overrides default blocklist with wildcards' {
+    export LIMA_SHELLENV_ALLOW="SSH_*,CUSTOM*"
+    export LIMA_SHELLENV_BLOCK="+*TOKEN"
+    export SSH_AUTH_SOCK=ssh_auth_sock
+    export SSH_CONNECTION=ssh_connection
+    export CUSTOM_VAR=custom_var
+    export MY_TOKEN=my_token
+    export UNRELATED=unrelated
+    run -0 limactl shell --preserve-env "$NAME" printenv
+    
+    assert_line SSH_AUTH_SOCK=ssh_auth_sock
+    assert_line SSH_CONNECTION=ssh_connection
+    assert_line CUSTOM_VAR=custom_var
+    refute_line --regexp '^MY_TOKEN='
+    assert_line UNRELATED=unrelated
+}
+
+@test 'invalid characters in patterns cause fatal errors' {
+    export LIMA_SHELLENV_BLOCK="FOO-BAR"
+    run ! limactl shell --preserve-env "$NAME" printenv
+    assert_output --partial "Invalid LIMA_SHELLENV_BLOCK pattern"
+    assert_output --partial "contains invalid character"
 }
 
 @test 'limactl info includes the default block list' {

hack/test-preserve-env.sh

@@ -339,10 +339,6 @@ if [[ -n ${CHECKS["mount-home"]} ]]; then
 	"${scriptdir}"/test-mount-home.sh "$NAME"
 fi
 
-if [[ -n ${CHECKS["preserve-env"]} ]]; then
-	"${scriptdir}"/test-preserve-env.sh "$NAME"
-fi
-
 if [[ -n ${CHECKS["ssh-over-vsock"]} ]]; then
 	if [[ "$(limactl ls "${NAME}" --yq .vmType)" == "vz" ]]; then
 		INFO "Testing SSH over vsock"

hack/test-templates.sh

@@ -4,7 +4,9 @@
 package envutil
 
 import (
+	"fmt"
 	"os"
+	"regexp"
 	"slices"
 	"strings"
 
@@ -42,27 +44,55 @@ var defaultBlockList = []string{
 	"_*", // Variables starting with underscore are typically internal
 }
 
+func validatePattern(pattern string) error {
+	invalidChar := regexp.MustCompile(`([^a-zA-Z0-9_*])`)
+	if matches := invalidChar.FindStringSubmatch(pattern); matches != nil {
+		invalidCharacter := matches[1]
+		pos := strings.Index(pattern, invalidCharacter)
+		return fmt.Errorf("pattern %q contains invalid character %q at position %d",
+			pattern, invalidCharacter, pos)
+	}
+	return nil
+}
+
 // getBlockList returns the list of environment variable patterns to be blocked.
-// The second return value indicates whether the list was explicitly set via LIMA_SHELLENV_BLOCK.
-func getBlockList() ([]string, bool) {
+func getBlockList() []string {
 	blockEnv := os.Getenv("LIMA_SHELLENV_BLOCK")
 	if blockEnv == "" {
-		return defaultBlockList, false
+		return defaultBlockList
 	}
-	after, found := strings.CutPrefix(blockEnv, "+")
-	if !found {
-		return parseEnvList(blockEnv), true
+
+	shouldAppend := strings.HasPrefix(blockEnv, "+")
+	patterns := parseEnvList(strings.TrimPrefix(blockEnv, "+"))
+
+	for _, pattern := range patterns {
+		if err := validatePattern(pattern); err != nil {
+			logrus.Fatalf("Invalid LIMA_SHELLENV_BLOCK pattern: %v", err)
+		}
+	}
+
+	if shouldAppend {
+		return slices.Concat(defaultBlockList, patterns)
 	}
-	return slices.Concat(defaultBlockList, parseEnvList(after)), true
+	return patterns
 }
 
 // getAllowList returns the list of environment variable patterns to be allowed.
-// The second return value indicates whether the list was explicitly set via LIMA_SHELLENV_ALLOW.
-func getAllowList() ([]string, bool) {
-	if allowEnv := os.Getenv("LIMA_SHELLENV_ALLOW"); allowEnv != "" {
-		return parseEnvList(allowEnv), true
+func getAllowList() []string {
+	allowEnv := os.Getenv("LIMA_SHELLENV_ALLOW")
+	if allowEnv == "" {
+		return nil
 	}
-	return nil, false
+
+	patterns := parseEnvList(allowEnv)
+
+	for _, pattern := range patterns {
+		if err := validatePattern(pattern); err != nil {
+			logrus.Fatalf("Invalid LIMA_SHELLENV_ALLOW pattern: %v", err)
+		}
+	}
+
+	return patterns
 }
 
 func parseEnvList(envList string) []string {
@@ -82,8 +112,14 @@ func matchesPattern(name, pattern string) bool {
 		return true
 	}
 
-	prefix, found := strings.CutSuffix(pattern, "*")
-	return found && strings.HasPrefix(name, prefix)
+	regexPattern := strings.ReplaceAll(pattern, "*", ".*")
+	regexPattern = "^" + regexPattern + "$"
+
+	match, err := regexp.MatchString(regexPattern, name)
+	if err != nil {
+		return false
+	}
+	return match
 }
 
 func matchesAnyPattern(name string, patterns []string) bool {
@@ -96,17 +132,10 @@ func matchesAnyPattern(name string, patterns []string) bool {
 // It returns a slice of environment variables that are not blocked by the current configuration.
 // The filtering is controlled by LIMA_SHELLENV_BLOCK and LIMA_SHELLENV_ALLOW environment variables.
 func FilterEnvironment() []string {
-	allowList, isAllowListSet := getAllowList()
-	blockList, isBlockListSet := getBlockList()
-
-	if isBlockListSet && isAllowListSet {
-		logrus.Warn("Both LIMA_SHELLENV_BLOCK and LIMA_SHELLENV_ALLOW are set. Block list will be ignored.")
-		blockList = nil
-	}
 	return filterEnvironmentWithLists(
 		os.Environ(),
-		allowList,
-		blockList,
+		getAllowList(),
+		getBlockList(),
 	)
 }
 
@@ -121,10 +150,7 @@ func filterEnvironmentWithLists(env, allowList, blockList []string) []string {
 
 		name := parts[0]
 
-		if len(allowList) > 0 {
-			if !matchesAnyPattern(name, allowList) {
-				continue
-			}
+		if len(allowList) > 0 && matchesAnyPattern(name, allowList) {
 			filtered = append(filtered, envVar)
 			continue
 		}