4 files changed, 138 insertions, 62 deletions
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 <Windows.h>
-#include <Versionhelpers.h>
-#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);