From 1d240a84ee3e223de1559c6445b11615bdcc67c9 Mon Sep 17 00:00:00 2001
From: Eric Kohl <eric.kohl@reactos.org>
Date: Sat, 13 Jun 2015 12:46:42 +0000
Subject: [PATCH] [USETUP] Allocate a disk layout buffer that is large enough
 to keep all partitions of the disk.

CORE-5818 resolve

svn path=/trunk/; revision=68122
---
 reactos/base/setup/usetup/partlist.c | 37 ++++++++++++++++++++++++----
 1 file changed, 32 insertions(+), 5 deletions(-)

diff --git a/reactos/base/setup/usetup/partlist.c b/reactos/base/setup/usetup/partlist.c
index 381111b0e74..bfdc4ecac30 100644
--- a/reactos/base/setup/usetup/partlist.c
+++ b/reactos/base/setup/usetup/partlist.c
@@ -1151,12 +1151,9 @@ AddDiskToList(
 
     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 @@ AddDiskToList(
                                    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