Refactor env handling: WithAllowedEnv for efficient os.Getenv() lookups

- Rename FilterByAllowlist -> WithAllowedEnv(envSet, extraAllowlist)
- Uses os.Getenv() for each allowed var instead of parsing os.Environ()
- Starts with envSet, adds allowed inherited vars (if not already set)
- cachedEnv now contains PICOCLAW_* + allowed inherited vars
This commit is contained in:
Keith Patrick 2026-03-09 03:07:57 +00:00
parent 7b82cac4df
commit 09b42753bf
2 changed files with 33 additions and 32 deletions

View file

@ -24,7 +24,7 @@ type ExecTool struct {
allowPatterns []*regexp.Regexp allowPatterns []*regexp.Regexp
customAllowPatterns []*regexp.Regexp customAllowPatterns []*regexp.Regexp
restrictToWorkspace bool restrictToWorkspace bool
cachedEnv map[string]string // cached sanitized env map from os.Environ() at init cachedEnv map[string]string // cached sanitized env from os.Getenv() at init
} }
var ( var (
@ -151,14 +151,8 @@ func NewExecToolWithConfig(workingDir string, restrict bool, config *config.Conf
// Ensure PICOCLAW_* vars are set for child processes // Ensure PICOCLAW_* vars are set for child processes
envSet = shell.WithPicoclawEnvVars(envSet, workingDir) envSet = shell.WithPicoclawEnvVars(envSet, workingDir)
// Build cached env: filter inherited env by allowlist, then merge with config envSet // Build cached env: start with envSet (PICOCLAW_*), then add allowed inherited vars
filteredBase := shell.FilterByAllowlist(os.Environ(), envAllowlist) cachedEnv := shell.WithAllowedEnv(envSet, envAllowlist)
// Pre-merge envSet into cachedEnv so PICOCLAW_* vars are preserved
// (MergeEnvVars returns []string, so we merge maps manually)
cachedEnv := filteredBase
for k, v := range envSet {
cachedEnv[k] = v
}
return &ExecTool{ return &ExecTool{
workingDir: workingDir, workingDir: workingDir,

View file

@ -78,48 +78,55 @@ var windowsEnvAllowlist = map[string]bool{
"HOMEPATH": true, "HOMEPATH": true,
} }
// FilterByAllowlist filters the inherited environment to only allowlisted variables. // WithAllowedEnv builds a map of allowed environment variables by looking them up.
// This should only be used at init time to create the cached environment. // This is more efficient than filtering os.Environ() with string parsing.
func FilterByAllowlist(baseEnv []string, extraAllowlist []string) map[string]string { // It starts with the provided env map, then adds allowed inherited vars (if not set).
if baseEnv == nil { // extraAllowlist adds to the default allowlist.
baseEnv = os.Environ() func WithAllowedEnv(envSet map[string]string, extraAllowlist []string) map[string]string {
// Start with provided envSet map
result := envSet
if result == nil {
result = make(map[string]string)
} }
allowed := make(map[string]bool, len(DefaultEnvAllowlist)+len(extraAllowlist)+len(windowsEnvAllowlist)) // Add default allowlist (only if not already set)
for k := range DefaultEnvAllowlist { for k := range DefaultEnvAllowlist {
allowed[envKey(k)] = true if _, exists := result[k]; !exists {
if val := os.Getenv(k); val != "" {
result[k] = val
}
}
} }
// Add Windows-specific vars
if runtime.GOOS == "windows" { if runtime.GOOS == "windows" {
for k := range windowsEnvAllowlist { for k := range windowsEnvAllowlist {
allowed[envKey(k)] = true if _, exists := result[k]; !exists {
if val := os.Getenv(k); val != "" {
result[k] = val
}
}
} }
} }
// Add extra allowlist from config
for _, k := range extraAllowlist { for _, k := range extraAllowlist {
allowed[envKey(k)] = true if _, exists := result[k]; !exists {
if val := os.Getenv(k); val != "" {
result[k] = val
}
}
} }
vars := make(map[string]string) return result
for _, entry := range baseEnv {
k, v, ok := strings.Cut(entry, "=")
if !ok {
continue
}
norm := envKey(k)
if allowed[norm] || isAllowedPrefix(norm) {
vars[norm] = v
}
}
return vars
} }
// MergeEnvVars merges multiple env sources into a final []string for exec.Cmd.Env. // MergeEnvVars merges multiple env sources into a final []string for exec.Cmd.Env.
// baseEnv is NOT filtered - it's assumed to already be sanitized (e.g., cachedEnv). // baseEnv is the cached map from AllowedEnv.
// envSet provides explicit key=value pairs (config, not filtered). // envSet provides explicit key=value pairs (config, not filtered).
// extraEnv provides additional key=value pairs from LLM (filtered by blocklist). // extraEnv provides additional key=value pairs from LLM (filtered by blocklist).
func MergeEnvVars(baseEnv map[string]string, envSet, extraEnv map[string]string) []string { func MergeEnvVars(baseEnv map[string]string, envSet, extraEnv map[string]string) []string {
vars := make(map[string]string, len(baseEnv)+len(envSet)+len(extraEnv)) vars := make(map[string]string, len(baseEnv)+len(envSet)+len(extraEnv))
// Start with base env (already sanitized) // Start with base env (already filtered)
for k, v := range baseEnv { for k, v := range baseEnv {
vars[envKey(k)] = v vars[envKey(k)] = v
} }