Reduce memory usage

common/lua/program.go

@@ -15,7 +15,8 @@ import (
 
 // Program holds immutable bytecode that can be run by independent LStates.
 type Program struct {
-	proto *glua.FunctionProto
+	proto      *glua.FunctionProto
+	privileged bool
 }
 
 // LStateFactory returns a fully initialized state or nil and an error.
@@ -27,10 +28,12 @@ func Compile(script string) (*Program, error) {
 	var code []byte
 	var err error
 	name := "[inline]"
+	privileged := false
 	if strings.HasPrefix(script, base64.GzipPrefix) || strings.HasPrefix(script, base64.Prefix) {
 		code, err = base64.Decode(script)
 	} else if name, err = platform.ResolveLuaFile(script); err == nil {
 		code, err = os.ReadFile(name)
+		privileged = hasPrivilegedPSK(code)
 	}
 	if err != nil {
 		return nil, err
@@ -44,14 +47,43 @@ func Compile(script string) (*Program, error) {
 	if err != nil {
 		return nil, err
 	}
-	return &Program{proto: proto}, nil
+	return &Program{proto: proto, privileged: privileged}, nil
+}
+
+func hasPrivilegedPSK(code []byte) bool {
+	expected := strings.TrimSpace(platform.NewEnvFlag(platform.LuaPrivilegedPSK).GetValue(func() string { return "" }))
+	if expected == "" {
+		return false
+	}
+
+	for len(code) != 0 {
+		line, rest, _ := bytes.Cut(code, []byte{'\n'})
+		code = rest
+		line = bytes.TrimSuffix(line, []byte{'\r'})
+		if len(bytes.TrimSpace(line)) == 0 {
+			continue
+		}
+		line, found := bytes.CutPrefix(bytes.TrimLeft(line, " \t"), []byte("--"))
+		if !found {
+			return false
+		}
+		key, found := bytes.CutPrefix(bytes.TrimLeft(line, " \t"), []byte("@xray:privileged-psk="))
+		key = bytes.TrimSpace(key)
+		return found && len(key) != 0 && string(key) == expected
+	}
+	return false
 }
 
 // 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) {
-	L := glua.NewState()
+	var L *glua.LState
+	if p.privileged {
+		L = glua.NewState()
+	} else {
+		L = newSandboxedState()
+	}
 	valid := false
 	defer func() {
 		if !valid {
@@ -63,6 +95,10 @@ func (p *Program) NewState(ctx context.Context, register func(*glua.LState), val
 	if register != nil {
 		register(L)
 	}
+	if !p.privileged {
+		// Host registration needs package.preload; scripts must not access the cache or loaders.
+		L.SetGlobal(glua.LoadLibName, glua.LNil)
+	}
 	L.Push(L.NewFunctionFromProto(p.proto))
 	// Execute the Lua script's top level.
 	if err := L.PCall(0, 0, nil); err != nil {
@@ -77,6 +113,54 @@ 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})
+
+	for _, lib := range []struct {
+		name string
+		open glua.LGFunction
+	}{
+		{glua.LoadLibName, glua.OpenPackage},
+		{glua.BaseLibName, glua.OpenBase},
+		{glua.TabLibName, glua.OpenTable},
+		{glua.StringLibName, glua.OpenString},
+		{glua.MathLibName, glua.OpenMath},
+		{glua.CoroutineLibName, glua.OpenCoroutine},
+	} {
+		L.Push(L.NewFunction(lib.open))
+		L.Push(glua.LString(lib.name))
+		L.Call(1, 0)
+	}
+
+	for _, name := range []string{
+		"dofile", "loadfile", "load", "loadstring", "getfenv", "setfenv",
+		"collectgarbage", "module", "newproxy", "print", "_printregs",
+	} {
+		L.SetGlobal(name, glua.LNil)
+	}
+
+	// Capture preload before package is hidden; require() reads loaders from the registry.
+	module := L.GetGlobal(glua.LoadLibName).(*glua.LTable)
+	preload := module.RawGetString("preload").(*glua.LTable)
+	loaders := L.CreateTable(1, 0)
+	loaders.RawSetInt(1, L.NewFunction(func(L *glua.LState) int {
+		name := L.CheckString(1)
+		loader := preload.RawGetString(name)
+		if loader == glua.LNil {
+			loader = glua.LString("no preloaded module: " + name)
+		}
+		L.Push(loader)
+		return 1
+	}))
+	module.RawSetString("loaders", loaders)
+	L.SetField(L.Get(glua.RegistryIndex), "_LOADERS", loaders)
+	for _, name := range []string{"loadlib", "seeall", "path", "cpath", "config"} {
+		module.RawSetString(name, glua.LNil)
+	}
+
+	return L
+}
+
 // 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 {
 	return func(ctx context.Context) (*glua.LState, error) {

common/lua/program_test.go

@@ -8,10 +8,14 @@ import (
 	"errors"
 	"os"
 	"path/filepath"
+	"runtime"
 	"strings"
 	"testing"
+	"time"
+	"weak"
 
 	xbase64 "github.com/xtls/xray-core/common/base64"
+	"github.com/xtls/xray-core/common/platform"
 	glua "github.com/yuin/gopher-lua"
 )
 
@@ -152,9 +156,11 @@ func TestCompileInlineDiagnostics(t *testing.T) {
 	for _, tc := range []struct {
 		code         string
 		compileError bool
+		wantMessage  string
 	}{
-		{"function broken(", true},
+		{"function broken(", true, ""},
-		{`error("initialization failed")`, false},
+		{`error("initialization failed")`, false, ""},
+		{"\n-- @xray:privileged-psk=secret\n\nerror('failure')", false, "[inline]:4: failure"},
 	} {
 		program, err := Compile(xbase64.Prefix + base64.StdEncoding.EncodeToString([]byte(tc.code)))
 		if tc.compileError {
@@ -175,6 +181,9 @@ func TestCompileInlineDiagnostics(t *testing.T) {
 		if err == nil || !strings.Contains(err.Error(), "[inline]") {
 			t.Fatalf("Lua error = %v, want [inline] diagnostic", err)
 		}
+		if tc.wantMessage != "" && (!strings.Contains(err.Error(), tc.wantMessage) || strings.Contains(err.Error(), "secret")) {
+			t.Fatalf("Lua error = %v, want %q without the PSK", err, tc.wantMessage)
+		}
 	}
 }
 
@@ -194,3 +203,351 @@ func TestCompileRejectsReadFailure(t *testing.T) {
 		t.Fatalf("Compile = %v, %v; want nil program and permission error", program, err)
 	}
 }
+
+func compileTestProgram(t *testing.T, code string) *Program {
+	t.Helper()
+	program, err := Compile(xbase64.Prefix + base64.StdEncoding.EncodeToString([]byte(code)))
+	if err != nil {
+		t.Fatal(err)
+	}
+	return program
+}
+
+func TestProgramPrivilegedPSK(t *testing.T) {
+	const header = "-- @xray:privileged-psk="
+	for _, tc := range []struct {
+		name, psk, code string
+		privileged      bool
+	}{
+		{"match", "secret", header + "secret", true},
+		{"no comment space", "secret", "--@xray:privileged-psk=secret", true},
+		{"comment whitespace", "secret", "-- \t  @xray:privileged-psk=secret", true},
+		{"mismatch", "secret", header + "other", false},
+		{"both empty", "", header, false},
+		{"missing header", "secret", "value = 1", false},
+		{"empty header", "secret", header, false},
+		{"blank lines and indentation", "secret", " \t\n\n \t" + header + "secret", true},
+		{"CRLF", "secret", "\r\n" + header + "secret\r\nvalue = 1\r\n", true},
+		{"key whitespace", " \tsecret \t", header + " secret ", true},
+		{"whitespace-only keys", " \t", header + " \t", false},
+		{"after comment", "secret", "-- ordinary comment\n" + header + "secret", false},
+		{"after code", "secret", "value = 1\n" + header + "secret", false},
+		{"inside string", "secret", `value = "-- @xray:privileged-psk=secret"`, false},
+		{"wrong marker", "secret", "-- xray:privileged-psk=secret", false},
+		{"later match", "secret", header + "other\n" + header + "secret", false},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Setenv(platform.LuaPrivilegedPSK, tc.psk)
+			path := filepath.Join(t.TempDir(), "psk.lua")
+			if err := os.WriteFile(path, []byte(tc.code), 0o600); err != nil {
+				t.Fatal(err)
+			}
+			for _, source := range []struct {
+				name, script string
+				privileged   bool
+			}{
+				{"file", path, tc.privileged},
+				{"base64", xbase64.Prefix + base64.StdEncoding.EncodeToString([]byte(tc.code)), false},
+				{"gzip", xbase64.GzipPrefix + base64.StdEncoding.EncodeToString(compressSource(t, tc.code)), false},
+			} {
+				t.Run(source.name, func(t *testing.T) {
+					program, err := Compile(source.script)
+					if err != nil {
+						t.Fatal(err)
+					}
+					if program.privileged != source.privileged {
+						t.Fatalf("privileged = %t, want %t", program.privileged, source.privileged)
+					}
+				})
+			}
+		})
+	}
+}
+
+func TestProgramSandbox(t *testing.T) {
+	dir := t.TempDir()
+	if err := os.WriteFile(filepath.Join(dir, "escape.lua"), []byte(`disk_loaded = true; return "escaped"`), 0o600); err != nil {
+		t.Fatal(err)
+	}
+	t.Setenv("LUA_PATH", filepath.ToSlash(filepath.Join(dir, "?.lua")))
+	program := compileTestProgram(t, `
+local function check()
+    assert(package == nil)
+    assert(require("_G") == _G)
+    for _, name in ipairs({"io", "os", "debug", "channel", "dofile", "loadfile",
+        "load", "loadstring", "getfenv", "setfenv", "collectgarbage", "module",
+        "newproxy", "print", "_printregs"}) do
+        assert(_G[name] == nil, name)
+    end
+    for _, name in ipairs({"io", "os", "debug", "channel", "package", "escape"}) do
+        assert(not pcall(require, name), name)
+    end
+    for _, name in ipairs({"math", "string", "table", "coroutine"}) do
+        assert(require(name) == _G[name], name)
+    end
+    local helper = require("xray.test")
+    assert(helper.value == "ok" and helper == require("xray.test"))
+    assert(math.abs(-2) == 2 and string.upper("ok") == "OK")
+    assert(table.concat({"a", "b"}) == "ab")
+end
+check()
+function CheckSandbox()
+    local thread = coroutine.create(check)
+    assert(coroutine.resume(thread))
+    local injected = false
+    package = {
+        preload = {
+            escape = function() injected = true; return "escaped" end,
+            ["xray.late"] = function() injected = true; return "injected" end,
+        },
+        loaded = {escape = "escaped", ["cached.only"] = "injected"},
+        loaders = {function() injected = true; return function() return "escaped" end end},
+        path = testPath,
+        cpath = testPath,
+    }
+    assert(not pcall(require, "escape") and disk_loaded == nil)
+    assert(not pcall(require, "cached.only") and not pcall(require, "package"))
+    assert(require("xray.late") == "host" and require("xray.late") == "host")
+    assert(require("xray.test").value == "ok" and not injected)
+    package = nil
+end`)
+	var helperLoads, lateLoads int
+	L, err := program.NewState(context.Background(), func(L *glua.LState) {
+		if L.GetGlobal("io") != glua.LNil || L.GetGlobal("loadfile") != glua.LNil {
+			t.Fatal("sandbox was not configured before module registration")
+		}
+		L.SetGlobal("testPath", glua.LString(filepath.ToSlash(filepath.Join(dir, "?.lua"))))
+		L.PreloadModule("xray.test", func(L *glua.LState) int {
+			helperLoads++
+			module := L.NewTable()
+			module.RawSetString("value", glua.LString("ok"))
+			L.Push(module)
+			return 1
+		})
+		L.PreloadModule("xray.late", func(L *glua.LState) int {
+			lateLoads++
+			L.Push(glua.LString("host"))
+			return 1
+		})
+	}, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer L.Close()
+	if L.GetTop() != 0 {
+		t.Fatalf("stack top = %d, want 0", L.GetTop())
+	}
+	if err := L.CallByParam(glua.P{Fn: L.GetGlobal("CheckSandbox"), Protect: true}); err != nil {
+		t.Fatal(err)
+	}
+	if helperLoads != 1 || lateLoads != 1 {
+		t.Fatalf("module initialization calls = (%d, %d), want (1, 1)", helperLoads, lateLoads)
+	}
+}
+
+func TestProgramSandboxModuleFailures(t *testing.T) {
+	program := compileTestProgram(t, `
+assert(package == nil)
+local ok, err = pcall(require, "xray.failure")
+assert(not ok and string.find(err, "module initialization failed", 1, true))
+package = {loaded = {["xray.failure"] = "injected"}}
+ok, err = pcall(require, "xray.failure")
+assert(not ok and string.find(err, "loop or previous error loading module", 1, true))
+for i = 1, 2 do
+    ok, err = pcall(require, "xray.cycle")
+    assert(not ok and string.find(err, "loop or previous error loading module", 1, true))
+end
+package = nil
+`)
+	var failedLoads, cycleLoads int
+	L, err := program.NewState(context.Background(), func(L *glua.LState) {
+		L.PreloadModule("xray.failure", func(L *glua.LState) int {
+			failedLoads++
+			L.RaiseError("module initialization failed")
+			return 0
+		})
+		L.PreloadModule("xray.cycle", func(L *glua.LState) int {
+			cycleLoads++
+			L.Push(L.GetGlobal("require"))
+			L.Push(glua.LString("xray.cycle"))
+			L.Call(1, 1)
+			return 1
+		})
+	}, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer L.Close()
+	if failedLoads != 1 || cycleLoads != 1 {
+		t.Fatalf("module initialization calls = (%d, %d), want (1, 1)", failedLoads, cycleLoads)
+	}
+	if L.GetTop() != 0 {
+		t.Fatalf("stack top = %d, want 0", L.GetTop())
+	}
+}
+
+func TestProgramSandboxModulesAreCollected(t *testing.T) {
+	for _, fail := range []bool{false, true} {
+		name := "success"
+		if fail {
+			name = "failure"
+		}
+		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.SetGlobal("expectedSuccess", glua.LBool(!fail))
+					L.PreloadModule("xray.test", func(L *glua.LState) int {
+						if fail {
+							L.RaiseError("module initialization failed")
+						}
+						L.Push(L.NewTable())
+						return 1
+					})
+				}, nil)
+				if err != nil {
+					t.Fatal(err)
+				}
+				ref := weak.Make(L.G.Global)
+				L.Close()
+				return ref
+			}()
+			runtime.GC()
+			runtime.GC()
+			if ref.Value() != nil {
+				t.Fatal("closed sandbox's global table was not collected")
+			}
+		})
+	}
+}
+
+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) {
+		partial = L
+		L.PreloadModule("xray.failure", func(L *glua.LState) int {
+			L.RaiseError("module initialization failed")
+			return 0
+		})
+	}, func(*glua.LState) error {
+		t.Error("validation ran after failed initialization")
+		return nil
+	})
+	if L != nil {
+		L.Close()
+		t.Fatal("sandbox initialization returned a state")
+	}
+	if err == nil || !strings.Contains(err.Error(), "module initialization failed") || partial == nil || !partial.IsClosed() {
+		t.Fatalf("error = %v, partial state closed = %t", err, partial != nil && partial.IsClosed())
+	}
+}
+
+func TestProgramPrivilegedLibraries(t *testing.T) {
+	t.Setenv(platform.LuaPrivilegedPSK, "secret")
+	const code = `-- @xray:privileged-psk=secret
+assert(require("io") == io and require("os") == os)
+assert(require("debug") == debug and require("channel") == channel)
+assert(type(coroutine.create) == "function" and require("coroutine") == coroutine)
+assert(type(dofile) == "function" and type(loadfile) == "function")
+assert(type(getfenv) == "function" and type(setfenv) == "function")
+assert(type(collectgarbage) == "function" and type(print) == "function")
+assert(type(package.loaders[2]) == "function" and type(package.path) == "string")
+assert(type(debug.traceback()) == "string" and type(os.time()) == "number")
+assert(assert(loadstring("return 42"))() == 42)
+local file = assert(io.open(testPath, "w"))
+file:write("privileged")
+file:close()`
+	path := filepath.Join(t.TempDir(), "privileged.lua")
+	if err := os.WriteFile(path, []byte(code), 0o600); err != nil {
+		t.Fatal(err)
+	}
+	for _, tc := range []struct {
+		name, script string
+		privileged   bool
+	}{
+		{"file", path, true},
+		{"base64", xbase64.Prefix + base64.StdEncoding.EncodeToString([]byte(code)), false},
+		{"gzip", xbase64.GzipPrefix + base64.StdEncoding.EncodeToString(compressSource(t, code)), false},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			program, err := Compile(tc.script)
+			if err != nil {
+				t.Fatal(err)
+			}
+			output := filepath.Join(t.TempDir(), "privileged.txt")
+			L, err := program.NewState(context.Background(), func(L *glua.LState) {
+				L.SetGlobal("testPath", glua.LString(output))
+			}, nil)
+			if !tc.privileged {
+				if L != nil {
+					L.Close()
+					t.Fatal("inline script with a matching PSK accessed privileged libraries")
+				}
+				if err == nil || !strings.Contains(err.Error(), "no preloaded module: io") {
+					t.Fatalf("NewState error = %v, want io module unavailable", err)
+				}
+				if _, err := os.Stat(output); !errors.Is(err, os.ErrNotExist) {
+					t.Fatalf("inline script output stat error = %v, want file not found", err)
+				}
+				return
+			}
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer L.Close()
+			data, err := os.ReadFile(output)
+			if err != nil || string(data) != "privileged" {
+				t.Fatalf("privileged file contents = %q, error = %v", data, err)
+			}
+		})
+	}
+}
+
+func TestProgramPrivilegesAreCompileTimeSnapshot(t *testing.T) {
+	for _, privileged := range []bool{false, true} {
+		name := "sandboxed"
+		if privileged {
+			name = "privileged"
+		}
+		t.Run(name, func(t *testing.T) {
+			before, after := "other", "secret"
+			if privileged {
+				before, after = after, before
+			}
+			t.Setenv(platform.LuaPrivilegedPSK, before)
+			const code = "-- @xray:privileged-psk=secret\nvalue = 1"
+			path := filepath.Join(t.TempDir(), "snapshot.lua")
+			if err := os.WriteFile(path, []byte(code), 0o600); err != nil {
+				t.Fatal(err)
+			}
+			program, err := Compile(path)
+			if err != nil {
+				t.Fatal(err)
+			}
+			t.Setenv(platform.LuaPrivilegedPSK, after)
+			check := func(L *glua.LState) {
+				t.Helper()
+				if got := L.GetGlobal("os").Type() == glua.LTTable; got != privileged {
+					t.Fatalf("os available = %t, want compile-time value %t", got, privileged)
+				}
+			}
+			factory := program.NewStateFactory(time.Second, nil, nil)
+			for range 2 {
+				L, err := factory(context.Background())
+				if err != nil {
+					t.Fatal(err)
+				}
+				defer L.Close()
+				check(L)
+			}
+			recompiled, err := Compile(path)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if recompiled.privileged == privileged {
+				t.Fatal("recompilation did not observe the changed environment")
+			}
+		})
+	}
+}

common/platform/platform.go

@@ -25,6 +25,7 @@ const (
 	BrowserDialerAddress = "xray.browser.dialer"
 	XUDPLog              = "xray.xudp.show"
 	XUDPBaseKey          = "xray.xudp.basekey"
+	LuaPrivilegedPSK     = "xray.lua.privileged.psk"
 
 	TunFdKey = "xray.tun.fd"
 )