1 files changed, 124 insertions, 22 deletions
diff --git a/src/Driver/Ntdriver.c b/src/Driver/Ntdriver.c
index 2d025fa5..6f068a8f 100644
--- a/src/Driver/Ntdriver.c
+++ b/src/Driver/Ntdriver.c
@@ -148,12 +148,38 @@ static KeAreAllApcsDisabledFn KeAreAllApcsDisabledPtr = NULL;
 static KeSetSystemGroupAffinityThreadFn KeSetSystemGroupAffinityThreadPtr = NULL;
 static KeQueryActiveGroupCountFn KeQueryActiveGroupCountPtr = NULL;
 static KeQueryActiveProcessorCountExFn KeQueryActiveProcessorCountExPtr = NULL;
+int EncryptionIoRequestCount = 0;
+int EncryptionItemCount = 0;
+int EncryptionFragmentSize = 0;
 
 POOL_TYPE ExDefaultNonPagedPoolType = NonPagedPool;
 ULONG ExDefaultMdlProtection = 0;
 
 PDEVICE_OBJECT VirtualVolumeDeviceObjects[MAX_MOUNTED_VOLUME_DRIVE_NUMBER + 1];
 
+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;
+}
+
 BOOL IsUefiBoot ()
 {
 	BOOL bStatus = FALSE;
@@ -1690,7 +1716,7 @@ NTSTATUS ProcessVolumeDeviceControlIrp (PDEVICE_OBJECT DeviceObject, PEXTENSION
 				ULONG ulNewInputLength = 0;
 				BOOL bForwardIoctl = FALSE;
 
-				if (inputLength >= minSizeGeneric && inputLength >= minSizedataSet && inputLength >= minSizeParameter)
+				if (((ULONGLONG) inputLength) >= minSizeGeneric && ((ULONGLONG) inputLength) >= minSizedataSet && ((ULONGLONG) inputLength) >= minSizeParameter)
 				{
 					if (bEntireSet)
 					{
@@ -1702,36 +1728,53 @@ NTSTATUS ProcessVolumeDeviceControlIrp (PDEVICE_OBJECT DeviceObject, PEXTENSION
 						}
 						else
 						{
-							DWORD dwDataSetOffset = ALIGN_VALUE (inputLength, sizeof(DEVICE_DATA_SET_RANGE));
+							DWORD dwDataSetOffset;
 							DWORD dwDataSetLength = sizeof(DEVICE_DATA_SET_RANGE);
 
-							Dump ("ProcessVolumeDeviceControlIrp: IOCTL_STORAGE_MANAGE_DATA_SET_ATTRIBUTES - DEVICE_DSM_FLAG_ENTIRE_DATA_SET_RANGE set. Setting data range to all volume.\n");
-
-							ulNewInputLength = dwDataSetOffset + dwDataSetLength;
-							pNewSetAttrs = (PDEVICE_MANAGE_DATA_SET_ATTRIBUTES) TCalloc (ulNewInputLength);
-							if (pNewSetAttrs)
+							if (AlignValue (inputLength,  sizeof(DEVICE_DATA_SET_RANGE), &dwDataSetOffset))
 							{
-								PDEVICE_DATA_SET_RANGE pRange = (PDEVICE_DATA_SET_RANGE) (((unsigned char*) pNewSetAttrs) + dwDataSetOffset);
+								Dump ("ProcessVolumeDeviceControlIrp: IOCTL_STORAGE_MANAGE_DATA_SET_ATTRIBUTES - DEVICE_DSM_FLAG_ENTIRE_DATA_SET_RANGE set. Setting data range to all volume.\n");
 
-								memcpy (pNewSetAttrs, pInputAttrs, inputLength);
+								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;
+										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;
+										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;
+										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;
+								}
 							}
 							else
 							{
-								Dump ("ProcessVolumeDeviceControlIrp: IOCTL_STORAGE_MANAGE_DATA_SET_ATTRIBUTES - Failed to allocate memory.\n");
-								Irp->IoStatus.Status = STATUS_INSUFFICIENT_RESOURCES;
+								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;
 								Irp->IoStatus.Information = 0;
 							}
 						}
@@ -2843,6 +2886,18 @@ NTSTATUS ProcessMainDeviceControlIrp (PDEVICE_OBJECT DeviceObject, PEXTENSION Ex
 		}
 		break;
 
+	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;
+
 	default:
 		return TCCompleteIrp (Irp, STATUS_INVALID_DEVICE_REQUEST, 0);
 	}
@@ -3253,6 +3308,7 @@ LPWSTR TCTranslateCode (ULONG ulCode)
 		TC_CASE_RET_NAME (VC_IOCTL_GET_DRIVE_GEOMETRY_EX);
 		TC_CASE_RET_NAME (VC_IOCTL_EMERGENCY_CLEAR_ALL_KEYS);
 		TC_CASE_RET_NAME (VC_IOCTL_IS_RAM_ENCRYPTION_ENABLED);
+		TC_CASE_RET_NAME (VC_IOCTL_ENCRYPTION_QUEUE_PARAMS);
 
 		TC_CASE_RET_NAME (IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS);
 
@@ -4572,7 +4628,7 @@ void EnsureNullTerminatedString (wchar_t *str, size_t maxSizeInBytes)
 void *AllocateMemoryWithTimeout (size_t size, int retryDelay, int timeout)
 {
 	LARGE_INTEGER waitInterval;
-	waitInterval.QuadPart = retryDelay * -10000;
+	waitInterval.QuadPart = ((LONGLONG)retryDelay) * -10000;
 
 	ASSERT (KeGetCurrentIrql() <= APC_LEVEL);
 	ASSERT (retryDelay > 0 && retryDelay <= timeout);
@@ -4755,6 +4811,52 @@ NTSTATUS ReadRegistryConfigFlags (BOOL driverEntry)
 		TCfree (data);
 	}
 
+	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;
+		
+		
+	}
+
+
 	return status;
 }