/* Legal Notice: Some portions of the source code contained in this file were derived from the source code of TrueCrypt 7.1a, which is Copyright (c) 2003-2012 TrueCrypt Developers Association and which is governed by the TrueCrypt License 3.0, also from the source code of Encryption for the Masses 2.02a, which is Copyright (c) 1998-2000 Paul Le Roux and which is governed by the 'License Agreement for Encryption for the Masses' and also from the source code of extcv, which is Copyright (c) 2009-2010 Kih-Oskh or Copyright (c) 2012-2013 Josef Schneider Modifications and additions to the original source code (contained in this file) and all other portions of this file are Copyright (c) 2013-2016 IDRIX and are governed by the Apache License 2.0 the full text of which is contained in the file License.txt included in VeraCrypt binary and source code distribution packages. */ #include "Tcdefs.h" #include #include #include #include #include #include #include #include #include "Apidrvr.h" #include "Volumes.h" #include "Crypto.h" #include "Dlgcode.h" #include "Language.h" #include "Pkcs5.h" #include "Random.h" #include "Progress.h" #include "InitDataArea.h" #include "ExpandVolume.h" #include "Resource.h" #ifndef SRC_POS #define SRC_POS (__FUNCTION__ ":" TC_TO_STRING(__LINE__)) #endif #define DEBUG_EXPAND_VOLUME #ifdef DEBUG_EXPAND_VOLUME #define DebugAddProgressDlgStatus AddProgressDlgStatus #else #define DebugAddProgressDlgStatus(a,b) #endif HWND hCurPage; /* Handle to window with progress bar (used by FormatNoFs)*/ int nPbar; /* Control ID of progress bar (used by FormatNoFs) */ volatile BOOL bVolTransformThreadCancel = FALSE; /* TRUE if the user cancels/pauses volume expansion */ // internal functions static int UpdateVolumeHeaderHostSize (wchar_t *lpszVolume, Password *pVolumePassword, HWND hwndDlg, uint64 newHostSize, uint64 *pDataSize, BOOL initFreeSpace); static int FsctlExtendVolume(wchar_t * szVolume, LONGLONG nTotalSectors ); /* MountVolTemp Mounts a trucrypt volume temporarily (using any free drive number) Parameters: hwndDlg : HWND [in] handle to parent window volumePath : char * [in] Pointer to a string that contains the volume path driveNo : int * [out] returns the drive number (0='A',...) password : Password * [in] Pointer to the volume password Return value: int with Truecrypt error code (ERR_SUCCESS on success) */ int MountVolTemp (HWND hwndDlg, wchar_t *volumePath, int *driveNo, Password *password, int pkcs5, int pim) { MountOptions mountOptions; ZeroMemory (&mountOptions, sizeof (mountOptions)); *driveNo = GetLastAvailableDrive (); if (*driveNo == -1) { *driveNo = -2; return ERR_NO_FREE_DRIVES; } mountOptions.ReadOnly = FALSE; mountOptions.Removable = ConfigReadInt ("MountVolumesRemovable", FALSE); mountOptions.ProtectHiddenVolume = FALSE; mountOptions.PreserveTimestamp = bPreserveTimestamp; mountOptions.PartitionInInactiveSysEncScope = FALSE; mountOptions.UseBackupHeader = FALSE; if (MountVolume (hwndDlg, *driveNo, volumePath, password, pkcs5, pim, FALSE, FALSE, FALSE, TRUE, &mountOptions, FALSE, FALSE) < 1) { *driveNo = -3; return ERR_VOL_MOUNT_FAILED; } return 0; } /* FsctlExtendVolume Expands a volume by sending the FSCTL_EXTEND_VOLUME ioctl command to the volume Parameters: szVolume : char * [in] Pointer to a string that contains the volume GUID nTotalSectors : LONGLONG [in] specifies the total size of the volume, in sectors Return value: int with Truecrypt error code (ERR_SUCCESS on success) Remarks: only supported by NTFS and RAW file systems */ static int FsctlExtendVolume(wchar_t * szVolume, LONGLONG nTotalSectors ) { HANDLE hDevice; // handle to the volume to be extended BOOL bResult; // results flag DWORD nbytes; // discard results DWORD dwError; int nStatus = ERR_OS_ERROR; hDevice = CreateFile(szVolume, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL); if (hDevice == INVALID_HANDLE_VALUE) goto error; bResult = DeviceIoControl(hDevice, FSCTL_EXTEND_VOLUME, &nTotalSectors, sizeof(nTotalSectors), NULL, 0, &nbytes, (LPOVERLAPPED) NULL); if (bResult) nStatus = ERR_SUCCESS; error: dwError = GetLastError (); if (hDevice != INVALID_HANDLE_VALUE) CloseHandle (hDevice); SetLastError (dwError); return nStatus; } BOOL GetFileSystemType(const wchar_t *szFileName, enum EV_FileSystem *pFS) { wchar_t szFS[256]; wchar_t root[MAX_PATH]; *pFS = EV_FS_TYPE_RAW; if (!GetVolumePathName (szFileName, root, ARRAYSIZE (root))) return FALSE; if ( GetVolumeInformation (root, NULL, 0, NULL, NULL, NULL, szFS, ARRAYSIZE(szFS)) ) { if (!wcsncmp (szFS, L"NTFS", 4)) *pFS = EV_FS_TYPE_NTFS; else if (!wcsncmp (szFS, L"FAT", 3)) // FAT16, FAT32 *pFS = EV_FS_TYPE_FAT; else if (!_wcsnicmp (szFS, L"exFAT", 5)) // exFAT *pFS = EV_FS_TYPE_EXFAT; else *pFS = EV_FS_TYPE_RAW; } else { return FALSE; } return TRUE; } /* QueryVolumeInfo Retrieves the free disk space and file size limit on the truecrypt volume host Parameters: hwndDlg : HWND [in] handle to parent window lpszVolume : char * [in] Pointer to a string that contains the volume path pHostSizeFree : uint64 * [out] returns the free space available on the host (always zero for devices) pSizeLimitFS : uint64 * [out] returns the file size limit of the host file system Return value: int with TrueCrypt error code (ERR_SUCCESS on success) */ int QueryVolumeInfo (HWND hwndDlg, const wchar_t *lpszVolume, uint64 * pHostSizeFree, uint64 * pSizeLimitFS ) { int nStatus = ERR_OS_ERROR; wchar_t szDiskFile[TC_MAX_PATH], root[MAX_PATH]; BOOL bDevice; enum EV_FileSystem fs; *pSizeLimitFS = (uint64)-1; CreateFullVolumePath (szDiskFile, sizeof(szDiskFile), lpszVolume, &bDevice); if (bDevice) { *pHostSizeFree=0; return ERR_SUCCESS; } if (!GetVolumePathName (szDiskFile, root, ARRAYSIZE (root))) { nStatus = ERR_OS_ERROR; goto error; } if( ! GetDiskFreeSpaceEx (root,(PULARGE_INTEGER)pHostSizeFree,NULL,NULL) ) { nStatus = ERR_OS_ERROR; goto error; } if ( ! GetFileSystemType(root,&fs) ) { nStatus = ERR_OS_ERROR; goto error; } /* file size limits FAT16 / FAT32 : 4 GB minus 1 byte (2^32 bytes minus 1 byte) exFAT: 128 PiB − 1 byte NTFS : Architecturally : 16 exabytes minus 1 KB (26^4 bytes minus 1 KB) Implementation (Windows Server 2008): 16 terabytes minus 64 KB (2^44 bytes minus 64 KB) */ switch (fs) { case EV_FS_TYPE_NTFS: *pSizeLimitFS = 16 * BYTES_PER_TB - 64 * BYTES_PER_KB; break; case EV_FS_TYPE_EXFAT: *pSizeLimitFS = 128 * BYTES_PER_PB - 1; break; case EV_FS_TYPE_FAT: *pSizeLimitFS = 4 * BYTES_PER_GB - 1; break; default: *pSizeLimitFS = (uint64)-1; } nStatus = ERR_SUCCESS; error: return nStatus; } BOOL GetNtfsNumberOfSectors(wchar_t * rootPath, uint64 * pNumberOfSectors, DWORD *pBytesPerSector) { HANDLE hDevice; BOOL bResult; DWORD nbytes, dwError; size_t len; NTFS_VOLUME_DATA_BUFFER ntfsvdb; wchar_t szVolumeGUID[128]; // get volume name if (!GetVolumeNameForVolumeMountPoint(rootPath,szVolumeGUID,ARRAYSIZE(szVolumeGUID))) { return FALSE; } // strip trailing backslash from volume GUID (otherwise it means root dir) len = wcslen(szVolumeGUID); if (len>0) --len; if (szVolumeGUID[len]==L'\\') szVolumeGUID[len]=0; hDevice = CreateFile(szVolumeGUID, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL); if (hDevice == INVALID_HANDLE_VALUE) return (FALSE); bResult = DeviceIoControl(hDevice, FSCTL_GET_NTFS_VOLUME_DATA, NULL, 0, &ntfsvdb, sizeof(ntfsvdb), &nbytes, (LPOVERLAPPED) NULL); if (bResult) { if (pNumberOfSectors) *pNumberOfSectors = ntfsvdb.NumberSectors.QuadPart; if (pBytesPerSector) *pBytesPerSector = ntfsvdb.BytesPerSector; } dwError = GetLastError (); CloseHandle(hDevice); SetLastError (dwError); return (bResult); } uint64 GetVolumeDataAreaSize (uint64 volumeSize, BOOL legacyVolume) { uint64 reservedSize; if (legacyVolume) reservedSize = TC_VOLUME_HEADER_SIZE_LEGACY; else reservedSize = TC_TOTAL_VOLUME_HEADERS_SIZE; if (volumeSize < reservedSize) return 0; return volumeSize - reservedSize; } uint64 GetVolumeSizeByDataAreaSize (uint64 dataAreaSize, BOOL legacyVolume) { uint64 reservedSize; if (legacyVolume) reservedSize = TC_VOLUME_HEADER_SIZE_LEGACY; else reservedSize = TC_TOTAL_VOLUME_HEADERS_SIZE; return dataAreaSize + reservedSize; } int ExtendFileSystem (HWND hwndDlg , wchar_t *lpszVolume, Password *pVolumePassword, int VolumePkcs5, int VolumePim, uint64 newDataAreaSize) { wchar_t szVolumeGUID[128]; int driveNo = -1; wchar_t rootPath[] = L"A:\\"; enum EV_FileSystem fs; DWORD dwError; int nStatus = ERR_SUCCESS; DWORD BytesPerSector; // mount and resize file system DebugAddProgressDlgStatus (hwndDlg, L"Mounting volume ...\r\n"); nStatus=MountVolTemp(hwndDlg, lpszVolume, &driveNo, pVolumePassword, VolumePkcs5, VolumePim); if (nStatus!=ERR_SUCCESS) { driveNo = -1; goto error; } rootPath[0] += driveNo; if ( !GetFileSystemType(rootPath,&fs) ) { dwError = GetLastError(); if (dwError == ERROR_UNRECOGNIZED_VOLUME) { // raw volume with unrecognized file system -> return with no error nStatus = ERR_SUCCESS; goto error; } nStatus = ERR_OS_ERROR; goto error; } if (fs != EV_FS_TYPE_RAW && fs != EV_FS_TYPE_NTFS ) { // FsctlExtendVolume only supports NTFS and RAW -> return with no error nStatus = ERR_SUCCESS; goto error; } // Get volume GUID if (!GetVolumeNameForVolumeMountPoint(rootPath,szVolumeGUID,ARRAYSIZE(szVolumeGUID))) { nStatus = ERR_OS_ERROR; goto error; } else { // strip trailing backslash from volume GUID (otherwise it means root dir) size_t len = wcslen(szVolumeGUID); if (len>0) --len; if (szVolumeGUID[len]==L'\\') szVolumeGUID[len]=0; } // Get Sector Size if ( !GetNtfsNumberOfSectors(rootPath, NULL, &BytesPerSector) ) { nStatus = ERR_OS_ERROR; goto error; } DebugAddProgressDlgStatus (hwndDlg, L"Extending file system ...\r\n"); // extend volume nStatus = FsctlExtendVolume(szVolumeGUID, newDataAreaSize/BytesPerSector ); error: dwError = GetLastError(); if (driveNo>=0) { DebugAddProgressDlgStatus (hwndDlg, L"Unmounting volume ...\r\n"); UnmountVolume (hwndDlg, driveNo, TRUE); } SetLastError (dwError); return nStatus; } /* ExpandVolume Sets the volume size in the volume header (and backup header) to a larger value, and resizes the filesystem within the volume (only NTFS supported) Parameters: hwndDlg : HWND [in] handle to progress dialog lpszVolume : char * [in] Pointer to a string that contains the path to the truecrypt volume pVolumePassword : Password * [in] Pointer to the volume password newHostSize : uint64 [in] new value of the volume host size (can be zero for devices, which means the volume should use all space of the host device) initFreeSpace : BOOL [in] if true, the new volume space will be initalized with random data Return value: int with Truecrypt error code (ERR_SUCCESS on success) Remarks: a lot of code is from TrueCrypt 'Common\Password.c' :: ChangePwd() */ static int ExpandVolume (HWND hwndDlg, wchar_t *lpszVolume, Password *pVolumePassword, int VolumePkcs5, int VolumePim, uint64 newHostSize, BOOL initFreeSpace) { int nDosLinkCreated = 1, nStatus = ERR_OS_ERROR; wchar_t szDiskFile[TC_MAX_PATH], szCFDevice[TC_MAX_PATH]; wchar_t szDosDevice[TC_MAX_PATH]; char buffer[TC_VOLUME_HEADER_EFFECTIVE_SIZE]; PCRYPTO_INFO cryptoInfo = NULL, ci = NULL; void *dev = INVALID_HANDLE_VALUE; DWORD dwError; BOOL bDevice; uint64 hostSize=0, newDataAreaSize, currentVolSize; DWORD HostSectorSize; FILETIME ftCreationTime; FILETIME ftLastWriteTime; FILETIME ftLastAccessTime; BOOL bTimeStampValid = FALSE; LARGE_INTEGER headerOffset; BOOL backupHeader; byte *wipeBuffer = NULL; uint32 workChunkSize = TC_VOLUME_HEADER_GROUP_SIZE; if (pVolumePassword->Length == 0) return -1; WaitCursor (); CreateFullVolumePath (szDiskFile, sizeof(szDiskFile), lpszVolume, &bDevice); if (bDevice == FALSE) { wcscpy (szCFDevice, szDiskFile); } else { nDosLinkCreated = FakeDosNameForDevice (szDiskFile, szDosDevice, sizeof(szDosDevice), szCFDevice, sizeof(szCFDevice), FALSE); if (nDosLinkCreated != 0) // note: nStatus == ERR_OS_ERROR goto error; } dev = CreateFile (szCFDevice, GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL); if (dev == INVALID_HANDLE_VALUE) goto error; if (bDevice) { /* This is necessary to determine the hidden volume header offset */ if (dev == INVALID_HANDLE_VALUE) { goto error; } else { PARTITION_INFORMATION diskInfo; DWORD dwResult; BOOL bResult; bResult = GetPartitionInfo (lpszVolume, &diskInfo); if (bResult) { hostSize = diskInfo.PartitionLength.QuadPart; HostSectorSize = TC_SECTOR_SIZE_FILE_HOSTED_VOLUME; //TO DO: get the real host disk sector size } else { DISK_GEOMETRY driveInfo; bResult = DeviceIoControl (dev, IOCTL_DISK_GET_DRIVE_GEOMETRY, NULL, 0, &driveInfo, sizeof (driveInfo), &dwResult, NULL); if (!bResult) goto error; hostSize = driveInfo.Cylinders.QuadPart * driveInfo.BytesPerSector * driveInfo.SectorsPerTrack * driveInfo.TracksPerCylinder; HostSectorSize = driveInfo.BytesPerSector; } if (hostSize == 0) { nStatus = ERR_VOL_SIZE_WRONG; goto error; } } } else { LARGE_INTEGER fileSize; if (!GetFileSizeEx (dev, &fileSize)) { nStatus = ERR_OS_ERROR; goto error; } hostSize = fileSize.QuadPart; HostSectorSize = TC_SECTOR_SIZE_FILE_HOSTED_VOLUME; //TO DO: get the real host disk sector size } if (Randinit ()) { if (CryptoAPILastError == ERROR_SUCCESS) nStatus = ERR_RAND_INIT_FAILED; else nStatus = ERR_CAPI_INIT_FAILED; goto error; } if (!bDevice && bPreserveTimestamp) { /* Remember the container modification/creation date and time, (used to reset file date and time of file-hosted volumes after password change (or attempt to), in order to preserve plausible deniability of hidden volumes (last password change time is stored in the volume header). */ if (GetFileTime ((HANDLE) dev, &ftCreationTime, &ftLastAccessTime, &ftLastWriteTime) == 0) { bTimeStampValid = FALSE; MessageBoxW (hwndDlg, GetString ("GETFILETIME_FAILED_PW"), lpszTitle, MB_OK | MB_ICONEXCLAMATION); } else bTimeStampValid = TRUE; } // Seek the volume header headerOffset.QuadPart = TC_VOLUME_HEADER_OFFSET; if (!SetFilePointerEx ((HANDLE) dev, headerOffset, NULL, FILE_BEGIN)) { nStatus = ERR_OS_ERROR; goto error; } /* Read in volume header */ nStatus = _lread ((HFILE) dev, buffer, sizeof (buffer)); if (nStatus != sizeof (buffer)) { // Windows may report EOF when reading sectors from the last cluster of a device formatted as NTFS memset (buffer, 0, sizeof (buffer)); } /* Try to decrypt the header */ nStatus = ReadVolumeHeader (FALSE, buffer, pVolumePassword, VolumePkcs5, VolumePim, FALSE, &cryptoInfo, NULL); if (nStatus == ERR_CIPHER_INIT_WEAK_KEY) nStatus = 0; // We can ignore this error here if (nStatus != 0) { cryptoInfo = NULL; goto error; } if (cryptoInfo->HeaderFlags & TC_HEADER_FLAG_ENCRYPTED_SYSTEM) { nStatus = ERR_SYS_HIDVOL_HEAD_REENC_MODE_WRONG; goto error; } if (bDevice && newHostSize == 0) { // this means we shall take all host space as new volume size newHostSize = hostSize; } if ( newHostSize % cryptoInfo->SectorSize != 0 || newHostSize > TC_MAX_VOLUME_SIZE || (bDevice && newHostSize > hostSize) ) { // 1. must be multiple of sector size // 2. truecrypt volume size limit // 3. for devices volume size can't be larger than host size cryptoInfo = NULL; nStatus = ERR_PARAMETER_INCORRECT; goto error; } newDataAreaSize = GetVolumeDataAreaSize (newHostSize, cryptoInfo->LegacyVolume); if (cryptoInfo->LegacyVolume) { if (bDevice) { if (initFreeSpace) { // unsupported cryptoInfo = NULL; nStatus = ERR_PARAMETER_INCORRECT; goto error; } else { // note: dummy value (only used for parameter checks) cryptoInfo->VolumeSize.Value = newDataAreaSize - TC_MINVAL_FS_EXPAND; } } else { cryptoInfo->VolumeSize.Value = GetVolumeDataAreaSize (hostSize, TRUE); } } currentVolSize = GetVolumeSizeByDataAreaSize (cryptoInfo->VolumeSize.Value, cryptoInfo->LegacyVolume); if ( newDataAreaSize < cryptoInfo->VolumeSize.Value + TC_MINVAL_FS_EXPAND ) { // shrinking a volume or enlarging by less then TC_MINVAL_FS_EXPAND is not allowed cryptoInfo = NULL; nStatus = ERR_PARAMETER_INCORRECT; goto error; } InitProgressBar ( newHostSize, currentVolSize, FALSE, FALSE, FALSE, TRUE); if (bVolTransformThreadCancel) { SetLastError(0); nStatus = ERR_USER_ABORT; goto error; } if (!bDevice) { LARGE_INTEGER liNewSize; liNewSize.QuadPart=(LONGLONG)newHostSize; // Preallocate the file if (!SetFilePointerEx (dev, liNewSize, NULL, FILE_BEGIN) || !SetEndOfFile (dev) || SetFilePointer (dev, 0, NULL, FILE_BEGIN) != 0) { nStatus = ERR_OS_ERROR; goto error; } } if (initFreeSpace) { uint64 startSector; int64 num_sectors; // fill new space with random data startSector = currentVolSize/HostSectorSize ; num_sectors = (newHostSize/HostSectorSize) - startSector; if (bDevice && !StartFormatWriteThread()) { nStatus = ERR_OS_ERROR; goto error; } DebugAddProgressDlgStatus(hwndDlg, L"Writing random data to new space ...\r\n"); SetFormatSectorSize(HostSectorSize); nStatus = FormatNoFs (hwndDlg, startSector, num_sectors, dev, cryptoInfo, FALSE); dwError = GetLastError(); StopFormatWriteThread(); SetLastError (dwError); } else { UpdateProgressBar(newHostSize); } if (nStatus != ERR_SUCCESS) { dwError = GetLastError(); DebugAddProgressDlgStatus(hwndDlg, L"Error: failed to write random data ...\r\n"); if ( !bDevice ) { // restore original size of the container file LARGE_INTEGER liOldSize; liOldSize.QuadPart=(LONGLONG)hostSize; if (!SetFilePointerEx (dev, liOldSize, NULL, FILE_BEGIN) || !SetEndOfFile (dev)) { DebugAddProgressDlgStatus(hwndDlg, L"Warning: failed to restore original size of the container file\r\n"); } } SetLastError (dwError); goto error; } RandSetHashFunction (cryptoInfo->pkcs5); // Re-encrypt the volume header forn non-legacy volumes: backup header first backupHeader = TRUE; headerOffset.QuadPart = TC_VOLUME_HEADER_OFFSET + newHostSize - TC_VOLUME_HEADER_GROUP_SIZE; /* note: updating the header is not neccessary for legay volumes */ while ( !cryptoInfo->LegacyVolume ) { if (backupHeader) DebugAddProgressDlgStatus(hwndDlg, L"Writing re-encrypted backup header ...\r\n"); else DebugAddProgressDlgStatus(hwndDlg, L"Writing re-encrypted primary header ...\r\n"); // Prepare new volume header nStatus = CreateVolumeHeaderInMemory (hwndDlg, FALSE, buffer, cryptoInfo->ea, cryptoInfo->mode, pVolumePassword, cryptoInfo->pkcs5, VolumePim, (char*)(cryptoInfo->master_keydata), &ci, newDataAreaSize, 0, // hiddenVolumeSize cryptoInfo->EncryptedAreaStart.Value, newDataAreaSize, cryptoInfo->RequiredProgramVersion, cryptoInfo->HeaderFlags, cryptoInfo->SectorSize, FALSE ); // use slow poll if (ci != NULL) crypto_close (ci); if (nStatus != 0) goto error; if (!SetFilePointerEx ((HANDLE) dev, headerOffset, NULL, FILE_BEGIN)) { nStatus = ERR_OS_ERROR; goto error; } if (!WriteEffectiveVolumeHeader (bDevice, dev, buffer)) { nStatus = ERR_OS_ERROR; goto error; } if ( ( backupHeader && !initFreeSpace ) || ( bDevice && !cryptoInfo->LegacyVolume && !cryptoInfo->hiddenVolume && cryptoInfo->HeaderVersion == 4 // BUG in TrueCrypt: doing this only for v4 make no sense && (cryptoInfo->HeaderFlags & TC_HEADER_FLAG_NONSYS_INPLACE_ENC) != 0 && (cryptoInfo->HeaderFlags & ~TC_HEADER_F
/*
 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 SHA256:
				derive_key_sha256 (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;
}