1 files changed, 507 insertions, 0 deletions
diff --git a/Common/EncryptionThreadPool.c b/Common/EncryptionThreadPool.c
new file mode 100644
index 0000000..cb5b1ab
--- /dev/null
+++ b/Common/EncryptionThreadPool.c
@@ -0,0 +1,507 @@
+/*
+ Copyright (c) 2008-2010 TrueCrypt Developers Association. All rights reserved.
+
+ Governed by the TrueCrypt License 3.0 the full text of which is contained in
+ the file License.txt included in TrueCrypt binary and source code distribution
+ packages.
+*/
+
+#include "EncryptionThreadPool.h"
+#include "Pkcs5.h"
+#ifdef DEVICE_DRIVER
+#include "Driver/Ntdriver.h"
+#endif
+
+#define TC_ENC_THREAD_POOL_MAX_THREAD_COUNT 64
+#define TC_ENC_THREAD_POOL_QUEUE_SIZE (TC_ENC_THREAD_POOL_MAX_THREAD_COUNT * 2)
+
+#ifdef DEVICE_DRIVER
+
+#define TC_THREAD_HANDLE PKTHREAD
+#define TC_THREAD_PROC VOID
+
+#define TC_SET_EVENT(EVENT) KeSetEvent (&EVENT, IO_DISK_INCREMENT, FALSE)
+#define TC_CLEAR_EVENT(EVENT) KeClearEvent (&EVENT)
+
+#define TC_MUTEX FAST_MUTEX
+#define TC_ACQUIRE_MUTEX(MUTEX) ExAcquireFastMutex (MUTEX)
+#define TC_RELEASE_MUTEX(MUTEX) ExReleaseFastMutex (MUTEX)
+
+#else // !DEVICE_DRIVER
+
+#define TC_THREAD_HANDLE HANDLE
+#define TC_THREAD_PROC unsigned __stdcall
+
+#define TC_SET_EVENT(EVENT) SetEvent (EVENT)
+#define TC_CLEAR_EVENT(EVENT) ResetEvent (EVENT)
+
+#define TC_MUTEX HANDLE
+#define TC_ACQUIRE_MUTEX(MUTEX) WaitForSingleObject (*(MUTEX), INFINITE)
+#define TC_RELEASE_MUTEX(MUTEX) ReleaseMutex (*(MUTEX))
+
+#endif // !DEVICE_DRIVER
+
+
+typedef enum
+{
+	WorkItemFree,
+	WorkItemReady,
+	WorkItemBusy
+} WorkItemState;
+
+
+typedef struct EncryptionThreadPoolWorkItemStruct
+{
+	WorkItemState State;
+	EncryptionThreadPoolWorkType Type;
+
+	TC_EVENT ItemCompletedEvent;
+
+	struct EncryptionThreadPoolWorkItemStruct *FirstFragment;
+	LONG OutstandingFragmentCount;
+
+	union
+	{
+		struct
+		{
+			PCRYPTO_INFO CryptoInfo;
+			byte *Data;
+			UINT64_STRUCT StartUnitNo;
+			uint32 UnitCount;
+
+		} Encryption;
+
+		struct
+		{
+			TC_EVENT *CompletionEvent;
+			LONG *CompletionFlag;
+			char *DerivedKey;
+			int IterationCount;
+			TC_EVENT *NoOutstandingWorkItemEvent;
+			LONG *OutstandingWorkItemCount;
+			char *Password;
+			int PasswordLength;
+			int Pkcs5Prf;
+			char *Salt;
+
+		} KeyDerivation;
+	};
+
+} EncryptionThreadPoolWorkItem;
+
+
+static volatile BOOL ThreadPoolRunning = FALSE;
+static volatile BOOL StopPending = FALSE;
+
+static uint32 ThreadCount;
+static TC_THREAD_HANDLE ThreadHandles[TC_ENC_THREAD_POOL_MAX_THREAD_COUNT];
+
+static EncryptionThreadPoolWorkItem WorkItemQueue[TC_ENC_THREAD_POOL_QUEUE_SIZE];
+
+static volatile int EnqueuePosition;
+static volatile int DequeuePosition;
+
+static TC_MUTEX EnqueueMutex;
+static TC_MUTEX DequeueMutex;
+
+static TC_EVENT WorkItemReadyEvent;
+static TC_EVENT WorkItemCompletedEvent;
+
+
+static WorkItemState GetWorkItemState (EncryptionThreadPoolWorkItem *workItem)
+{
+	return InterlockedExchangeAdd ((LONG *) &workItem->State, 0);
+}
+
+
+static void SetWorkItemState (EncryptionThreadPoolWorkItem *workItem, WorkItemState newState)
+{
+	InterlockedExchange ((LONG *) &workItem->State, (LONG) newState);
+}
+
+
+static TC_THREAD_PROC EncryptionThreadProc (void *threadArg)
+{
+	EncryptionThreadPoolWorkItem *workItem;
+
+	while (!StopPending)
+	{
+		TC_ACQUIRE_MUTEX (&DequeueMutex);
+
+		workItem = &WorkItemQueue[DequeuePosition++];
+
+		if (DequeuePosition >= TC_ENC_THREAD_POOL_QUEUE_SIZE)
+			DequeuePosition = 0;
+
+		while (!StopPending && GetWorkItemState (workItem) != WorkItemReady)
+		{
+			TC_WAIT_EVENT (WorkItemReadyEvent);
+		}
+
+		SetWorkItemState (workItem, WorkItemBusy);
+
+		TC_RELEASE_MUTEX (&DequeueMutex);
+
+		if (StopPending)
+			break;
+
+		switch (workItem->Type)
+		{
+		case DecryptDataUnitsWork:
+			DecryptDataUnitsCurrentThread (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);
+			break;
+
+		case DeriveKeyWork:
+			switch (workItem->KeyDerivation.Pkcs5Prf)
+			{
+			case RIPEMD160:
+				derive_key_ripemd160 (workItem->KeyDerivation.Password, workItem->KeyDerivation.PasswordLength, workItem->KeyDerivation.Salt, PKCS5_SALT_SIZE,
+					workItem->KeyDerivation.IterationCount, workItem->KeyDerivation.DerivedKey, GetMaxPkcs5OutSize());
+				break;
+
+			case SHA512:
+				derive_key_sha512 (workItem->KeyDerivation.Password, workItem->KeyDerivation.PasswordLength, workItem->KeyDerivation.Salt, PKCS5_SALT_SIZE,
+					workItem->KeyDerivation.IterationCount, workItem->KeyDerivation.DerivedKey, GetMaxPkcs5OutSize());
+				break;
+
+			case WHIRLPOOL:
+				derive_key_whirlpool (workItem->KeyDerivation.Password, workItem->KeyDerivation.PasswordLength, workItem->KeyDerivation.Salt, PKCS5_SALT_SIZE,
+					workItem->KeyDerivation.IterationCount, workItem->KeyDerivation.DerivedKey, GetMaxPkcs5OutSize());
+				break;
+
+			case SHA1:
+				derive_key_sha1 (workItem->KeyDerivation.Password, workItem->KeyDerivation.PasswordLength, workItem->KeyDerivation.Salt, PKCS5_SALT_SIZE,
+					workItem->KeyDerivation.IterationCount, workItem->KeyDerivation.DerivedKey, GetMaxPkcs5OutSize());
+				break;
+
+			default:		
+				TC_THROW_FATAL_EXCEPTION;
+			} 
+
+			InterlockedExchange (workItem->KeyDerivation.CompletionFlag, TRUE);
+			TC_SET_EVENT (*workItem->KeyDerivation.CompletionEvent);
+			
+			if (InterlockedDecrement (workItem->KeyDerivation.OutstandingWorkItemCount) == 0)
+				TC_SET_EVENT (*workItem->KeyDerivation.NoOutstandingWorkItemEvent);
+
+			SetWorkItemState (workItem, WorkItemFree);
+			TC_SET_EVENT (WorkItemCompletedEvent);
+			continue;
+
+		default:
+			TC_THROW_FATAL_EXCEPTION;
+		}
+
+		if (workItem != workItem->FirstFragment)
+		{
+			SetWorkItemState (workItem, WorkItemFree);
+			TC_SET_EVENT (WorkItemCompletedEvent);
+		}
+
+		if (InterlockedDecrement (&workItem->FirstFragment->OutstandingFragmentCount) == 0)
+			TC_SET_EVENT (workItem->FirstFragment->ItemCompletedEvent);
+	}
+
+#ifdef DEVICE_DRIVER
+	PsTerminateSystemThread (STATUS_SUCCESS);
+#else
+	_endthreadex (0);
+    return 0;
+#endif
+}
+
+
+BOOL EncryptionThreadPoolStart (size_t encryptionFreeCpuCount)
+{
+	size_t cpuCount, i;
+
+	if (ThreadPoolRunning)
+		return TRUE;
+
+#ifdef DEVICE_DRIVER
+	cpuCount = GetCpuCount();
+#else
+	{
+		SYSTEM_INFO sysInfo;
+		GetSystemInfo (&sysInfo);
+		cpuCount = sysInfo.dwNumberOfProcessors;
+	}
+#endif
+
+	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;
+
+	StopPending = FALSE;
+	DequeuePosition = 0;
+	EnqueuePosition = 0;
+
+#ifdef DEVICE_DRIVER
+	KeInitializeEvent (&WorkItemReadyEvent, SynchronizationEvent, FALSE);
+	KeInitializeEvent (&WorkItemCompletedEvent, SynchronizationEvent, FALSE);
+#else
+	WorkItemReadyEvent = CreateEvent (NULL, FALSE, FALSE, NULL);
+	if (!WorkItemReadyEvent)
+		return FALSE;
+	
+	WorkItemCompletedEvent = CreateEvent (NULL, FALSE, FALSE, NULL);
+	if (!WorkItemCompletedEvent)
+		return FALSE;
+#endif
+	
+#ifdef DEVICE_DRIVER
+	ExInitializeFastMutex (&DequeueMutex);
+	ExInitializeFastMutex (&EnqueueMutex);
+#else
+	DequeueMutex = CreateMutex (NULL, FALSE, NULL);
+	if (!DequeueMutex)
+		return FALSE;
+
+	EnqueueMutex = CreateMutex (NULL, FALSE, NULL);
+	if (!EnqueueMutex)
+		return FALSE;
+#endif
+
+	memset (WorkItemQueue, 0, sizeof (WorkItemQueue));
+
+	for (i = 0; i < sizeof (WorkItemQueue) / sizeof (WorkItemQueue[0]); ++i)
+	{
+		WorkItemQueue[i].State = WorkItemFree;
+
+#ifdef DEVICE_DRIVER
+		KeInitializeEvent (&WorkItemQueue[i].ItemCompletedEvent, SynchronizationEvent, FALSE);
+#else
+		WorkItemQueue[i].ItemCompletedEvent = CreateEvent (NULL, FALSE, FALSE, NULL);
+		if (!WorkItemQueue[i].ItemCompletedEvent)
+		{
+			EncryptionThreadPoolStop();
+			return FALSE;
+		}
+#endif
+	}
+
+	for (ThreadCount = 0; ThreadCount < cpuCount; ++ThreadCount)
+	{
+#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;
+		}
+	}
+
+	ThreadPoolRunning = TRUE;
+	return TRUE;
+}
+
+
+void EncryptionThreadPoolStop ()
+{
+	size_t i;
+
+	if (!ThreadPoolRunning)
+		return;
+
+	StopPending = TRUE;
+	TC_SET_EVENT (WorkItemReadyEvent);
+
+	for (i = 0; i < ThreadCount; ++i)
+	{
+#ifdef DEVICE_DRIVER
+		TCStopThread (ThreadHandles[i], &WorkItemReadyEvent);
+#else
+		TC_WAIT_EVENT (ThreadHandles[i]);
+#endif
+	}
+
+	ThreadCount = 0;
+
+#ifndef DEVICE_DRIVER
+	CloseHandle (DequeueMutex);
+	CloseHandle (EnqueueMutex);
+
+	CloseHandle (WorkItemReadyEvent);
+	CloseHandle (WorkItemCompletedEvent);
+
+	for (i = 0; i < sizeof (WorkItemQueue) / sizeof (WorkItemQueue[0]); ++i)
+	{
+		if (WorkItemQueue[i].ItemCompletedEvent)
+			CloseHandle (WorkItemQueue[i].ItemCompletedEvent);
+	}
+#endif
+
+	ThreadPoolRunning = FALSE;
+}
+
+
+void EncryptionThreadPoolBeginKeyDerivation (TC_EVENT *completionEvent, TC_EVENT *noOutstandingWorkItemEvent, LONG *completionFlag, LONG *outstandingWorkItemCount, int pkcs5Prf, char *password, int passwordLength, char *salt, int iterationCount, char *derivedKey)
+{
+	EncryptionThreadPoolWorkItem *workItem;
+
+	if (!ThreadPoolRunning)
+		TC_THROW_FATAL_EXCEPTION;
+
+	TC_ACQUIRE_MUTEX (&EnqueueMutex);
+
+	workItem = &WorkItemQueue[EnqueuePosition++];
+	if (EnqueuePosition >= TC_ENC_THREAD_POOL_QUEUE_SIZE)
+		EnqueuePosition = 0;
+
+	while (GetWorkItemState (workItem) != WorkItemFree)
+	{
+		TC_WAIT_EVENT (WorkItemCompletedEvent);
+	}
+
+	workItem->Type = DeriveKeyWork;
+	workItem->KeyDerivation.CompletionEvent = completionEvent;
+	workItem->KeyDerivation.CompletionFlag = completionFlag;
+	workItem->KeyDerivation.DerivedKey = derivedKey;
+	workItem->KeyDerivation.IterationCount = iterationCount;
+	workItem->KeyDerivation.NoOutstandingWorkItemEvent = noOutstandingWorkItemEvent;
+	workItem->KeyDerivation.OutstandingWorkItemCount = outstandingWorkItemCount;
+	workItem->KeyDerivation.Password = password;
+	workItem->KeyDerivation.PasswordLength = passwordLength;
+	workItem->KeyDerivation.Pkcs5Prf = pkcs5Prf;
+	workItem->KeyDerivation.Salt = salt;
+
+	InterlockedIncrement (outstandingWorkItemCount);
+	TC_CLEAR_EVENT (*noOutstandingWorkItemEvent);
+
+	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)
+{
+	uint32 fragmentCount;
+	uint32 unitsPerFragment;
+	uint32 remainder;
+
+	byte *fragmentData;
+	uint64 fragmentStartUnitNo;
+
+	EncryptionThreadPoolWorkItem *workItem;
+	EncryptionThreadPoolWorkItem *firstFragmentWorkItem;
+	
+	if (unitCount == 0)
+		return;
+	
+	if (!ThreadPoolRunning || unitCount == 1)
+	{
+		switch (type)
+		{
+		case DecryptDataUnitsWork:
+			DecryptDataUnitsCurrentThread (data, startUnitNo, unitCount, cryptoInfo);
+			break;
+
+		case EncryptDataUnitsWork:
+			EncryptDataUnitsCurrentThread (data, startUnitNo, unitCount, cryptoInfo);
+			break;
+
+		default:
+			TC_THROW_FATAL_EXCEPTION;
+		}
+
+		return;
+	}
+
+	if (unitCount <= ThreadCount)
+	{
+		fragmentCount = unitCount;
+		unitsPerFragment = 1;
+		remainder = 0;
+	}
+	else
+	{
+		/* Note that it is not efficient to divide the data into fragments smaller than a few hundred bytes.
+		The reason is that the overhead associated with thread handling would in most cases make a multi-threaded 
+		process actually slower than a single-threaded process. */
+
+		fragmentCount = ThreadCount;
+		unitsPerFragment = unitCount / ThreadCount;
+		remainder = unitCount % ThreadCount;
+
+		if (remainder > 0)
+			++unitsPerFragment;
+	}
+	
+	fragmentData = data;
+	fragmentStartUnitNo = startUnitNo->Value;
+
+	TC_ACQUIRE_MUTEX (&EnqueueMutex);
+	firstFragmentWorkItem = &WorkItemQueue[EnqueuePosition];
+
+	while (GetWorkItemState (firstFragmentWorkItem) != WorkItemFree)
+	{
+		TC_WAIT_EVENT (WorkItemCompletedEvent);
+	}
+
+	firstFragmentWorkItem->OutstandingFragmentCount = fragmentCount;
+
+	while (fragmentCount-- > 0)
+	{
+		workItem = &WorkItemQueue[EnqueuePosition++];
+		if (EnqueuePosition >= TC_ENC_THREAD_POOL_QUEUE_SIZE)
+			EnqueuePosition = 0;
+
+		while (GetWorkItemState (workItem) != WorkItemFree)
+		{
+			TC_WAIT_EVENT (WorkItemCompletedEvent);
+		}
+
+		workItem->Type = type;
+		workItem->FirstFragment = firstFragmentWorkItem;
+
+		workItem->Encryption.CryptoInfo = cryptoInfo;
+		workItem->Encryption.Data = fragmentData;
+		workItem->Encryption.UnitCount = unitsPerFragment;
+		workItem->Encryption.StartUnitNo.Value = fragmentStartUnitNo;
+
+ 		fragmentData += unitsPerFragment * ENCRYPTION_DATA_UNIT_SIZE;
+		fragmentStartUnitNo += unitsPerFragment;
+
+		if (remainder > 0 && --remainder == 0)
+			--unitsPerFragment;
+
+		SetWorkItemState (workItem, WorkItemReady);
+		TC_SET_EVENT (WorkItemReadyEvent);
+	}
+
+	TC_RELEASE_MUTEX (&EnqueueMutex);
+
+	TC_WAIT_EVENT (firstFragmentWorkItem->ItemCompletedEvent);
+	SetWorkItemState (firstFragmentWorkItem, WorkItemFree);
+	TC_SET_EVENT (WorkItemCompletedEvent);
+}
+
+
+size_t GetEncryptionThreadCount ()
+{
+	return ThreadPoolRunning ? ThreadCount : 0;
+}
+
+
+size_t GetMaxEncryptionThreadCount ()
+{
+	return TC_ENC_THREAD_POOL_MAX_THREAD_COUNT;
+}
+
+
+BOOL IsEncryptionThreadPoolRunning ()
+{
+	return ThreadPoolRunning;
+}