[+] Aurora::Process::GetBinaryClassPath

[+] Aurora::Process::SetBinaryClassPath
[+] Aurora::Process::AddBinaryClassPath
[+] Aurora::Process::LoadModuleEx
[+] Aurora::Process::GetProcHandle
[+] Aurora::Process::GetProcAddressEx
[+] Aurora::ProcessConfig
[*] Cleanup a bit

Include/Aurora/Process/Process.hpp

@@ -29,6 +29,7 @@ namespace Aurora::Process
         eModulePathSystemDir, /// /usr/lib/, windir/system32
         eModulePathUserDir,   /// /usr/local/lib
         eProcessDirectory,    /// 
+        eClassPath,
         //eOSSpecified,       /// LD_LIBRARY_PATH + /etc/ld.so.conf, AddDllDirectory
         eSpecified
     ));
@@ -40,18 +41,19 @@ namespace Aurora::Process
     // Ubuntu 2018 (mandatory 2021+): https://lists.ubuntu.com/archives/ubuntu-devel-announce/2018-November/001253.html https://wiki.debian.org/UsrMerge
     // Fedora (2011?): https://fedoraproject.org/wiki/Features/UsrMove
     
-    static AuList<EModulePath> kUserOverloadableSearchPath  = {EModulePath::eSpecified, EModulePath::eModulePathCWD, EModulePath::eProcessDirectory, EModulePath::eModulePathUserDir, EModulePath::eModulePathSystemDir};
+    static AuList<EModulePath> kUserOverloadableSearchPath  = { EModulePath::eSpecified, EModulePath::eModulePathCWD, EModulePath::eClassPath, EModulePath::eProcessDirectory, EModulePath::eModulePathUserDir, EModulePath::eModulePathSystemDir };
-    static AuList<EModulePath> kAdminOverloadableSearchPath = {EModulePath::eSpecified, EModulePath::eModulePathSystemDir, EModulePath::eProcessDirectory, EModulePath::eModulePathCWD, EModulePath::eModulePathUserDir};
+    static AuList<EModulePath> kAdminOverloadableSearchPath = { EModulePath::eSpecified, EModulePath::eModulePathSystemDir, EModulePath::eProcessDirectory, EModulePath::eClassPath, EModulePath::eModulePathCWD, EModulePath::eModulePathUserDir };
-    static AuList<EModulePath> kSpecifiedPath               = {EModulePath::eSpecified};
+    static AuList<EModulePath> kSpecifiedPath               = { EModulePath::eSpecified, EModulePath::eClassPath };
-    static AuList<EModulePath> kCWDPath                     = {EModulePath::eModulePathCWD};
+    static AuList<EModulePath> kSpecifiedPathStrict         = { EModulePath::eSpecified };
-    static AuList<EModulePath> kProcPath                    = {EModulePath::eProcessDirectory};
+    static AuList<EModulePath> kCWDPath                     = { EModulePath::eModulePathCWD };
+    static AuList<EModulePath> kProcPath                    = { EModulePath::eProcessDirectory };
 
     struct ModuleLoadRequest
     {
-        inline ModuleLoadRequest(const AuList<AuString> &directories, const AuString &mod) : mod(mod), specifiedSearchPaths(&directories), searchPath(&kSpecifiedPath)
+        inline ModuleLoadRequest(const AuList<AuString> &directories, const AuString &mod) : mod(mod), specifiedSearchPaths(&directories), searchPath(&kSpecifiedPathStrict)
         {}
 
-        inline ModuleLoadRequest(const AuList<AuString> &directories, const AuString &mod, const AuString &version) : mod(mod), version(version), specifiedSearchPaths(&directories), searchPath(&kSpecifiedPath)
+        inline ModuleLoadRequest(const AuList<AuString> &directories, const AuString &mod, const AuString &version) : mod(mod), version(version), specifiedSearchPaths(&directories), searchPath(&kSpecifiedPathStrict)
         {}
         
         inline ModuleLoadRequest(const AuString &mod) : mod(mod), searchPath(&kAdminOverloadableSearchPath) 
@@ -79,6 +81,17 @@ namespace Aurora::Process
     */
     AUKN_SYM bool LoadModule(const ModuleLoadRequest &request);
 
+    AUKN_SYM void *LoadModuleEx(const ModuleLoadRequest &request);
 
-    AUKN_SYM AuMach GetProcAddress(AuString mod, AuString symbol);
+    AUKN_SYM void *GetProcHandle(const AuString &name);
+
+    AUKN_SYM AuMach GetProcAddress(const AuString &mod, const AuString &symbol);
+
+    AUKN_SYM AuMach GetProcAddressEx(void *pHandle, const AuString &symbol);
+
+    AUKN_SYM AuList<AuString> GetBinaryClassPath();
+
+    AUKN_SYM bool SetBinaryClassPath(const AuList<AuString> &list, bool preloadAll = false);
+    
+    AUKN_SYM bool AddBinaryClassPath(const AuString &dir, bool preloadAll = false);
 }

Include/Aurora/Runtime.hpp

@@ -409,6 +409,13 @@ namespace Aurora
         AuUInt8 uSignalGAIOWorkerThreadDone { 28 };
     };
 
+    struct ProcessConfig
+    {
+        bool bAlwaysPreloadEntireClassPath { false };
+        bool bForcePreload { false };
+        bool bEnablePreload { true };
+    };
+
     struct RuntimeStartInfo
     {
         ConsoleConfig console;
@@ -419,6 +426,7 @@ namespace Aurora
         AuAlignTo<64, DebugConfig> debug;
         AuAlignTo<32, ThreadingConfig> threadingConfig;
         AuAlignTo<32, LinuxConfig> linuxConfig;
+        AuAlignTo<32, ProcessConfig> processConfig;
         AuAlignTo<32, DummyConfig> padding;
     };
 

Source/AuProcAddresses.NT.cpp

@@ -137,7 +137,9 @@ namespace Aurora
         ADD_GET_PROC(Kernel32, VerSetConditionMask)
         ADD_GET_PROC(Kernel32, QueryPerformanceCounter)
         ADD_GET_PROC(Kernel32, QueryPerformanceFrequency)
-
+        ADD_GET_PROC(Kernel32, RemoveDllDirectory)
+        ADD_GET_PROC(Kernel32, AddDllDirectory)
+            
         ADD_GET_PROC_BI2(Kernel32, PSAPILegacy, K32GetProcessMemoryInfo, GetProcessMemoryInfo)
             
         ADD_GET_PROC(Sync, WaitOnAddress)

Source/AuProcAddresses.NT.hpp

@@ -239,6 +239,14 @@ namespace Aurora
     );
 #endif
 
+    inline BOOL(__stdcall *pRemoveDllDirectory)(
+        PVOID                       Cookie
+    );
+
+    inline PVOID(__stdcall *pAddDllDirectory)(
+        PCWSTR                      NewDirectory
+    );
+
     inline INT(__stdcall *pGetAddrInfoExCancel)(
         LPHANDLE                    lpHandle
     );

Source/Process/AuProcess.cpp

@@ -40,6 +40,7 @@
 namespace Aurora::Process
 {
     static AuFutexMutex gSpinLock;
+    static AuList<AuString> gClassPath;
     static AuHashMap<AuString, void *> gModuleHandles;
     static const bool kIsMainSigned = false;
 
@@ -273,12 +274,12 @@ namespace Aurora::Process
     }
     #endif
 
-    static bool LoadModule(const AuString &name, const AuString &path)
+    static void *LoadModule(const AuString &name, const AuString &path)
     {
     #if defined(AURORA_IS_MODERNNT_DERIVED)
         if (!pLoadLibraryW)
         {
-            return false;
+            return {};
         }
 
         auto handle = pLoadLibraryW(Locale::ConvertFromUTF8(path).c_str());
@@ -296,11 +297,12 @@ namespace Aurora::Process
         }
     #endif
 
-        gModuleHandles.insert(AuMakePair(name, (void *)handle));
+        auto pRet = (void *)handle;
-        return true;
+        gModuleHandles.insert(AuMakePair(name, pRet));
+        return pRet;
     }
 
-    static bool TryLoadModule(const AuString &path,
+    static void *TryLoadModule(const AuString &path,
     #if defined(AURORA_IS_MODERNNT_DERIVED)
                               const AuString &abs,
     #endif
@@ -320,7 +322,7 @@ namespace Aurora::Process
         {
             SysPushErrorNested("Couldn't open existing file. Race exploit?");
             fail = true;
-            return false;
+            return {};
         }
 
         auto absPath = AuIOFS::NormalizePathRet(abs);
@@ -334,7 +336,7 @@ namespace Aurora::Process
                 SysPushErrorNested("Couldn't verify file {}", path);
                 fail = true;
                 AuWin32CloseHandle(mitigateTimeOfUse);
-                return false;
+                return {};
             }
         #else
             AuLogWarn("Can't verify {} on this platform", path);
@@ -350,13 +352,13 @@ namespace Aurora::Process
             {
                 SysPushErrorNested("Couldn't open existing file. Race exploit?");
                 fail = true;
-                return false;
+                return {};
             }
 
             if ((sb.st_mode & S_IXUSR) == 0)
             {
                 fail = true;
-                return false;
+                return {};
             }
         }
     #endif
@@ -377,7 +379,7 @@ namespace Aurora::Process
         return returnValue;
     }
 
-    static bool TryLoadModule(const AuString &path, const AuString &auDll, const AuString &genericDll, const ModuleLoadRequest &request, bool &fail)
+    static void *TryLoadModule(const AuString &path, const AuString &auDll, const AuString &genericDll, const ModuleLoadRequest &request, bool &fail)
     {
         AuString a = path + "/" + auDll;
         AuString b = path + "/" + genericDll;
@@ -412,34 +414,61 @@ namespace Aurora::Process
         return {};
     }
 
-    static bool TryLoadModule(EModulePath path, const AuString &auDll, const AuString &genericDll, const ModuleLoadRequest &request)
+    static void *TryLoadModule(EModulePath path, const AuString &auDll, const AuString &genericDll, const ModuleLoadRequest &request)
     {
         AuString pathA, pathB;
+        AuList<AuString> arrayPaths;
 
         switch (path)
         {
+        case EModulePath::eClassPath:
+        {
+            AU_LOCK_GUARD(gSpinLock);
+            arrayPaths = gClassPath;
+            break;
+        }
         case EModulePath::eModulePathCWD:
-            if (!Process::GetWorkingDirectory(pathA)) return false;
+        {
+            if (!Process::GetWorkingDirectory(pathA))
+            {
+                return {};
+            }
 
             break;
+        }
         case EModulePath::eProcessDirectory:
-            if (!Process::GetProcDirectory(pathA)) return false;
+        {
+            if (!Process::GetProcDirectory(pathA))
+            {
+                return {};
+            }
 
             break;
+        }
         case EModulePath::eModulePathSystemDir:
+        {
             pathA = AuIOFS::GetSystemLibPath().value_or(AuString {});
             pathB = AuIOFS::GetSystemLibPath2().value_or(AuString {});
 
-            if (pathA.empty()) return false;
+            if (pathA.empty())
+            {
+                return {};
+            }
 
             break;
+        }
         case EModulePath::eModulePathUserDir:
+        {
             pathA = AuIOFS::GetUserLibPath().value_or(AuString {});
             pathB = AuIOFS::GetUserLibPath2().value_or(AuString {});
-            
+
-            if (pathA.empty()) return false;
+            if (pathA.empty())
+            {
+                return {};
+            }
 
             break;
+        }
         //case EModulePath::eOSSpecified:
         //
         //    break;
@@ -452,27 +481,40 @@ namespace Aurora::Process
 
         if (pathA.size())
         {
-            if (TryLoadModule(pathA, auDll, genericDll, request, fail))
+            if (auto pRet = TryLoadModule(pathA, auDll, genericDll, request, fail))
             {
-                return true;
+                return pRet;
             }
         
             if (fail)
             {
-                return false;
+                return {};
             }
         }
 
         if (pathB.size())
         {
-            if (TryLoadModule(pathB, auDll, genericDll, request, fail))
+            if (auto pRet = TryLoadModule(pathB, auDll, genericDll, request, fail))
             {
-                return true;
+                return pRet;
             }
 
             if (fail)
             {
-                return false;
+                return {};
+            }
+        }
+
+        for (const auto &dir : arrayPaths)
+        {
+            if (auto pRet = TryLoadModule(dir, auDll, genericDll, request, fail))
+            {
+                return pRet;
+            }
+
+            if (fail)
+            {
+                return {};
             }
         }
 
@@ -480,14 +522,14 @@ namespace Aurora::Process
         {
             for (const auto &val : *request.specifiedSearchPaths)
             {
-                if (TryLoadModule(val, auDll, genericDll, request, fail))
+                if (auto pRet = TryLoadModule(val, auDll, genericDll, request, fail))
                 {
-                    return true;
+                    return pRet;
                 }
                 
                 if (fail)
                 {
-                    return false;
+                    return {};
                 }
             }
         }
@@ -496,13 +538,18 @@ namespace Aurora::Process
     }
 
     AUKN_SYM bool LoadModule(const ModuleLoadRequest &request)
+    {
+        return bool(LoadModuleEx(request));
+    }
+
+    AUKN_SYM void *LoadModuleEx(const ModuleLoadRequest &request)
     {
         AU_LOCK_GUARD(gSpinLock);
         
         auto h = gModuleHandles.find(request.mod);
         if (h != gModuleHandles.end())
         {
-            return true;
+            return h->second;
         }
 
         auto au   = request.mod + ConstructAuDllSuffix();
@@ -528,37 +575,47 @@ namespace Aurora::Process
         auto searchPath = request.searchPath ? request.searchPath : &kUserOverloadableSearchPath;
         for (EModulePath path : *searchPath)
         {
-            if (TryLoadModule(path, au, base, request))
+            if (auto pRet = TryLoadModule(path, au, base, request))
             {
-                return true;
+                return pRet;
             }
         }
 
-        return false;
+        return {};
     }
 
-    AUKN_SYM AuMach GetProcAddress(AuString mod, AuString symbol)
+    AUKN_SYM AuMach GetProcAddressEx(void *pHandle, const AuString &symbol)
     {
-        AU_LOCK_GUARD(gSpinLock);
-
-        auto h = gModuleHandles.find(mod);
-        if (h == gModuleHandles.end())
-        {
-            SysPushErrorGen("Module {} is not loaded", mod);
-            return {};
-        }
-        auto handle = h->second;
-        
     #if defined(AURORA_IS_MODERNNT_DERIVED)
         if (!pGetProcAddress)
         {
             return {};
         }
 
-        auto ret = reinterpret_cast<AuMach>(pGetProcAddress(reinterpret_cast<HMODULE>(handle), symbol.c_str()));
+        auto ret = reinterpret_cast<AuMach>(pGetProcAddress(reinterpret_cast<HMODULE>(pHandle), symbol.c_str()));
     #else
-        auto ret = reinterpret_cast<AuMach>(dlsym(handle, symbol.c_str()));
+        auto ret = reinterpret_cast<AuMach>(dlsym(hapHandlendle, symbol.c_str()));
     #endif
+        return ret;
+    }
+
+    AUKN_SYM void *GetProcHandle(const AuString &name)
+    {
+        AU_LOCK_GUARD(gSpinLock);
+
+        auto h = gModuleHandles.find(name);
+        if (h == gModuleHandles.end())
+        {
+            SysPushErrorGen("Module {} is not loaded", name);
+            return {};
+        }
+
+        return h->second;
+    }
+
+    AUKN_SYM AuMach GetProcAddress(const AuString &mod, const AuString &symbol)
+    {
+        auto ret = GetProcAddressEx(GetProcHandle(mod), symbol);
 
         if (!ret)
         {
@@ -590,6 +647,177 @@ namespace Aurora::Process
     #endif
     }
 
+    static void PreloadDLLsDoOnce(const AuList<AuString> &dirs)
+    {
+        AU_DEBUG_MEMCRUNCH;
+
+        AuList<void *> cookieList;
+
+    #if defined(AURORA_PLATFORM_WIN32)
+        for (const auto dir : dirs)
+        {
+            void *pCookie {};
+
+            if (pRemoveDllDirectory)
+            {
+                if (pAddDllDirectory)
+                {
+                    pCookie = pAddDllDirectory(AuLocale::ConvertFromUTF8(dir).c_str());
+                }
+            }
+
+            cookieList.push_back(pCookie);
+        }
+    #endif
+
+        for (const auto dir : dirs)
+        {
+            AuList<AuString> files;
+            AuList<AuString> sharedObjects;
+
+            if (!AuFS::FilesInDirectory(dir, files))
+            {
+                SysPushErrorIO("Couldn't readdir: {}", dir);
+                continue;
+            }
+
+            auto endingPattern = ConstructAuDllSuffix();
+            auto endingAlt = GetPlatformExt(Build::kCurrentPlatform);
+
+            for (const auto &file : files)
+            {
+                if (!(file.ends_with(endingPattern.c_str()) ||
+                      file.ends_with(endingAlt)))
+                {
+                    continue;
+                }
+
+                sharedObjects.push_back(file);
+            }
+
+            {
+                AuUInt uSuccess {};
+                do
+                {
+                    uSuccess = 0;
+
+                    for (auto itr = sharedObjects.begin();
+                         itr != sharedObjects.end();  )
+                    {
+                        auto &sharedDLL = *itr;
+
+                        AuLogDbg("[attempt] Loading shared object: {}", sharedDLL);
+
+                        AuString a = dir + "/" + sharedDLL;
+
+                    #if defined(AURORA_IS_MODERNNT_DERIVED)
+                        AuString aAbs = dir + "/" + sharedDLL + ".";
+                    #endif
+
+                        bool fail {};
+                        ModuleLoadRequest request { AuList<AuString>{ dir }, sharedDLL };
+
+                        if (TryLoadModule(a,
+                        #if defined(AURORA_IS_MODERNNT_DERIVED)
+                                          aAbs,
+                        #endif
+                                          request,
+                                          fail
+                            ))
+                        {
+                            if (!fail)
+                            {
+                                AuLogDbg("[attempt] Loaded shared object: {}!", sharedDLL);
+                                uSuccess++;
+                                itr = sharedObjects.erase(itr);
+                            }
+                            else
+                            {
+                                itr++;
+                            }
+                        }
+                        else
+                        {
+                            itr++;
+                        }
+                    }
+                }
+                while (uSuccess);
+            }
+        }
+
+    #if defined(AURORA_PLATFORM_WIN32)
+        for (const auto pCookie : AuExchange(cookieList, {}))
+        {
+            if (pRemoveDllDirectory)
+            {
+                if (pCookie)
+                {
+                    pRemoveDllDirectory(pCookie);
+                }
+            }
+        }
+    #endif
+    }
+
+    static void PreloadDLLs(const AuString &dir)
+    {
+        AU_DEBUG_MEMCRUNCH;
+        if (gRuntimeConfig.processConfig.bAlwaysPreloadEntireClassPath)
+        {
+            PreloadDLLsDoOnce(gClassPath);
+        }
+        else
+        {
+            PreloadDLLsDoOnce({ dir });
+        }
+    }
+
+    AUKN_SYM bool SetBinaryClassPath(const AuList<AuString> &list, bool preloadAll)
+    {
+        AU_DEBUG_MEMCRUNCH;
+        AU_LOCK_GUARD(gSpinLock);
+        AuExchange(gClassPath, list);
+
+        if (((preloadAll && gRuntimeConfig.processConfig.bEnablePreload)) ||
+             (gRuntimeConfig.processConfig.bForcePreload))
+        {
+            if (gRuntimeConfig.processConfig.bAlwaysPreloadEntireClassPath)
+            {
+                PreloadDLLsDoOnce(gClassPath);
+            }
+            else
+            {
+                PreloadDLLsDoOnce(list);
+            }
+        }
+
+        return true;
+    }
+
+    AUKN_SYM bool AddBinaryClassPath(const AuString &dir, bool preloadAll)
+    {
+        AU_LOCK_GUARD(gSpinLock);
+
+        if (!AuTryInsert(gClassPath, dir))
+        {
+            return false;
+        }
+
+        if (((preloadAll && gRuntimeConfig.processConfig.bEnablePreload)) ||
+             (gRuntimeConfig.processConfig.bForcePreload))
+        {
+            PreloadDLLs(dir);
+        }
+
+        return true;
+    }
+
+    AUKN_SYM AuList<AuString> GetBinaryClassPath()
+    {
+        AU_LOCK_GUARD(gSpinLock);
+        return gClassPath;
+    }
 
     void InitProcess()
     {