From 65449d4f03be0ad05b560b56fbbaf680ccde14b9 Mon Sep 17 00:00:00 2001 From: Meo597 <197331664+Meo597@users.noreply.github.com> Date: Thu, 15 Oct 2026 11:18:27 +0800 Subject: [PATCH] lua: require explicit registry sizing for program states to reduce memory usage --- app/dns/script.go | 1 + app/router/script.go | 1 + common/lua/program.go | 24 +++++++++++++++++------- common/lua/program_test.go | 38 +++++++++++++++++++------------------- 4 files changed, 38 insertions(+), 26 deletions(-) diff --git a/app/dns/script.go b/app/dns/script.go index c465d175a..c4cae5197 100644 --- a/app/dns/script.go +++ b/app/dns/script.go @@ -30,6 +30,7 @@ func newScriptEngine(script string) (*scriptEngine, error) { func (e *scriptEngine) start(server *DNS) error { pool, err := xlua.NewPool(server.ctx, scriptExecutionTimeout, e.program.NewStateFactory( scriptExecutionTimeout*20, + 256, func(L *lua.LState) { geodata.RegisterLua(L) log.RegisterLua(L) diff --git a/app/router/script.go b/app/router/script.go index 2893e0624..e2659cb87 100644 --- a/app/router/script.go +++ b/app/router/script.go @@ -31,6 +31,7 @@ func newScriptEngine(script string) (*scriptEngine, error) { func (e *scriptEngine) start(router *Router) error { pool, err := xlua.NewPool(router.ctx, scriptExecutionTimeout, e.program.NewStateFactory( scriptExecutionTimeout*20, + 256, func(L *lua.LState) { geodata.RegisterLua(L) log.RegisterLua(L) diff --git a/common/lua/program.go b/common/lua/program.go index da64ac220..51a41fa78 100644 --- a/common/lua/program.go +++ b/common/lua/program.go @@ -59,12 +59,19 @@ func Compile(script string) (*Program, error) { // NewState creates a state, runs register, executes the program under ctx, and // runs validate. It removes the initialization context before returning a state // owned by the caller. -func (p *Program) NewState(ctx context.Context, register func(*glua.LState), validate func(*glua.LState) error) (*glua.LState, error) { +// Choose registrySize for the script type's expected stack usage, balancing +// per-state memory usage against the number of states retained in a pool. +func (p *Program) NewState(ctx context.Context, registrySize int, register func(*glua.LState), validate func(*glua.LState) error) (*glua.LState, error) { + options := glua.Options{ + RegistrySize: registrySize, + RegistryMaxSize: glua.RegistrySize, + RegistryGrowStep: registrySize, + } var L *glua.LState if p.privileged { - L = glua.NewState() + L = glua.NewState(options) } else { - L = newSandboxedState() + L = newSandboxedState(options) } valid := false defer func() { @@ -95,8 +102,9 @@ func (p *Program) NewState(ctx context.Context, register func(*glua.LState), val return L, nil } -func newSandboxedState() *glua.LState { - L := glua.NewState(glua.Options{SkipOpenLibs: true}) +func newSandboxedState(options glua.Options) *glua.LState { + options.SkipOpenLibs = true + L := glua.NewState(options) for _, lib := range []struct { name string @@ -239,10 +247,12 @@ func redactMemoryAddresses(message string) string { } // NewStateFactory returns a factory that gives each state an initialization timeout. -func (p *Program) NewStateFactory(initTimeout time.Duration, register func(*glua.LState), validate func(*glua.LState) error) LStateFactory { +// Choose registrySize for the script type's expected stack usage, balancing +// per-state memory usage against the number of states retained in a pool. +func (p *Program) NewStateFactory(initTimeout time.Duration, registrySize int, register func(*glua.LState), validate func(*glua.LState) error) LStateFactory { return func(ctx context.Context) (*glua.LState, error) { initCtx, cancel := context.WithTimeout(ctx, initTimeout) defer cancel() - return p.NewState(initCtx, register, validate) + return p.NewState(initCtx, registrySize, register, validate) } } diff --git a/common/lua/program_test.go b/common/lua/program_test.go index c18b69466..6a24cbeb9 100644 --- a/common/lua/program_test.go +++ b/common/lua/program_test.go @@ -33,13 +33,13 @@ func TestProgramStatesAreIndependent(t *testing.T) { if program.proto.SourceName != path { t.Fatalf("program source = %q, want %q", program.proto.SourceName, path) } - first, err := program.NewState(context.Background(), nil, nil) + first, err := program.NewState(context.Background(), 256, nil, nil) if err != nil { t.Fatal(err) } defer first.Close() first.SetGlobal("value", glua.LNumber(42)) - second, err := program.NewState(context.Background(), nil, nil) + second, err := program.NewState(context.Background(), 256, nil, nil) if err != nil { t.Fatal(err) } @@ -60,7 +60,7 @@ func TestProgramInitializationObservesCancellation(t *testing.T) { } ctx, cancel := context.WithCancel(context.Background()) cancel() - state, err := program.NewState(ctx, nil, nil) + state, err := program.NewState(ctx, 256, nil, nil) if err == nil || state != nil { if state != nil { state.Close() @@ -80,7 +80,7 @@ func TestNewStateClosesFailedValidation(t *testing.T) { } wantErr := errors.New("invalid script") var checked *glua.LState - L, err := program.NewState(context.Background(), nil, func(L *glua.LState) error { + L, err := program.NewState(context.Background(), 256, nil, func(L *glua.LState) error { checked = L return wantErr }) @@ -122,7 +122,7 @@ func TestCompileInline(t *testing.T) { if program.proto.SourceName != "[inline]" { t.Fatalf("program source = %q, want [inline]", program.proto.SourceName) } - state, err := program.NewState(context.Background(), nil, nil) + state, err := program.NewState(context.Background(), 256, nil, nil) if err != nil { t.Fatal(err) } @@ -175,7 +175,7 @@ func TestCompileInlineDiagnostics(t *testing.T) { if err != nil { t.Fatal(err) } - state, initErr := program.NewState(context.Background(), nil, nil) + state, initErr := program.NewState(context.Background(), 256, nil, nil) if state != nil { state.Close() t.Fatal("NewState returned a state after failed initialization") @@ -370,7 +370,7 @@ func TestProgramPrivilegedSignatureFixture(t *testing.T) { if program.privileged != signed { t.Fatalf("privileged = %t, want %t", program.privileged, signed) } - L, err := program.NewState(context.Background(), nil, nil) + L, err := program.NewState(context.Background(), 256, nil, nil) if err != nil { t.Fatal(err) } @@ -501,7 +501,7 @@ func TestProgramInlineSignaturesCannotAuthorize(t *testing.T) { if err != nil || program == nil || program.privileged { t.Fatalf("inline signed source Compile = %v, %v; want sandbox without verification", program, err) } - L, err := program.NewState(context.Background(), nil, nil) + L, err := program.NewState(context.Background(), 256, nil, nil) if err != nil { t.Fatal(err) } @@ -533,7 +533,7 @@ func TestProgramSignedDiagnostics(t *testing.T) { if err != nil { t.Fatal(err) } - L, err := program.NewState(context.Background(), nil, nil) + L, err := program.NewState(context.Background(), 256, nil, nil) if L != nil { L.Close() t.Fatal("failing signed script returned a state") @@ -600,7 +600,7 @@ function CheckSandbox() package = nil end`) var helperLoads, lateLoads int - L, err := program.NewState(context.Background(), func(L *glua.LState) { + L, err := program.NewState(context.Background(), 256, func(L *glua.LState) { if L.GetGlobal("io") != glua.LNil || L.GetGlobal("loadfile") != glua.LNil { t.Fatal("sandbox was not configured before module registration") } @@ -659,7 +659,7 @@ assert(tostring("hello") == "hello") string.__tostring = nil `) var metamethodCalls int - L, err := program.NewState(context.Background(), func(L *glua.LState) { + L, err := program.NewState(context.Background(), 256, func(L *glua.LState) { L.SetGlobal("testIPv4", luar.New(L, net.ParseIP("192.0.2.1"))) L.SetGlobal("testIPv6", luar.New(L, net.ParseIP("2001:db8::1"))) PushUserData(L, net.ParseIP("198.51.100.1")) @@ -733,7 +733,7 @@ for _, format in ipairs({string.format, require("string").format, ("").format}) end assert(("%s"):format({}) == "") `) - L, err := program.NewState(context.Background(), func(L *glua.LState) { + L, err := program.NewState(context.Background(), 256, func(L *glua.LState) { L.SetGlobal("testIP", luar.New(L, net.ParseIP("192.0.2.1"))) PushUserData(L, bytes.NewBufferString("0xdeadbeef")) L.SetGlobal("testNativeStringer", L.Get(-1)) @@ -852,7 +852,7 @@ ok, a = xpcall(function() error(businessError, 0) end, testPanic) assert(not ok and a == "{
}") `) var sent []string - L, err := program.NewState(context.Background(), func(L *glua.LState) { + L, err := program.NewState(context.Background(), 256, func(L *glua.LState) { callable := L.NewTable() metatable := L.NewTable() metatable.RawSetString("__call", L.NewFunction(func(L *glua.LState) int { @@ -912,7 +912,7 @@ end `) nativeErrors := make(map[string]error) originalMessages := make(map[string]string) - L, err := program.NewState(context.Background(), func(L *glua.LState) { + L, err := program.NewState(context.Background(), 256, func(L *glua.LState) { luaError := &glua.ApiError{Object: L.NewFunction(func(*glua.LState) int { return 0 })} stringerError := xerrors.New("DNS query failed").Base(luaError) messages := L.NewTable() @@ -964,7 +964,7 @@ end package = nil `) var failedLoads, cycleLoads int - L, err := program.NewState(context.Background(), func(L *glua.LState) { + L, err := program.NewState(context.Background(), 256, func(L *glua.LState) { L.PreloadModule("xray.failure", func(L *glua.LState) int { failedLoads++ L.RaiseError("module initialization failed") @@ -999,7 +999,7 @@ func TestProgramSandboxModulesAreCollected(t *testing.T) { t.Run(name, func(t *testing.T) { ref := func() weak.Pointer[glua.LTable] { program := compileTestProgram(t, `payload = string.rep("x", 1024 * 1024); assert(pcall(require, "xray.test") == expectedSuccess)`) - L, err := program.NewState(context.Background(), func(L *glua.LState) { + L, err := program.NewState(context.Background(), 256, func(L *glua.LState) { L.SetGlobal("expectedSuccess", glua.LBool(!fail)) L.PreloadModule("xray.test", func(L *glua.LState) int { if fail { @@ -1028,7 +1028,7 @@ func TestProgramSandboxModulesAreCollected(t *testing.T) { func TestProgramSandboxClosesFailedInitialization(t *testing.T) { program := compileTestProgram(t, `require("xray.failure")`) var partial *glua.LState - L, err := program.NewState(context.Background(), func(L *glua.LState) { + L, err := program.NewState(context.Background(), 256, func(L *glua.LState) { partial = L L.PreloadModule("xray.failure", func(L *glua.LState) int { L.RaiseError("module initialization failed") @@ -1093,7 +1093,7 @@ file:close()` t.Fatal(err) } output := filepath.Join(t.TempDir(), "privileged.txt") - L, err := program.NewState(context.Background(), func(L *glua.LState) { + L, err := program.NewState(context.Background(), 256, func(L *glua.LState) { L.SetGlobal("testPath", glua.LString(output)) }, nil) if !tc.privileged { @@ -1156,7 +1156,7 @@ func TestProgramPrivilegesAreCompileTimeSnapshot(t *testing.T) { t.Fatalf("os.execute available = %t, want compile-time value %t", got, privileged) } } - factory := program.NewStateFactory(time.Second, nil, nil) + factory := program.NewStateFactory(time.Second, 256, nil, nil) for range 2 { L, err := factory(context.Background()) if err != nil {