diff options
Diffstat (limited to 'src/Common/EncryptionThreadPool.c')
-rw-r--r-- | src/Common/EncryptionThreadPool.c | 192 |
1 files changed, 152 insertions, 40 deletions
diff --git a/src/Common/EncryptionThreadPool.c b/src/Common/EncryptionThreadPool.c index 0dbaec5b..62b2cae8 100644 --- a/src/Common/EncryptionThreadPool.c +++ b/src/Common/EncryptionThreadPool.c @@ -20,6 +20,12 @@ #define TC_ENC_THREAD_POOL_MAX_THREAD_COUNT 256 //64 #define TC_ENC_THREAD_POOL_QUEUE_SIZE (TC_ENC_THREAD_POOL_MAX_THREAD_COUNT * 2) +#define TC_ENC_THREAD_POOL_LEGACY_MAX_THREAD_COUNT 64 +#define TC_ENC_THREAD_POOL_LEGACY_QUEUE_SIZE (TC_ENC_THREAD_POOL_LEGACY_MAX_THREAD_COUNT * 2) + +static volatile size_t ThreadPoolCount = TC_ENC_THREAD_POOL_LEGACY_MAX_THREAD_COUNT; +static volatile int ThreadQueueSize = TC_ENC_THREAD_POOL_LEGACY_QUEUE_SIZE; + #ifdef DEVICE_DRIVER #define TC_THREAD_HANDLE PKTHREAD @@ -102,6 +108,18 @@ typedef struct EncryptionThreadPoolWorkItemStruct char *Salt; } KeyDerivation; + + struct + { + TC_EVENT *KeyDerivationCompletedEvent; + TC_EVENT *NoOutstandingWorkItemEvent; + LONG *outstandingWorkItemCount; + void* keyInfoBuffer; + int keyInfoBufferSize; + void* keyDerivationWorkItems; + int keyDerivationWorkItemsSize; + + } ReadVolumeHeaderFinalization; }; } EncryptionThreadPoolWorkItem; @@ -178,19 +196,22 @@ static void SetWorkItemState (EncryptionThreadPoolWorkItem *workItem, WorkItemSt static TC_THREAD_PROC EncryptionThreadProc (void *threadArg) { EncryptionThreadPoolWorkItem *workItem; + if (threadArg) + { #ifdef DEVICE_DRIVER - SetThreadCpuGroupAffinity ((USHORT) *(WORD*)(threadArg)); + 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); - } + 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 + } while (!StopPending) @@ -199,7 +220,7 @@ static TC_THREAD_PROC EncryptionThreadProc (void *threadArg) workItem = &WorkItemQueue[DequeuePosition++]; - if (DequeuePosition >= TC_ENC_THREAD_POOL_QUEUE_SIZE) + if (DequeuePosition >= ThreadQueueSize) DequeuePosition = 0; while (!StopPending && GetWorkItemState (workItem) != WorkItemReady) @@ -227,8 +248,8 @@ static TC_THREAD_PROC EncryptionThreadProc (void *threadArg) case DeriveKeyWork: switch (workItem->KeyDerivation.Pkcs5Prf) { - case RIPEMD160: - derive_key_ripemd160 (workItem->KeyDerivation.Password, workItem->KeyDerivation.PasswordLength, workItem->KeyDerivation.Salt, PKCS5_SALT_SIZE, + case BLAKE2S: + derive_key_blake2s (workItem->KeyDerivation.Password, workItem->KeyDerivation.PasswordLength, workItem->KeyDerivation.Salt, PKCS5_SALT_SIZE, workItem->KeyDerivation.IterationCount, workItem->KeyDerivation.DerivedKey, GetMaxPkcs5OutSize()); break; @@ -266,6 +287,37 @@ static TC_THREAD_PROC EncryptionThreadProc (void *threadArg) TC_SET_EVENT (WorkItemCompletedEvent); continue; + case ReadVolumeHeaderFinalizationWork: + TC_WAIT_EVENT (*(workItem->ReadVolumeHeaderFinalization.NoOutstandingWorkItemEvent)); + + if (workItem->ReadVolumeHeaderFinalization.keyDerivationWorkItems) + { + burn (workItem->ReadVolumeHeaderFinalization.keyDerivationWorkItems, workItem->ReadVolumeHeaderFinalization.keyDerivationWorkItemsSize); +#if !defined(DEVICE_DRIVER) + VirtualUnlock (workItem->ReadVolumeHeaderFinalization.keyDerivationWorkItems, workItem->ReadVolumeHeaderFinalization.keyDerivationWorkItemsSize); +#endif + TCfree (workItem->ReadVolumeHeaderFinalization.keyDerivationWorkItems); + } + + if (workItem->ReadVolumeHeaderFinalization.keyInfoBuffer) + { + burn (workItem->ReadVolumeHeaderFinalization.keyInfoBuffer, workItem->ReadVolumeHeaderFinalization.keyInfoBufferSize); +#if !defined(DEVICE_DRIVER) + VirtualUnlock (workItem->ReadVolumeHeaderFinalization.keyInfoBuffer, workItem->ReadVolumeHeaderFinalization.keyInfoBufferSize); +#endif + TCfree (workItem->ReadVolumeHeaderFinalization.keyInfoBuffer); + } + +#if !defined(DEVICE_DRIVER) + CloseHandle (*(workItem->ReadVolumeHeaderFinalization.KeyDerivationCompletedEvent)); + CloseHandle (*(workItem->ReadVolumeHeaderFinalization.NoOutstandingWorkItemEvent)); +#endif + TCfree (workItem->ReadVolumeHeaderFinalization.KeyDerivationCompletedEvent); + TCfree (workItem->ReadVolumeHeaderFinalization.NoOutstandingWorkItemEvent); + TCfree (workItem->ReadVolumeHeaderFinalization.outstandingWorkItemCount); + SetWorkItemState (workItem, WorkItemFree); + TC_SET_EVENT (WorkItemCompletedEvent); + continue; default: TC_THROW_FATAL_EXCEPTION; } @@ -288,13 +340,11 @@ static TC_THREAD_PROC EncryptionThreadProc (void *threadArg) #endif } +#ifndef DEVICE_DRIVER -BOOL EncryptionThreadPoolStart (size_t encryptionFreeCpuCount) +size_t GetCpuCount (WORD* pGroupCount) { - size_t cpuCount = 0, i = 0; -#ifdef DEVICE_DRIVER - cpuCount = GetCpuCount(); -#else + size_t cpuCount = 0; SYSTEM_INFO sysInfo; GetActiveProcessorGroupCountFn GetActiveProcessorGroupCountPtr = (GetActiveProcessorGroupCountFn) GetProcAddress (GetModuleHandle (L"Kernel32.dll"), "GetActiveProcessorGroupCount"); GetActiveProcessorCountFn GetActiveProcessorCountPtr = (GetActiveProcessorCountFn) GetProcAddress (GetModuleHandle (L"Kernel32.dll"), "GetActiveProcessorCount"); @@ -307,25 +357,47 @@ BOOL EncryptionThreadPoolStart (size_t encryptionFreeCpuCount) totalProcessors += (size_t) GetActiveProcessorCountPtr(j); } cpuCount = totalProcessors; + if (pGroupCount) + *pGroupCount = groupCount; } else { GetSystemInfo(&sysInfo); - cpuCount = sysInfo.dwNumberOfProcessors; + cpuCount = (size_t) sysInfo.dwNumberOfProcessors; + if (pGroupCount) + *pGroupCount = 1; } + + return cpuCount; +} + #endif + +BOOL EncryptionThreadPoolStart (size_t encryptionFreeCpuCount) +{ + size_t cpuCount = 0, i = 0; + WORD groupCount = 1; + + cpuCount = GetCpuCount(&groupCount); + if (ThreadPoolRunning) return TRUE; + if (groupCount > 1) + { + ThreadPoolCount = TC_ENC_THREAD_POOL_MAX_THREAD_COUNT; + ThreadQueueSize = TC_ENC_THREAD_POOL_QUEUE_SIZE; + } + if (cpuCount > encryptionFreeCpuCount) cpuCount -= encryptionFreeCpuCount; if (cpuCount < 2) return TRUE; - if (cpuCount > TC_ENC_THREAD_POOL_MAX_THREAD_COUNT) - cpuCount = TC_ENC_THREAD_POOL_MAX_THREAD_COUNT; + if (cpuCount > ThreadPoolCount) + cpuCount = ThreadPoolCount; StopPending = FALSE; DequeuePosition = 0; @@ -377,32 +449,38 @@ BOOL EncryptionThreadPoolStart (size_t encryptionFreeCpuCount) for (ThreadCount = 0; ThreadCount < cpuCount; ++ThreadCount) { + WORD* pThreadArg = NULL; + if (groupCount > 1) + { #ifdef DEVICE_DRIVER - ThreadProcessorGroups[ThreadCount] = GetCpuGroup ((size_t) ThreadCount); + ThreadProcessorGroups[ThreadCount] = GetCpuGroup ((size_t) ThreadCount); #else - // Determine which processor group to bind the thread to. - if (GetActiveProcessorGroupCountPtr && GetActiveProcessorCountPtr) - { - WORD j, groupCount = GetActiveProcessorGroupCountPtr(); - uint32 totalProcessors = 0U; - for (j = 0U; j < groupCount; j++) + GetActiveProcessorCountFn GetActiveProcessorCountPtr = (GetActiveProcessorCountFn) GetProcAddress (GetModuleHandle (L"Kernel32.dll"), "GetActiveProcessorCount"); + // Determine which processor group to bind the thread to. + if (GetActiveProcessorCountPtr) { - totalProcessors += (uint32) GetActiveProcessorCountPtr(j); - if (totalProcessors >= ThreadCount) + WORD j; + uint32 totalProcessors = 0U; + for (j = 0U; j < groupCount; j++) { - ThreadProcessorGroups[ThreadCount] = j; - break; + totalProcessors += (uint32) GetActiveProcessorCountPtr(j); + if (totalProcessors >= ThreadCount) + { + ThreadProcessorGroups[ThreadCount] = j; + break; + } } } - } - else - ThreadProcessorGroups[ThreadCount] = 0; + else + ThreadProcessorGroups[ThreadCount] = 0; #endif + pThreadArg = &ThreadProcessorGroups[ThreadCount]; + } #ifdef DEVICE_DRIVER - if (!NT_SUCCESS(TCStartThread(EncryptionThreadProc, (void*)(&ThreadProcessorGroups[ThreadCount]), &ThreadHandles[ThreadCount]))) + if (!NT_SUCCESS(TCStartThread(EncryptionThreadProc, (void*) pThreadArg, &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*) pThreadArg, 0, NULL))) #endif { @@ -465,7 +543,7 @@ void EncryptionThreadPoolBeginKeyDerivation (TC_EVENT *completionEvent, TC_EVENT TC_ACQUIRE_MUTEX (&EnqueueMutex); workItem = &WorkItemQueue[EnqueuePosition++]; - if (EnqueuePosition >= TC_ENC_THREAD_POOL_QUEUE_SIZE) + if (EnqueuePosition >= ThreadQueueSize) EnqueuePosition = 0; while (GetWorkItemState (workItem) != WorkItemFree) @@ -493,6 +571,40 @@ void EncryptionThreadPoolBeginKeyDerivation (TC_EVENT *completionEvent, TC_EVENT TC_RELEASE_MUTEX (&EnqueueMutex); } +void EncryptionThreadPoolBeginReadVolumeHeaderFinalization (TC_EVENT *keyDerivationCompletedEvent, TC_EVENT *noOutstandingWorkItemEvent, LONG* outstandingWorkItemCount, + void* keyInfoBuffer, int keyInfoBufferSize, + void* keyDerivationWorkItems, int keyDerivationWorkItemsSize) +{ + EncryptionThreadPoolWorkItem *workItem; + + if (!ThreadPoolRunning) + TC_THROW_FATAL_EXCEPTION; + + TC_ACQUIRE_MUTEX (&EnqueueMutex); + + workItem = &WorkItemQueue[EnqueuePosition++]; + if (EnqueuePosition >= ThreadQueueSize) + EnqueuePosition = 0; + + while (GetWorkItemState (workItem) != WorkItemFree) + { + TC_WAIT_EVENT (WorkItemCompletedEvent); + } + + workItem->Type = ReadVolumeHeaderFinalizationWork; + workItem->ReadVolumeHeaderFinalization.NoOutstandingWorkItemEvent = noOutstandingWorkItemEvent; + workItem->ReadVolumeHeaderFinalization.KeyDerivationCompletedEvent = keyDerivationCompletedEvent; + workItem->ReadVolumeHeaderFinalization.keyInfoBuffer = keyInfoBuffer; + workItem->ReadVolumeHeaderFinalization.keyInfoBufferSize = keyInfoBufferSize; + workItem->ReadVolumeHeaderFinalization.keyDerivationWorkItems = keyDerivationWorkItems; + workItem->ReadVolumeHeaderFinalization.keyDerivationWorkItemsSize = keyDerivationWorkItemsSize; + workItem->ReadVolumeHeaderFinalization.outstandingWorkItemCount = outstandingWorkItemCount; + + SetWorkItemState (workItem, WorkItemReady); + TC_SET_EVENT (WorkItemReadyEvent); + TC_RELEASE_MUTEX (&EnqueueMutex); +} + void EncryptionThreadPoolDoWork (EncryptionThreadPoolWorkType type, byte *data, const UINT64_STRUCT *startUnitNo, uint32 unitCount, PCRYPTO_INFO cryptoInfo) { @@ -564,7 +676,7 @@ void EncryptionThreadPoolDoWork (EncryptionThreadPoolWorkType type, byte *data, while (fragmentCount-- > 0) { workItem = &WorkItemQueue[EnqueuePosition++]; - if (EnqueuePosition >= TC_ENC_THREAD_POOL_QUEUE_SIZE) + if (EnqueuePosition >= ThreadQueueSize) EnqueuePosition = 0; while (GetWorkItemState (workItem) != WorkItemFree) @@ -580,7 +692,7 @@ void EncryptionThreadPoolDoWork (EncryptionThreadPoolWorkType type, byte *data, workItem->Encryption.UnitCount = unitsPerFragment; workItem->Encryption.StartUnitNo.Value = fragmentStartUnitNo; - fragmentData += unitsPerFragment * ENCRYPTION_DATA_UNIT_SIZE; + fragmentData += ((uint64)unitsPerFragment) * ENCRYPTION_DATA_UNIT_SIZE; fragmentStartUnitNo += unitsPerFragment; if (remainder > 0 && --remainder == 0) @@ -606,7 +718,7 @@ size_t GetEncryptionThreadCount () size_t GetMaxEncryptionThreadCount () { - return TC_ENC_THREAD_POOL_MAX_THREAD_COUNT; + return ThreadPoolCount; } |