1 files changed, 239 insertions, 27 deletions
diff --git a/src/Common/EncryptionThreadPool.c b/src/Common/EncryptionThreadPool.c
index 618e8474..79f1c890 100644
--- a/src/Common/EncryptionThreadPool.c
+++ b/src/Common/EncryptionThreadPool.c
@@ -16,9 +16,16 @@
 #include "Driver/Ntdriver.h"
 #endif
 
-#define TC_ENC_THREAD_POOL_MAX_THREAD_COUNT 64
+//Increasing the maximum number of threads 
+#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
@@ -43,6 +50,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
 
 
@@ -69,7 +88,7 @@ typedef struct EncryptionThreadPoolWorkItemStruct
 		struct
 		{
 			PCRYPTO_INFO CryptoInfo;
-			byte *Data;
+			uint8 *Data;
 			UINT64_STRUCT StartUnitNo;
 			uint32 UnitCount;
 
@@ -89,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;
@@ -99,6 +130,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];
 
@@ -111,6 +143,43 @@ static TC_MUTEX DequeueMutex;
 static TC_EVENT WorkItemReadyEvent;
 static TC_EVENT WorkItemCompletedEvent;
 
+#if defined(_WIN64)
+void EncryptDataUnitsCurrentThreadEx (unsigned __int8 *buf, const UINT64_STRUCT *structUnitNo, TC_LARGEST_COMPILER_UINT nbrUnits, PCRYPTO_INFO ci)
+{
+	if (IsRamEncryptionEnabled())
+	{
+		CRYPTO_INFO tmpCI;
+		memcpy (&tmpCI, ci, sizeof (CRYPTO_INFO));
+		VcUnprotectKeys (&tmpCI, VcGetEncryptionID (ci));
+
+		EncryptDataUnitsCurrentThread (buf, structUnitNo, nbrUnits, &tmpCI);
+
+		burn (&tmpCI, sizeof(CRYPTO_INFO));
+	}
+	else
+		EncryptDataUnitsCurrentThread (buf, structUnitNo, nbrUnits, ci);
+}
+
+void DecryptDataUnitsCurrentThreadEx (unsigned __int8 *buf, const UINT64_STRUCT *structUnitNo, TC_LARGEST_COMPILER_UINT nbrUnits, PCRYPTO_INFO ci)
+{
+	if (IsRamEncryptionEnabled())
+	{
+		CRYPTO_INFO tmpCI;
+		memcpy (&tmpCI, ci, sizeof (CRYPTO_INFO));
+		VcUnprotectKeys (&tmpCI, VcGetEncryptionID (ci));
+
+		DecryptDataUnitsCurrentThread (buf, structUnitNo, nbrUnits, &tmpCI);
+
+		burn (&tmpCI, sizeof(CRYPTO_INFO));
+	}
+	else
+		DecryptDataUnitsCurrentThread (buf, structUnitNo, nbrUnits, ci);
+}
+
+#else
+#define EncryptDataUnitsCurrentThreadEx EncryptDataUnitsCurrentThread
+#define DecryptDataUnitsCurrentThreadEx DecryptDataUnitsCurrentThread
+#endif
 
 static WorkItemState GetWorkItemState (EncryptionThreadPoolWorkItem *workItem)
 {
@@ -127,6 +196,23 @@ static void SetWorkItemState (EncryptionThreadPoolWorkItem *workItem, WorkItemSt
 static TC_THREAD_PROC EncryptionThreadProc (void *threadArg)
 {
 	EncryptionThreadPoolWorkItem *workItem;
+	if (threadArg)
+	{
+#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
+	}
+
 
 	while (!StopPending)
 	{
@@ -134,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)
@@ -152,18 +238,18 @@ static TC_THREAD_PROC EncryptionThreadProc (void *threadArg)
 		switch (workItem->Type)
 		{
 		case DecryptDataUnitsWork:
-			DecryptDataUnitsCurrentThread (workItem->Encryption.Data, &workItem->Encryption.StartUnitNo, workItem->Encryption.UnitCount, workItem->Encryption.CryptoInfo);
+			DecryptDataUnitsCurrentThreadEx (workItem->Encryption.Data, &workItem->Encryption.StartUnitNo, workItem->Encryption.UnitCount, workItem->Encryption.CryptoInfo);
 			break;
 
 		case EncryptDataUnitsWork:
-			EncryptDataUnitsCurrentThread (workItem->Encryption.Data, &workItem->Encryption.StartUnitNo, workItem->Encryption.UnitCount, workItem->Encryption.CryptoInfo);
+			EncryptDataUnitsCurrentThreadEx (workItem->Encryption.Data, &workItem->Encryption.StartUnitNo, workItem->Encryption.UnitCount, workItem->Encryption.CryptoInfo);
 			break;
 
 		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;
 
@@ -201,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;
 		}
@@ -223,23 +340,55 @@ static TC_THREAD_PROC EncryptionThreadProc (void *threadArg)
 #endif
 }
 
+#ifndef DEVICE_DRIVER
+
+size_t GetCpuCount (WORD* pGroupCount)
+{
+	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");
+	if (GetActiveProcessorGroupCountPtr && GetActiveProcessorCountPtr)
+	{
+		WORD j, groupCount = GetActiveProcessorGroupCountPtr();
+		size_t totalProcessors = 0;
+		for (j = 0; j < groupCount; ++j)
+		{
+			totalProcessors += (size_t) GetActiveProcessorCountPtr(j);
+		}
+		cpuCount = totalProcessors;
+		if (pGroupCount)
+			*pGroupCount = groupCount;
+	}
+	else
+	{
+		GetSystemInfo(&sysInfo);
+		cpuCount = (size_t) sysInfo.dwNumberOfProcessors;
+		if (pGroupCount)
+			*pGroupCount = 1;
+	}
+
+	return cpuCount;
+}
+
+#endif
+
 
 BOOL EncryptionThreadPoolStart (size_t encryptionFreeCpuCount)
 {
-	size_t cpuCount, i;
+	size_t cpuCount = 0, i = 0;
+	WORD groupCount = 1;
+
+	cpuCount = GetCpuCount(&groupCount);
 
 	if (ThreadPoolRunning)
 		return TRUE;
 
-#ifdef DEVICE_DRIVER
-	cpuCount = GetCpuCount();
-#else
+	if (groupCount > 1)
 	{
-		SYSTEM_INFO sysInfo;
-		GetSystemInfo (&sysInfo);
-		cpuCount = sysInfo.dwNumberOfProcessors;
+		ThreadPoolCount = TC_ENC_THREAD_POOL_MAX_THREAD_COUNT;
+		ThreadQueueSize = TC_ENC_THREAD_POOL_QUEUE_SIZE;
 	}
-#endif
 
 	if (cpuCount > encryptionFreeCpuCount)
 		cpuCount -= encryptionFreeCpuCount;
@@ -247,8 +396,8 @@ BOOL EncryptionThreadPoolStart (size_t 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;
@@ -300,11 +449,40 @@ BOOL EncryptionThreadPoolStart (size_t encryptionFreeCpuCount)
 
 	for (ThreadCount = 0; ThreadCount < cpuCount; ++ThreadCount)
 	{
+		WORD* pThreadArg = NULL;
+		if (groupCount > 1)
+		{
 #ifdef DEVICE_DRIVER
-		if (!NT_SUCCESS (TCStartThread (EncryptionThreadProc, NULL, &ThreadHandles[ThreadCount])))
+			ThreadProcessorGroups[ThreadCount] = GetCpuGroup ((size_t) ThreadCount);
 #else
-		if (!(ThreadHandles[ThreadCount] = (HANDLE) _beginthreadex (NULL, 0, EncryptionThreadProc, NULL, 0, NULL)))
+			GetActiveProcessorCountFn GetActiveProcessorCountPtr = (GetActiveProcessorCountFn) GetProcAddress (GetModuleHandle (L"Kernel32.dll"), "GetActiveProcessorCount");
+			// Determine which processor group to bind the thread to.
+			if (GetActiveProcessorCountPtr)
+			{
+				WORD j;
+				uint32 totalProcessors = 0U;
+				for (j = 0U; j < groupCount; j++)
+				{
+					totalProcessors += (uint32) GetActiveProcessorCountPtr(j);
+					if (totalProcessors >= ThreadCount)
+					{
+						ThreadProcessorGroups[ThreadCount] = j;
+						break;
+					}
+				}
+			}
+			else
+				ThreadProcessorGroups[ThreadCount] = 0;
 #endif
+			pThreadArg = &ThreadProcessorGroups[ThreadCount];
+		}
+
+#ifdef DEVICE_DRIVER
+		if (!NT_SUCCESS(TCStartThread(EncryptionThreadProc, (void*) pThreadArg, &ThreadHandles[ThreadCount])))
+#else
+		if (!(ThreadHandles[ThreadCount] = (HANDLE)_beginthreadex(NULL, 0, EncryptionThreadProc, (void*) pThreadArg, 0, NULL)))
+#endif
+
 		{
 			EncryptionThreadPoolStop();
 			return FALSE;
@@ -365,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)
@@ -393,14 +571,48 @@ 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)
+void EncryptionThreadPoolDoWork (EncryptionThreadPoolWorkType type, uint8 *data, const UINT64_STRUCT *startUnitNo, uint32 unitCount, PCRYPTO_INFO cryptoInfo)
 {
 	uint32 fragmentCount;
 	uint32 unitsPerFragment;
 	uint32 remainder;
 
-	byte *fragmentData;
+	uint8 *fragmentData;
 	uint64 fragmentStartUnitNo;
 
 	EncryptionThreadPoolWorkItem *workItem;
@@ -414,11 +626,11 @@ void EncryptionThreadPoolDoWork (EncryptionThreadPoolWorkType type, byte *data,
 		switch (type)
 		{
 		case DecryptDataUnitsWork:
-			DecryptDataUnitsCurrentThread (data, startUnitNo, unitCount, cryptoInfo);
+			DecryptDataUnitsCurrentThreadEx (data, startUnitNo, unitCount, cryptoInfo);
 			break;
 
 		case EncryptDataUnitsWork:
-			EncryptDataUnitsCurrentThread (data, startUnitNo, unitCount, cryptoInfo);
+			EncryptDataUnitsCurrentThreadEx (data, startUnitNo, unitCount, cryptoInfo);
 			break;
 
 		default:
@@ -464,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)
@@ -480,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)
@@ -506,7 +718,7 @@ size_t GetEncryptionThreadCount ()
 
 size_t GetMaxEncryptionThreadCount ()
 {
-	return TC_ENC_THREAD_POOL_MAX_THREAD_COUNT;
+	return ThreadPoolCount;
 }