From b443f83c774eb741f610b71bec80632da8782f51 Mon Sep 17 00:00:00 2001 From: praydog Date: Wed, 3 Apr 2024 02:53:06 -0700 Subject: [PATCH] More efficient way of dealing with modules in current dir --- src/Main.cpp | 9 +++- src/REFramework.cpp | 116 ++++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 120 insertions(+), 5 deletions(-) diff --git a/src/Main.cpp b/src/Main.cpp index e54eb5205..4971c3b2a 100644 --- a/src/Main.cpp +++ b/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 } } diff --git a/src/REFramework.cpp b/src/REFramework.cpp index b3ddcd3fb..767a78ec3 100644 --- a/src/REFramework.cpp +++ b/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 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"); });