lua: require explicit registry sizing for program states to reduce memory usage

This commit is contained in:
Meo597
2026-10-15 11:18:27 +08:00
parent d5bfb0a5c7
commit 65449d4f03
4 changed files with 38 additions and 26 deletions
+1
View File
@@ -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)
+1
View File
@@ -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)
+17 -7
View File
@@ -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)
}
}
+19 -19
View File
@@ -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({}) == "<opaque>")
`)
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 == "{<address> <address>}")
`)
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 {