[PARTINFO] Further improvements to the Partition Information Tool.
authorHermès Bélusca-Maïto <hermes.belusca-maito@reactos.org>
Fri, 8 Jan 2021 00:26:37 +0000 (01:26 +0100)
committerHermès Bélusca-Maïto <hermes.belusca-maito@reactos.org>
Fri, 8 Jan 2021 00:37:05 +0000 (01:37 +0100)
- Migrate all the IOCTL calls to NtDeviceIoControlFile() calls in
  order to know the actual returned data and Status values.

- Add support for displaying information retrieved from:
  * IOCTL_DISK_GET_DRIVE_GEOMETRY_EX ,
  * IOCTL_DISK_GET_PARTITION_INFO and IOCTL_DISK_GET_PARTITION_INFO_EX
    for "partition 0" (aka. the whole disk).

- Compactify the drive layouts after getting these from
  IOCTL_DISK_GET_DRIVE_LAYOUT(_EX).

- Fix the HexDump() when displaying a buffer whose size is not a
  multiple of 16.

modules/rosapps/applications/sysutils/utils/partinfo/partinfo.c

index d835e5b..38cd7f7 100644 (file)
@@ -3,7 +3,7 @@
  * LICENSE:     GPL-2.0+ (https://spdx.org/licenses/GPL-2.0+)
  * PURPOSE:     Displays disk and partition information for MBR and GPT disks.
  * COPYRIGHT:   Copyright 2001-2002 Eric Kohl
- *              Copyright 2020 Hermes Belusca-Maito
+ *              Copyright 2020-2021 Hermes Belusca-Maito
  */
 
 /* INCLUDES *****************************************************************/
 #include <windows.h>
 #include <ntndk.h>
 
+/* The maximum information a DISK_GEOMETRY_EX dynamic structure can contain */
+typedef struct _DISK_GEOMETRY_EX_INTERNAL
+{
+    DISK_GEOMETRY Geometry;
+    LARGE_INTEGER DiskSize;
+    DISK_PARTITION_INFO Partition;
+    DISK_DETECTION_INFO Detection;
+} DISK_GEOMETRY_EX_INTERNAL, *PDISK_GEOMETRY_EX_INTERNAL;
+
+#define DRIVE_LAYOUT_INFO_ENTRY_SIZE \
+    RTL_FIELD_SIZE(DRIVE_LAYOUT_INFORMATION, PartitionEntry[0])
+
+#define DRIVE_LAYOUT_INFO_SIZE(n) \
+    (FIELD_OFFSET(DRIVE_LAYOUT_INFORMATION, PartitionEntry) + \
+        ((n) * DRIVE_LAYOUT_INFO_ENTRY_SIZE))
+
+#define DRIVE_LAYOUT_INFOEX_ENTRY_SIZE \
+    RTL_FIELD_SIZE(DRIVE_LAYOUT_INFORMATION_EX, PartitionEntry[0])
+
+#define DRIVE_LAYOUT_INFOEX_SIZE(n) \
+    (FIELD_OFFSET(DRIVE_LAYOUT_INFORMATION_EX, PartitionEntry) + \
+        ((n) * DRIVE_LAYOUT_INFOEX_ENTRY_SIZE))
+
+
+/* Rounding up for number; not the same as ROUND_UP() with power-of-two sizes */
+#define ROUND_UP_NUM(num, up)   ((((num) + (up) - 1) / (up)) * (up))
+
+
 // #define DUMP_DATA
-#define DUMP_SIZE_INFO
 
 
 /* FORMATTING HELPERS *******************************************************/
@@ -27,6 +54,12 @@ static PCSTR PartitionStyleNames[] = {"MBR", "GPT", "RAW", "Unknown"};
           ? PartitionStyleNames[(PartStyle)] \
           : PartitionStyleNames[_countof(PartitionStyleNames)-1] )
 
+static PCSTR DetectTypeNames[] = { "None", "DetectInt13", "DetectExInt13", "Unknown" };
+#define DETECT_TYPE_NAME(DetectType) \
+    ( ((DetectType) <= DetectExInt13)   \
+          ? DetectTypeNames[(DetectType)] \
+          : DetectTypeNames[_countof(DetectTypeNames)-1] )
+
 #define GUID_FORMAT_STR "%08x-%04x-%04x-%02x%02x-%02x%02x%02x%02x%02x%02x"
 #define GUID_ELEMENTS(Guid) \
     (Guid)->Data1, (Guid)->Data2, (Guid)->Data3, \
@@ -54,16 +87,20 @@ void HexDump(
       offset += 16;
     }
 
-    ptr = (PUCHAR)((ULONG_PTR)buffer + offset);
-    printf("%08lx ", offset);
-    while (offset < size)
+    if (offset < size)
     {
-        printf(" %02hx", *ptr);
-        offset++;
-        ptr++;
+        ptr = (PUCHAR)((ULONG_PTR)buffer + offset);
+        printf("%08lx ", offset);
+        while (offset < size)
+        {
+            printf(" %02hx", *ptr);
+            offset++;
+            ptr++;
+        }
+        printf("\n");
     }
 
-    printf("\n\n\n");
+    printf("\n");
 }
 #endif
 
@@ -77,16 +114,29 @@ int main(int argc, char *argv[])
     NTSTATUS Status;
     ULONG ulDrive;
     HANDLE hDisk;
-    DWORD dwRead;
-    DWORD dwLastError;
-    DWORD i;
+    ULONG i;
     SYSTEM_DEVICE_INFORMATION DeviceInfo;
-    DISK_GEOMETRY DiskGeometry;
+    IO_STATUS_BLOCK Iosb;
     ULONG BufferSize;
-    PDRIVE_LAYOUT_INFORMATION LayoutBuffer;
-    PDRIVE_LAYOUT_INFORMATION_EX LayoutBufferEx;
+    union
+    {
+        DISK_GEOMETRY Info;
+        DISK_GEOMETRY_EX_INTERNAL InfoEx;
+    } DiskGeometry;
+    PDISK_GEOMETRY pDiskGeometry;
+    union
+    {
+        PARTITION_INFORMATION Info;
+        PARTITION_INFORMATION_EX InfoEx;
+    } PartInfo;
+    PPARTITION_INFORMATION pPartInfo;
+    PPARTITION_INFORMATION_EX pPartInfoEx;
+    union
+    {
+        PDRIVE_LAYOUT_INFORMATION Info;
+        PDRIVE_LAYOUT_INFORMATION_EX InfoEx;
+    } LayoutBuffer;
     PVOID ptr;
-    GUID Guid;
     WCHAR DriveName[40];
 
     if (argc != 2)
@@ -144,90 +194,304 @@ int main(int argc, char *argv[])
         return 0;
     }
 
+    printf("Drive number: %lu\n\n", ulDrive);
+
 
     /*
      * Get the drive geometry.
      */
-    if (!DeviceIoControl(hDisk,
-                         IOCTL_DISK_GET_DRIVE_GEOMETRY,
-                         NULL,
-                         0,
-                         &DiskGeometry,
-                         sizeof(DiskGeometry),
-                         &dwRead,
-                         NULL))
+    Status = NtDeviceIoControlFile(hDisk,
+                                   NULL,
+                                   NULL,
+                                   NULL,
+                                   &Iosb,
+                                   IOCTL_DISK_GET_DRIVE_GEOMETRY,
+                                   NULL,
+                                   0,
+                                   &DiskGeometry.Info,
+                                   sizeof(DiskGeometry.Info));
+    if (!NT_SUCCESS(Status))
     {
-        printf("DeviceIoControl(IOCTL_DISK_GET_DRIVE_GEOMETRY) failed! Error: %lu\n",
-               GetLastError());
-        CloseHandle(hDisk);
-        return 0;
+        printf("NtDeviceIoControlFile(IOCTL_DISK_GET_DRIVE_GEOMETRY) failed! (Status 0x%lx)\n", Status);
     }
+    else
+    {
+        pDiskGeometry = &DiskGeometry.Info;
 
 #ifdef DUMP_DATA
-    HexDump(&DiskGeometry, dwRead);
+        HexDump(&DiskGeometry.Info, (ULONG)Iosb.Information);
 #endif
-    printf("Drive number: %lu\n\n", ulDrive);
+        printf("IOCTL_DISK_GET_DRIVE_GEOMETRY\n"
+               "Cylinders: %I64u\nMediaType: 0x%x\nTracksPerCylinder: %lu\n"
+               "SectorsPerTrack: %lu\nBytesPerSector: %lu\n",
+               pDiskGeometry->Cylinders.QuadPart,
+               pDiskGeometry->MediaType,
+               pDiskGeometry->TracksPerCylinder,
+               pDiskGeometry->SectorsPerTrack,
+               pDiskGeometry->BytesPerSector);
+    }
+
+
+    /*
+     * Get the extended drive geometry.
+     */
+    printf("\n");
+    Status = NtDeviceIoControlFile(hDisk,
+                                   NULL,
+                                   NULL,
+                                   NULL,
+                                   &Iosb,
+                                   IOCTL_DISK_GET_DRIVE_GEOMETRY_EX,
+                                   NULL,
+                                   0,
+                                   &DiskGeometry.InfoEx,
+                                   sizeof(DiskGeometry.InfoEx));
+    if (!NT_SUCCESS(Status))
+    {
+        printf("NtDeviceIoControlFile(IOCTL_DISK_GET_DRIVE_GEOMETRY_EX) failed! (Status 0x%lx)\n", Status);
+    }
+    else
+    {
+        // PDISK_DETECTION_INFO DiskDetectInfo = &DiskGeometry.InfoEx.Detection; // DiskGeometryGetDetect(&DiskGeometry.InfoEx);
+        // PDISK_PARTITION_INFO DiskPartInfo = &DiskGeometry.InfoEx.Partition; // DiskGeometryGetPartition(&DiskGeometry.InfoEx);
+        pDiskGeometry = &DiskGeometry.InfoEx.Geometry;
+
+#ifdef DUMP_DATA
+        HexDump(&DiskGeometry.InfoEx, (ULONG)Iosb.Information);
+#endif
+        printf("IOCTL_DISK_GET_DRIVE_GEOMETRY_EX\n"
+               "Cylinders: %I64u\nMediaType: 0x%x\nTracksPerCylinder: %lu\n"
+               "SectorsPerTrack: %lu\nBytesPerSector: %lu\n"
+               "DiskSize: %I64u\n",
+               pDiskGeometry->Cylinders.QuadPart,
+               pDiskGeometry->MediaType,
+               pDiskGeometry->TracksPerCylinder,
+               pDiskGeometry->SectorsPerTrack,
+               pDiskGeometry->BytesPerSector,
+               DiskGeometry.InfoEx.DiskSize.QuadPart);
+
+        printf("SizeOfDetectInfo: %lu  DetectionType: %lu (%s)\n",
+               DiskGeometry.InfoEx.Detection.SizeOfDetectInfo,
+               DiskGeometry.InfoEx.Detection.DetectionType,
+               DETECT_TYPE_NAME(DiskGeometry.InfoEx.Detection.DetectionType));
+        switch (DiskGeometry.InfoEx.Detection.DetectionType)
+        {
+        case DetectInt13:
+        {
+            PDISK_INT13_INFO pInt13 = &DiskGeometry.InfoEx.Detection.Int13;
+            printf("  DriveSelect: %u  MaxCylinders: %lu  SectorsPerTrack: %u  MaxHeads: %u  NumberDrives: %u\n",
+                   pInt13->DriveSelect, pInt13->MaxCylinders,
+                   pInt13->SectorsPerTrack, pInt13->MaxHeads,
+                   pInt13->NumberDrives);
+            break;
+        }
+
+        case DetectExInt13:
+        {
+            PDISK_EX_INT13_INFO pExInt13 = &DiskGeometry.InfoEx.Detection.ExInt13;
+            printf("  ExBufferSize: %u  ExFlags: %u\n"
+                   "  ExCylinders: %lu  ExHeads: %lu  ExSectorsPerTrack: %lu\n"
+                   "  ExSectorsPerDrive: %I64u  ExSectorSize: %u  ExReserved: %u\n",
+                   pExInt13->ExBufferSize, pExInt13->ExFlags,
+                   pExInt13->ExCylinders,  pExInt13->ExHeads,
+                   pExInt13->ExSectorsPerTrack, pExInt13->ExSectorsPerDrive,
+                   pExInt13->ExSectorSize, pExInt13->ExReserved);
+            break;
+        }
+
+        case DetectNone:
+        default:
+            break;
+        }
+
+        printf("SizeOfPartitionInfo: %lu  PartitionStyle: %lu [%s]\n",
+               DiskGeometry.InfoEx.Partition.SizeOfPartitionInfo,
+               DiskGeometry.InfoEx.Partition.PartitionStyle,
+               PARTITION_STYLE_NAME(DiskGeometry.InfoEx.Partition.PartitionStyle));
+
+        if (DiskGeometry.InfoEx.Partition.PartitionStyle == PARTITION_STYLE_MBR)
+        {
+            printf("  Signature: 0x%08lx  Checksum 0x%08lx\n",
+                   DiskGeometry.InfoEx.Partition.Mbr.Signature,
+                   DiskGeometry.InfoEx.Partition.Mbr.CheckSum);
+        }
+        else if (DiskGeometry.InfoEx.Partition.PartitionStyle == PARTITION_STYLE_GPT)
+        {
+            printf("  DiskId: {" GUID_FORMAT_STR "}\n",
+                   GUID_ELEMENTS(&DiskGeometry.InfoEx.Partition.Gpt.DiskId));
+        }
+        else
+        {
+            /* Unknown */
+            printf("\n");
+        }
+    }
 
-    printf("IOCTL_DISK_GET_DRIVE_GEOMETRY\n"
-           "Cylinders: %I64u\nMediaType: %x\nTracksPerCylinder: %lu\n"
-           "SectorsPerTrack: %lu\nBytesPerSector: %lu\n\n",
-           DiskGeometry.Cylinders.QuadPart,
-           DiskGeometry.MediaType,
-           DiskGeometry.TracksPerCylinder,
-           DiskGeometry.SectorsPerTrack,
-           DiskGeometry.BytesPerSector);
-
-#if 0 // TODO!
-    /* Get extended drive geometry */
-    // IOCTL_DISK_GET_DRIVE_GEOMETRY_EX
+
+    /*
+     * Display partition 0 (i.e. the whole disk) information.
+     */
+    printf("\n");
+    Status = NtDeviceIoControlFile(hDisk,
+                                   NULL,
+                                   NULL,
+                                   NULL,
+                                   &Iosb,
+                                   IOCTL_DISK_GET_PARTITION_INFO,
+                                   NULL,
+                                   0,
+                                   &PartInfo.Info,
+                                   sizeof(PartInfo.Info));
+    if (!NT_SUCCESS(Status))
+    {
+        printf("NtDeviceIoControlFile(IOCTL_DISK_GET_PARTITION_INFO) failed! (Status 0x%lx)\n", Status);
+    }
+    else
+    {
+        pPartInfo = &PartInfo.Info;
+
+#ifdef DUMP_DATA
+        HexDump(&PartInfo.Info, (ULONG)Iosb.Information);
 #endif
+        printf("IOCTL_DISK_GET_PARTITION_INFO\n"
+               "nr: %ld boot: %1x type: 0x%x (%s) start: 0x%016I64x count: 0x%016I64x hidden: 0x%lx\n",
+               pPartInfo->PartitionNumber,
+               pPartInfo->BootIndicator,
+               pPartInfo->PartitionType,
+               pPartInfo->RecognizedPartition ? "Recognized" : "Not recognized",
+               pPartInfo->StartingOffset.QuadPart,
+               pPartInfo->PartitionLength.QuadPart,
+               pPartInfo->HiddenSectors);
+    }
+
+    printf("\n");
+    Status = NtDeviceIoControlFile(hDisk,
+                                   NULL,
+                                   NULL,
+                                   NULL,
+                                   &Iosb,
+                                   IOCTL_DISK_GET_PARTITION_INFO_EX,
+                                   NULL,
+                                   0,
+                                   &PartInfo.InfoEx,
+                                   sizeof(PartInfo.InfoEx));
+    if (!NT_SUCCESS(Status))
+    {
+        printf("NtDeviceIoControlFile(IOCTL_DISK_GET_PARTITION_INFO_EX) failed! (Status 0x%lx)\n", Status);
+    }
+    else
+    {
+        pPartInfoEx = &PartInfo.InfoEx;
+
+#ifdef DUMP_DATA
+        HexDump(&PartInfo.InfoEx, (ULONG)Iosb.Information);
+#endif
+        printf("IOCTL_DISK_GET_PARTITION_INFO_EX\n");
+
+        if (pPartInfoEx->PartitionStyle == PARTITION_STYLE_MBR)
+        {
+            printf("nr: %ld [%s] boot: %1x type: 0x%x (%s) start: 0x%016I64x count: 0x%016I64x hidden: 0x%lx\n",
+                   pPartInfoEx->PartitionNumber,
+                   PARTITION_STYLE_NAME(pPartInfoEx->PartitionStyle),
+                   pPartInfoEx->Mbr.BootIndicator,
+                   pPartInfoEx->Mbr.PartitionType,
+                   pPartInfoEx->Mbr.RecognizedPartition ? "Recognized" : "Not recognized",
+                   pPartInfoEx->StartingOffset.QuadPart,
+                   pPartInfoEx->PartitionLength.QuadPart,
+                   pPartInfoEx->Mbr.HiddenSectors);
+#if (NTDDI_VERSION >= NTDDI_WINBLUE)
+            printf("    PartitionId: {" GUID_FORMAT_STR "}\n",
+                   GUID_ELEMENTS(&pPartInfoEx->Mbr.PartitionId));
+#endif
+        }
+        else if (pPartInfoEx->PartitionStyle == PARTITION_STYLE_GPT)
+        {
+            printf("nr: %ld [%s]\n"
+                   "    type : {" GUID_FORMAT_STR "}\n"
+                   "    id   : {" GUID_FORMAT_STR "}\n"
+                   "    attrs: 0x%016I64x\n"
+                   "    name : '%.*S'\n"
+                   "    start: 0x%016I64x count: 0x%016I64x\n",
+                   pPartInfoEx->PartitionNumber,
+                   PARTITION_STYLE_NAME(pPartInfoEx->PartitionStyle),
+                   GUID_ELEMENTS(&pPartInfoEx->Gpt.PartitionType),
+                   GUID_ELEMENTS(&pPartInfoEx->Gpt.PartitionId),
+                   pPartInfoEx->Gpt.Attributes,
+                   (int)_countof(pPartInfoEx->Gpt.Name),
+                   pPartInfoEx->Gpt.Name,
+                   pPartInfoEx->StartingOffset.QuadPart,
+                   pPartInfoEx->PartitionLength.QuadPart);
+        }
+    }
 
 
     /*
      * Retrieve the legacy partition layout
      */
+    printf("\n");
 
-    /* Allocate a layout buffer with 4 partition entries first */
-    BufferSize = sizeof(DRIVE_LAYOUT_INFORMATION) +
-                       ((4 - ANYSIZE_ARRAY) * sizeof(PARTITION_INFORMATION));
-    LayoutBuffer = malloc(BufferSize);
-    if (!LayoutBuffer)
+    /* Allocate a layout buffer with initially 4 partition entries (or 16 for NEC PC-98) */
+    BufferSize = DRIVE_LAYOUT_INFO_SIZE(IsNEC_98 ? 16 : 4);
+    LayoutBuffer.Info = RtlAllocateHeap(RtlGetProcessHeap(),
+                                        HEAP_ZERO_MEMORY,
+                                        BufferSize);
+    if (!LayoutBuffer.Info)
     {
         printf("Out of memory!");
-        CloseHandle(hDisk);
-        return 0;
+        goto Quit;
     }
-    memset(LayoutBuffer, 0, BufferSize);
 
-    /* Keep looping while the drive layout buffer is too small */
+    /*
+     * Keep looping while the drive layout buffer is too small.
+     * Iosb.Information or PartitionCount only contain actual info only
+     * once NtDeviceIoControlFile(IOCTL_DISK_GET_DRIVE_LAYOUT) succeeds.
+     */
     for (;;)
     {
-        if (DeviceIoControl(hDisk,
-                            IOCTL_DISK_GET_DRIVE_LAYOUT,
-                            NULL,
-                            0,
-                            LayoutBuffer,
-                            BufferSize,
-                            &dwRead,
-                            NULL))
+        Status = NtDeviceIoControlFile(hDisk,
+                                       NULL,
+                                       NULL,
+                                       NULL,
+                                       &Iosb,
+                                       IOCTL_DISK_GET_DRIVE_LAYOUT,
+                                       NULL,
+                                       0,
+                                       LayoutBuffer.Info,
+                                       BufferSize);
+        if (NT_SUCCESS(Status))
         {
-            dwLastError = ERROR_SUCCESS;
+            /* We succeeded; compactify the layout structure but keep the
+             * number of partition entries rounded up to a multiple of 4. */
+#if 1
+            BufferSize = DRIVE_LAYOUT_INFO_SIZE(ROUND_UP_NUM(LayoutBuffer.Info->PartitionCount, 4));
+#else
+            BufferSize = (ULONG)Iosb.Information;
+#endif
+            ptr = RtlReAllocateHeap(RtlGetProcessHeap(),
+                                    HEAP_REALLOC_IN_PLACE_ONLY,
+                                    LayoutBuffer.Info, BufferSize);
+            if (!ptr)
+            {
+                printf("Compactification failed; keeping original structure.\n");
+            }
+            else
+            {
+                LayoutBuffer.Info = ptr;
+            }
+            Status = STATUS_SUCCESS;
             break;
         }
 
-        dwLastError = GetLastError();
-        if (dwLastError != ERROR_INSUFFICIENT_BUFFER)
+        if (Status != STATUS_BUFFER_TOO_SMALL)
         {
-            printf("DeviceIoControl(IOCTL_DISK_GET_DRIVE_LAYOUT) failed! Error: %lu\n",
-                   dwLastError);
+            printf("NtDeviceIoControlFile(IOCTL_DISK_GET_DRIVE_LAYOUT) failed! (Status 0x%lx)\n", Status);
 
             /* Bail out if any other error than "invalid function" has been emitted.
              * This happens for example when calling it on GPT disks. */
-            if (dwLastError != ERROR_INVALID_FUNCTION)
+            if (Status != STATUS_INVALID_DEVICE_REQUEST)
             {
-                free(LayoutBuffer);
-                CloseHandle(hDisk);
-                return 0;
+                RtlFreeHeap(RtlGetProcessHeap(), 0, LayoutBuffer.Info);
+                goto Quit;
             }
             else
             {
@@ -236,45 +500,46 @@ int main(int argc, char *argv[])
             }
         }
 
-        /* Reallocate the buffer */
-        BufferSize += 4 * sizeof(PARTITION_INFORMATION);
-        ptr = realloc(LayoutBuffer, BufferSize);
+        /* Reallocate the buffer by chunks of 4 entries */
+        BufferSize += 4 * DRIVE_LAYOUT_INFO_ENTRY_SIZE;
+        ptr = RtlReAllocateHeap(RtlGetProcessHeap(),
+                                HEAP_ZERO_MEMORY,
+                                LayoutBuffer.Info, BufferSize);
         if (!ptr)
         {
             printf("Out of memory!");
-            free(LayoutBuffer);
-            CloseHandle(hDisk);
-            return 0;
+            RtlFreeHeap(RtlGetProcessHeap(), 0, LayoutBuffer.Info);
+            goto Quit;
         }
-        LayoutBuffer = ptr;
-        memset(LayoutBuffer, 0, BufferSize);
+        LayoutBuffer.Info = ptr;
     }
 
-    if (dwLastError == ERROR_SUCCESS)
+    if (Status == STATUS_SUCCESS)
     {
 #ifdef DUMP_DATA
-        HexDump(LayoutBuffer, dwRead);
+        HexDump(LayoutBuffer.Info, (ULONG)Iosb.Information);
 #endif
-
         printf("IOCTL_DISK_GET_DRIVE_LAYOUT\n"
                "Partitions: %lu  Signature: 0x%08lx\n",
-               LayoutBuffer->PartitionCount,
-               LayoutBuffer->Signature);
+               LayoutBuffer.Info->PartitionCount,
+               LayoutBuffer.Info->Signature);
 
-        for (i = 0; i < LayoutBuffer->PartitionCount; i++)
+        for (i = 0; i < LayoutBuffer.Info->PartitionCount; i++)
         {
-            printf(" %ld: nr: %ld boot: %1x type: %x start: 0x%016I64x count: 0x%016I64x hidden: 0x%lx\n",
+            pPartInfo = &LayoutBuffer.Info->PartitionEntry[i];
+
+            printf(" %ld: nr: %ld boot: %1x type: 0x%x (%s) start: 0x%016I64x count: 0x%016I64x hidden: 0x%lx\n",
                    i,
-                   LayoutBuffer->PartitionEntry[i].PartitionNumber,
-                   LayoutBuffer->PartitionEntry[i].BootIndicator,
-                   LayoutBuffer->PartitionEntry[i].PartitionType,
-                   LayoutBuffer->PartitionEntry[i].StartingOffset.QuadPart,
-                   LayoutBuffer->PartitionEntry[i].PartitionLength.QuadPart,
-                   LayoutBuffer->PartitionEntry[i].HiddenSectors);
+                   pPartInfo->PartitionNumber,
+                   pPartInfo->BootIndicator,
+                   pPartInfo->PartitionType,
+                   pPartInfo->RecognizedPartition ? "Recognized" : "Not recognized",
+                   pPartInfo->StartingOffset.QuadPart,
+                   pPartInfo->PartitionLength.QuadPart,
+                   pPartInfo->HiddenSectors);
         }
-
-        free(LayoutBuffer);
     }
+    RtlFreeHeap(RtlGetProcessHeap(), 0, LayoutBuffer.Info);
 
 
     /*
@@ -282,46 +547,67 @@ int main(int argc, char *argv[])
      */
     printf("\n");
 
-    /* Allocate a layout buffer with 4 partition entries first */
-    BufferSize = sizeof(DRIVE_LAYOUT_INFORMATION_EX) +
-                       ((4 - ANYSIZE_ARRAY) * sizeof(PARTITION_INFORMATION_EX));
-    LayoutBufferEx = malloc(BufferSize);
-    if (!LayoutBufferEx)
+    /* Allocate a layout buffer with initially 4 partition entries (or 16 for NEC PC-98) */
+    BufferSize = DRIVE_LAYOUT_INFOEX_SIZE(IsNEC_98 ? 16 : 4);
+    LayoutBuffer.InfoEx = RtlAllocateHeap(RtlGetProcessHeap(),
+                                          HEAP_ZERO_MEMORY,
+                                          BufferSize);
+    if (!LayoutBuffer.InfoEx)
     {
         printf("Out of memory!");
-        CloseHandle(hDisk);
-        return 0;
+        goto Quit;
     }
-    memset(LayoutBufferEx, 0, BufferSize);
 
-    /* Keep looping while the drive layout buffer is too small */
+    /*
+     * Keep looping while the drive layout buffer is too small.
+     * Iosb.Information or PartitionCount only contain actual info only
+     * once NtDeviceIoControlFile(IOCTL_DISK_GET_DRIVE_LAYOUT_EX) succeeds.
+     */
     for (;;)
     {
-        if (DeviceIoControl(hDisk,
-                            IOCTL_DISK_GET_DRIVE_LAYOUT_EX,
-                            NULL,
-                            0,
-                            LayoutBufferEx,
-                            BufferSize,
-                            &dwRead,
-                            NULL))
+        Status = NtDeviceIoControlFile(hDisk,
+                                       NULL,
+                                       NULL,
+                                       NULL,
+                                       &Iosb,
+                                       IOCTL_DISK_GET_DRIVE_LAYOUT_EX,
+                                       NULL,
+                                       0,
+                                       LayoutBuffer.InfoEx,
+                                       BufferSize);
+        if (NT_SUCCESS(Status))
         {
-            dwLastError = ERROR_SUCCESS;
+            /* We succeeded; compactify the layout structure but keep the
+             * number of partition entries rounded up to a multiple of 4. */
+#if 1
+            BufferSize = DRIVE_LAYOUT_INFOEX_SIZE(ROUND_UP_NUM(LayoutBuffer.InfoEx->PartitionCount, 4));
+#else
+            BufferSize = (ULONG)Iosb.Information;
+#endif
+            ptr = RtlReAllocateHeap(RtlGetProcessHeap(),
+                                    HEAP_REALLOC_IN_PLACE_ONLY,
+                                    LayoutBuffer.InfoEx, BufferSize);
+            if (!ptr)
+            {
+                printf("Compactification failed; keeping original structure.\n");
+            }
+            else
+            {
+                LayoutBuffer.InfoEx = ptr;
+            }
+            Status = STATUS_SUCCESS;
             break;
         }
 
-        dwLastError = GetLastError();
-        if (dwLastError != ERROR_INSUFFICIENT_BUFFER)
+        if (Status != STATUS_BUFFER_TOO_SMALL)
         {
-            printf("DeviceIoControl(IOCTL_DISK_GET_DRIVE_LAYOUT_EX) failed! Error: %lu\n",
-                   dwLastError);
+            printf("NtDeviceIoControlFile(IOCTL_DISK_GET_DRIVE_LAYOUT_EX) failed! (Status 0x%lx)\n", Status);
 
             /* Bail out if any other error than "invalid function" has been emitted */
-            if (dwLastError != ERROR_INVALID_FUNCTION)
+            if (Status != STATUS_INVALID_DEVICE_REQUEST)
             {
-                free(LayoutBufferEx);
-                CloseHandle(hDisk);
-                return 0;
+                RtlFreeHeap(RtlGetProcessHeap(), 0, LayoutBuffer.InfoEx);
+                goto Quit;
             }
             else
             {
@@ -330,74 +616,77 @@ int main(int argc, char *argv[])
             }
         }
 
-        /* Reallocate the buffer */
-        BufferSize += 4 * sizeof(PARTITION_INFORMATION_EX);
-        ptr = realloc(LayoutBufferEx, BufferSize);
+        /* Reallocate the buffer by chunks of 4 entries */
+        BufferSize += 4 * DRIVE_LAYOUT_INFOEX_ENTRY_SIZE;
+        ptr = RtlReAllocateHeap(RtlGetProcessHeap(),
+                                HEAP_ZERO_MEMORY,
+                                LayoutBuffer.InfoEx, BufferSize);
         if (!ptr)
         {
             printf("Out of memory!");
-            free(LayoutBufferEx);
-            CloseHandle(hDisk);
-            return 0;
+            RtlFreeHeap(RtlGetProcessHeap(), 0, LayoutBuffer.InfoEx);
+            goto Quit;
         }
-        LayoutBufferEx = ptr;
-        memset(LayoutBufferEx, 0, BufferSize);
+        LayoutBuffer.InfoEx = ptr;
     }
 
-    if (dwLastError == ERROR_SUCCESS)
+    if (Status == STATUS_SUCCESS)
     {
 #ifdef DUMP_DATA
-        HexDump(LayoutBufferEx, dwRead);
+        HexDump(LayoutBuffer.InfoEx, (ULONG)Iosb.Information);
 #endif
-
         printf("IOCTL_DISK_GET_DRIVE_LAYOUT_EX\n"
-               "PartitionStyle: [%s]\n",
-               PARTITION_STYLE_NAME(LayoutBufferEx->PartitionStyle));
+               "PartitionStyle: %lu [%s]\n",
+               LayoutBuffer.InfoEx->PartitionStyle,
+               PARTITION_STYLE_NAME(LayoutBuffer.InfoEx->PartitionStyle));
 
-        if (LayoutBufferEx->PartitionStyle == PARTITION_STYLE_MBR)
+        if (LayoutBuffer.InfoEx->PartitionStyle == PARTITION_STYLE_MBR)
         {
             printf("Partitions: %lu  Signature: 0x%08lx",
-                   LayoutBufferEx->PartitionCount,
-                   LayoutBufferEx->Mbr.Signature);
+                   LayoutBuffer.InfoEx->PartitionCount,
+                   LayoutBuffer.InfoEx->Mbr.Signature);
 #if (NTDDI_VERSION >= NTDDI_WIN10_RS1)
             printf("  Checksum 0x%08lx\n",
-                   LayoutBufferEx->Mbr.CheckSum);
+                   LayoutBuffer.InfoEx->Mbr.CheckSum);
 #else
             printf("\n");
 #endif
 
-            for (i = 0; i < LayoutBufferEx->PartitionCount; i++)
+            for (i = 0; i < LayoutBuffer.InfoEx->PartitionCount; i++)
             {
-                printf(" %ld: nr: %ld [%s] boot: %1x type: %x start: 0x%016I64x count: 0x%016I64x hidden: 0x%lx\n",
+                pPartInfoEx = &LayoutBuffer.InfoEx->PartitionEntry[i];
+
+                printf(" %ld: nr: %ld [%s] boot: %1x type: 0x%x (%s) start: 0x%016I64x count: 0x%016I64x hidden: 0x%lx\n",
                        i,
-                       LayoutBufferEx->PartitionEntry[i].PartitionNumber,
-                       PARTITION_STYLE_NAME(LayoutBufferEx->PartitionEntry[i].PartitionStyle),
-                       LayoutBufferEx->PartitionEntry[i].Mbr.BootIndicator,
-                       LayoutBufferEx->PartitionEntry[i].Mbr.PartitionType,
-                       LayoutBufferEx->PartitionEntry[i].StartingOffset.QuadPart,
-                       LayoutBufferEx->PartitionEntry[i].PartitionLength.QuadPart,
-                       LayoutBufferEx->PartitionEntry[i].Mbr.HiddenSectors);
+                       pPartInfoEx->PartitionNumber,
+                       PARTITION_STYLE_NAME(pPartInfoEx->PartitionStyle),
+                       pPartInfoEx->Mbr.BootIndicator,
+                       pPartInfoEx->Mbr.PartitionType,
+                       pPartInfoEx->Mbr.RecognizedPartition ? "Recognized" : "Not recognized",
+                       pPartInfoEx->StartingOffset.QuadPart,
+                       pPartInfoEx->PartitionLength.QuadPart,
+                       pPartInfoEx->Mbr.HiddenSectors);
 #if (NTDDI_VERSION >= NTDDI_WINBLUE)
-                Guid = LayoutBufferEx->PartitionEntry[i].Mbr.PartitionId;
                 printf("    PartitionId: {" GUID_FORMAT_STR "}\n",
-                       GUID_ELEMENTS(&Guid));
+                       GUID_ELEMENTS(&pPartInfoEx->Mbr.PartitionId));
 #endif
             }
         }
-        else if (LayoutBufferEx->PartitionStyle == PARTITION_STYLE_GPT)
+        else if (LayoutBuffer.InfoEx->PartitionStyle == PARTITION_STYLE_GPT)
         {
-            Guid = LayoutBufferEx->Gpt.DiskId;
             printf("Partitions: %lu  MaxPartitionCount: %lu\n"
                    "DiskId: {" GUID_FORMAT_STR "}\n"
                    "StartingUsableOffset: 0x%016I64x  UsableLength: 0x%016I64x\n",
-                   LayoutBufferEx->PartitionCount,
-                   LayoutBufferEx->Gpt.MaxPartitionCount,
-                   GUID_ELEMENTS(&Guid),
-                   LayoutBufferEx->Gpt.StartingUsableOffset.QuadPart,
-                   LayoutBufferEx->Gpt.UsableLength.QuadPart);
+                   LayoutBuffer.InfoEx->PartitionCount,
+                   LayoutBuffer.InfoEx->Gpt.MaxPartitionCount,
+                   GUID_ELEMENTS(&LayoutBuffer.InfoEx->Gpt.DiskId),
+                   LayoutBuffer.InfoEx->Gpt.StartingUsableOffset.QuadPart,
+                   LayoutBuffer.InfoEx->Gpt.UsableLength.QuadPart);
 
-            for (i = 0; i < LayoutBufferEx->PartitionCount; i++)
+            for (i = 0; i < LayoutBuffer.InfoEx->PartitionCount; i++)
             {
+                pPartInfoEx = &LayoutBuffer.InfoEx->PartitionEntry[i];
+
                 printf(" %ld: nr: %ld [%s]\n"
                        "    type : {" GUID_FORMAT_STR "}\n"
                        "    id   : {" GUID_FORMAT_STR "}\n"
@@ -405,22 +694,21 @@ int main(int argc, char *argv[])
                        "    name : '%.*S'\n"
                        "    start: 0x%016I64x count: 0x%016I64x\n",
                        i,
-                       LayoutBufferEx->PartitionEntry[i].PartitionNumber,
-                       PARTITION_STYLE_NAME(LayoutBufferEx->PartitionEntry[i].PartitionStyle),
-                       GUID_ELEMENTS(&LayoutBufferEx->PartitionEntry[i].Gpt.PartitionType),
-                       GUID_ELEMENTS(&LayoutBufferEx->PartitionEntry[i].Gpt.PartitionId),
-                       LayoutBufferEx->PartitionEntry[i].Gpt.Attributes,
-                       (int)_countof(LayoutBufferEx->PartitionEntry[i].Gpt.Name),
-                       LayoutBufferEx->PartitionEntry[i].Gpt.Name,
-                       LayoutBufferEx->PartitionEntry[i].StartingOffset.QuadPart,
-                       LayoutBufferEx->PartitionEntry[i].PartitionLength.QuadPart);
+                       pPartInfoEx->PartitionNumber,
+                       PARTITION_STYLE_NAME(pPartInfoEx->PartitionStyle),
+                       GUID_ELEMENTS(&pPartInfoEx->Gpt.PartitionType),
+                       GUID_ELEMENTS(&pPartInfoEx->Gpt.PartitionId),
+                       pPartInfoEx->Gpt.Attributes,
+                       (int)_countof(pPartInfoEx->Gpt.Name),
+                       pPartInfoEx->Gpt.Name,
+                       pPartInfoEx->StartingOffset.QuadPart,
+                       pPartInfoEx->PartitionLength.QuadPart);
             }
         }
-
-        free(LayoutBufferEx);
     }
+    RtlFreeHeap(RtlGetProcessHeap(), 0, LayoutBuffer.InfoEx);
 
+Quit:
     CloseHandle(hDisk);
-
     return 0;
 }