diff --git a/README.md b/README.md index 8bf30abbc..007515bb0 100644 --- a/README.md +++ b/README.md @@ -195,7 +195,7 @@ Sed selects the longest regex match at the leftmost position, including for in-p For simple commands, rshell expands only enough words to identify the command, then applies command, elevation, and mode policy before expanding the remaining arguments, inline assignments, or redirects. A rejected command therefore cannot trigger command substitutions or redirect side effects through its unused arguments. Once a command is authorized, redirects retain normal shell semantics: they are established before execution and remain effective even if the command later exits nonzero. -Shell expansion is bounded independently of the execution timeout: a command may receive at most 16,384 expanded arguments totaling 10 MiB, and one `Run` may produce at most 64 MiB of expanded argument, assignment, redirect, and heredoc text across all loops, subshells, and pipeline stages. Individual variable values remain capped at 1 MiB and heredocs at 10 MiB. +Shell expansion is bounded independently of the execution timeout: a command may receive at most 16,384 expanded arguments totaling 10 MiB, and one `Run` may produce at most 64 MiB of expanded argument, assignment, redirect, and heredoc text across all loops, subshells, and pipeline stages. The argument cap is enforced before an unquoted variable or command substitution is field-split, not after the fields are built. Individual variable values remain capped at 1 MiB and heredocs at 10 MiB. The variables a script holds at once are capped at 1 MiB of values plus 1 MiB of names and fixed per-variable overhead (about 16k variables), shared by every concurrently running subshell and pipeline stage; a subshell that overrides or clears an inherited variable is still charged for the copy its parent holds. Redirect targets undergo shell expansion and must resolve to exactly one field; ambiguous targets fail before the file is opened. Quoting preserves literal spaces and wildcard characters. diff --git a/SHELL_FEATURES.md b/SHELL_FEATURES.md index 05752f4c7..9fa56f15c 100644 --- a/SHELL_FEATURES.md +++ b/SHELL_FEATURES.md @@ -230,7 +230,8 @@ the substituted command. - ✅ Single quotes: `'literal'` - ✅ Double quotes: `"with $expansion"` - ✅ Globbing: `*`, `?`, `[abc]`, `[a-z]`, `[!a]` -- ✅ Bounded expansion: at most 16,384 expanded arguments and 10 MiB of expanded argument text per command; at most 64 MiB of expanded argument, assignment, redirect, and heredoc text per `Run`, shared across loops, subshells, and pipeline stages; the execution context is checked between produced fields +- ✅ Bounded expansion: at most 16,384 expanded arguments and 10 MiB of expanded argument text per command; at most 64 MiB of expanded argument, assignment, redirect, and heredoc text per `Run`, shared across loops, subshells, and pipeline stages; the execution context is checked between produced fields; the argument cap is enforced on an upper bound before an unquoted variable is field-split and before an unquoted command substitution's output is split, so the limit fires before the fields are built +- ✅ Bounded variable storage: each value is capped at 1 MiB; the values a script holds at once are capped at 1 MiB in total, and variable names plus a fixed 64-byte per-variable overhead at a further 1 MiB (about 16k variables), even when the values are empty; both caps are shared by every subshell and pipeline stage that is alive at the same time, and storage is returned when a subshell or stage finishes; overriding or clearing an inherited variable in a subshell does not refund the copy its parent still holds - ✅ Line continuation: `\` at end of line - ✅ Comments: `# text` - ❌ Extended globbing: `@(pat)`, `*(pat)`, etc. diff --git a/interp/api.go b/interp/api.go index f3c7c48b7..11b70bc32 100644 --- a/interp/api.go +++ b/interp/api.go @@ -225,6 +225,12 @@ type runnerState struct { lastExpandExit exitStatus // used to surface exit statuses while expanding fields + // activeFields is the field collector currently expanding a word on this + // runner, if any. cmdSubst uses it to charge the output of unquoted + // command substitutions against the field budget before the expander + // splits it. Subshells do not inherit it. + activeFields *fieldCollector + // startTime is captured once at the beginning of Run() and passed to // all builtin invocations so they share a consistent time reference. startTime time.Time @@ -569,9 +575,10 @@ func (r *Runner) Reset() { // count against the user-visible MaxTotalVarsBytes cap. Those values are // small and bounded; only the variables that a script itself creates or // modifies should count. ALLOWED_PATHS is operator-configured and - // typically a few hundred bytes, so this is safe. + // typically a few hundred bytes, so this is safe. This also starts the + // run-wide budget that subshells and pipeline stages share. if ov, ok := r.writeEnv.(*overlayEnviron); ok { - ov.totalBytes = 0 + ov.resetAccounting() } r.didReset = true @@ -1104,6 +1111,15 @@ func (r *Runner) subshell(background bool) *Runner { return r2 } +// releaseVarStorage returns a finished subshell's variable storage to the +// run-wide budget shared with its parent and siblings. Call it once the +// subshell from [Runner.subshell] has stopped running. +func (r *Runner) releaseVarStorage() { + if ov, ok := r.writeEnv.(*overlayEnviron); ok { + ov.release() + } +} + // classifyRunOutcome maps the error Run() is about to return to a stable, // low-cardinality enum tag on the run span. The outcome answers "did the // interpreter do its job", not "did the script exit zero" — any script diff --git a/interp/internal_errors_test.go b/interp/internal_errors_test.go index 2915673c6..043a96333 100644 --- a/interp/internal_errors_test.go +++ b/interp/internal_errors_test.go @@ -266,7 +266,7 @@ func TestExpandErrStorageExhaustedAborts(t *testing.T) { // expandEnv.Set wraps the error with the variable name before it reaches // expandErr, producing "name: variable storage limit exceeded (N bytes total)". - storageErr := fmt.Errorf("UNSET_VAR: %w", &errTotalVarStorageExceeded{total: 1048577}) + storageErr := fmt.Errorf("UNSET_VAR: %w", &errTotalVarStorageExceeded{total: varStorage{values: 1048577}}) r.expandErr(storageErr) assert.True(t, r.exit.exiting, "expandErr must set exiting=true for storage-cap errors so the script aborts") diff --git a/interp/memory_accounting_test.go b/interp/memory_accounting_test.go new file mode 100644 index 000000000..e3aec4123 --- /dev/null +++ b/interp/memory_accounting_test.go @@ -0,0 +1,356 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2026-present Datadog, Inc. + +package interp_test + +import ( + "bytes" + "context" + "fmt" + "io" + "runtime" + "strings" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "mvdan.cc/sh/v3/syntax" + + "github.com/DataDog/rshell/interp" +) + +// maxAttackAllocBytes bounds the total allocation of the memory-exhaustion +// scripts below. Without the fixes each one allocates hundreds of megabytes +// or more; bounded, each stays within a few megabytes of real work. +const maxAttackAllocBytes = 64 << 20 + +// runMeasured runs script with exactly opts (no commands are allowed unless +// opts allow them) and returns its output and the bytes allocated while it +// ran. The tests using it are not parallel, so no other test allocates +// concurrently. +func runMeasured(t *testing.T, script string, stdin io.Reader, stderr io.Writer, opts ...interp.RunnerOption) (stdout string, totalAlloc uint64) { + t.Helper() + prog, err := syntax.NewParser().Parse(strings.NewReader(script), "") + require.NoError(t, err) + var outBuf bytes.Buffer + runner, err := interp.New(append([]interp.RunnerOption{ + interp.StdIO(stdin, &outBuf, stderr), + interp.MaxExecutionTime(30 * time.Second), + }, opts...)...) + require.NoError(t, err) + defer runner.Close() + + runtime.GC() + var before, after runtime.MemStats + runtime.ReadMemStats(&before) + _ = runner.Run(context.Background(), prog) + runtime.ReadMemStats(&after) + return outBuf.String(), after.TotalAlloc - before.TotalAlloc +} + +// spacedValueScript assigns x a 512 KiB value made of single-character words +// separated by spaces: 262,144 fields when split. +const spacedValueScript = "x='a '\n" + + "for i in 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18; do x=\"$x$x\"; done\n" + +// TestVULN4760BraceExpansionBoundedBeforeCommandGate pins that a compact brace +// range neither materializes millions of fields before the command gate nor +// overruns the brace limit for an allowed command. +func TestVULN4760BraceExpansionBoundedBeforeCommandGate(t *testing.T) { + for name, tc := range map[string]struct { + script string + stderr string + }{ + "command not allowed": {"not_allowed {1..3000000}\n", "not_allowed: command not allowed"}, + "command allowed": {"echo {1..3000000}\n", "brace expansion would exceed"}, + } { + t.Run(name, func(t *testing.T) { + var stderr bytes.Buffer + start := time.Now() + _, alloc := runMeasured(t, tc.script, nil, &stderr, + interp.AllowedCommands([]string{"rshell:echo"})) + + assert.Contains(t, stderr.String(), tc.stderr) + assert.Less(t, alloc, uint64(maxAttackAllocBytes)) + assert.Less(t, time.Since(start), 5*time.Second) + }) + } +} + +// TestVULN5664UnquotedSplitBoundedBeforeExpansion pins that the field cap is +// enforced on an upper bound before an unquoted word is split, not after the +// expander has built every field. No command is allowed: the word is the +// command name. +func TestVULN5664UnquotedSplitBoundedBeforeExpansion(t *testing.T) { + for name, word := range map[string]string{ + "command": strings.Repeat("$x", 20) + "\n", + "for-loop": "for v in " + strings.Repeat("$x", 20) + "; do :; done\n", + } { + t.Run(name, func(t *testing.T) { + var stderr bytes.Buffer + stdout, alloc := runMeasured(t, spacedValueScript+word+"echo AFTER\n", nil, &stderr) + + assert.Contains(t, stderr.String(), "expansion exceeds maximum field count") + assert.NotContains(t, stdout, "AFTER") + assert.Less(t, alloc, uint64(maxAttackAllocBytes)) + }) + } +} + +// TestVULN5664UnquotedCommandSubstitutionSplitBounded covers the same +// allocate-then-check gap for unquoted command substitutions: their output is +// charged against the field cap before the expander splits it. +func TestVULN5664UnquotedCommandSubstitutionSplitBounded(t *testing.T) { + var stderr bytes.Buffer + stdout, alloc := runMeasured(t, + spacedValueScript+`echo $(echo "$x")$(echo "$x")$(echo "$x")$(echo "$x")$(echo "$x") >/dev/null`+"\necho AFTER\n", + nil, &stderr, interp.AllowedCommands([]string{"rshell:echo"})) + + assert.Contains(t, stderr.String(), "expansion exceeds maximum field count") + assert.NotContains(t, stdout, "AFTER") + assert.Less(t, alloc, uint64(maxAttackAllocBytes)) +} + +// TestVULN5630EmptyVariablesChargedAgainstStorageCap pins that variable names +// and entries count against MaxTotalVarsBytes, so empty-valued variables +// created by read cannot grow storage without bound. +func TestVULN5630EmptyVariablesChargedAgainstStorageCap(t *testing.T) { + var stderr bytes.Buffer + stdout, _ := runMeasured(t, + "for p in a b c d e f; do read $p{a..y}{a..y}{a..z} <&2); echo AFTER\n", + nil, &stderr, interp.AllowedCommands([]string{"rshell:printf", "rshell:echo", "rshell:true"})) + assert.Contains(t, stderr.String(), "expansion exceeds maximum field count") + assert.NotContains(t, stdout, "AFTER") + assert.NotContains(t, stdout+stderr.String(), "REACHED") +} + +// signalWriter collects output and closes found once it contains substr. +type signalWriter struct { + mu sync.Mutex + buf bytes.Buffer + substr string + found chan struct{} + once sync.Once +} + +func (w *signalWriter) Write(p []byte) (int, error) { + w.mu.Lock() + defer w.mu.Unlock() + w.buf.Write(p) + if strings.Contains(w.buf.String(), w.substr) { + w.once.Do(func() { close(w.found) }) + } + return len(p), nil +} + +func (w *signalWriter) String() string { + w.mu.Lock() + defer w.mu.Unlock() + return w.buf.String() +} + +// TestVULN1675ConcurrentPipelineStagesShareStorageCap pins that concurrent +// pipeline stages draw on one run-wide variable budget. Each stage builds a +// 512 KiB variable and then blocks in cat, so all of them are live at once; +// three of them exceed MaxTotalVarsBytes together. +func TestVULN1675ConcurrentPipelineStagesShareStorageCap(t *testing.T) { + stage := `{ X=AAAAAAAAAAAAAAAA; X=$X$X; X=$X$X; X=$X$X; X=$X$X; X=$X$X; ` + + `X=$X$X; X=$X$X; X=$X$X; X=$X$X; X=$X$X; X=$X$X; X=$X$X; X=$X$X; X=$X$X; X=$X$X; cat; }` + script := strings.Join([]string{stage, stage, stage}, " | ") + "\n" + + // The first stage's cat holds the pipeline open until the cap fires (or, + // without the fix, until the timeout below), keeping every stage live. + stdinR, stdinW := io.Pipe() + stderr := &signalWriter{substr: "variable storage limit exceeded", found: make(chan struct{})} + go func() { + select { + case <-stderr.found: + case <-time.After(5 * time.Second): + } + stdinW.Close() + }() + runMeasured(t, script, stdinR, stderr, interp.AllowedCommands([]string{"rshell:cat"})) + + assert.Contains(t, stderr.String(), "variable storage limit exceeded") +} + +// TestMemoryAccountingNormalWorkloads checks that ordinary scripts stay within +// the limits: storage released by finished subshells and pipeline stages is +// reusable, and moderate splitting and variable counts are unaffected. +func TestMemoryAccountingNormalWorkloads(t *testing.T) { + big := "X=AAAAAAAAAAAAAAAA; X=$X$X; X=$X$X; X=$X$X; X=$X$X; X=$X$X; " + + "X=$X$X; X=$X$X; X=$X$X; X=$X$X; X=$X$X; X=$X$X; X=$X$X; X=$X$X; X=$X$X; X=$X$X" + words := strings.TrimSpace(strings.Repeat("w ", 2000)) + tests := map[string]struct { + script string + want string + }{ + "sequential pipelines release storage": { + script: "for i in {1..40}; do { " + big + "; echo ok; } | cat >/dev/null; done; echo DONE\n", + want: "DONE\n", + }, + "command substitutions release storage": { + script: "for i in {1..40}; do y=$(" + big + "; echo ok); done; echo $y DONE\n", + want: "ok DONE\n", + }, + "subshells release storage": { + script: "for i in {1..40}; do ( " + big + " ); done; echo DONE\n", + want: "DONE\n", + }, + "parent storage is visible to every stage": { + script: big + "\n{ Y=a; echo $Y; } | { Z=b; cat; echo $Z; }\n", + want: "a\nb\n", + }, + "thousands of variables": { + script: "for i in {1..5000}; do read v$i </dev/null; Y=$PAYLOAD; echo OK; }\n", + want: "OK\n", + }, + } + for name, tc := range tests { + t.Run(name, func(t *testing.T) { + var stderr bytes.Buffer + stdout, _ := runMeasured(t, tc.script, nil, &stderr, env, + interp.AllowedCommands([]string{"rshell:echo", "rshell:cat", "rshell:true"})) + assert.Equal(t, tc.want, stdout, "stderr: %s", stderr.String()) + if tc.stderr == "" { + assert.Empty(t, stderr.String()) + } else { + assert.Contains(t, stderr.String(), tc.stderr) + } + }) + } +} diff --git a/interp/runner_exec.go b/interp/runner_exec.go index 8b5cdf79c..138bb28a7 100644 --- a/interp/runner_exec.go +++ b/interp/runner_exec.go @@ -210,6 +210,11 @@ func (r *Runner) callExpr(ctx context.Context, cm *syntax.CallExpr, redirs []*sy type restoreVar struct { name string vr expand.Variable + // env and snap, when the environment is an overlay, are that + // overlay and its own binding before the assignment, restored + // exactly instead of vr. + env *overlayEnviron + snap bindingSnapshot } type inlineAssignment struct { name string @@ -224,6 +229,10 @@ func (r *Runner) callExpr(ctx context.Context, cm *syntax.CallExpr, redirs []*sy if isCd && (restore.name == "PWD" || restore.name == "OLDPWD") { continue } + if restore.env != nil { + restore.env.restoreBinding(restore.snap) + continue + } r.setVarRestore(restore.name, restore.vr) } }() @@ -238,9 +247,16 @@ func (r *Runner) callExpr(ctx context.Context, cm *syntax.CallExpr, redirs []*sy func() { // Earlier assignments are visible while expanding later values, // but redirects expand against the original environment. + // The temporary overlay's storage is returned to the shared + // budget on every exit; the assignments are charged again when + // they are installed for the command. previousEnv := r.writeEnv - r.writeEnv = newOverlayEnviron(previousEnv, false) - defer func() { r.writeEnv = previousEnv }() + tempEnv := newOverlayEnviron(previousEnv, false) + r.writeEnv = tempEnv + defer func() { + tempEnv.release() + r.writeEnv = previousEnv + }() for _, as := range cm.Assigns { name := as.Name.Value @@ -267,7 +283,11 @@ func (r *Runner) callExpr(ctx context.Context, cm *syntax.CallExpr, redirs []*sy seenRestore := map[string]bool{} for _, assignment := range assignments { if !seenRestore[assignment.name] { - restores = append(restores, restoreVar{assignment.name, assignment.prev}) + restore := restoreVar{name: assignment.name, vr: assignment.prev} + if ov, ok := r.writeEnv.(*overlayEnviron); ok { + restore.env, restore.snap = ov, ov.snapshotBinding(assignment.name) + } + restores = append(restores, restore) seenRestore[assignment.name] = true } r.setVar(assignment.name, assignment.vr) @@ -288,6 +308,7 @@ func (r *Runner) cmd(ctx context.Context, cm syntax.Command) { // clear the flag that suppresses nested pipeline spans. r2.inPipeline = false r2.stmts(ctx, cm.Stmts) + r2.releaseVarStorage() r.exit = r2.exit if !r.exit.limitExit { r.exit.exiting = false @@ -375,6 +396,9 @@ func (r *Runner) cmd(ctx context.Context, cm syntax.Command) { }() rLeft.exit.fatal(fmt.Errorf("internal error")) } + // Release the producer's storage before closing the + // pipe, so a consumer that has seen EOF can reuse it. + rLeft.releaseVarStorage() pw.Close() wg.Done() }() @@ -382,6 +406,7 @@ func (r *Runner) cmd(ctx context.Context, cm syntax.Command) { rLeft.exit.exiting = false }() rRight.stmt(ctx, cm.Y) + rRight.releaseVarStorage() r.exit = rRight.exit if !r.exit.limitExit { r.exit.exiting = false diff --git a/interp/runner_expand.go b/interp/runner_expand.go index ffe1e2a4b..920f34abe 100644 --- a/interp/runner_expand.go +++ b/interp/runner_expand.go @@ -128,13 +128,21 @@ func (r *Runner) cmdSubst(w io.Writer, cs *syntax.CmdSubst) error { return nil } } - _, err = io.Copy(w, io.LimitReader(f, maxCmdSubstOutput)) + var buf bytes.Buffer + _, err = io.Copy(&buf, io.LimitReader(f, maxCmdSubstOutput)) var exitCode uint8 if err != nil { exitCode = 1 } r.lastExpandExit = exitStatus{code: exitCode} r.lastExit = r.lastExpandExit + if err != nil { + return err + } + if err := r.chargeSplitSubstOutput(cs, buf.Bytes()); err != nil { + return err + } + _, err = w.Write(buf.Bytes()) return err } @@ -148,6 +156,7 @@ func (r *Runner) cmdSubst(w io.Writer, cs *syntax.CmdSubst) error { // escape the subshell — only the diagnostic is suppressed. r2.inLoop = r.inLoop r2.stmts(r.ectx, cs.Stmts) + r2.releaseVarStorage() r2.exit.exiting = false r.lastExpandExit = r2.exit r.lastExit = r.lastExpandExit @@ -165,6 +174,9 @@ func (r *Runner) cmdSubst(w io.Writer, cs *syntax.CmdSubst) error { r.exit.limitExit = true return nil } + if err := r.chargeSplitSubstOutput(cs, buf.Bytes()); err != nil { + return err + } _, err := w.Write(buf.Bytes()) return err } @@ -287,6 +299,41 @@ type fieldCollector struct { maxFields int fields []string bytes int64 + + // The expander builds every field of a word (of one brace alternative) + // before yielding the first, so the field cap must be enforced on an + // upper bound before expansion rather than on the yielded fields alone. + // wordSources bounds the fields the current word's literal parts and + // unquoted parameters can produce; splitSubsts are its unquoted command + // substitutions, whose output is only known once they run and is charged + // to substSources by cmdSubst. + wordSources int + splitSubsts []splitSubst + substSources int +} + +// splitSubst is an unquoted command substitution of the word being expanded, +// with what the static estimate assumed about its neighbours, so that its +// output can be charged with the same field continuation the expander uses. +type splitSubst struct { + cs *syntax.CmdSubst + // inBrace marks a substitution inside a brace alternative. Only the + // largest alternative's sources are counted, so its output is charged + // in full, without continuation credits. + inBrace bool + // openBefore is whether a field is open before the substitution, as far + // as the static walk knows; prev is the index of an immediately + // preceding substitution, whose actual output decides that instead. + openBefore bool + prev int + // deferredNext is whether the part after the substitution starts a new + // field unless the substitution leaves one open. The static walk does + // not count it; it is charged with the substitution's output, which runs + // before the expander reaches that part. + deferredNext bool + + // ran and openAfter record the substitution's output once it has run. + ran, openAfter bool } func (r *Runner) newFieldCollector(maxFields int) *fieldCollector { @@ -307,7 +354,16 @@ func (c *fieldCollector) setCommandPrefixFields(n int) bool { return true } +func (c *fieldCollector) fieldLimitError() error { + return &expansionLimitError{message: fmt.Sprintf( + "expansion exceeds maximum field count (%d)", c.maxFields)} +} + func (c *fieldCollector) add(words ...*syntax.Word) bool { + previous := c.r.activeFields + c.r.activeFields = c + defer func() { c.r.activeFields = previous }() + ifs := c.r.fieldSplitIFS() for _, word := range words { if c.r.stop(c.r.ectx) { return false @@ -320,6 +376,13 @@ func (c *fieldCollector) add(words ...*syntax.Word) bool { c.r.expandErr(err) return false } + c.splitSubsts = c.splitSubsts[:0] + c.wordSources = c.r.wordFieldSources(estimatedWord.Parts, ifs, &c.splitSubsts) + c.substSources = 0 + if c.wordSources > c.maxFields-len(c.fields) { + c.r.expandErr(c.fieldLimitError()) + return false + } for field, err := range expand.FieldsSeq(c.r.ecfg, word) { if err != nil { @@ -332,9 +395,11 @@ func (c *fieldCollector) add(words ...*syntax.Word) bool { if c.r.stop(c.r.ectx) { return false } + // A yield means the current brace alternative has been fully + // expanded; later alternatives run their substitutions afresh. + c.substSources = 0 if len(c.fields) >= c.maxFields { - c.r.expandErr(&expansionLimitError{message: fmt.Sprintf( - "expansion exceeds maximum field count (%d)", c.maxFields)}) + c.r.expandErr(c.fieldLimitError()) return false } fieldBytes := int64(len(field)) @@ -354,6 +419,182 @@ func (c *fieldCollector) add(words ...*syntax.Word) bool { return true } +// fieldSplitIFS returns the IFS value the expander splits with, mirroring +// mvdan.cc/sh/v3/expand's default when IFS is unset. +func (r *Runner) fieldSplitIFS() string { + if vr := r.lookupVar("IFS"); vr.IsSet() { + return vr.String() + } + return " \t\n" +} + +// wordFieldSources returns an upper bound on the number of fields one brace +// alternative of a word with these parts can produce, not counting the output +// of unquoted command substitutions, which are appended to splitSubsts. +// +// Every field the expander emits starts with a distinct "source": either an +// unsplit part (a literal or a quoted part) or a run of non-IFS characters in +// a split value, unless it continues the field the preceding part left open, +// since adjacent parts concatenate. Counting sources therefore bounds the +// field count without expanding the word. The one field an all-empty word can +// still produce requires a quoted part, which is itself a source. +func (r *Runner) wordFieldSources(parts []syntax.WordPart, ifs string, splitSubsts *[]splitSubst) int { + sources, _ := r.partsFieldSources(parts, ifs, false, false, splitSubsts) + return sources +} + +// partsFieldSources counts the field sources of parts, given whether a field +// is already open before them, and reports whether one is open after them. +// When that is unknown it reports false, which can only overcount the +// sources that follow. +func (r *Runner) partsFieldSources(parts []syntax.WordPart, ifs string, open, inBrace bool, splitSubsts *[]splitSubst) (int, bool) { + sources := 0 + // afterSubst is the index of the substitution the previous part was, or + // -1 when the previous part was something else. + afterSubst := -1 + for _, part := range parts { + counted := false + switch part := part.(type) { + case *syntax.ParamExp: + if part.Param != nil { + value := r.lookupVar(part.Param.Value).String() + n, closes, opens := splitFieldSources(value, ifs) + if n > 0 && open && !closes { + n-- // the first run continues the open field + } + counted = n > 0 && !closes + sources += n + if value != "" { + open = opens + } + } + case *syntax.CmdSubst: + *splitSubsts = append(*splitSubsts, splitSubst{ + cs: part, inBrace: inBrace, openBefore: open, prev: afterSubst, + }) + afterSubst = len(*splitSubsts) - 1 + open = false + continue + case *syntax.BraceExp: + // FieldsSeq expands each alternative as a separate word, so only + // the largest alternative is built at once. + largest, allOpen := 0, true + for _, elem := range part.Elems { + n, elemOpen := r.partsFieldSources(elem.Parts, ifs, open, true, splitSubsts) + largest = max(largest, n) + allOpen = allOpen && elemOpen + } + sources += largest + open = allOpen + default: + counted = !open + if !open { + sources++ + } + open = true + } + if afterSubst >= 0 && !inBrace && counted { + (*splitSubsts)[afterSubst].deferredNext = true + sources-- + } + afterSubst = -1 + } + return sources, open +} + +// splitFieldSources counts the runs of non-IFS characters in s and reports +// whether s starts with an IFS character (closing any open field) and whether +// it ends inside a run (leaving a field open). +func splitFieldSources(s, ifs string) (sources int, closes, opens bool) { + if ifs == "" { + return min(len(s), 1), false, s != "" + } + inField := false + for i, c := range s { + isIFS := strings.ContainsRune(ifs, c) + if i == 0 { + closes = isIFS + } + if !isIFS && !inField { + sources++ + } + inField = !isIFS + } + return sources, closes, inField +} + +// chargeSplitSubstOutput charges the output of an unquoted command +// substitution against the field budget of the word being expanded, before +// the expander splits it into fields. +func (r *Runner) chargeSplitSubstOutput(cs *syntax.CmdSubst, out []byte) error { + c := r.activeFields + if c == nil { + return nil + } + idx := -1 + for i := range c.splitSubsts { + if c.splitSubsts[i].cs == cs { + idx = i + } + } + if idx < 0 { + return nil + } + c.substSources += c.substFieldSources(idx, normalizeSubstOutput(out), r.fieldSplitIFS()) + if c.wordSources+c.substSources > c.maxFields-len(c.fields) { + return c.fieldLimitError() + } + return nil +} + +// normalizeSubstOutput returns the text the expander splits for a command +// substitution's output, mirroring mvdan.cc/sh/v3/expand: NUL bytes are +// removed and trailing newlines trimmed. Removing a NUL can join the bytes of +// a multi-byte IFS character, so counting the raw output could undercount. +func normalizeSubstOutput(out []byte) string { + end := len(out) + for end > 0 && (out[end-1] == '\n' || out[end-1] == 0) { + end-- + } + var sb strings.Builder + sb.Grow(end) + for _, b := range out[:end] { + if b != 0 { + sb.WriteByte(b) + } + } + return sb.String() +} + +// substFieldSources counts the field sources the idx-th substitution's +// normalized output adds, with the field continuation the expander applies +// across adjacent parts: a first run continues a field left open before it, +// and a deferred next part is a new source only if the output does not leave +// a field open. +func (c *fieldCollector) substFieldSources(idx int, out, ifs string) int { + s := &c.splitSubsts[idx] + n, closes, opens := splitFieldSources(out, ifs) + if s.inBrace { + return n + } + before := s.openBefore + if s.prev >= 0 && c.splitSubsts[s.prev].ran { + before = c.splitSubsts[s.prev].openAfter + } + if n > 0 && before && !closes { + n-- + } + after := before + if out != "" { + after = opens + } + s.ran, s.openAfter = true, after + if s.deferredNext && !after { + n++ + } + return n +} + func (r *Runner) fields(words ...*syntax.Word) []string { collector := r.newFieldCollector(MaxExpandedArgumentsPerCommand) collector.add(words...) diff --git a/interp/var_budget_test.go b/interp/var_budget_test.go new file mode 100644 index 000000000..97065f1bb --- /dev/null +++ b/interp/var_budget_test.go @@ -0,0 +1,148 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2026-present Datadog, Inc. + +package interp + +import ( + "bytes" + "context" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "mvdan.cc/sh/v3/expand" + "mvdan.cc/sh/v3/syntax" +) + +func strVar(s string) expand.Variable { + return expand.Variable{Set: true, Kind: expand.String, Str: s} +} + +// TestVarBudgetSharedByConcurrentOverlays checks the budget directly: two +// sibling overlays cannot together exceed the cap, and a released sibling's +// storage becomes available again. +func TestVarBudgetSharedByConcurrentOverlays(t *testing.T) { + root := &overlayEnviron{} + root.resetAccounting() + require.NoError(t, root.Set("BASE", strVar(strings.Repeat("b", 100<<10)))) + + left := newOverlayEnviron(root, true) + right := newOverlayEnviron(root, true) + half := strings.Repeat("x", 500<<10) + require.NoError(t, left.Set("X", strVar(half))) + err := right.Set("X", strVar(half)) + var storageErr *errTotalVarStorageExceeded + require.ErrorAs(t, err, &storageErr, "siblings must share one budget") + + left.release() + require.NoError(t, right.Set("X", strVar(half)), "released storage must be reusable") + + // A subshell of a stage is charged against its ancestors' and its + // siblings' storage through the shared budget. + nested := newOverlayEnviron(right, false) + assert.ErrorAs(t, nested.Set("Y", strVar(half)), &storageErr) + require.NoError(t, nested.Set("Y", strVar(strings.Repeat("y", 100<<10)))) + nested.release() + right.release() + assert.Equal(t, root.charged, root.budget.live) +} + +// TestVarBudgetReturnsToRootAfterRun checks that every subshell and pipeline +// stage releases what it charged, so only the root's own storage remains. +func TestVarBudgetReturnsToRootAfterRun(t *testing.T) { + script := `A=1 +x=$(B=2; echo $B) +( C=3 ) +{ D=4; echo $D; } | { E=5; cat; } | { F=$(G=6; echo $G); cat; } +for i in 1 2 3; do H=$i | cat; done +I=7 J=$I true +` + prog, err := syntax.NewParser().Parse(strings.NewReader(script), "") + require.NoError(t, err) + var out bytes.Buffer + r, err := New(StdIO(nil, &out, &out), AllowedCommands([]string{"rshell:echo", "rshell:cat", "rshell:true"})) + require.NoError(t, err) + t.Cleanup(func() { r.Close() }) + require.NoError(t, r.Run(context.Background(), prog), out.String()) + + root := r.writeEnv.(*overlayEnviron) + assert.Equal(t, root.charged, root.budget.live) +} + +// TestVarBudgetInheritedStorageNotRefunded checks that hiding or replacing a +// variable an ancestor still holds does not refund the ancestor's storage. +func TestVarBudgetInheritedStorageNotRefunded(t *testing.T) { + big := strings.Repeat("a", 600<<10) + other := strings.Repeat("b", 600<<10) + for name, background := range map[string]bool{"subshell": false, "pipeline stage": true} { + t.Run(name, func(t *testing.T) { + var storageErr *errTotalVarStorageExceeded + + root := &overlayEnviron{} + root.resetAccounting() + require.NoError(t, root.Set("A", strVar(big))) + + replace := newOverlayEnviron(root, background) + assert.ErrorAs(t, replace.Set("A", strVar(other)), &storageErr, + "replacing A must not refund the parent's copy") + replace.release() + + clear := newOverlayEnviron(root, background) + require.NoError(t, clear.Set("A", strVar(""))) + assert.ErrorAs(t, clear.Set("B", strVar(other)), &storageErr, + "clearing A must not refund the parent's copy") + clear.release() + + assert.Equal(t, root.charged, root.budget.live) + }) + } +} + +// TestVarBudgetTemporaryAssignmentsRestoreExactly checks that the binding +// before FOO=val cmd comes back exactly: a variable that did not exist leaves +// no entry behind (an undeclared entry would cost nothing yet still occupy +// the map), an inherited value is not copied into the child overlay, and the +// overlay owns exactly what it owned before. +func TestVarBudgetTemporaryAssignmentsRestoreExactly(t *testing.T) { + script := `A=x +NEW1= true +NEW2=v true +A=y true +( A=z true; INNER= true; echo "$A" ) +` + prog, err := syntax.NewParser().Parse(strings.NewReader(script), "") + require.NoError(t, err) + var out bytes.Buffer + r, err := New(StdIO(nil, &out, &out), AllowedCommands([]string{"rshell:echo", "rshell:true"})) + require.NoError(t, err) + t.Cleanup(func() { r.Close() }) + require.NoError(t, r.Run(context.Background(), prog), out.String()) + assert.Equal(t, "x\n", out.String()) + + root := r.writeEnv.(*overlayEnviron) + for _, name := range []string{"NEW1", "NEW2", "INNER"} { + _, ok := root.values[name] + assert.False(t, ok, "%s left an entry", name) + _, ok = root.owned[name] + assert.False(t, ok, "%s left owned storage", name) + } + assert.Equal(t, strVar("x").Str, root.values["A"].Str) + assert.Equal(t, varStorageCost("A", root.values["A"]), root.owned["A"]) + assert.Equal(t, root.charged, root.budget.live) + + // A child overlay restoring an inherited binding hides nothing and owns + // nothing once the command ends. + child := newOverlayEnviron(root, false) + snap := child.snapshotBinding("A") + require.NoError(t, child.Set("A", strVar(strings.Repeat("c", 1024)))) + child.restoreBinding(snap) + _, ok := child.values["A"] + assert.False(t, ok) + assert.Equal(t, varStorage{}, child.charged) + assert.Equal(t, "x", child.Get("A").Str) + child.release() + assert.Equal(t, root.charged, root.budget.live) +} diff --git a/interp/var_size_test.go b/interp/var_size_test.go index 6537cf6a3..57943c56b 100644 --- a/interp/var_size_test.go +++ b/interp/var_size_test.go @@ -89,17 +89,17 @@ func TestSubshellTotalVarStorageDoubleCount(t *testing.T) { // by newOverlayEnviron AND charged again as a full "new" write (oldBytes=0), // causing the delta to be inflated by len(prev.Str) and incorrectly hitting the // cap even when real memory use was well within bounds. +// +// The parent keeps its copy while the subshell runs, so the override itself is +// charged once on top of it: 3 × 340 KiB fits, while charging any of the three +// values twice would not. func TestNonBackgroundSubshellVarOverrideTracking(t *testing.T) { - // 512 KiB in parent; override with the same 512 KiB value inside a ( ) subshell. - // Real subshell memory = 512 KiB — well within the 1 MiB cap. - // Before the fix: totalBytes seed = 512 KiB, delta = +512 KiB (oldBytes=0), - // sum = 1 MiB, then Y=z would push it over → cap hit. - value512K := strings.Repeat("x", 512*1024) - script := fmt.Sprintf("X=%s\n( X=%s; Y=z; echo OK )\necho DONE\n", value512K, value512K) + value340K := strings.Repeat("x", 340*1024) + script := fmt.Sprintf("X=%s\n( X=%s; Y=%s; echo OK )\necho DONE\n", value340K, value340K, value340K) stdout, _, code := runScript(t, script, "") - assert.Contains(t, stdout, "OK", "subshell override should succeed: real storage is 512 KiB, within the 1 MiB cap") + assert.Contains(t, stdout, "OK", "subshell override should succeed: live storage is 1020 KiB, within the 1 MiB cap") assert.Contains(t, stdout, "DONE") assert.Equal(t, 0, code) } diff --git a/interp/vars.go b/interp/vars.go index b7957234d..81665b2e9 100644 --- a/interp/vars.go +++ b/interp/vars.go @@ -12,6 +12,7 @@ import ( "runtime" "strconv" "strings" + "sync" "mvdan.cc/sh/v3/expand" "mvdan.cc/sh/v3/syntax" @@ -22,64 +23,107 @@ import ( const MaxVarBytes = 1 << 20 // 1 MiB // MaxTotalVarsBytes is the maximum total size in bytes of all variable values -// combined. Assignments that would push the total over this limit are rejected. +// a script holds at once. The budget is shared by every subshell and pipeline +// stage of a run that is alive at the same time. Assignments that would push +// the total over this limit are rejected. const MaxTotalVarsBytes = 1 << 20 // 1 MiB +// MaxTotalVarEntryBytes is the maximum total of variable names plus a fixed +// per-variable overhead (varEntryOverheadBytes) a script holds at once. It +// bounds the number of variables, which MaxTotalVarsBytes alone does not: +// empty values cost nothing against it. With short names this allows roughly +// 16k variables. It is shared like MaxTotalVarsBytes. +const MaxTotalVarEntryBytes = 1 << 20 // 1 MiB + +// varEntryOverheadBytes is charged against MaxTotalVarEntryBytes for every +// stored variable in addition to its name. It approximates the map entry and +// expand.Variable bookkeeping, so that empty-valued variables are not free. +const varEntryOverheadBytes = 64 + +// varStorage is variable storage charged against the run-wide caps: value +// bytes against MaxTotalVarsBytes, and name bytes plus per-variable overhead +// against MaxTotalVarEntryBytes. +type varStorage struct { + values int + entries int +} + +func (a varStorage) add(b varStorage) varStorage { + return varStorage{values: a.values + b.values, entries: a.entries + b.entries} +} + +func (a varStorage) sub(b varStorage) varStorage { + return varStorage{values: a.values - b.values, entries: a.entries - b.entries} +} + +func (a varStorage) grows(from varStorage) bool { + return a.values > from.values || a.entries > from.entries +} + +func (a varStorage) exceedsCaps() bool { + return a.values > MaxTotalVarsBytes || a.entries > MaxTotalVarEntryBytes +} + +// varStorageCost returns what one stored variable is charged. Undeclared +// variables have no entry and cost nothing. +func varStorageCost(name string, vr expand.Variable) varStorage { + if !vr.Declared() { + return varStorage{} + } + return varStorage{values: len(vr.Str), entries: len(name) + varEntryOverheadBytes} +} + +// varStorageBudget is the run-wide variable storage budget shared by every +// overlay derived from the same root environment. live is the sum of the +// storage each live overlay owns (see overlayEnviron.owned). Subshells and +// pipeline stages run on their own overlays, and pipeline stages run +// concurrently, so without a shared budget each one could hold the full caps +// at the same time. +type varStorageBudget struct { + mu sync.Mutex + live varStorage +} + // errTotalVarStorageExceeded is returned by overlayEnviron.Set when the total // variable storage cap is exceeded. It is a distinct type so that setVar can // treat it as a fatal error and abort the script. type errTotalVarStorageExceeded struct { - total int + total varStorage } func (e *errTotalVarStorageExceeded) Error() string { - return fmt.Sprintf("variable storage limit exceeded (%d bytes total)", e.total) + if e.total.values > MaxTotalVarsBytes { + return fmt.Sprintf("variable storage limit exceeded (%d bytes total)", e.total.values) + } + return fmt.Sprintf("variable storage limit exceeded (%d bytes of names and entries total)", e.total.entries) } func newOverlayEnviron(parent expand.Environ, background bool) *overlayEnviron { oenv := &overlayEnviron{} if !background { oenv.parent = parent - // Seed totalBytes from the parent's tracked counter so that the cap is - // enforced consistently across subshell nesting levels (preventing a - // script from resetting the counter to zero at each nesting level and - // allocating O(depth × MaxTotalVarsBytes) before hitting any limit). - // - // When the parent is an *overlayEnviron we use its tracked counter - // directly, which reflects only script-assigned variables (Env() - // variables are excluded from the cap via the Reset() zero-out, so they - // do not appear in pov.totalBytes). Summing via parent.Each() instead - // would also count inherited Env() variables, producing false cap hits - // for legitimate callers that provide a large Env() variable. - // - // When the parent is not an *overlayEnviron (fallback) we must sum - // manually since there is no pre-computed counter. - if pov, ok := parent.(*overlayEnviron); ok { - oenv.totalBytes = pov.totalBytes - } else { - parent.Each(func(_ string, vr expand.Variable) bool { - oenv.totalBytes += len(vr.Str) - return true - }) - } } else { + // The copied entries share their storage with the parent, which keeps + // owning (and being charged for) them while this overlay is alive. oenv.values = make(map[string]expand.Variable) maps.Insert(oenv.values, parent.Each) - // Seed totalBytes from the parent's tracked counter rather than - // re-summing all variables (which would include inherited env vars - // provided via Env() that the main runner intentionally excludes from - // the cap). Using the parent's counter keeps background-subshell - // accounting consistent with the main runner and prevents false positives - // when the caller provides a large Env() variable: an assignment that - // succeeds in the parent must also succeed in a pipeline/background subshell. - if pov, ok := parent.(*overlayEnviron); ok { - oenv.totalBytes = pov.totalBytes - } else { - // Fallback for non-overlayEnviron parents: sum all values. - for _, vr := range oenv.values { - oenv.totalBytes += len(vr.Str) - } - } + } + if pov, ok := parent.(*overlayEnviron); ok { + // Join the parent's run-wide budget. The parent's storage stays + // charged there for as long as it is alive, so variables this overlay + // inherits are counted exactly once, and hiding or replacing them here + // does not refund storage the parent still holds. Variables provided + // through Env() are not owned by any overlay and are never counted. + oenv.budget = pov.sharedBudget() + } else { + // Fallback for non-overlayEnviron parents, which have no budget: + // count every inherited entry. + var base varStorage + parent.Each(func(name string, vr expand.Variable) bool { + base = base.add(varStorageCost(name, vr)) + return true + }) + oenv.budget = &varStorageBudget{live: base} } return oenv } @@ -90,14 +134,79 @@ type overlayEnviron struct { // which we can safely reuse without data races, such as non-background subshells. parent expand.Environ values map[string]expand.Variable - // totalBytes tracks the total script-assigned variable storage counted - // against MaxTotalVarsBytes. For background subshells, where all parent - // variables are copied into [values], this equals the sum of len(value.Str) - // over [values]. For non-background subshells (e.g. ( ) or $( )), the - // counter is seeded from the parent's counter so parent-inherited bytes - // are included even before any variable is written to [values]; those - // inherited variables are NOT present in [values] themselves. - totalBytes int + + // budget is the run-wide budget shared with every other overlay derived + // from the same root. It is created lazily for standalone overlays. + budget *varStorageBudget + // owned is the storage, per variable, that this overlay itself assigned + // and holds in budget, counted against MaxTotalVarsBytes and + // MaxTotalVarEntryBytes (see varStorageCost). Entries inherited from a + // parent, or copied from it for background subshells, are owned by the + // parent and absent here. charged is the sum of owned. + owned map[string]varStorage + charged varStorage +} + +// sharedBudget returns o's run-wide budget, creating one seeded with o's own +// storage if o has none yet. It is only called while o has no live +// children, so the lazy creation cannot race. +func (o *overlayEnviron) sharedBudget() *varStorageBudget { + if o.budget == nil { + o.budget = &varStorageBudget{live: o.charged} + } + return o.budget +} + +// own records that o now holds vr for name, replacing what o previously owned +// for it, and updates the run-wide budget. When enforce is set and o's +// storage for name grows, the change is rejected if the storage held by every +// live overlay (ancestors, concurrent pipeline stages and their subshells) +// would exceed a cap. The check and the update happen under the budget lock +// so concurrent stages cannot both pass against the same snapshot. +func (o *overlayEnviron) own(name string, vr expand.Variable, enforce bool) error { + newCost := varStorageCost(name, vr) + oldCost := o.owned[name] + delta := newCost.sub(oldCost) + budget := o.sharedBudget() + budget.mu.Lock() + defer budget.mu.Unlock() + if enforce && newCost.grows(oldCost) { + if after := budget.live.add(delta); after.exceedsCaps() { + return &errTotalVarStorageExceeded{total: after} + } + } + budget.live = budget.live.add(delta) + o.charged = o.charged.add(delta) + if newCost == (varStorage{}) { + delete(o.owned, name) + } else { + if o.owned == nil { + o.owned = make(map[string]varStorage) + } + o.owned[name] = newCost + } + return nil +} + +// release returns o's owned storage to the run-wide budget. It must be called +// once the subshell, pipeline stage or temporary overlay using o has +// finished; o must not be written afterwards. +func (o *overlayEnviron) release() { + if o.budget != nil { + o.budget.mu.Lock() + o.budget.live = o.budget.live.sub(o.charged) + o.budget.mu.Unlock() + } + o.owned = nil + o.charged = varStorage{} +} + +// resetAccounting forgets what o owns, so its current variables are not +// counted, and starts a fresh run-wide budget rooted at o. +func (o *overlayEnviron) resetAccounting() { + o.budget = &varStorageBudget{} + o.owned = nil + o.charged = varStorage{} } func (o *overlayEnviron) Get(name string) expand.Variable { @@ -137,51 +246,35 @@ func (o *overlayEnviron) Set(name string, vr expand.Variable) error { if prev.ReadOnly { return fmt.Errorf("readonly variable") } + // Only storage o owns is refunded; a parent's entry stays charged + // to the parent. if prev.Local { + // The unset local retains its slot but no value. vr.Local = true - // Subtract old value from total; the unset local retains its slot but no value. - // Use prev.Str unconditionally: non-background subshells seed totalBytes from - // the parent, so parent-inherited vars are already counted even when !inOverlay. - o.totalBytes -= len(prev.Str) - if o.totalBytes < 0 { - o.totalBytes = 0 // defensive guard against invariant violation - } + _ = o.own(name, vr, false) o.values[name] = vr return nil } - // Same reasoning as above: subtract unconditionally so parent-seeded bytes are - // correctly released when a parent-inherited variable is deleted in a subshell. - o.totalBytes -= len(prev.Str) - if o.totalBytes < 0 { - o.totalBytes = 0 // defensive guard against invariant violation - } + _ = o.own(name, expand.Variable{}, false) delete(o.values, name) return nil } - // modifying the entire variable — enforce total storage cap. - // Use the previous value's byte count unconditionally: for non-background - // subshells, newOverlayEnviron seeds totalBytes by summing parent variables, - // so the parent's bytes are already counted. If we only credited oldBytes - // when inOverlay (i.e. the variable was already in our overlay), we would - // double-charge the parent's contribution on the first override, erroneously - // inflating totalBytes by len(prev.Str). - oldBytes := len(prev.Str) - newBytes := len(vr.Str) - delta := newBytes - oldBytes - if delta > 0 && o.totalBytes+delta > MaxTotalVarsBytes { - return &errTotalVarStorageExceeded{total: o.totalBytes + delta} - } - o.totalBytes += delta + // modifying the entire variable — enforce total storage cap. Only what + // o previously owned for name is credited: overriding a parent's entry + // leaves the parent's storage charged, since the parent still holds it. vr.Local = prev.Local || vr.Local + if err := o.own(name, vr, true); err != nil { + return err + } o.values[name] = vr return nil } // setUncapped sets a variable without enforcing MaxTotalVarsBytes. It is -// used by setVarRestore to restore inline command variables (e.g. FOO=val cmd) -// to their original values after the command returns, even when the script -// has filled storage close to the cap during command execution. totalBytes -// is still updated so that subsequent cap checks use an accurate baseline. +// used by setVarRestore to roll back a variable (e.g. OLDPWD after a failed +// cd) even when the script has filled storage close to the cap. The owned +// storage is still updated so that subsequent cap checks use an accurate +// baseline. func (o *overlayEnviron) setUncapped(name string, vr expand.Variable) { if o.values == nil { o.values = make(map[string]expand.Variable) @@ -190,16 +283,57 @@ func (o *overlayEnviron) setUncapped(name string, vr expand.Variable) { if !inOverlay && o.parent != nil { prev = o.parent.Get(name) } - oldBytes := len(prev.Str) - newBytes := len(vr.Str) - o.totalBytes += newBytes - oldBytes - if o.totalBytes < 0 { - o.totalBytes = 0 - } vr.Local = prev.Local || vr.Local + _ = o.own(name, vr, false) o.values[name] = vr } +// bindingSnapshot records o's own binding for a name: its entry, if any, and +// the storage o owns for it. +type bindingSnapshot struct { + name string + vr expand.Variable + inOverlay bool + owned varStorage +} + +// snapshotBinding records o's binding for name so restoreBinding can put it +// back exactly once a temporary assignment (FOO=val cmd) ends. +func (o *overlayEnviron) snapshotBinding(name string) bindingSnapshot { + vr, inOverlay := o.values[name] + return bindingSnapshot{name: name, vr: vr, inOverlay: inOverlay, owned: o.owned[name]} +} + +// restoreBinding puts back a binding recorded by snapshotBinding, without +// enforcing the storage caps. A name o had no entry for is deleted rather +// than set, so an inherited value shows through again without o taking +// ownership of a copy, and a previously absent name leaves no entry behind. +// o's ownership returns to what it was, whatever the command assigned. +func (o *overlayEnviron) restoreBinding(snap bindingSnapshot) { + budget := o.sharedBudget() + budget.mu.Lock() + delta := snap.owned.sub(o.owned[snap.name]) + budget.live = budget.live.add(delta) + budget.mu.Unlock() + o.charged = o.charged.add(delta) + if snap.owned == (varStorage{}) { + delete(o.owned, snap.name) + } else { + if o.owned == nil { + o.owned = make(map[string]varStorage) + } + o.owned[snap.name] = snap.owned + } + if snap.inOverlay { + if o.values == nil { + o.values = make(map[string]expand.Variable) + } + o.values[snap.name] = snap.vr + } else { + delete(o.values, snap.name) + } +} + func (o *overlayEnviron) Each(f func(name string, vr expand.Variable) bool) { // Emit our own overrides first so they take precedence. for name, vr := range o.values { @@ -210,7 +344,7 @@ func (o *overlayEnviron) Each(f func(name string, vr expand.Variable) bool) { // Then emit parent variables, skipping any that we already overrode above. // Without this guard, a variable that exists in both the parent environment // and our overlay (because the script re-assigned it) would be emitted twice, - // causing newOverlayEnviron to seed totalBytes at 2× the real storage. + // and anything summing Each() would count it twice. if o.parent != nil { o.parent.Each(func(name string, vr expand.Variable) bool { if _, override := o.values[name]; override {