From 1f87e6412766019d62fa45bfe8489604083a411e Mon Sep 17 00:00:00 2001
From: Mounir IDRASSI <mounir.idrassi@idrix.fr>
Date: Sat, 13 Feb 2016 15:40:15 +0100
Subject: Windows: handle rare case where size of disk can't be retrieved using
 IOCTL_DISK_GET_PARTITION_INFO_EX. use IOCTL_DISK_GET_DRIVE_GEOMETRY in such
 case and add workaround to display correctly partitions in dialog even if
 disk size is not available.

---
 src/Common/Dlgcode.c | 43 ++++++++++++++++++++++++++++++++++++++++++-
 src/Common/Dlgcode.h |  1 +
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/src/Common/Dlgcode.c b/src/Common/Dlgcode.c
index e3e70985..da37795d 100644
--- a/src/Common/Dlgcode.c
+++ b/src/Common/Dlgcode.c
@@ -3422,7 +3422,9 @@ BOOL CALLBACK RawDevicesDlgProc (HWND hwndDlg, UINT msg, WPARAM wParam, LPARAM l
 				// Path
 				if (!device.IsPartition || device.DynamicVolume)
 				{
-					if (!device.Floppy && device.Size == 0)
+					if (!device.Floppy && (device.Size == 0) 
+						&& (device.IsPartition || device.Partitions.empty() || device.Partitions[0].Size == 0)
+						)
 						continue;
 
 					if (line > 1)
@@ -7709,6 +7711,33 @@ BOOL GetDriveGeometry (const wchar_t *deviceName, PDISK_GEOMETRY diskGeometry)
 		return FALSE;
 }
 
+BOOL GetPhysicalDriveGeometry (int driveNumber, PDISK_GEOMETRY diskGeometry)
+{
+	HANDLE hDev;
+	BOOL bResult = FALSE;
+	TCHAR devicePath[MAX_PATH];
+
+	StringCchPrintfW (devicePath, ARRAYSIZE (devicePath), L"\\\\.\\PhysicalDrive%d", driveNumber);
+
+	if ((hDev = CreateFileW (devicePath, 0, 0, NULL, OPEN_EXISTING, 0, NULL)) != INVALID_HANDLE_VALUE)
+	{
+		DWORD bytesRead = 0;
+
+		ZeroMemory (diskGeometry, sizeof (DISK_GEOMETRY));
+
+		if (	DeviceIoControl (hDev, IOCTL_DISK_GET_DRIVE_GEOMETRY, NULL, 0, diskGeometry, sizeof (DISK_GEOMETRY), &bytesRead, NULL)
+			&& (bytesRead == sizeof (DISK_GEOMETRY)) 
+			&& diskGeometry->BytesPerSector)
+		{
+			bResult = TRUE;
+		}
+
+		CloseHandle (hDev);
+	}
+
+	return bResult;
+}
+
 
 // Returns drive letter number assigned to device (-1 if none)
 int GetDiskDeviceDriveLetter (PWSTR deviceName)
@@ -10812,6 +10841,18 @@ std::vector <HostDevice> GetAvailableHostDevices (bool noDeviceProperties, bool
 				device.Bootable = partInfo.BootIndicator ? true : false;
 				device.Size = partInfo.PartitionLength.QuadPart;
 			}
+			else
+			{
+				// retrieve size using DISK_GEOMETRY
+				DISK_GEOMETRY deviceGeometry = {0};
+				if (	GetDriveGeometry (devPath, &deviceGeometry)
+						||	((partNumber == 0) && GetPhysicalDriveGeometry (devNumber, &deviceGeometry))
+					)
+				{
+					device.Size = deviceGeometry.Cylinders.QuadPart * (LONGLONG) deviceGeometry.BytesPerSector 
+						* (LONGLONG) deviceGeometry.SectorsPerTrack * (LONGLONG) deviceGeometry.TracksPerCylinder;
+				}
+			}
 
 			device.HasUnencryptedFilesystem = (detectUnencryptedFilesystems && openTest.FilesystemDetected) ? true : false;
 
diff --git a/src/Common/Dlgcode.h b/src/Common/Dlgcode.h
index efaf935b..20199924 100644
--- a/src/Common/Dlgcode.h
+++ b/src/Common/Dlgcode.h
@@ -452,6 +452,7 @@ void OpenOnlineHelp ();
 BOOL GetPartitionInfo (const wchar_t *deviceName, PPARTITION_INFORMATION rpartInfo);
 BOOL GetDeviceInfo (const wchar_t *deviceName, DISK_PARTITION_INFO_STRUCT *info);
 BOOL GetDriveGeometry (const wchar_t *deviceName, PDISK_GEOMETRY diskGeometry);
+BOOL GetPhysicalDriveGeometry (int driveNumber, PDISK_GEOMETRY diskGeometry);
 BOOL IsVolumeDeviceHosted (const wchar_t *lpszDiskFile);
 int CompensateXDPI (int val);
 int CompensateYDPI (int val);
-- 
cgit v1.2.3