More efficient way of dealing with modules in current dir

src/Main.cpp

@@ -15,6 +15,7 @@
 
 HMODULE g_dinput = 0;
 std::mutex g_load_mutex{};
+extern bool g_success_made_ldr_notification;
 
 void failed() {
     MessageBox(0, "REFramework: Unable to load the original dinput8.dll. Please report this to the developer.", "REFramework", 0);
@@ -75,11 +76,15 @@ void startup_thread(HMODULE reframework_module) {
 
 #if defined(MHRISE)
         if (our_dll) {
-            utility::spoof_module_paths_in_exe_dir();
+            if (!g_success_made_ldr_notification) {
+                utility::spoof_module_paths_in_exe_dir();
+            }
             utility::unlink(*our_dll);
         }
 #elif defined (DD2)
-        utility::spoof_module_paths_in_exe_dir();
+        if (!g_success_made_ldr_notification) {
+            utility::spoof_module_paths_in_exe_dir();
+        }
 #endif
     }
 }

src/REFramework.cpp

@@ -137,6 +137,81 @@ void REFramework::hook_monitor() {
     }
 }
 
+typedef struct _LDR_DLL_UNLOADED_NOTIFICATION_DATA {
+    ULONG Flags;                    //Reserved.
+    PCUNICODE_STRING FullDllName;   //The full path name of the DLL module.
+    PCUNICODE_STRING BaseDllName;   //The base file name of the DLL module.
+    PVOID DllBase;                  //A pointer to the base address for the DLL in memory.
+    ULONG SizeOfImage;              //The size of the DLL image, in bytes.
+} LDR_DLL_UNLOADED_NOTIFICATION_DATA, *PLDR_DLL_UNLOADED_NOTIFICATION_DATA;
+
+typedef struct _LDR_DLL_LOADED_NOTIFICATION_DATA {
+    ULONG Flags;                    //Reserved.
+    PCUNICODE_STRING FullDllName;   //The full path name of the DLL module.
+    PCUNICODE_STRING BaseDllName;   //The base file name of the DLL module.
+    PVOID DllBase;                  //A pointer to the base address for the DLL in memory.
+    ULONG SizeOfImage;              //The size of the DLL image, in bytes.
+} LDR_DLL_LOADED_NOTIFICATION_DATA, *PLDR_DLL_LOADED_NOTIFICATION_DATA;
+
+typedef union _LDR_DLL_NOTIFICATION_DATA {
+    LDR_DLL_LOADED_NOTIFICATION_DATA Loaded;
+    LDR_DLL_UNLOADED_NOTIFICATION_DATA Unloaded;
+} LDR_DLL_NOTIFICATION_DATA, *PLDR_DLL_NOTIFICATION_DATA;
+
+using PLDR_DLL_NOTIFICATION_FUNCTION = void (*)(
+  ULONG                       NotificationReason,
+  PLDR_DLL_NOTIFICATION_DATA NotificationData,
+  PVOID                       Context
+);
+
+using LdrRegisterDllNotification_t = NTSTATUS (*) (
+    ULONG                          Flags,
+    PLDR_DLL_NOTIFICATION_FUNCTION NotificationFunction,
+    PVOID                          Context,
+    PVOID                          *Cookie
+);
+
+#define LDR_DLL_NOTIFICATION_REASON_LOADED 1
+#define LDR_DLL_NOTIFICATION_REASON_UNLOADED 2
+
+std::optional<std::filesystem::path> g_current_game_path{};
+bool g_success_made_ldr_notification{false};
+
+void CALLBACK ldr_notification_callback(
+    ULONG                       NotificationReason,
+    PLDR_DLL_NOTIFICATION_DATA NotificationData,
+    PVOID                       Context
+) 
+try {
+    // From what I can tell, the PEB entries get filled in by the time this is called
+    // so we're good.
+    if (NotificationReason == LDR_DLL_NOTIFICATION_REASON_LOADED) {
+        if (NotificationData->Loaded.BaseDllName != nullptr && NotificationData->Loaded.BaseDllName->Buffer != nullptr) {
+            std::wstring base_dll_name = NotificationData->Loaded.BaseDllName->Buffer;
+            spdlog::info("LdrRegisterDllNotification: Loaded: {}", utility::narrow(base_dll_name));
+        }
+
+        if (g_current_game_path && NotificationData->Loaded.FullDllName != nullptr && NotificationData->Loaded.FullDllName->Buffer != nullptr) {
+            std::wstring full_dll_name = NotificationData->Loaded.FullDllName->Buffer;
+            std::filesystem::path full_dll_path = full_dll_name;
+
+            if (full_dll_path.parent_path() == *g_current_game_path) {
+                spdlog::info("LdrRegisterDllNotification: DLL loaded from game directory: {}", utility::narrow(full_dll_name));
+
+#if defined(DD2) || defined(MHRISE)
+                utility::spoof_module_paths_in_exe_dir();
+#endif
+            }
+        } else {
+            spdlog::info("LdrRegisterDllNotification: DLL loaded from unknown location");
+        }
+    }
+} catch (const std::exception& e) {
+    spdlog::error("ldr_notification_callback: Exception occurred: {}", e.what());
+} catch(...) {
+    spdlog::error("ldr_notification_callback: Unknown exception occurred");
+}
+
 typedef NTSTATUS (WINAPI* PFN_LdrLockLoaderLock)(ULONG Flags, ULONG *State, ULONG_PTR *Cookie);
 typedef NTSTATUS (WINAPI* PFN_LdrUnlockLoaderLock)(ULONG Flags, ULONG_PTR Cookie);
 
@@ -162,6 +237,12 @@ REFramework::REFramework(HMODULE reframework_module)
     spdlog::info("Game Module Addr: {:x}", (uintptr_t)m_game_module);
     spdlog::info("Game Module Size: {:x}", module_size);
 
+    if (auto current_game_path = utility::get_module_pathw(m_game_module); current_game_path.has_value()) {
+        g_current_game_path = *current_game_path;
+        g_current_game_path = g_current_game_path->parent_path();
+        spdlog::info("Current game path: {}", utility::narrow(g_current_game_path->c_str()));
+    }
+
     // preallocate some memory for minhook to mitigate failures (temporarily at least... this should in theory fail when too many hooks are made)
     // but, 64 slots should be enough for now. 
     // so... TODO: modify minhook to use absolute jumps when failing to allocate memory nearby
@@ -217,6 +298,28 @@ REFramework::REFramework(HMODULE reframework_module)
         } else {
             spdlog::info("RtlGetVersion() not found");
         }
+
+        // Do this at least once before setting up our callback.
+#if defined(DD2) || defined(MHRISE)
+        utility::spoof_module_paths_in_exe_dir();
+#endif
+
+        // Register our LdrRegisterDllNotification callback
+        spdlog::info("Registering LdrRegisterDllNotification callback...");
+        const auto ldr_register_dll_notification = (LdrRegisterDllNotification_t)GetProcAddress(ntdll, "LdrRegisterDllNotification");
+
+        if (ldr_register_dll_notification != nullptr) {
+            PVOID cookie = nullptr;
+            g_success_made_ldr_notification = NT_SUCCESS(ldr_register_dll_notification(0, ldr_notification_callback, nullptr, &cookie));
+
+            if (g_success_made_ldr_notification) {
+                spdlog::info("LdrRegisterDllNotification callback registered successfully");
+            } else {
+                spdlog::info("LdrRegisterDllNotification callback failed to register");
+            }
+        } else {
+            spdlog::info("LdrRegisterDllNotification not found");
+        }
     } else {
         spdlog::info("ntdll.dll not found");
     }
@@ -251,7 +354,10 @@ REFramework::REFramework(HMODULE reframework_module)
     utility::load_module_from_current_directory(L"openxr_loader.dll");
     LoadLibraryA("dxgi.dll");
     LoadLibraryA("d3d11.dll");
-    utility::spoof_module_paths_in_exe_dir();
+
+    if (!g_success_made_ldr_notification) {
+        utility::spoof_module_paths_in_exe_dir();
+    }
 #endif
 
 #if defined(RE8)
@@ -1567,7 +1673,9 @@ bool REFramework::initialize_game_data() {
 
         try {
 #if defined(MHRISE) || defined(DD2)
-            utility::spoof_module_paths_in_exe_dir();
+            if (!g_success_made_ldr_notification) {
+                utility::spoof_module_paths_in_exe_dir();
+            }
 #endif
             reframework::initialize_sdk();
 
@@ -1634,7 +1742,9 @@ bool REFramework::initialize_game_data() {
         }
 
 #if defined(MHRISE) || defined(DD2)
-        utility::spoof_module_paths_in_exe_dir();
+        if (!g_success_made_ldr_notification) {
+            utility::spoof_module_paths_in_exe_dir();
+        }
 #endif
         spdlog::info("Game data initialization thread finished");
     });