[fix] apply the Machine step budget and print func to load()ed modules

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

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

run_test.go

@@ -10,13 +10,102 @@ import (
 	"reflect"
 	"strings"
 	"testing"
+	"testing/fstest"
 	"time"
 
 	"github.com/1set/starlet"
 	"github.com/1set/starlight/convert"
 	"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