diff options
| author | Mounir IDRASSI <mounir.idrassi@idrix.fr> | 2020-02-10 02:20:46 +0100 |
|---|---|---|
| committer | Mounir IDRASSI <mounir.idrassi@idrix.fr> | 2020-02-10 02:23:25 +0100 |
| commit | 94d3a1919c8eee4d7bfd7280ee4964477dee02e3 (patch) | |
| tree | 6b745e402e302dfaa9df976a51f569e0f62619ac /src/Driver/Ntdriver.c | |
| parent | c1d670fd752b96abac09b919d7ca1bff4ebafcbc (diff) | |
| download | VeraCrypt-94d3a1919c8eee4d7bfd7280ee4964477dee02e3.tar.gz VeraCrypt-94d3a1919c8eee4d7bfd7280ee4964477dee02e3.zip | |
Windows: Implement support for processor groups in the driver and fix build issues caused by previous implementation
Diffstat (limited to 'src/Driver/Ntdriver.c')
| -rw-r--r-- | src/Driver/Ntdriver.c | 66 |
1 files changed, 58 insertions, 8 deletions
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) { |