multica/server/internal/daemon/execenv/codex_shell_env.go

package execenv

import (
	"bytes"
	"fmt"
	"log/slog"
	"os"
	"regexp"
	"sort"
	"strings"

	"github.com/pelletier/go-toml/v2"
	"github.com/pelletier/go-toml/v2/unstable"
)

// Codex filters environment variables before running shell tool subprocesses.
// Its default secret guard drops names containing KEY, SECRET, or TOKEN, so a
// daemon-spawned Codex process can have MULTICA_TOKEN while `multica issue ...`
// inside a shell tool does not. The daemon therefore owns the shell policy in
// each task's isolated Codex home.
const (
	shellEnvironmentPolicyKey  = "shell_environment_policy"
	multicaShellEnvBeginMarker = "# BEGIN multica-managed shell-environment (do not edit; regenerated by daemon)"
	multicaShellEnvEndMarker   = "# END multica-managed shell-environment"
)

var shellEnvBlockRe = regexp.MustCompile(
	`(?ms)^` + regexp.QuoteMeta(multicaShellEnvBeginMarker) +
		`.*?^` + regexp.QuoteMeta(multicaShellEnvEndMarker) + `\n*`)

type codexShellEnvironmentPolicy struct {
	Inherit               string   `toml:"inherit"`
	IgnoreDefaultExcludes bool     `toml:"ignore_default_excludes"`
	IncludeOnly           []string `toml:"include_only"`
}

type tomlByteRange struct {
	start int
	end   int
}

// CodexShellEnvAllowlist returns exact (not glob) environment names for the
// managed shell policy.
//
// Inherited process variables keep Codex's documented default filtering for
// names containing KEY, SECRET, or TOKEN. Credential-looking explicit values
// are included only when their names also appear in authorizedExplicit, which
// daemon.go derives solely from the current agent's blocklist-checked
// custom_env. Inherited MULTICA_* variables are always dropped because they
// belong to the daemon process, not necessarily to this task. Explicit
// MULTICA_* values are safe to include because daemon.go blocklists that
// namespace from agent custom_env and constructs those values from the current
// task.
func CodexShellEnvAllowlist(inherited []string, explicit map[string]string, authorizedExplicit []string) []string {
	authorized := make(map[string]struct{}, len(authorizedExplicit))
	for _, key := range authorizedExplicit {
		if key != "" {
			authorized[strings.ToUpper(key)] = struct{}{}
		}
	}

	// Codex's glob matching is case-insensitive. De-duplicate on the same basis
	// so Windows Path/PATH aliases cannot create ambiguous policy entries.
	allowed := make(map[string]string, len(inherited)+len(explicit))
	add := func(key string, isExplicit bool) {
		if key == "" {
			return
		}
		upper := strings.ToUpper(key)
		if strings.HasPrefix(upper, "MULTICA_") {
			if !isExplicit {
				return
			}
		} else if codexDefaultExcludesEnvKey(upper) {
			if !isExplicit {
				return
			}
			if _, ok := authorized[upper]; !ok {
				return
			}
		}
		allowed[upper] = key
	}

	for _, entry := range inherited {
		key, _, ok := strings.Cut(entry, "=")
		if !ok {
			continue
		}
		add(key, false)
	}
	for key := range explicit {
		add(key, true)
	}

	keys := make([]string, 0, len(allowed))
	for _, key := range allowed {
		keys = append(keys, key)
	}
	sort.Slice(keys, func(i, j int) bool {
		return strings.ToUpper(keys[i]) < strings.ToUpper(keys[j])
	})
	return keys
}

func codexDefaultExcludesEnvKey(upperKey string) bool {
	return strings.Contains(upperKey, "KEY") ||
		strings.Contains(upperKey, "SECRET") ||
		strings.Contains(upperKey, "TOKEN")
}

func renderMulticaShellEnvBlock(includeOnly []string) (string, error) {
	policy, err := toml.Marshal(codexShellEnvironmentPolicy{
		Inherit:               "all",
		IgnoreDefaultExcludes: true,
		IncludeOnly:           includeOnly,
	})
	if err != nil {
		return "", fmt.Errorf("render shell_environment_policy: %w", err)
	}

	var b strings.Builder
	b.WriteString(multicaShellEnvBeginMarker)
	b.WriteByte('\n')
	b.WriteString("[")
	b.WriteString(shellEnvironmentPolicyKey)
	b.WriteString("]\n")
	b.Write(bytes.TrimRight(policy, "\n"))
	b.WriteByte('\n')
	b.WriteString(multicaShellEnvEndMarker)
	b.WriteByte('\n')
	return b.String(), nil
}

// stripUserShellEnvPolicy removes complete TOML expressions that define a
// shell_environment_policy, including root dotted/inline values, quoted table
// forms, nested set tables, and profile overlays. It deliberately preserves
// comments and every unrelated expression byte-for-byte.
func stripUserShellEnvPolicy(content []byte) ([]byte, error) {
	if len(bytes.TrimSpace(content)) == 0 {
		return content, nil
	}

	// The iterative parser exposes source ranges but does not enforce all TOML
	// semantic rules (for example duplicate definitions). Validate with the
	// stable decoder first so a malformed user config is never partially
	// rewritten into a different malformed config.
	var decoded map[string]any
	if err := toml.Unmarshal(content, &decoded); err != nil {
		return nil, fmt.Errorf("parse config.toml before shell environment update: %w", err)
	}

	var parser unstable.Parser
	parser.Reset(content)
	var currentTable []string
	ranges := make([]tomlByteRange, 0, 8)
	for parser.NextExpression() {
		expr := parser.Expression()
		keys := tomlNodeKeys(expr)
		switch expr.Kind {
		case unstable.Table, unstable.ArrayTable:
			currentTable = keys
			if isCodexShellEnvPolicyPath(keys) {
				r, err := tomlTableHeaderLineRange(content, expr)
				if err != nil {
					return nil, err
				}
				ranges = append(ranges, r)
			}
		case unstable.KeyValue:
			fullKey := make([]string, 0, len(currentTable)+len(keys))
			fullKey = append(fullKey, currentTable...)
			fullKey = append(fullKey, keys...)
			if isCodexShellEnvPolicyPath(fullKey) {
				start := int(expr.Raw.Offset)
				ranges = append(ranges, tomlByteRange{start: start, end: start + int(expr.Raw.Length)})
			}
		}
	}
	if err := parser.Error(); err != nil {
		return nil, fmt.Errorf("parse config.toml shell environment expressions: %w", err)
	}
	return removeTOMLByteRanges(content, ranges)
}

func tomlNodeKeys(node *unstable.Node) []string {
	it := node.Key()
	keys := make([]string, 0, 4)
	for it.Next() {
		keys = append(keys, string(it.Node().Data))
	}
	return keys
}

func isCodexShellEnvPolicyPath(keys []string) bool {
	if len(keys) > 0 && keys[0] == shellEnvironmentPolicyKey {
		return true
	}
	return len(keys) >= 3 && keys[0] == "profiles" && keys[2] == shellEnvironmentPolicyKey
}

func tomlTableHeaderLineRange(content []byte, table *unstable.Node) (tomlByteRange, error) {
	it := table.Key()
	if !it.Next() {
		return tomlByteRange{}, errorsForInvalidShellEnvTOML("table header has no key")
	}
	offset := int(it.Node().Raw.Offset)
	if offset < 0 || offset > len(content) {
		return tomlByteRange{}, errorsForInvalidShellEnvTOML("table key range is outside config")
	}
	start := bytes.LastIndexByte(content[:offset], '\n') + 1
	end := len(content)
	if newline := bytes.IndexByte(content[offset:], '\n'); newline >= 0 {
		end = offset + newline + 1
	}
	return tomlByteRange{start: start, end: end}, nil
}

func errorsForInvalidShellEnvTOML(detail string) error {
	return fmt.Errorf("parse config.toml shell environment expressions: %s", detail)
}

func removeTOMLByteRanges(content []byte, ranges []tomlByteRange) ([]byte, error) {
	if len(ranges) == 0 {
		return append([]byte(nil), content...), nil
	}
	sort.Slice(ranges, func(i, j int) bool {
		if ranges[i].start == ranges[j].start {
			return ranges[i].end < ranges[j].end
		}
		return ranges[i].start < ranges[j].start
	})

	merged := make([]tomlByteRange, 0, len(ranges))
	for _, r := range ranges {
		if r.start < 0 || r.end < r.start || r.end > len(content) {
			return nil, errorsForInvalidShellEnvTOML("expression range is outside config")
		}
		if len(merged) > 0 && r.start <= merged[len(merged)-1].end {
			if r.end > merged[len(merged)-1].end {
				merged[len(merged)-1].end = r.end
			}
			continue
		}
		merged = append(merged, r)
	}

	var out bytes.Buffer
	previous := 0
	for _, r := range merged {
		out.Write(content[previous:r.start])
		previous = r.end
	}
	out.Write(content[previous:])
	return out.Bytes(), nil
}

// EnsureCodexShellEnvPolicyConfig replaces every user/profile shell policy in
// the task-local config with one validated daemon-managed policy. The source
// config in the user's shared Codex home is never modified.
func EnsureCodexShellEnvPolicyConfig(configPath string, includeOnly []string, logger *slog.Logger) error {
	data, err := os.ReadFile(configPath)
	if err != nil && !os.IsNotExist(err) {
		return fmt.Errorf("read config.toml: %w", err)
	}

	withoutManaged := []byte(shellEnvBlockRe.ReplaceAllString(string(data), ""))
	stripped, err := stripUserShellEnvPolicy(withoutManaged)
	if err != nil {
		return err
	}
	block, err := renderMulticaShellEnvBlock(includeOnly)
	if err != nil {
		return err
	}

	base := strings.TrimRight(string(stripped), " \t\r\n")
	updated := block
	if base != "" {
		updated = base + "\n\n" + block
	}
	var validated map[string]any
	if err := toml.Unmarshal([]byte(updated), &validated); err != nil {
		return fmt.Errorf("validate generated config.toml shell environment policy: %w", err)
	}

	if updated == string(data) {
		return nil
	}
	if err := os.WriteFile(configPath, []byte(updated), 0o644); err != nil {
		return fmt.Errorf("write config.toml: %w", err)
	}
	if logger != nil {
		logger.Debug("codex shell environment: wrote managed shell_environment_policy",
			"config_path", configPath,
			"include_count", len(includeOnly),
		)
	}
	return nil
}