move affinity changes to common

ThirteenAG · Nov 13, 2024 · f10a645 · f10a645
1 parent 8932d30
commit f10a645
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 81 deletions.
diff --git a/includes/stdafx.h b/includes/stdafx.h
@@ -673,6 +673,45 @@ class RegistryWrapper
     }
 };
 
+namespace AffinityChanges
+{
+    static inline DWORD_PTR gameThreadAffinity = 0;
+    static inline bool Init()
+    {
+        DWORD_PTR processAffinity, systemAffinity;
+        if (!GetProcessAffinityMask(GetCurrentProcess(), &processAffinity, &systemAffinity))
+        {
+            return false;
+        }
+
+        DWORD_PTR otherCoresAff = (processAffinity - 1) & processAffinity;
+        if (otherCoresAff == 0) // Only one core is available for the game
+        {
+            return false;
+        }
+        gameThreadAffinity = processAffinity & ~otherCoresAff;
+
+        SetThreadAffinityMask(GetCurrentThread(), gameThreadAffinity);
+
+        return true;
+    }
+
+    static inline HANDLE WINAPI CreateThread_GameThread(LPSECURITY_ATTRIBUTES lpThreadAttributes, SIZE_T dwStackSize, LPTHREAD_START_ROUTINE lpStartAddress,
+        PVOID lpParameter, DWORD dwCreationFlags, LPDWORD lpThreadId)
+    {
+        HANDLE hThread = CreateThread(lpThreadAttributes, dwStackSize, lpStartAddress, lpParameter, dwCreationFlags | CREATE_SUSPENDED, lpThreadId);
+        if (hThread != nullptr)
+        {
+            SetThreadAffinityMask(hThread, gameThreadAffinity);
+            if ((dwCreationFlags & CREATE_SUSPENDED) == 0) // Resume only if the game didn't pass CREATE_SUSPENDED
+            {
+                ResumeThread(hThread);
+            }
+        }
+        return hThread;
+    }
+}
+
 namespace WindowedModeWrapper
 {
     static bool bBorderlessWindowed = true;

diff --git a/source/NFSUnderground2.WidescreenFix/dllmain.cpp b/source/NFSUnderground2.WidescreenFix/dllmain.cpp
@@ -161,16 +161,6 @@ void LoadDatFile()
     }
 }
 
-static constexpr DWORD AffinityMask = 1;
-HANDLE WINAPI CustomCreateThread(LPSECURITY_ATTRIBUTES lpThreadAttributes, SIZE_T dwStackSize, LPTHREAD_START_ROUTINE lpStartAddress,
-    LPVOID lpParameter, DWORD dwCreationFlags, LPDWORD lpThreadId)
-{
-    HANDLE hThread = CreateThread(lpThreadAttributes, dwStackSize, lpStartAddress, lpParameter, dwCreationFlags, lpThreadId);
-    if (hThread)
-        SetThreadAffinityMask(hThread, AffinityMask);
-    return hThread;
-}
-
 namespace NOSTrailFix
 {
     static uint32_t NOSTrailFrameDelay = 1;
@@ -880,34 +870,11 @@ void Init()
 
     if (bSingleCoreAffinity)
     {
-        HINSTANCE					hInstance = GetModuleHandle(nullptr);
-        PIMAGE_NT_HEADERS			ntHeader = (PIMAGE_NT_HEADERS)((DWORD_PTR)hInstance + ((PIMAGE_DOS_HEADER)hInstance)->e_lfanew);
-        PIMAGE_IMPORT_DESCRIPTOR	pImports = (PIMAGE_IMPORT_DESCRIPTOR)((DWORD_PTR)hInstance + ntHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress);
-
-        // Find KERNEL32.DLL
-        for (; pImports->Name != 0; pImports++)
+        if (AffinityChanges::Init())
         {
-            if (!_stricmp((const char*)((DWORD_PTR)hInstance + pImports->Name), "KERNEL32.DLL"))
-            {
-                if (pImports->OriginalFirstThunk != 0)
-                {
-                    PIMAGE_IMPORT_BY_NAME* pFunctions = (PIMAGE_IMPORT_BY_NAME*)((DWORD_PTR)hInstance + pImports->OriginalFirstThunk);
-                    for (ptrdiff_t j = 0; pFunctions[j] != nullptr; j++)
-                    {
-                        if (!strcmp((const char*)((DWORD_PTR)hInstance + pFunctions[j]->Name), "CreateThread"))
-                        {
-                            // Overwrite the address with the address to a custom CreateThread
-                            DWORD dwProtect[2];
-                            DWORD_PTR* pAddress = &((DWORD_PTR*)((DWORD_PTR)hInstance + pImports->FirstThunk))[j];
-                            VirtualProtect(pAddress, sizeof(DWORD_PTR), PAGE_EXECUTE_READWRITE, &dwProtect[0]);
-                            *pAddress = (DWORD_PTR)CustomCreateThread;
-                            VirtualProtect(pAddress, sizeof(DWORD_PTR), dwProtect[0], &dwProtect[1]);
-                            SetThreadAffinityMask(GetCurrentThread(), AffinityMask);
-                            break;
-                        }
-                    }
-                }
-            }
+            IATHook::Replace(GetModuleHandleA(NULL), "KERNEL32.DLL",
+                std::forward_as_tuple("CreateThread", AffinityChanges::CreateThread_GameThread)
+            );
         }
     }
 
@@ -1072,7 +1039,7 @@ CEXP void InitializeASI()
 {
     std::call_once(CallbackHandler::flag, []()
     {
-        CallbackHandler::RegisterCallback(Init, hook::pattern("C7 00 80 02 00 00 C7 01 E0 01 00 00"));
+        CallbackHandler::RegisterCallbackAtGetSystemTimeAsFileTime(Init, hook::pattern("C7 00 80 02 00 00 C7 01 E0 01 00 00"));
     });
 }
 

diff --git a/source/SplinterCellDoubleAgent.WidescreenFix/dllmain.cpp b/source/SplinterCellDoubleAgent.WidescreenFix/dllmain.cpp
@@ -7,45 +7,6 @@
 #pragma comment(lib, "winmm.lib") // needed for timeBeginPeriod()/timeEndPeriod()
 #include <LEDEffects.h>
 
-namespace AffinityChanges
-{
-    DWORD_PTR gameThreadAffinity = 0;
-    static bool Init()
-    {
-        DWORD_PTR processAffinity, systemAffinity;
-        if (!GetProcessAffinityMask(GetCurrentProcess(), &processAffinity, &systemAffinity))
-        {
-            return false;
-        }
-
-        DWORD_PTR otherCoresAff = (processAffinity - 1) & processAffinity;
-        if (otherCoresAff == 0) // Only one core is available for the game
-        {
-            return false;
-        }
-        gameThreadAffinity = processAffinity & ~otherCoresAff;
-
-        SetThreadAffinityMask(GetCurrentThread(), gameThreadAffinity);
-
-        return true;
-    }
-
-    static HANDLE WINAPI CreateThread_GameThread(LPSECURITY_ATTRIBUTES lpThreadAttributes, SIZE_T dwStackSize, LPTHREAD_START_ROUTINE lpStartAddress,
-        PVOID lpParameter, DWORD dwCreationFlags, LPDWORD lpThreadId)
-    {
-        HANDLE hThread = CreateThread(lpThreadAttributes, dwStackSize, lpStartAddress, lpParameter, dwCreationFlags | CREATE_SUSPENDED, lpThreadId);
-        if (hThread != nullptr)
-        {
-            SetThreadAffinityMask(hThread, gameThreadAffinity);
-            if ((dwCreationFlags & CREATE_SUSPENDED) == 0) // Resume only if the game didn't pass CREATE_SUSPENDED
-            {
-                ResumeThread(hThread);
-            }
-        }
-        return hThread;
-    }
-}
-
 uint32_t curAmmoInClip = 1;
 uint32_t curClipCapacity = 1;
 void AmmoInClip()
@@ -880,10 +841,12 @@ void InitEngine()
 
     if (bSingleCoreAffinity)
     {
-        AffinityChanges::Init();
-        IATHook::Replace(GetModuleHandle(L"Engine"), "kernel32.dll",
-            std::forward_as_tuple("CreateThread", AffinityChanges::CreateThread_GameThread)
-        );
+        if (AffinityChanges::Init())
+        {
+            IATHook::Replace(GetModuleHandle(L"Engine"), "kernel32.dll",
+                std::forward_as_tuple("CreateThread", AffinityChanges::CreateThread_GameThread)
+            );
+        }
     }
 }