From 9de3b4eb23873d8cc74fa952d1f83bc9ae93ff1c Mon Sep 17 00:00:00 2001 From: Kevin Tang <73975146+vt128@users.noreply.github.com> Date: Sun, 12 Jul 2026 23:07:52 +0800 Subject: [PATCH] [fix] apply the Machine step budget and print func to load()ed modules MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A module executed by load() ran in a bare starlark.Thread that inherited none of the Machine's execution context, so: - the step budget did not apply — a top-level loop in a load()ed source module ran unbounded, escaping the DoS guard (the budget only armed the main thread); - print() in a loaded module went to stdout via a hardcoded fmt.Println, bypassing the Machine's print func. The load cache now builds its child thread from the Machine (newLoadThread): the same print func, an independent copy of the step budget (with the MaxStepsExceeded panic), and the current run's context local. The budget is per-thread — enough to stop one runaway loaded module, which is the DoS the bare thread let through; it is not a shared aggregate counter. Because a loaded module can now panic (OnMaxSteps), the load cache's get() recovers around doLoad: it readies the entry with the error, drops it so a retry re-executes, and re-raises so runInternal maps it to a typed error. Without this the half-built cache entry was never readied, and a second Run of a module that hit the budget blocked forever on its ready channel. Not covered here (separate follow-ups, noted honestly rather than half-done): lazy loaders built with MakeModuleLoaderFromString/Reader/File still run on their own bare thread and inherit no budget/print (fixing needs an API change); and a context timeout cancels only the main thread, so a loaded module's pure-Starlark loop is bounded by the step budget, not the deadline. Test: TestMachine_LoadInheritsMachineContext — a loaded module hits the step budget, its print() reaches the Machine print func, and a second run after a budget panic returns promptly instead of deadlocking on the load cache. Co-Authored-By: Claude Opus 4.8 (1M context) --- cache.go | 53 +++++++++++++++++++++++++++---- run.go | 35 ++++++++++++++++++++- run_test.go | 89 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 170 insertions(+), 7 deletions(-) diff --git a/cache.go b/cache.go index 3e07ce6..27fa1fb 100644 --- a/cache.go +++ b/cache.go @@ -28,6 +28,11 @@ type cache struct { execOpts *syntax.FileOptions loadMod func(s string) (starlark.StringDict, error) // load from built-in module first readFile func(s string) ([]byte, error) // and then from file system + // newThread builds the thread that executes a loaded module, carrying the + // Machine's execution context (print func, step budget, run context). When + // nil the load path falls back to a bare thread — a module loaded that way + // would inherit none of the Machine's guards. + newThread func(load func(*starlark.Thread, string) (starlark.StringDict, error)) *starlark.Thread } type entry struct { @@ -77,7 +82,34 @@ func (c *cache) get(cc *cycleChecker, module string) (starlark.StringDict, error c.cacheMu.Unlock() e.setOwner(cc) + // doLoad can panic: a loaded module can exhaust the step budget, which + // OnMaxSteps raises as a panic. Without recovery the half-built entry + // would never be readied or removed, so every later load of the same + // module blocks forever on <-e.ready. Ready the entry (with the error), + // drop it so a retry re-executes, then re-raise so runInternal maps it + // to a typed error. A `loaded` sentinel — not `recover() != nil` — + // decides whether doLoad panicked, because under the go 1.19 floor + // recover() returns nil for panic(nil), which would skip the cleanup. + loaded := false + defer func() { + if loaded { + return + } + r := recover() + if e.err == nil { + if err, ok := r.(error); ok { + e.err = err + } else { + e.err = fmt.Errorf("panic loading module %q: %v", module, r) + } + } + e.setOwner(nil) + close(e.ready) + c.remove(module) + panic(r) + }() e.globals, e.err = c.doLoad(cc, module) + loaded = true e.setOwner(nil) // Broadcast that the entry is now ready. @@ -87,12 +119,21 @@ func (c *cache) get(cc *cycleChecker, module string) (starlark.StringDict, error } func (c *cache) doLoad(cc *cycleChecker, module string) (starlark.StringDict, error) { - thread := &starlark.Thread{ - Print: func(_ *starlark.Thread, msg string) { fmt.Println(msg) }, - Load: func(_ *starlark.Thread, module string) (starlark.StringDict, error) { - // Tunnel the cycle-checker state for this "thread of loading". - return c.get(cc, module) - }, + // Tunnel the cycle-checker state for this "thread of loading". + loadFn := func(_ *starlark.Thread, module string) (starlark.StringDict, error) { + return c.get(cc, module) + } + var thread *starlark.Thread + if c.newThread != nil { + // inherit the Machine's print func, step budget, and run context so a + // loaded module cannot bypass them (e.g. a top-level loop escaping the + // DoS guard, or print() going to stdout instead of the print func) + thread = c.newThread(loadFn) + } else { + thread = &starlark.Thread{ + Print: func(_ *starlark.Thread, msg string) { fmt.Println(msg) }, + Load: loadFn, + } } // 1: load from built-in module, the first field returns nil if not found diff --git a/run.go b/run.go index 202b069..199a12a 100644 --- a/run.go +++ b/run.go @@ -274,7 +274,8 @@ func (m *Machine) prepareThread(extras StringAnyMap) (err error) { readFile: func(name string) ([]byte, error) { return readScriptFile(name, m.scriptFS) }, - globals: m.predeclared, + globals: m.predeclared, + newThread: m.newLoadThread, } m.thread = &starlark.Thread{ Name: "starlet", @@ -305,6 +306,38 @@ func (m *Machine) prepareThread(extras StringAnyMap) (err error) { return nil } +// newLoadThread builds the thread that runs a module executed by load(), +// mirroring the main thread's execution context: the same print func, an +// independent copy of the step budget (so a loaded module's work is bounded +// by the DoS guard instead of escaping it), and the current run's context +// local. The step budget is per-thread, not a shared aggregate counter, so a +// loaded module gets its own MaxSteps allowance — enough to stop a runaway +// loop, which is the DoS the bare thread let through. +func (m *Machine) newLoadThread(load func(*starlark.Thread, string) (starlark.StringDict, error)) *starlark.Thread { + t := &starlark.Thread{ + Name: "starlet:load", + Print: m.printFunc, + Load: load, + } + limit := m.maxSteps + if limit == 0 { + limit = math.MaxUint64 + } + t.SetMaxExecutionSteps(limit) + if lim := m.maxSteps; lim > 0 { + t.OnMaxSteps = func(*starlark.Thread) { + // recovered by the run/call recover and mapped to a typed error + panic(MaxStepsExceededError{Limit: lim}) + } + } + if m.thread != nil { + if ctx := m.thread.Local("context"); ctx != nil { + t.SetLocal("context", ctx) + } + } + return t +} + // applyStepBudget arms the thread with the configured step budget; it must // run before every execution. The Starlark runtime normalizes a zero limit // to "unlimited" only on a thread's very first use, so the translation to diff --git a/run_test.go b/run_test.go index 04a1125..f0b4a5b 100644 --- a/run_test.go +++ b/run_test.go @@ -10,6 +10,7 @@ import ( "reflect" "strings" "testing" + "testing/fstest" "time" "github.com/1set/starlet" @@ -17,6 +18,94 @@ import ( "go.starlark.net/starlark" ) +// TestMachine_LoadInheritsMachineContext: a module executed by load() ran in a +// bare thread that inherited none of the Machine's execution context, so the +// step budget did not apply (a top-level loop in a loaded module bypassed the +// whole DoS guard) and print() in a loaded module went to stdout instead of +// the Machine's print function. The load path now builds its child thread from +// the Machine's settings. +func TestMachine_LoadInheritsMachineContext(t *testing.T) { + t.Run("step budget applies to a loaded module", func(t *testing.T) { + sfs := fstest.MapFS{ + // a top-level comprehension that burns far more than the budget + "heavy.star": &fstest.MapFile{Data: []byte(`x = len([i for i in range(1000000)])`)}, + } + m := starlet.NewDefault() + m.SetScript("main.star", []byte(`load("heavy.star", "x"); y = x`), sfs) + m.SetMaxExecutionSteps(1000) + _, err := m.Run() + var me starlet.MaxStepsExceededError + if !errors.As(err, &me) { + t.Fatalf("expected the loaded module to hit the step budget, got: %v", err) + } + }) + + t.Run("print in a loaded module uses the Machine print func", func(t *testing.T) { + sfs := fstest.MapFS{ + "greet.star": &fstest.MapFile{Data: []byte(`print("hi from loaded")` + "\n" + `msg = 1`)}, + } + m := starlet.NewDefault() + m.SetScript("main.star", []byte(`load("greet.star", "msg"); z = msg`), sfs) + pf, cmp := getPrintCompareFunc(t) + m.SetPrintFunc(pf) + if _, err := m.Run(); err != nil { + t.Fatalf("run: %v", err) + } + cmp("hi from loaded\n") + }) + + t.Run("a budget panic in a loaded module does not poison the cache", func(t *testing.T) { + // The step-budget panic unwinds through the load cache; the abandoned + // entry must not leave a later load of the same module blocked on its + // ready channel forever. A second Run must return promptly, not hang. + sfs := fstest.MapFS{ + "heavy.star": &fstest.MapFile{Data: []byte(`x = len([i for i in range(1000000)])`)}, + } + m := starlet.NewDefault() + m.SetScript("main.star", []byte(`load("heavy.star", "x"); y = x`), sfs) + m.SetMaxExecutionSteps(1000) + + var me starlet.MaxStepsExceededError + if _, err := m.Run(); !errors.As(err, &me) { + t.Fatalf("first run: expected MaxStepsExceededError, got: %v", err) + } + done := make(chan error, 1) + go func() { _, err := m.Run(); done <- err }() + select { + case err := <-done: + if !errors.As(err, &me) { + t.Fatalf("second run: expected MaxStepsExceededError, got: %v", err) + } + case <-time.After(5 * time.Second): + t.Fatal("second run deadlocked: the load cache entry was poisoned by the budget panic") + } + }) + + t.Run("a panic(nil) loader does not poison the cache", func(t *testing.T) { + // The cache cleanup keys on a completion sentinel, not recover() != nil, + // so it survives panic(nil) too (recover() is nil for it under go 1.19). + // A second Run must not deadlock on the abandoned entry. + mods := starlet.ModuleLoaderMap{ + "boom": func() (starlark.StringDict, error) { panic(nil) }, + } + m := starlet.NewWithLoaders(nil, nil, mods) + m.SetScript("main.star", []byte(`load("boom", "x")`), nil) + + runOnce := func() { + defer func() { _ = recover() }() // the panic propagates back out; swallow it + _, _ = m.Run() + } + runOnce() + done := make(chan struct{}) + go func() { runOnce(); close(done) }() + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("second run deadlocked after a panic(nil) loader") + } + }) +} + func Test_DefaultMachine_Run_NoCode(t *testing.T) { m := starlet.NewDefault() // run with empty script