From 6122525ed348d747a0a0a4f46727f12a3b38d180 Mon Sep 17 00:00:00 2001
From: Jamie Reece Wilson <me@reece.sx>
Date: Sat, 20 Apr 2024 05:41:56 +0100
Subject: [PATCH] [*] Fix AuProcess deadlock and parity issue between some unix
 loaders and Windows [*] Slight UWP progress in AuProcessMap.NT.cpp (not
 really)

---
 Source/Process/AuProcess.cpp       | 85 +++++++++++++++++++++++++++---
 Source/Process/AuProcessMap.NT.cpp | 12 +++--
 2 files changed, 86 insertions(+), 11 deletions(-)

diff --git a/Source/Process/AuProcess.cpp b/Source/Process/AuProcess.cpp
index 71ab4187..f8b482e8 100644
--- a/Source/Process/AuProcess.cpp
+++ b/Source/Process/AuProcess.cpp
@@ -39,7 +39,7 @@
 
 namespace Aurora::Process
 {
-    static AuFutexMutex gSpinLock;
+    static AuCriticalSection gSpinLock;
     static AuList<AuString> gClassPath;
     static AuHashMap<AuString, void *> gModuleHandles;
     static const bool kIsMainSigned = false;
@@ -634,16 +634,75 @@ namespace Aurora::Process
 
     AUKN_SYM AuMach GetProcAddressEx(void *pHandle, const AuString &symbol)
     {
+        AuMach ret {};
+
     #if defined(AURORA_IS_MODERNNT_DERIVED)
         if (!pGetProcAddress)
         {
             return {};
         }
 
-        auto ret = reinterpret_cast<AuMach>(pGetProcAddress(reinterpret_cast<HMODULE>(pHandle), symbol.c_str()));
+        if (pHandle)
+        {
+            ret = reinterpret_cast<AuMach>(pGetProcAddress(reinterpret_cast<HMODULE>(pHandle), symbol.c_str()));
+        }
+        else
+        {
+        #if defined(AURORA_PLATFORM_WIN32)
+            ret = reinterpret_cast<AuMach>(pGetProcAddress(::GetModuleHandleW(nullptr), symbol.c_str()));
+
+            if (!ret)
+            {
+                HMODULE hHandle;
+                if (GetModuleHandleExW(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT,
+                                       reinterpret_cast<LPCWSTR>(&GetProcAddressEx), &hHandle))
+                {
+                    ret = reinterpret_cast<AuMach>(pGetProcAddress(hHandle, symbol.c_str()));
+                }
+            }
+
+            if (!ret)
+        #endif
+            {
+                AU_LOCK_GUARD(gSpinLock);
+
+                for (const auto &[string, hHandle] : gModuleHandles)
+                {
+                    ret = reinterpret_cast<AuMach>(pGetProcAddress(reinterpret_cast<HMODULE>(hHandle), symbol.c_str()));
+                    if (ret)
+                    {
+                        break;
+                    }
+                }
+            }
+        }
     #else
-        auto ret = reinterpret_cast<AuMach>(dlsym(pHandle, symbol.c_str()));
+        if (pHandle)
+        {
+            ret = reinterpret_cast<AuMach>(dlsym(pHandle, symbol.c_str()));
+        }
+        else
+        {
+        #if defined(RTLD_DEFAULT)
+            ret = reinterpret_cast<AuMach>(dlsym(RTLD_DEFAULT, symbol.c_str()));
+
+            if (!ret)
+        #endif
+            {
+                AU_LOCK_GUARD(gSpinLock);
+
+                for (const auto &[string, hHandle] : gModuleHandles)
+                {
+                    ret = reinterpret_cast<AuMach>(dlsym(hHandle, symbol.c_str()));
+                    if (ret)
+                    {
+                        break;
+                    }
+                }
+            }
+        }
     #endif
+
         return ret;
     }
 
@@ -663,14 +722,24 @@ namespace Aurora::Process
 
     AUKN_SYM AuMach GetProcAddress(const AuString &mod, const AuString &symbol)
     {
-        auto ret = GetProcAddressEx(GetProcHandle(mod), symbol);
-
-        if (!ret)
+        if (mod.empty())
         {
-            SysPushErrorGen("Couldn't resolve symbol in module {} of mangled name '{}'", mod, symbol);
+            return GetProcAddressEx(nullptr, symbol);
         }
 
-        return ret;
+        if (auto pHandle = GetProcHandle(mod))
+        {
+            auto ret = GetProcAddressEx(pHandle, symbol);
+
+            if (!ret)
+            {
+                SysPushErrorGen("Couldn't resolve symbol in module {} of mangled name '{}'", mod, symbol);
+            }
+
+            return ret;
+        }
+
+        return {};
     }
 
     AUKN_SYM void Exit(AuUInt32 exitcode)
diff --git a/Source/Process/AuProcessMap.NT.cpp b/Source/Process/AuProcessMap.NT.cpp
index 26f05f68..2dfd7805 100644
--- a/Source/Process/AuProcessMap.NT.cpp
+++ b/Source/Process/AuProcessMap.NT.cpp
@@ -216,14 +216,20 @@ namespace Aurora::Process
 
     bool MakeAwarePtr(AuUInt pointer)
     {
-        HMODULE handle;
-        if (!GetModuleHandleExW(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT, reinterpret_cast<LPCWSTR>(pointer), &handle))
+    #if defined(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS)
+        HMODULE hHandle;
+        if (!GetModuleHandleExW(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT,
+                                reinterpret_cast<LPCWSTR>(pointer),
+                                &hHandle))
         {
             return false;
         }
 
-        MakeAware(handle);
+        MakeAware(hHandle);
         return true;
+    #else
+        return false;
+    #endif
     }
 
     AuOptional<Section> LookupArbitrarySection(AuUInt address)