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:
parent
7b82cac4df
commit
09b42753bf
2 changed files with 33 additions and 32 deletions
|
|
@ -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,
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
}
|
}
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue