Add parse_pe function to check for PE headers and enhance os.isexec for Windows

4 files changed, 100 insertions, 2 deletions

diff --git a/core/src/xmake/engine.c b/core/src/xmake/engine.c
index a77cc4a9b..4a55df07f 100644
--- a/core/src/xmake/engine.c
+++ b/core/src/xmake/engine.c
@@ -177,6 +177,7 @@ tb_int_t xm_io_file_flush(lua_State *lua);
 tb_int_t xm_io_file_close(lua_State *lua);
 tb_int_t xm_io_file_convert(lua_State *lua);
 tb_int_t xm_io_file_isatty(lua_State *lua);
+tb_int_t xm_io_parse_pe(lua_State *lua);
 
 // the io/filelock functions
 tb_int_t xm_io_filelock_open(lua_State *lua);
@@ -497,6 +498,7 @@ static luaL_Reg const g_io_functions[] = {
     { "file_write", xm_io_file_write },
     { "file_flush", xm_io_file_flush },
     { "file_isatty", xm_io_file_isatty },
+    { "parse_pe", xm_io_parse_pe },
     { "file_close", xm_io_file_close },
     { "file_convert", xm_io_file_convert },
     { "file_rawfd", xm_io_file_rawfd },
diff --git a/core/src/xmake/io/parse_pe.c b/core/src/xmake/io/parse_pe.c
new file mode 100644
index 000000000..017bf989b
--- /dev/null
+++ b/core/src/xmake/io/parse_pe.c
@@ -0,0 +1,66 @@
+#include "prefix.h"
+
+tb_int_t xm_io_parse_pe(lua_State* lua) {
+    // check
+    tb_assert_and_check_return_val(lua, 0);
+
+    // get file path
+    tb_char_t const* path = luaL_checkstring(lua, 1);
+    tb_check_return_val(path, 0);
+
+    // open file
+    tb_file_ref_t file = tb_file_init(path, TB_FILE_MODE_RO);
+    if (!file) {
+        lua_pushnil(lua);
+        lua_pushfstring(lua, "cannot open file: %s", path);
+        return 2;
+    }
+
+    // read dos header
+    tb_byte_t buffer[1024];
+    if (tb_file_read(file, buffer, 0x40)) {
+        // check MZ
+        if (buffer[0] == 'M' && buffer[1] == 'Z') {
+            // get pe offset
+            tb_uint32_t pe_offset = 0;
+            // e_lfanew is at 0x3c
+            pe_offset = tb_bits_get_u32_le(buffer + 0x3c);
+            
+            // seek to pe header
+            if (tb_file_seek(file, pe_offset, TB_FILE_SEEK_BEG) == pe_offset) {
+                 if (tb_file_read(file, buffer, 24)) { // Signature(4) + FileHeader(20)
+                      // check PE signature
+                      if (buffer[0] == 'P' && buffer[1] == 'E' && buffer[2] == 0 && buffer[3] == 0) {
+                          // get machine
+                          tb_uint16_t machine = tb_bits_get_u16_le(buffer + 4);
+                          
+                          lua_newtable(lua);
+                          
+                          lua_pushstring(lua, "arch");
+                          tb_char_t const* arch = tb_null;
+                          switch (machine) {
+                              case 0x014c: arch = "x86"; break;
+                              case 0x8664: arch = "x64"; break;
+                              case 0xAA64: arch = "arm64"; break; 
+                              case 0x01c0: arch = "arm"; break; 
+                              case 0x01c4: arch = "arm"; break; 
+                          }
+                          if (arch) {
+                              lua_pushstring(lua, arch);
+                              lua_settable(lua, -3);
+                          } else {
+                               lua_pop(lua, 1); // pop "arch"
+                          }
+                           
+                          tb_file_exit(file);
+                          return 1; 
+                      }
+                 }
+            }
+        }
+    }
+
+    tb_file_exit(file);
+    lua_pushnil(lua);
+    return 1;
+}
diff --git a/xmake/core/base/os.lua b/xmake/core/base/os.lua
index 963a015c1..33a073773 100644
--- a/xmake/core/base/os.lua
+++ b/xmake/core/base/os.lua
@@ -1158,15 +1158,38 @@ function os.isexec(filepath)
     if os.isfile(filepath) then
         if os.host() == "windows" then
             local ext = path.extension(filepath):lower()
-            if ext == ".exe" or ext == ".cmd" or ext == ".bat" then
+            if ext == ".exe" or ext == ".cmd" or ext == ".bat" or ext == ".ps1" then
                 return true
             end
+            -- check for PE header
+            local peinfo = io.parse_pe(filepath)
+            if peinfo then
+                local pe_arch = peinfo.arch
+                -- x64 or x86
+                if pe_arch == "x86" or pe_arch == "x64" and os.is_arch("x86_64") then
+                    return true
+                else
+                    return false
+                end
+                -- arm64
+                if pe_arch == "arm64" and os.is_arch("arm64") then
+                    return true
+                else
+                    return false
+                end
+                -- arm32
+                if pe_arch == "arm" and os.is_arch("arm.*") then
+                    return true
+                else
+                    return false
+                end
+            end
         else
             return os._access(filepath, "x")
         end
     end
     if os.host() == "windows" then
-        for _, suffix in ipairs({".exe", ".cmd", ".bat"}) do
+        for _, suffix in ipairs({".exe", ".cmd", ".bat", ".ps1"}) do
             if os.isfile(filepath .. suffix) then
                 return true
             end
diff --git a/xmake/core/sandbox/modules/io.lua b/xmake/core/sandbox/modules/io.lua
index 0cafaa5f2..5657d02b9 100644
--- a/xmake/core/sandbox/modules/io.lua
+++ b/xmake/core/sandbox/modules/io.lua
@@ -372,6 +372,13 @@ function sandbox_io.tail(filepath, linecount, opt)
     io.tail(filepath, linecount, opt)
 end
 
+-- parse pe
+function sandbox_io.parse_pe(filepath)
+    assert(filepath)
+    filepath = vformat(filepath)
+    return io.parse_pe(filepath)
+end
+
 -- lazy loading stdfile
 setmetatable(sandbox_io, { __index = function (tbl, key)
         local val = rawget(tbl, key)