[ekohl] 68122: [USETUP] Allocate a disk layout buffer that is large enough to keep all partitions of the disk. CORE-5818 resolve

Author: ekohl
Date: Sat Jun 13 12:46:42 2015
New Revision: 68122

URL: http://svn.reactos.org/svn/reactos?rev=68122&view=rev
Log:
[USETUP]
Allocate a disk layout buffer that is large enough to keep all partitions of the disk.

CORE-5818 resolve

Modified:
    trunk/reactos/base/setup/usetup/partlist.c

Modified: trunk/reactos/base/setup/usetup/partlist.c
URL: http://svn.reactos.org/svn/reactos/trunk/reactos/base/setup/usetup/partlist....
==============================================================================
--- trunk/reactos/base/setup/usetup/partlist.c	[iso-8859-1] (original)
+++ trunk/reactos/base/setup/usetup/partlist.c	[iso-8859-1] Sat Jun 13 12:46:42 2015
@@ -1151,12 +1151,9 @@

InsertAscendingList(&List->DiskListHead, DiskEntry, DISKENTRY, ListEntry, DiskNumber);

-    /*
-     * Allocate a buffer for 26 logical drives (2 entries each == 52)
-     * plus the main partiton table (4 entries). Total 56 entries.
-     */
+    /* Allocate a layout buffer for 4 partition entries first */
     LayoutBufferSize = sizeof(DRIVE_LAYOUT_INFORMATION) +
-                       ((56 - ANYSIZE_ARRAY) * sizeof(PARTITION_INFORMATION));
+                       ((4 - ANYSIZE_ARRAY) * sizeof(PARTITION_INFORMATION));
     DiskEntry->LayoutBuffer = RtlAllocateHeap(ProcessHeap,
                                               HEAP_ZERO_MEMORY,
                                               LayoutBufferSize);
@@ -1175,6 +1172,36 @@
                                    0,
                                    DiskEntry->LayoutBuffer,
                                    LayoutBufferSize);
+    DPRINT("Status: 0x%08lx\n", Status);
+    DPRINT("PartitionCount: %lu\n", DiskEntry->LayoutBuffer->PartitionCount);
+
+    /* If we need more than 4 partition entries, reallocte the buffer and
+       retrieve the disk layout again */
+    if (!NT_SUCCESS(Status) && DiskEntry->LayoutBuffer->PartitionCount > 4)
+    {
+        LayoutBufferSize = sizeof(DRIVE_LAYOUT_INFORMATION) +
+                           ((DiskEntry->LayoutBuffer->PartitionCount - ANYSIZE_ARRAY) * sizeof(PARTITION_INFORMATION));
+        DiskEntry->LayoutBuffer = RtlReAllocateHeap(ProcessHeap,
+                                                    HEAP_ZERO_MEMORY,
+                                                    DiskEntry->LayoutBuffer,
+                                                    LayoutBufferSize);
+        if (DiskEntry->LayoutBuffer == NULL)
+        {
+            return;
+        }
+
+        Status = NtDeviceIoControlFile(FileHandle,
+                                       NULL,
+                                       NULL,
+                                       NULL,
+                                       &Iosb,
+                                       IOCTL_DISK_GET_DRIVE_LAYOUT,
+                                       NULL,
+                                       0,
+                                       DiskEntry->LayoutBuffer,
+                                       LayoutBufferSize);
+    }
+
     if (NT_SUCCESS(Status))
     {
 #ifdef DUMP_PARTITION_TABLE