refactor: enhance stdin handling by adding utf8 BOM removal and main function check

2 files changed, 12 insertions, 2 deletions

diff --git a/tests/modules/stdin/test.lua b/tests/modules/stdin/test.lua
index 1cef7a362..bcd9efca6 100644
--- a/tests/modules/stdin/test.lua
+++ b/tests/modules/stdin/test.lua
@@ -70,11 +70,13 @@ else
     if pwsh ~= "" then
         test_shell("pwsh_single", string.format('%s -c "echo \\"print(\'hello_pwsh\')\\" | %s l --stdin"', pwsh, xmake), "hello_pwsh")
         test_shell("pwsh_calc", string.format('%s -c "echo \\"local f = 1+1; print(f)\\" | %s l --stdin"', pwsh, xmake), "2")
-        test_shell("pwsh_multi", string.format('%s -c "echo \\"print(\'pline1\')\\nprint(\'pline2\')\\" | %s l --stdin"', pwsh, xmake), "pline1[\r\n]+pline2")
+        test_shell("pwsh_main", string.format('%s -c "echo \\"function main() print(\'in_pwsh_main\') end\\" | %s l --stdin"', pwsh, xmake), "in_pwsh_main")
+        test_shell("pwsh_multi", string.format('%s -c "echo \\"print(\'pline1\')\\" \\"print(\'pline2\')\\" | %s l --stdin"', pwsh, xmake), "pline1[\r\n]+pline2")
     end
 
     test_shell("sh_single", string.format('echo "print(\'hello_sh\')" | %s l --stdin', xmake), "hello_sh")
     test_shell("sh_calc", string.format('echo "local f = 1+1; print(f)" | %s l --stdin', xmake), "2")
+    test_shell("sh_main", string.format('echo "function main() print(\'in_sh_main\') end" | %s l --stdin', xmake), "in_sh_main")
     test_shell("sh_multi", string.format('printf "print(\'shell_line1\')\\nprint(\'shell_line2\')" | %s l --stdin', xmake), "shell_line1[\r\n]+shell_line2")
 end
 end
diff --git a/xmake/plugins/lua/main.lua b/xmake/plugins/lua/main.lua
index 0c0b44005..1cc13dbc3 100644
--- a/xmake/plugins/lua/main.lua
+++ b/xmake/plugins/lua/main.lua
@@ -63,6 +63,10 @@ function main()
         if script == "-" or from_stdin then
             local script_content = io.read("*a")
             if script_content then
+                -- remove utf8 bom
+                if script_content:startswith("\239\187\191") then
+                    script_content = script_content:sub(4)
+                end
                 import("core.base.tty")
                 local shell = tty.shell()
                 if shell == "cmd" or shell == "powershell" or shell == "pwsh" or os.host() == "windows" then
@@ -72,7 +76,11 @@ function main()
                     end
                     script_content = script_content:replace("\\n", "\n", {plain = true}):replace("\\r", "\r", {plain = true})
                 end
-                script_content = "function main(...)\n" .. script_content .. "\nend"
+
+                if not script_content:find("function main", 1, true) then
+                    script_content = "function main(...)\n" .. script_content .. "\nend"
+                end
+
                 script = os.tmpfile() .. ".lua"
                 io.writefile(script, script_content)
                 script_file_to_remove = script