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 +++++ src/Driver/Ntdriver.c | 66 ++++++++++++++++++--- src/Driver/Ntdriver.h | 2 + 4 files changed, 138 insertions(+), 62 deletions(-) (limited to 'src') 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 diff --git a/src/Driver/Ntdriver.c b/src/Driver/Ntdriver.c index 8caf7bfa..2ae17f5a 100644 --- a/src/Driver/Ntdriver.c +++ b/src/Driver/Ntdriver.c @@ -143,6 +143,9 @@ static KeSaveExtendedProcessorStateFn KeSaveExtendedProcessorStatePtr = NULL; static KeRestoreExtendedProcessorStateFn KeRestoreExtendedProcessorStatePtr = NULL; static ExGetFirmwareEnvironmentVariableFn ExGetFirmwareEnvironmentVariablePtr = NULL; static KeAreAllApcsDisabledFn KeAreAllApcsDisabledPtr = NULL; +static KeSetSystemGroupAffinityThreadFn KeSetSystemGroupAffinityThreadPtr = NULL; +static KeQueryActiveGroupCountFn KeQueryActiveGroupCountPtr = NULL; +static KeQueryActiveProcessorCountExFn KeQueryActiveProcessorCountExPtr = NULL; POOL_TYPE ExDefaultNonPagedPoolType = NonPagedPool; ULONG ExDefaultMdlProtection = 0; @@ -283,13 +286,20 @@ NTSTATUS DriverEntry (PDRIVER_OBJECT DriverObject, PUNICODE_STRING RegistryPath) } // KeSaveExtendedProcessorState/KeRestoreExtendedProcessorState are available starting from Windows 7 + // KeQueryActiveGroupCount/KeQueryActiveProcessorCountEx/KeSetSystemGroupAffinityThread are available starting from Windows 7 if ((OsMajorVersion > 6) || (OsMajorVersion == 6 && OsMinorVersion >= 1)) { - UNICODE_STRING saveFuncName, restoreFuncName; + UNICODE_STRING saveFuncName, restoreFuncName, groupCountFuncName, procCountFuncName, setAffinityFuncName; RtlInitUnicodeString(&saveFuncName, L"KeSaveExtendedProcessorState"); RtlInitUnicodeString(&restoreFuncName, L"KeRestoreExtendedProcessorState"); + RtlInitUnicodeString(&groupCountFuncName, L"KeQueryActiveGroupCount"); + RtlInitUnicodeString(&procCountFuncName, L"KeQueryActiveProcessorCountEx"); + RtlInitUnicodeString(&setAffinityFuncName, L"KeSetSystemGroupAffinityThread"); KeSaveExtendedProcessorStatePtr = (KeSaveExtendedProcessorStateFn) MmGetSystemRoutineAddress(&saveFuncName); KeRestoreExtendedProcessorStatePtr = (KeRestoreExtendedProcessorStateFn) MmGetSystemRoutineAddress(&restoreFuncName); + KeSetSystemGroupAffinityThreadPtr = (KeSetSystemGroupAffinityThreadFn) MmGetSystemRoutineAddress(&setAffinityFuncName); + KeQueryActiveGroupCountPtr = (KeQueryActiveGroupCountFn) MmGetSystemRoutineAddress(&groupCountFuncName); + KeQueryActiveProcessorCountExPtr = (KeQueryActiveProcessorCountExFn) MmGetSystemRoutineAddress(&procCountFuncName); } // ExGetFirmwareEnvironmentVariable is available starting from Windows 8 @@ -4488,16 +4498,27 @@ NTSTATUS TCCompleteDiskIrp (PIRP irp, NTSTATUS status, ULONG_PTR information) size_t GetCpuCount () { - KAFFINITY activeCpuMap = KeQueryActiveProcessors(); - size_t mapSize = sizeof (activeCpuMap) * 8; size_t cpuCount = 0; - - while (mapSize--) + if (KeQueryActiveGroupCountPtr && KeQueryActiveProcessorCountExPtr) + { + USHORT i, groupCount = KeQueryActiveGroupCountPtr (); + for (i = 0; i < groupCount; i++) + { + cpuCount += (size_t) KeQueryActiveProcessorCountExPtr (i); + } + } + else { - if (activeCpuMap & 1) - ++cpuCount; + KAFFINITY activeCpuMap = KeQueryActiveProcessors(); + size_t mapSize = sizeof (activeCpuMap) * 8; - activeCpuMap >>= 1; + while (mapSize--) + { + if (activeCpuMap & 1) + ++cpuCount; + + activeCpuMap >>= 1; + } } if (cpuCount == 0) @@ -4506,6 +4527,35 @@ size_t GetCpuCount () return cpuCount; } +USHORT GetCpuGroup (size_t index) +{ + if (KeQueryActiveGroupCountPtr && KeQueryActiveProcessorCountExPtr) + { + USHORT i, groupCount = KeQueryActiveGroupCountPtr (); + size_t cpuCount = 0; + for (i = 0; i < groupCount; i++) + { + cpuCount += (size_t) KeQueryActiveProcessorCountExPtr (i); + if (cpuCount >= index) + { + return i; + } + } + } + + return 0; +} + +void SetThreadCpuGroupAffinity (USHORT index) +{ + if (KeSetSystemGroupAffinityThreadPtr) + { + GROUP_AFFINITY groupAffinity = {0}; + groupAffinity.Mask = ~0ULL; + groupAffinity.Group = index; + KeSetSystemGroupAffinityThreadPtr (&groupAffinity, NULL); + } +} void EnsureNullTerminatedString (wchar_t *str, size_t maxSizeInBytes) { diff --git a/src/Driver/Ntdriver.h b/src/Driver/Ntdriver.h index 2e4d6555..25ee64e9 100644 --- a/src/Driver/Ntdriver.h +++ b/src/Driver/Ntdriver.h @@ -174,6 +174,8 @@ NTSTATUS TCCompleteDiskIrp (PIRP irp, NTSTATUS status, ULONG_PTR information); NTSTATUS ProbeRealDriveSize (PDEVICE_OBJECT driveDeviceObject, LARGE_INTEGER *driveSize); BOOL UserCanAccessDriveDevice (); size_t GetCpuCount (); +USHORT GetCpuGroup (size_t index); +void SetThreadCpuGroupAffinity (USHORT index); void EnsureNullTerminatedString (wchar_t *str, size_t maxSizeInBytes); void *AllocateMemoryWithTimeout (size_t size, int retryDelay, int timeout); BOOL IsDriveLetterAvailable (int nDosDriveNo, DeviceNamespaceType namespaceType); -- cgit v1.2.3