1 files changed, 61 insertions, 21 deletions
diff --git a/src/Driver/Ntdriver.c b/src/Driver/Ntdriver.c
index 2d025fa5..ca840fea 100644
--- a/src/Driver/Ntdriver.c
+++ b/src/Driver/Ntdriver.c
@@ -154,6 +154,29 @@ 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 +1713,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 +1725,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);
 
-								pRange->StartingOffset = (ULONGLONG) Extension->cryptoInfo->hiddenVolume ? Extension->cryptoInfo->hiddenVolumeOffset : Extension->cryptoInfo->volDataAreaOffset;
-								pRange->LengthInBytes = Extension->DiskLength;
+										memcpy (pNewSetAttrs, pInputAttrs, inputLength);
 
-								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;
+										pRange->StartingOffset = (ULONGLONG) Extension->cryptoInfo->hiddenVolume ? Extension->cryptoInfo->hiddenVolumeOffset : Extension->cryptoInfo->volDataAreaOffset;
+										pRange->LengthInBytes = Extension->DiskLength;
 
-								bForwardIoctl = TRUE;
+										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;
+								}
 							}
 							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;
 							}
 						}