diff --git a/inc/delay_imports.h b/inc/delay_imports.h index 19bbc23..5221ec4 100644 --- a/inc/delay_imports.h +++ b/inc/delay_imports.h @@ -1,20 +1,24 @@ #ifndef DELAY_IMPORTS_H #define DELAY_IMPORTS_H +#define STATUS_SUCCESS ((NTSTATUS)0x00000000L) #define ThreadQuerySetWin32StartAddress 9 typedef NTSTATUS(WINAPI* RTLVERIFYVERSIONINFOPROC)(PRTL_OSVERSIONINFOEXW, ULONG, ULONGLONG); -typedef ULONGLONG(WINAPI* VERSETCONDITIONMASKPROC)(ULONGLONG, DWORD, BYTE); typedef const char* (CDECL* WINE_GET_VERSIONPROC)(); typedef void (CDECL* WINE_GET_HOST_VERSIONPROC)(const char** sysname, const char** release); typedef NTSTATUS(WINAPI* NTQUERYINFORMATIONTHREADPROC)(HANDLE, LONG, PVOID, ULONG, PULONG); +typedef ULONGLONG(WINAPI* VERSETCONDITIONMASKPROC)(ULONGLONG, DWORD, BYTE); +typedef BOOL(WINAPI* GETMODULEHANDLEEXAPROC)(DWORD, LPCSTR, HMODULE*); + extern NTQUERYINFORMATIONTHREADPROC NtQueryInformationThread; extern RTLVERIFYVERSIONINFOPROC RtlVerifyVersionInfo; extern WINE_GET_VERSIONPROC wine_get_version; extern WINE_GET_HOST_VERSIONPROC wine_get_host_version; extern VERSETCONDITIONMASKPROC VerSetConditionMaskProc; +extern GETMODULEHANDLEEXAPROC GetModuleHandleExAProc; void imports_init(); diff --git a/src/delay_imports.c b/src/delay_imports.c index f224784..9becb1f 100644 --- a/src/delay_imports.c +++ b/src/delay_imports.c @@ -8,6 +8,7 @@ WINE_GET_VERSIONPROC wine_get_version; WINE_GET_HOST_VERSIONPROC wine_get_host_version; VERSETCONDITIONMASKPROC VerSetConditionMaskProc; +GETMODULEHANDLEEXAPROC GetModuleHandleExAProc; void imports_init() { @@ -24,5 +25,6 @@ void imports_init() if (mod) { VerSetConditionMaskProc = (VERSETCONDITIONMASKPROC)GetProcAddress(mod, "VerSetConditionMask"); + GetModuleHandleExAProc = (GETMODULEHANDLEEXAPROC)GetProcAddress(mod, "GetModuleHandleExA"); } } diff --git a/src/utils.c b/src/utils.c index a6649c4..47bd680 100644 --- a/src/utils.c +++ b/src/utils.c @@ -13,6 +13,7 @@ #include "utils.h" #include "config.h" #include "versionhelpers.h" +#include "delay_imports.h" /* @@ -103,7 +104,49 @@ void util_set_thread_affinity(DWORD tid) HANDLE thread = OpenThread(THREAD_QUERY_INFORMATION | THREAD_SET_INFORMATION, FALSE, tid); if (thread) { - SetThreadAffinityMask(thread, 1); + void* start = NULL; + NTSTATUS status = STATUS_PENDING; + + if (NtQueryInformationThread) + { + status = NtQueryInformationThread(thread, ThreadQuerySetWin32StartAddress, &start, sizeof(start), NULL); + } + + if (status == STATUS_SUCCESS && start && GetModuleHandleExAProc) + { + char game_exe_path[MAX_PATH] = { 0 }; + char game_dir[MAX_PATH] = { 0 }; + + if (GetModuleFileNameA(NULL, game_exe_path, sizeof(game_exe_path))) + { + _splitpath(game_exe_path, NULL, game_dir, NULL, NULL); + + char mod_path[MAX_PATH] = { 0 }; + char mod_dir[MAX_PATH] = { 0 }; + char mod_filename[MAX_PATH] = { 0 }; + HMODULE mod = NULL; + + if (GetModuleHandleExAProc( + GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT, start, &mod)) + { + if (GetModuleFileNameA(mod, mod_path, sizeof(mod_path))) + { + _splitpath(mod_path, NULL, mod_dir, mod_filename, NULL); + + if (_strnicmp(game_dir, mod_dir, strlen(game_dir)) == 0 || + _strcmpi(mod_filename, "WINMM") == 0) + { + SetThreadAffinityMask(thread, 1); + } + } + } + } + } + else + { + SetThreadAffinityMask(thread, 1); + } + CloseHandle(thread); } }