From 94d3a1919c8eee4d7bfd7280ee4964477dee02e3 Mon Sep 17 00:00:00 2001 From: Mounir IDRASSI Date: Mon, 10 Feb 2020 02:20:46 +0100 Subject: Windows: Implement support for processor groups in the driver and fix build issues caused by previous implementation --- src/Common/EncryptionThreadPool.c | 117 ++++++++++++++++++++------------------ src/Common/Tcdefs.h | 15 +++++ 2 files changed, 78 insertions(+), 54 deletions(-) (limited to 'src/Common') diff --git a/src/Common/EncryptionThreadPool.c b/src/Common/EncryptionThreadPool.c index c11182a2..0dbaec5b 100644 --- a/src/Common/EncryptionThreadPool.c +++ b/src/Common/EncryptionThreadPool.c @@ -10,11 +10,6 @@ code distribution packages. */ -#if defined(_WIN32) || defined(WIN32) || defined(_WIN64) || defined (WIN64) -#include -#include -#endif - #include "EncryptionThreadPool.h" #include "Pkcs5.h" #ifdef DEVICE_DRIVER @@ -49,6 +44,18 @@ #define TC_ACQUIRE_MUTEX(MUTEX) WaitForSingleObject (*(MUTEX), INFINITE) #define TC_RELEASE_MUTEX(MUTEX) ReleaseMutex (*(MUTEX)) +typedef BOOL (WINAPI *SetThreadGroupAffinityFn)( + HANDLE hThread, + const GROUP_AFFINITY *GroupAffinity, + PGROUP_AFFINITY PreviousGroupAffinity +); + +typedef WORD (WINAPI* GetActiveProcessorGroupCountFn)(); + +typedef DWORD (WINAPI *GetActiveProcessorCountFn)( + WORD GroupNumber +); + #endif // !DEVICE_DRIVER @@ -105,6 +112,7 @@ static volatile BOOL StopPending = FALSE; static uint32 ThreadCount; static TC_THREAD_HANDLE ThreadHandles[TC_ENC_THREAD_POOL_MAX_THREAD_COUNT]; +static WORD ThreadProcessorGroups[TC_ENC_THREAD_POOL_MAX_THREAD_COUNT]; static EncryptionThreadPoolWorkItem WorkItemQueue[TC_ENC_THREAD_POOL_QUEUE_SIZE]; @@ -117,10 +125,6 @@ static TC_MUTEX DequeueMutex; static TC_EVENT WorkItemReadyEvent; static TC_EVENT WorkItemCompletedEvent; -#if defined(_WIN32) || defined(WIN32) || defined(_WIN64) || defined (WIN64) -static uint32 totalProcessors; -#endif - #if defined(_WIN64) void EncryptDataUnitsCurrentThreadEx (unsigned __int8 *buf, const UINT64_STRUCT *structUnitNo, TC_LARGEST_COMPILER_UINT nbrUnits, PCRYPTO_INFO ci) { @@ -173,13 +177,21 @@ static void SetWorkItemState (EncryptionThreadPoolWorkItem *workItem, WorkItemSt static TC_THREAD_PROC EncryptionThreadProc (void *threadArg) { - -#if defined(_WIN32) || defined(WIN32) || defined(_WIN64) || defined(WIN64) - GROUP_AFFINITY groupAffinity = { ~0ULL, *(int*)(threadArg), { 0, 0, 0 } }; - BOOL value = SetThreadGroupAffinity(GetCurrentThread(), &groupAffinity, NULL); + EncryptionThreadPoolWorkItem *workItem; +#ifdef DEVICE_DRIVER + SetThreadCpuGroupAffinity ((USHORT) *(WORD*)(threadArg)); +#else + SetThreadGroupAffinityFn SetThreadGroupAffinityPtr = (SetThreadGroupAffinityFn) GetProcAddress (GetModuleHandle (L"kernel32.dll"), "SetThreadGroupAffinity"); + if (SetThreadGroupAffinityPtr && threadArg) + { + GROUP_AFFINITY groupAffinity = {0}; + groupAffinity.Mask = ~0ULL; + groupAffinity.Group = *(WORD*)(threadArg); + SetThreadGroupAffinityPtr(GetCurrentThread(), &groupAffinity, NULL); + } + #endif - EncryptionThreadPoolWorkItem *workItem; while (!StopPending) { @@ -279,36 +291,33 @@ static TC_THREAD_PROC EncryptionThreadProc (void *threadArg) BOOL EncryptionThreadPoolStart (size_t encryptionFreeCpuCount) { - size_t cpuCount = 0, i = 0, processors = 0, totalProcessors = 0; - int threadProcessorGroups[128]; -#if defined(_WIN32) || defined(WIN32) || defined(_WIN64) || defined (WIN64) - for (i = 0; i < GetActiveProcessorGroupCount(); ++i) - { - processors = GetActiveProcessorCount(i); - totalProcessors += processors; -} -#endif - - if (ThreadPoolRunning) - return TRUE; - -#if defined(_WIN32) || defined(WIN32) || defined(_WIN64) || defined(WIN64) - SYSTEM_INFO sysInfo; - GetSystemInfo(&sysInfo); - cpuCount = (IsWindows7OrGreater) ? totalProcessors : sysInfo.dwNumberOfProcessors; -#endif - -/* + size_t cpuCount = 0, i = 0; #ifdef DEVICE_DRIVER cpuCount = GetCpuCount(); #else + SYSTEM_INFO sysInfo; + GetActiveProcessorGroupCountFn GetActiveProcessorGroupCountPtr = (GetActiveProcessorGroupCountFn) GetProcAddress (GetModuleHandle (L"Kernel32.dll"), "GetActiveProcessorGroupCount"); + GetActiveProcessorCountFn GetActiveProcessorCountPtr = (GetActiveProcessorCountFn) GetProcAddress (GetModuleHandle (L"Kernel32.dll"), "GetActiveProcessorCount"); + if (GetActiveProcessorGroupCountPtr && GetActiveProcessorCountPtr) { - SYSTEM_INFO sysInfo; - GetSystemInfo (&sysInfo); + WORD j, groupCount = GetActiveProcessorGroupCountPtr(); + size_t totalProcessors = 0; + for (j = 0; j < groupCount; ++j) + { + totalProcessors += (size_t) GetActiveProcessorCountPtr(j); + } + cpuCount = totalProcessors; + } + else + { + GetSystemInfo(&sysInfo); cpuCount = sysInfo.dwNumberOfProcessors; } #endif -*/ + + if (ThreadPoolRunning) + return TRUE; + if (cpuCount > encryptionFreeCpuCount) cpuCount -= encryptionFreeCpuCount; @@ -368,34 +377,34 @@ BOOL EncryptionThreadPoolStart (size_t encryptionFreeCpuCount) for (ThreadCount = 0; ThreadCount < cpuCount; ++ThreadCount) { - -#if defined(_WIN32) || defined(WIN32) || defined(_WIN64) || defined(WIN64) +#ifdef DEVICE_DRIVER + ThreadProcessorGroups[ThreadCount] = GetCpuGroup ((size_t) ThreadCount); +#else // Determine which processor group to bind the thread to. - totalProcessors = 0U; - for (i = 0U; i < GetActiveProcessorGroupCount(); ++i) + if (GetActiveProcessorGroupCountPtr && GetActiveProcessorCountPtr) { - totalProcessors += GetActiveProcessorCount(i); - if (totalProcessors >= ThreadCount) + WORD j, groupCount = GetActiveProcessorGroupCountPtr(); + uint32 totalProcessors = 0U; + for (j = 0U; j < groupCount; j++) { - threadProcessorGroups[ThreadCount] = i; - break; + totalProcessors += (uint32) GetActiveProcessorCountPtr(j); + if (totalProcessors >= ThreadCount) + { + ThreadProcessorGroups[ThreadCount] = j; + break; + } } } + else + ThreadProcessorGroups[ThreadCount] = 0; #endif #ifdef DEVICE_DRIVER - if (!NT_SUCCESS(TCStartThread(EncryptionThreadProc, (void*)(&threadProcessorGroups[ThreadCount]), &ThreadHandles[ThreadCount]))) + if (!NT_SUCCESS(TCStartThread(EncryptionThreadProc, (void*)(&ThreadProcessorGroups[ThreadCount]), &ThreadHandles[ThreadCount]))) #else - if (!(ThreadHandles[ThreadCount] = (HANDLE)_beginthreadex(NULL, 0, EncryptionThreadProc, (void*)(&threadProcessorGroups[ThreadCount]), 0, NULL))) + if (!(ThreadHandles[ThreadCount] = (HANDLE)_beginthreadex(NULL, 0, EncryptionThreadProc, (void*)(&ThreadProcessorGroups[ThreadCount]), 0, NULL))) #endif -/* -#ifdef DEVICE_DRIVER - if (!NT_SUCCESS (TCStartThread (EncryptionThreadProc, NULL, &ThreadHandles[ThreadCount]))) -#else - if (!(ThreadHandles[ThreadCount] = (HANDLE) _beginthreadex (NULL, 0, EncryptionThreadProc, NULL, 0, NULL))) -#endif -*/ { EncryptionThreadPoolStop(); return FALSE; diff --git a/src/Common/Tcdefs.h b/src/Common/Tcdefs.h index c9a38a0e..c89cb779 100644 --- a/src/Common/Tcdefs.h +++ b/src/Common/Tcdefs.h @@ -263,6 +263,10 @@ extern ULONG AllocTag; typedef int BOOL; #endif +#ifndef WORD +typedef USHORT WORD; +#endif + #ifndef BOOLEAN typedef unsigned char BOOLEAN; #endif @@ -295,6 +299,17 @@ typedef NTSTATUS (NTAPI *ExGetFirmwareEnvironmentVariableFn) ( typedef BOOLEAN (NTAPI *KeAreAllApcsDisabledFn) (); +typedef void (NTAPI *KeSetSystemGroupAffinityThreadFn)( + PGROUP_AFFINITY Affinity, + PGROUP_AFFINITY PreviousAffinity +); + +typedef USHORT (NTAPI *KeQueryActiveGroupCountFn)(); + +typedef ULONG (NTAPI *KeQueryActiveProcessorCountExFn)( + USHORT GroupNumber +); + extern NTSTATUS NTAPI KeSaveExtendedProcessorState ( __in ULONG64 Mask, PXSTATE_SAVE XStateSave -- cgit v1.2.3