+#include <initguid.h>

+#include <Ntddstor.h>

+BOOL EraseKeysOnShutdown = TRUE; // by default, we erase encryption keys on system shutdown

+static KeQueryInterruptTimePreciseFn KeQueryInterruptTimePrecisePtr = NULL;

+static KeSetSystemGroupAffinityThreadFn KeSetSystemGroupAffinityThreadPtr = NULL;

+static KeQueryActiveGroupCountFn KeQueryActiveGroupCountPtr = NULL;

+static KeQueryActiveProcessorCountExFn KeQueryActiveProcessorCountExPtr = NULL;

+int EncryptionIoRequestCount = 0;

+int EncryptionItemCount = 0;

+int EncryptionFragmentSize = 0;

+BOOL AlignValue (ULONG ulValue, ULONG ulAlignment, ULONG *pulResult)

+{

+ BOOL bRet = FALSE;

+ HRESULT hr;

+ if (ulAlignment == 0)

+ {

+ *pulResult = ulValue;

+ bRet = TRUE;

+ }

+ else

+ {

+ ulAlignment -= 1;

+ hr = ULongAdd (ulValue, ulAlignment, &ulValue);

+ if (S_OK == hr)

+ {

+ *pulResult = ulValue & (~ulAlignment);

+ bRet = TRUE;

+ }

+ }

+

+ return bRet;

+}

+

+ if (KeQueryInterruptTimePrecisePtr)

+ {

+ iSeed.QuadPart = KeQueryInterruptTimePrecisePtr (&iSeed2.QuadPart);

+ WHIRLPOOL_add ((unsigned char *) &(iSeed.QuadPart), sizeof(iSeed.QuadPart), &tctx);

+ WHIRLPOOL_add ((unsigned char *) &(iSeed2.QuadPart), sizeof(iSeed2.QuadPart), &tctx);

+ }

+ else

+ {

+ iSeed.QuadPart = KeQueryInterruptTime ();

+ WHIRLPOOL_add ((unsigned char *) &(iSeed.QuadPart), sizeof(iSeed.QuadPart), &tctx);

+ }

+ // KeQueryActiveGroupCount/KeQueryActiveProcessorCountEx/KeSetSystemGroupAffinityThread are available starting from Windows 7

+ UNICODE_STRING saveFuncName, restoreFuncName, groupCountFuncName, procCountFuncName, setAffinityFuncName;

+ RtlInitUnicodeString(&groupCountFuncName, L"KeQueryActiveGroupCount");

+ RtlInitUnicodeString(&procCountFuncName, L"KeQueryActiveProcessorCountEx");

+ RtlInitUnicodeString(&setAffinityFuncName, L"KeSetSystemGroupAffinityThread");

+ KeSetSystemGroupAffinityThreadPtr = (KeSetSystemGroupAffinityThreadFn) MmGetSystemRoutineAddress(&setAffinityFuncName);

+ KeQueryActiveGroupCountPtr = (KeQueryActiveGroupCountFn) MmGetSystemRoutineAddress(&groupCountFuncName);

+ KeQueryActiveProcessorCountExPtr = (KeQueryActiveProcessorCountExFn) MmGetSystemRoutineAddress(&procCountFuncName);

+ // KeQueryInterruptTimePrecise is available starting from Windows 8.1

+ if ((OsMajorVersion > 6) || (OsMajorVersion == 6 && OsMinorVersion >= 3))

+ {

+ UNICODE_STRING funcName;

+ RtlInitUnicodeString(&funcName, L"KeQueryInterruptTimePrecise");

+ KeQueryInterruptTimePrecisePtr = (KeQueryInterruptTimePreciseFn) MmGetSystemRoutineAddress(&funcName);

+ }

+

+ // Add 1MB to the disk size to emulate the geometry of a real MBR disk

+ outputBuffer->DiskSize.QuadPart = Extension->DiskLength + BYTES_PER_MB;

+ outputBuffer->StartingOffset.QuadPart = BYTES_PER_MB; // Set offset to 1MB to emulate the partition offset on a real MBR disk

+ outputBuffer->StartingOffset.QuadPart = BYTES_PER_MB; // Set offset to 1MB to emulate the partition offset on a real MBR disk

+ outputBuffer->PartitionEntry->StartingOffset.QuadPart = BYTES_PER_MB; // Set offset to 1MB to emulate the partition offset on a real MBR disk

+ outputBuffer->PartitionEntry->StartingOffset.QuadPart = BYTES_PER_MB; // Set offset to 1MB to emulate the partition offset on a real MBR disk

+ DWORD dwBuffersize = min (pVerifyInformation->Length, 16 * PAGE_SIZE);

+ PVOID buffer = TCalloc (dwBuffersize);

+ LARGE_INTEGER offset;

+ DWORD dwRemainingBytes = pVerifyInformation->Length, dwReadCount;

+ while (dwRemainingBytes)

+ {

+ dwReadCount = min (dwBuffersize, dwRemainingBytes);

+ Irp->IoStatus.Status = ZwReadFile (Extension->hDeviceFile, NULL, NULL, NULL, &ioStatus, buffer, dwReadCount, &offset, NULL);

+ if (NT_SUCCESS (Irp->IoStatus.Status) && ioStatus.Information != dwReadCount)

+ {

+ Irp->IoStatus.Status = STATUS_INVALID_PARAMETER;

+ break;

+ }

+ else if (!NT_SUCCESS (Irp->IoStatus.Status))

+ break;

+

+ dwRemainingBytes -= dwReadCount;

+ offset.QuadPart += (ULONGLONG) dwReadCount;

+ }

+

+ burn (buffer, dwBuffersize);

+ TCfree (buffer);

+ extents->Extents[0].StartingOffset.QuadPart = BYTES_PER_MB; // Set offset to 1MB to emulate the partition offset on a real MBR disk

+ capacity->DiskLength.QuadPart = Extension->DiskLength + BYTES_PER_MB; // Add 1MB to the disk size to emulate the geometry of a real MBR disk

+ capacity->NumberOfBlocks.QuadPart = capacity->DiskLength.QuadPart / capacity->BlockLength;

+ if (((ULONGLONG) inputLength) >= minSizeGeneric && ((ULONGLONG) inputLength) >= minSizedataSet && ((ULONGLONG) inputLength) >= minSizeParameter)

+ DWORD dwDataSetOffset;

+ if (AlignValue (inputLength, sizeof(DEVICE_DATA_SET_RANGE), &dwDataSetOffset))

+ Dump ("ProcessVolumeDeviceControlIrp: IOCTL_STORAGE_MANAGE_DATA_SET_ATTRIBUTES - DEVICE_DSM_FLAG_ENTIRE_DATA_SET_RANGE set. Setting data range to all volume.\n");

+ if (S_OK == ULongAdd(dwDataSetOffset, dwDataSetLength, &ulNewInputLength))

+ {

+ pNewSetAttrs = (PDEVICE_MANAGE_DATA_SET_ATTRIBUTES) TCalloc (ulNewInputLength);

+ if (pNewSetAttrs)

+ {

+ PDEVICE_DATA_SET_RANGE pRange = (PDEVICE_DATA_SET_RANGE) (((unsigned char*) pNewSetAttrs) + dwDataSetOffset);

+ memcpy (pNewSetAttrs, pInputAttrs, inputLength);

+ pRange->StartingOffset = (ULONGLONG) Extension->cryptoInfo->hiddenVolume ? Extension->cryptoInfo->hiddenVolumeOffset : Extension->cryptoInfo->volDataAreaOffset;

+ pRange->LengthInBytes = Extension->DiskLength;

+ pNewSetAttrs->Size = sizeof(DEVICE_MANAGE_DATA_SET_ATTRIBUTES);

+ pNewSetAttrs->Action = action;

+ pNewSetAttrs->Flags = pInputAttrs->Flags & (~DEVICE_DSM_FLAG_ENTIRE_DATA_SET_RANGE);

+ pNewSetAttrs->ParameterBlockOffset = pInputAttrs->ParameterBlockOffset;

+ pNewSetAttrs->ParameterBlockLength = pInputAttrs->ParameterBlockLength;

+ pNewSetAttrs->DataSetRangesOffset = dwDataSetOffset;

+ pNewSetAttrs->DataSetRangesLength = dwDataSetLength;

+

+ bForwardIoctl = TRUE;

+ }

+ else

+ {

+ Dump ("ProcessVolumeDeviceControlIrp: IOCTL_STORAGE_MANAGE_DATA_SET_ATTRIBUTES - Failed to allocate memory.\n");

+ Irp->IoStatus.Status = STATUS_INSUFFICIENT_RESOURCES;

+ Irp->IoStatus.Information = 0;

+ }

+ }

+ else

+ {

+ Dump ("ProcessVolumeDeviceControlIrp: IOCTL_STORAGE_MANAGE_DATA_SET_ATTRIBUTES - DEVICE_DSM_FLAG_ENTIRE_DATA_SET_RANGE set but data range length computation overflowed.\n");

+ Irp->IoStatus.Status = STATUS_INVALID_PARAMETER;

+ Irp->IoStatus.Information = 0;

+ }

+ Dump ("ProcessVolumeDeviceControlIrp: IOCTL_STORAGE_MANAGE_DATA_SET_ATTRIBUTES - DEVICE_DSM_FLAG_ENTIRE_DATA_SET_RANGE set but data set offset computation overflowed.\n");

+ Irp->IoStatus.Status = STATUS_INVALID_PARAMETER;

+

+ case VC_IOCTL_ENCRYPTION_QUEUE_PARAMS:

+ if (ValidateIOBufferSize (Irp, sizeof (EncryptionQueueParameters), ValidateOutput))

+ {

+ EncryptionQueueParameters* pParams = (EncryptionQueueParameters*) Irp->AssociatedIrp.SystemBuffer;

+ pParams->EncryptionFragmentSize = EncryptionFragmentSize;

+ pParams->EncryptionIoRequestCount = EncryptionIoRequestCount;

+ pParams->EncryptionItemCount = EncryptionItemCount;

+ Irp->IoStatus.Information = sizeof (EncryptionQueueParameters);

+ Irp->IoStatus.Status = STATUS_SUCCESS;

+ }

+ break;

+

+ if (bDevice && pThreadBlock->mount->bPartitionInInactiveSysEncScope

+ && (!Extension->cryptoInfo->hiddenVolume)

+ && (Extension->cryptoInfo->EncryptedAreaLength.Value != Extension->cryptoInfo->VolumeSize.Value)

+ )

+ {

+ // Support partial encryption only in the case of system encryption

+ Extension->Queue.EncryptedAreaStart = 0;

+ Extension->Queue.EncryptedAreaEnd = Extension->cryptoInfo->EncryptedAreaLength.Value - 1;

+ if (Extension->Queue.CryptoInfo->EncryptedAreaLength.Value == 0)

+ {

+ Extension->Queue.EncryptedAreaStart = -1;

+ Extension->Queue.EncryptedAreaEnd = -1;

+ }

+ Extension->Queue.bSupportPartialEncryption = TRUE;

+ }

+ TC_CASE_RET_NAME (VC_IOCTL_ENCRYPTION_QUEUE_PARAMS);

+ else if (ulCode == IOCTL_DISK_GROW_PARTITION)

+ return (LPWSTR) _T ("IOCTL_DISK_GROW_PARTITION");

+ VOID (*PsDereferenceImpersonationTokenD) (PACCESS_TOKEN ImpersonationToken);

+ UNICODE_STRING name;

+ RtlInitUnicodeString (&name, L"PsDereferenceImpersonationToken");

+ PsDereferenceImpersonationTokenD = MmGetSystemRoutineAddress (&name);

+ if (!PsDereferenceImpersonationTokenD)

+ TC_BUG_CHECK (STATUS_NOT_IMPLEMENTED);

+# define PsDereferencePrimaryToken

+# define PsDereferenceImpersonationToken PsDereferenceImpersonationTokenD

+ SeDeleteClientSecurity (&Extension->SecurityClientContext);

+# undef PsDereferencePrimaryToken

+# undef PsDereferenceImpersonationToken

+ ULONG sectorSize;

+ status = GetDeviceSectorSize (driveDeviceObject, &sectorSize);

+ if (!NT_SUCCESS (status))

+ return status;

+

+ sectorBuffer = TCalloc (sectorSize);

+ for (offset.QuadPart = sysLength.QuadPart; ; offset.QuadPart += sectorSize)

+ status = TCReadDevice (driveDeviceObject, sectorBuffer, offset, sectorSize);

+ status = TCWriteDevice (driveDeviceObject, sectorBuffer, offset, sectorSize);

+size_t GetCpuCount (WORD* pGroupCount)

+ if (KeQueryActiveGroupCountPtr && KeQueryActiveProcessorCountExPtr)

+ {

+ USHORT i, groupCount = KeQueryActiveGroupCountPtr ();

+ for (i = 0; i < groupCount; i++)

+ {

+ cpuCount += (size_t) KeQueryActiveProcessorCountExPtr (i);

+ }

+ if (pGroupCount)

+ *pGroupCount = groupCount;

+ }

+ else

+ KAFFINITY activeCpuMap = KeQueryActiveProcessors();

+ size_t mapSize = sizeof (activeCpuMap) * 8;

+

+ while (mapSize--)

+ {

+ if (activeCpuMap & 1)

+ ++cpuCount;

+

+ activeCpuMap >>= 1;

+ }

+ if (pGroupCount)

+ *pGroupCount = 1;

+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);

+ }

+}

+ waitInterval.QuadPart = ((LONGLONG)retryDelay) * -10000;

+ if (driverEntry && NT_SUCCESS (TCReadRegistryKey (&name, VC_ENCRYPTION_IO_REQUEST_COUNT, &data)))

+ {

+ if (data->Type == REG_DWORD)

+ EncryptionIoRequestCount = *(uint32 *) data->Data;

+

+ TCfree (data);

+ }

+

+ if (driverEntry && NT_SUCCESS (TCReadRegistryKey (&name, VC_ENCRYPTION_ITEM_COUNT, &data)))

+ {

+ if (data->Type == REG_DWORD)

+ EncryptionItemCount = *(uint32 *) data->Data;

+

+ TCfree (data);

+ }

+

+ if (driverEntry && NT_SUCCESS (TCReadRegistryKey (&name, VC_ENCRYPTION_FRAGMENT_SIZE, &data)))

+ {

+ if (data->Type == REG_DWORD)

+ EncryptionFragmentSize = *(uint32 *) data->Data;

+

+ TCfree (data);

+ }

+

+ if (driverEntry)

+ {

+ if (EncryptionIoRequestCount < TC_ENC_IO_QUEUE_PREALLOCATED_IO_REQUEST_COUNT)

+ EncryptionIoRequestCount = TC_ENC_IO_QUEUE_PREALLOCATED_IO_REQUEST_COUNT;

+ else if (EncryptionIoRequestCount > TC_ENC_IO_QUEUE_PREALLOCATED_IO_REQUEST_MAX_COUNT)

+ EncryptionIoRequestCount = TC_ENC_IO_QUEUE_PREALLOCATED_IO_REQUEST_MAX_COUNT;

+

+ if ((EncryptionItemCount == 0) || (EncryptionItemCount > (EncryptionIoRequestCount / 2)))

+ EncryptionItemCount = EncryptionIoRequestCount / 2;

+

+ /* EncryptionFragmentSize value in registry is expressed in KiB */

+ /* Maximum allowed value for EncryptionFragmentSize is 2048 KiB */

+ EncryptionFragmentSize *= 1024;

+ if (EncryptionFragmentSize == 0)

+ EncryptionFragmentSize = TC_ENC_IO_QUEUE_MAX_FRAGMENT_SIZE;

+ else if (EncryptionFragmentSize > (8 * TC_ENC_IO_QUEUE_MAX_FRAGMENT_SIZE))

+ EncryptionFragmentSize = 8 * TC_ENC_IO_QUEUE_MAX_FRAGMENT_SIZE;

+

+

+ }

+

+ if (driverEntry && NT_SUCCESS (TCReadRegistryKey (&name, VC_ERASE_KEYS_SHUTDOWN, &data)))

+ {

+ if (data->Type == REG_DWORD)

+ {

+ if (*((uint32 *) data->Data))

+ EraseKeysOnShutdown = TRUE;

+ else

+ EraseKeysOnShutdown = FALSE;

+ }

+

+ TCfree (data);

+ }

+

+

+NTSTATUS NTAPI KeSaveExtendedProcessorStateVC (

+VOID NTAPI KeRestoreExtendedProcessorStateVC (