nihilist/VeraCrypt
1
0
Fork 0
mirror of https://github.com/veracrypt/VeraCrypt synced 2024-11-28 05:53:29 +01:00
Code Issues Packages Projects Releases Wiki Activity

Windows: Implement support for processor groups in the driver and fix build issues caused by previous implementation

Browse Source
This commit is contained in:
Mounir IDRASSI 2020-02-10 02:20:46 +01:00
parent c1d670fd75
commit 94d3a1919c
No known key found for this signature in database
GPG Key ID: 02C30AE90FAE4A6F
4 changed files with 139 additions and 63 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

119
src/Common/EncryptionThreadPool.c
View File

@ -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)
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)
{
processors = GetActiveProcessorCount(i);
totalProcessors += processors;
}
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 defined(_WIN32) || defined(WIN32) || defined(_WIN64) || defined(WIN64)
SYSTEM_INFO sysInfo;
GetSystemInfo(&sysInfo);
cpuCount = (IsWindows7OrGreater) ? totalProcessors : sysInfo.dwNumberOfProcessors;
#endif
/*
#ifdef DEVICE_DRIVER
cpuCount = GetCpuCount();
#else
{
SYSTEM_INFO sysInfo;
GetSystemInfo (&sysInfo);
cpuCount = sysInfo.dwNumberOfProcessors;
}
#endif
*/
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;

15
src/Common/Tcdefs.h
View File

@ -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

66
src/Driver/Ntdriver.c
View File

@ -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)
{
if (activeCpuMap & 1)
++cpuCount;
USHORT i, groupCount = KeQueryActiveGroupCountPtr ();
for (i = 0; i < groupCount; i++)
{
cpuCount += (size_t) KeQueryActiveProcessorCountExPtr (i);
}
}
else
{
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)
{

2
src/Driver/Ntdriver.h
View File

@ -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);