From 74b82118d5f77116ec5f4a1e2438cdc77cead40f Mon Sep 17 00:00:00 2001
From: Mounir IDRASSI <mounir.idrassi@idrix.fr>
Date: Mon, 15 May 2017 16:22:48 +0200
Subject: Windows: use IOCTL_DISK_GET_DRIVE_GEOMETRY_EX instead of the
 deprecated IOCTL_DISK_GET_DRIVE_GEOMETRY in order to get accurate disk size
 value.

---
 src/Driver/Ntdriver.c | 35 ++++++++++++++++++++++++++++++++---
 src/Driver/Ntvol.c    | 10 +++++-----
 2 files changed, 37 insertions(+), 8 deletions(-)

(limited to 'src/Driver')

diff --git a/src/Driver/Ntdriver.c b/src/Driver/Ntdriver.c
index eaa8fa8c..ea9d725b 100644
--- a/src/Driver/Ntdriver.c
+++ b/src/Driver/Ntdriver.c
@@ -1483,6 +1483,34 @@ NTSTATUS ProcessMainDeviceControlIrp (PDEVICE_OBJECT DeviceObject, PEXTENSION Ex
 		}
 		break;
 
+	case VC_IOCTL_GET_DRIVE_GEOMETRY_EX:
+		if (ValidateIOBufferSize (Irp, sizeof (DISK_GEOMETRY_EX_STRUCT), ValidateInputOutput))
+		{
+			DISK_GEOMETRY_EX_STRUCT *g = (DISK_GEOMETRY_EX_STRUCT *) Irp->AssociatedIrp.SystemBuffer;
+			{
+				NTSTATUS ntStatus;
+				DISK_GEOMETRY_EX geo = {0};
+
+				memset (g, 0, sizeof (DISK_GEOMETRY_EX_STRUCT));
+
+				EnsureNullTerminatedString (g->deviceName, sizeof (g->deviceName));
+
+				ntStatus = TCDeviceIoControl (g->deviceName,
+					IOCTL_DISK_GET_DRIVE_GEOMETRY_EX,
+					NULL, 0, &geo, sizeof (geo));
+
+				if (NT_SUCCESS(ntStatus))
+				{
+					memcpy (&g->diskGeometry, &geo.Geometry, sizeof (DISK_GEOMETRY));
+					g->DiskSize.QuadPart = geo.DiskSize.QuadPart;
+				}
+
+				Irp->IoStatus.Information = sizeof (DISK_GEOMETRY_EX_STRUCT);
+				Irp->IoStatus.Status = ntStatus;
+			}
+		}
+		break;
+
 	case TC_IOCTL_PROBE_REAL_DRIVE_SIZE:
 		if (ValidateIOBufferSize (Irp, sizeof (ProbeRealDriveSizeRequest), ValidateInputOutput))
 		{
@@ -2125,6 +2153,7 @@ LPWSTR TCTranslateCode (ULONG ulCode)
 		TC_CASE_RET_NAME (TC_IOCTL_START_DECOY_SYSTEM_WIPE);
 		TC_CASE_RET_NAME (TC_IOCTL_WIPE_PASSWORD_CACHE);
 		TC_CASE_RET_NAME (TC_IOCTL_WRITE_BOOT_DRIVE_SECTOR);
+		TC_CASE_RET_NAME (VC_IOCTL_GET_DRIVE_GEOMETRY_EX);
 
 		TC_CASE_RET_NAME (IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS);
 
@@ -3421,14 +3450,14 @@ NTSTATUS WriteRegistryConfigFlags (uint32 flags)
 NTSTATUS GetDeviceSectorSize (PDEVICE_OBJECT deviceObject, ULONG *bytesPerSector)
 {
 	NTSTATUS status;
-	DISK_GEOMETRY geometry;
+	DISK_GEOMETRY_EX geometry;
 
-	status = SendDeviceIoControlRequest (deviceObject, IOCTL_DISK_GET_DRIVE_GEOMETRY, NULL, 0, &geometry, sizeof (geometry));
+	status = SendDeviceIoControlRequest (deviceObject, IOCTL_DISK_GET_DRIVE_GEOMETRY_EX, NULL, 0, &geometry, sizeof (geometry));
 
 	if (!NT_SUCCESS (status))
 		return status;
 
-	*bytesPerSector = geometry.BytesPerSector;
+	*bytesPerSector = geometry.Geometry.BytesPerSector;
 	return STATUS_SUCCESS;
 }
 
diff --git a/src/Driver/Ntvol.c b/src/Driver/Ntvol.c
index cfcc45e6..b3e08747 100644
--- a/src/Driver/Ntvol.c
+++ b/src/Driver/Ntvol.c
@@ -82,7 +82,7 @@ NTSTATUS TCOpenVolume (PDEVICE_OBJECT DeviceObject,
 		PARTITION_INFORMATION pi;
 		PARTITION_INFORMATION_EX pix;
 		LARGE_INTEGER diskLengthInfo;
-		DISK_GEOMETRY dg;
+		DISK_GEOMETRY_EX dg;
 		STORAGE_PROPERTY_QUERY storagePropertyQuery = {0};
 		STORAGE_ACCESS_ALIGNMENT_DESCRIPTOR storageDescriptor = {0};
 
@@ -94,12 +94,12 @@ NTSTATUS TCOpenVolume (PDEVICE_OBJECT DeviceObject,
 		if (!NT_SUCCESS (ntStatus))
 			goto error;
 
-		ntStatus = TCSendHostDeviceIoControlRequest (DeviceObject, Extension, IOCTL_DISK_GET_DRIVE_GEOMETRY, (char *) &dg, sizeof (dg));
+		ntStatus = TCSendHostDeviceIoControlRequest (DeviceObject, Extension, IOCTL_DISK_GET_DRIVE_GEOMETRY_EX, (char *) &dg, sizeof (dg));
 		if (!NT_SUCCESS (ntStatus))
 			goto error;
 
-		lDiskLength.QuadPart = dg.Cylinders.QuadPart * dg.SectorsPerTrack * dg.TracksPerCylinder * dg.BytesPerSector;
-		Extension->HostBytesPerSector = dg.BytesPerSector;
+		lDiskLength.QuadPart = dg.DiskSize.QuadPart;
+		Extension->HostBytesPerSector = dg.Geometry.BytesPerSector;
 
 		storagePropertyQuery.PropertyId = StorageAccessAlignmentProperty;
 		storagePropertyQuery.QueryType = PropertyStandardQuery;
@@ -113,7 +113,7 @@ NTSTATUS TCOpenVolume (PDEVICE_OBJECT DeviceObject,
 		}
 		else
 		{
-			Extension->HostBytesPerPhysicalSector = dg.BytesPerSector;
+			Extension->HostBytesPerPhysicalSector = dg.Geometry.BytesPerSector;
 		}
 
 		// Drive geometry is used only when IOCTL_DISK_GET_PARTITION_INFO fails
-- 
cgit v1.2.3