From 1aae252fc3d5354c6f96dd66f816db7dad889e4e Mon Sep 17 00:00:00 2001
From: Meo597 <197331664+Meo597@users.noreply.github.com>
Date: Tue, 13 Oct 2026 11:38:06 +0800
Subject: [PATCH] Lua: Restrict PSK privileges to local Lua files

---
 common/lua/program.go      |  4 +--
 common/lua/program_test.go | 59 ++++++++++++++++++++++++++++++++------
 2 files changed, 53 insertions(+), 10 deletions(-)

diff --git a/common/lua/program.go b/common/lua/program.go
index 506e1d3b4..11247ee11 100644
--- a/common/lua/program.go
+++ b/common/lua/program.go
@@ -28,17 +28,17 @@ 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
 	}
 
-	privileged := hasPrivilegedPSK(code)
-
 	chunk, err := parse.Parse(bytes.NewReader(code), name)
 	if err != nil {
 		return nil, err
diff --git a/common/lua/program_test.go b/common/lua/program_test.go
index e8795585d..8ee85ac48 100644
--- a/common/lua/program_test.go
+++ b/common/lua/program_test.go
@@ -238,9 +238,27 @@ func TestProgramPrivilegedPSK(t *testing.T) {
 	} {
 		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)
+			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)
+					}
+				})
 			}
 		})
 	}
@@ -446,10 +464,11 @@ file:close()`
 	}
 	for _, tc := range []struct {
 		name, script string
+		privileged   bool
 	}{
-		{"file", path},
-		{"base64", xbase64.Prefix + base64.StdEncoding.EncodeToString([]byte(code))},
-		{"gzip", xbase64.GzipPrefix + base64.StdEncoding.EncodeToString(compressSource(t, code))},
+		{"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)
@@ -460,6 +479,19 @@ file:close()`
 			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)
 			}
@@ -485,7 +517,14 @@ func TestProgramPrivilegesAreCompileTimeSnapshot(t *testing.T) {
 			}
 			t.Setenv(platform.LuaPrivilegedPSK, before)
 			const code = "-- @xray:privileged-psk=secret\nvalue = 1"
-			program := compileTestProgram(t, code)
+			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()
@@ -502,7 +541,11 @@ func TestProgramPrivilegesAreCompileTimeSnapshot(t *testing.T) {
 				defer L.Close()
 				check(L)
 			}
-			if recompiled := compileTestProgram(t, code); recompiled.privileged == privileged {
+			recompiled, err := Compile(path)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if recompiled.privileged == privileged {
 				t.Fatal("recompilation did not observe the changed environment")
 			}
 		})