Lua: Add sandbox mechanism

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.
@@ -36,6 +37,8 @@ func Compile(script string) (*Program, error) {
 		return nil, err
 	}
 
+	privileged := hasPrivilegedPSK(code)
+
 	chunk, err := parse.Parse(bytes.NewReader(code), name)
 	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 {
@@ -77,6 +109,45 @@ 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)
+	}
+
+	// require() reads loaders from the registry.
+	module := L.GetGlobal(glua.LoadLibName).(*glua.LTable)
+	loaders := L.NewTable()
+	loaders.RawSetInt(1, module.RawGetString("loaders").(*glua.LTable).RawGetInt(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

@@ -10,8 +10,10 @@ import (
 	"path/filepath"
 	"strings"
 	"testing"
+	"time"
 
 	xbase64 "github.com/xtls/xray-core/common/base64"
+	"github.com/xtls/xray-core/common/platform"
 	glua "github.com/yuin/gopher-lua"
 )
 
@@ -152,9 +154,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 +179,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 +201,210 @@ 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)
+			program := compileTestProgram(t, tc.code)
+			if program.privileged != tc.privileged {
+				t.Fatalf("privileged = %t, want %t", program.privileged, tc.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(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", "escape"}) do
+        assert(package.loaded[name] == nil and package.preload[name] == nil, name)
+        assert(not pcall(require, name), name)
+    end
+    for name in pairs(package) do
+        assert(name == "preload" or name == "loaded" or name == "loaders", name)
+    end
+    assert(type(package.loaders[1]) == "function" and package.loaders[2] == nil)
+    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))
+    package.path = testPath
+    package.cpath = testPath
+    assert(not pcall(require, "escape") and disk_loaded == nil)
+    package.path = nil
+    package.cpath = nil
+end`)
+	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 {
+			module := L.NewTable()
+			module.RawSetString("value", glua.LString("ok"))
+			L.Push(module)
+			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)
+	}
+}
+
+func TestProgramSandboxClosesFailedInitialization(t *testing.T) {
+	program := compileTestProgram(t, `require("io")`)
+	var partial *glua.LState
+	L, err := program.NewState(context.Background(), func(L *glua.LState) {
+		partial = L
+	}, 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 || 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(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
+	}{
+		{"file", path},
+		{"base64", xbase64.Prefix + base64.StdEncoding.EncodeToString([]byte(code))},
+		{"gzip", xbase64.GzipPrefix + base64.StdEncoding.EncodeToString(compressSource(t, code))},
+	} {
+		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 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"
+			program := compileTestProgram(t, code)
+			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)
+			}
+			if recompiled := compileTestProgram(t, code); 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"
 )