[reactos] 02/03: [SETUPLIB][USETUP] Minor improvements.

[Hermès Bélusca-Maïto]

https://git.reactos.org/?p=reactos.git;a=commitdiff;h=29cc1843d074106649aa6…

commit 29cc1843d074106649aa6b0155f9ff347c7abb45
Author:     Hermès Bélusca-Maïto <hermes.belusca-maito(a)reactos.org>
AuthorDate: Sun Feb 24 19:04:57 2019 +0100
Commit:     Hermès Bélusca-Maïto <hermes.belusca-maito(a)reactos.org>
CommitDate: Sun Mar 10 15:41:20 2019 +0100

    [SETUPLIB][USETUP] Minor improvements.
    
    - Make some partitioning functions not dependent on the selected
"CurrentPartition".
    - Add some sanity checks.
    - Improve some of the "is-partitioned" checks.
---
 base/setup/lib/setuplib.c       |   2 +
 base/setup/lib/utils/osdetect.c |   2 +
 base/setup/lib/utils/partlist.c | 201 +++++++++++++++++++++++-----------------
 base/setup/lib/utils/partlist.h |  18 +++-
 base/setup/usetup/usetup.c      |  26 +++++-
 5 files changed, 155 insertions(+), 94 deletions(-)

diff --git a/base/setup/lib/setuplib.c b/base/setup/lib/setuplib.c
index 9e4d22ee37..0501fe72ef 100644
--- a/base/setup/lib/setuplib.c
+++ b/base/setup/lib/setuplib.c
@@ -630,6 +630,8 @@ InitDestinationPaths(
     WCHAR PathBuffer[MAX_PATH];
     NTSTATUS Status;
 
+    ASSERT(PartEntry->IsPartitioned && PartEntry->PartitionNumber != 0);
+
     /* Create 'pSetupData->DestinationRootPath' string */
     RtlFreeUnicodeString(&pSetupData->DestinationRootPath);
     Status = RtlStringCchPrintfW(PathBuffer, ARRAYSIZE(PathBuffer),
diff --git a/base/setup/lib/utils/osdetect.c b/base/setup/lib/utils/osdetect.c
index 8dd35cd55b..d191c58998 100644
--- a/base/setup/lib/utils/osdetect.c
+++ b/base/setup/lib/utils/osdetect.c
@@ -697,6 +697,8 @@ FindNTOSInstallations(
     ULONG Version;
     WCHAR PathBuffer[MAX_PATH];
 
+    ASSERT(PartEntry->IsPartitioned && PartEntry->PartitionNumber != 0);
+
     /* Set PartitionRootPath */
     RtlStringCchPrintfW(PathBuffer, ARRAYSIZE(PathBuffer),
                         L"\\Device\\Harddisk%lu\\Partition%lu\\",
diff --git a/base/setup/lib/utils/partlist.c b/base/setup/lib/utils/partlist.c
index 5c20c4bf36..269bb99b2f 100644
--- a/base/setup/lib/utils/partlist.c
+++ b/base/setup/lib/utils/partlist.c
@@ -342,9 +342,10 @@ AssignDriveLetters(
             if (PartEntry->IsPartitioned &&
                 !IsContainerPartition(PartEntry->PartitionType))
             {
+                ASSERT(PartEntry->PartitionType != PARTITION_ENTRY_UNUSED);
+
                 if (IsRecognizedPartition(PartEntry->PartitionType) ||
-                    (PartEntry->PartitionType == PARTITION_ENTRY_UNUSED &&
-                     PartEntry->SectorCount.QuadPart != 0LL))
+                    PartEntry->SectorCount.QuadPart != 0LL)
                 {
                     if (Letter <= L'Z')
                     {
@@ -373,9 +374,10 @@ AssignDriveLetters(
 
             if (PartEntry->IsPartitioned)
             {
+                ASSERT(PartEntry->PartitionType != PARTITION_ENTRY_UNUSED);
+
                 if (IsRecognizedPartition(PartEntry->PartitionType) ||
-                    (PartEntry->PartitionType == PARTITION_ENTRY_UNUSED &&
-                     PartEntry->SectorCount.QuadPart != 0LL))
+                    PartEntry->SectorCount.QuadPart != 0LL)
                 {
                     if (Letter <= L'Z')
                     {
@@ -903,6 +905,7 @@ AddPartitionToDisk(
     else if (IsRecognizedPartition(PartEntry->PartitionType))
     {
         ASSERT(PartitionInfo->RecognizedPartition);
+        ASSERT(PartEntry->IsPartitioned && PartEntry->PartitionNumber !=
0);
 
         /* Open the volume, ignore any errors */
         RtlStringCchPrintfW(PathBuffer, ARRAYSIZE(PathBuffer),
@@ -2384,7 +2387,7 @@ UpdateDiskLayout(
 
     if (DiskEntry->DiskStyle == PARTITION_STYLE_GPT)
     {
-        DPRINT("GPT-partitioned disk detected, not currently supported by
SETUP!\n");
+        DPRINT1("GPT-partitioned disk detected, not currently supported by
SETUP!\n");
         return;
     }
 
@@ -2405,14 +2408,16 @@ UpdateDiskLayout(
 
         if (PartEntry->IsPartitioned)
         {
+            ASSERT(PartEntry->PartitionType != PARTITION_ENTRY_UNUSED);
+
             PartitionInfo = &DiskEntry->LayoutBuffer->PartitionEntry[Index];
             PartEntry->PartitionIndex = Index;
 
-            /* Reset the current partition number only for newly-created partitions */
+            /* Reset the current partition number only for newly-created (unmounted)
partitions */
             if (PartEntry->New)
                 PartEntry->PartitionNumber = 0;
 
-            PartEntry->OnDiskPartitionNumber =
(!IsContainerPartition(PartEntry->PartitionType)) ? PartitionNumber : 0;
+            PartEntry->OnDiskPartitionNumber =
(!IsContainerPartition(PartEntry->PartitionType) ? PartitionNumber : 0);
 
             if (!IsSamePrimaryLayoutEntry(PartitionInfo, DiskEntry, PartEntry))
             {
@@ -2447,17 +2452,19 @@ UpdateDiskLayout(
 
         if (PartEntry->IsPartitioned)
         {
+            ASSERT(PartEntry->PartitionType != PARTITION_ENTRY_UNUSED);
+
             PartitionInfo = &DiskEntry->LayoutBuffer->PartitionEntry[Index];
             PartEntry->PartitionIndex = Index;
 
-            DPRINT1("Updating logical partition entry %lu\n", Index);
-
-            /* Reset the current partition number only for newly-created partitions */
+            /* Reset the current partition number only for newly-created (unmounted)
partitions */
             if (PartEntry->New)
                 PartEntry->PartitionNumber = 0;
 
             PartEntry->OnDiskPartitionNumber = PartitionNumber;
 
+            DPRINT1("Updating logical partition entry %lu\n", Index);
+
             PartitionInfo->StartingOffset.QuadPart =
PartEntry->StartSector.QuadPart * DiskEntry->BytesPerSector;
             PartitionInfo->PartitionLength.QuadPart =
PartEntry->SectorCount.QuadPart * DiskEntry->BytesPerSector;
             PartitionInfo->HiddenSectors = DiskEntry->SectorAlignment;
@@ -2569,7 +2576,10 @@ GetPrevUnpartitionedEntry(
                                           PARTENTRY,
                                           ListEntry);
         if (PrevPartEntry->IsPartitioned == FALSE)
+        {
+            ASSERT(PrevPartEntry->PartitionType == PARTITION_ENTRY_UNUSED);
             return PrevPartEntry;
+        }
     }
 
     return NULL;
@@ -2601,7 +2611,10 @@ GetNextUnpartitionedEntry(
                                           PARTENTRY,
                                           ListEntry);
         if (NextPartEntry->IsPartitioned == FALSE)
+        {
+            ASSERT(NextPartEntry->PartitionType == PARTITION_ENTRY_UNUSED);
             return NextPartEntry;
+        }
     }
 
     return NULL;
@@ -2610,39 +2623,36 @@ GetNextUnpartitionedEntry(
 BOOLEAN
 CreatePrimaryPartition(
     IN PPARTLIST List,
+    IN PPARTENTRY SelectedEntry,
     IN ULONGLONG SectorCount,
     IN BOOLEAN AutoCreate)
 {
     ERROR_NUMBER Error;
-    PDISKENTRY DiskEntry;
     PPARTENTRY PartEntry;
 
     DPRINT1("CreatePrimaryPartition(%I64u)\n", SectorCount);
 
     if (List == NULL ||
-        List->CurrentDisk == NULL ||
-        List->CurrentPartition == NULL ||
-        List->CurrentPartition->IsPartitioned)
+        SelectedEntry == NULL ||
+        SelectedEntry->DiskEntry == NULL ||
+        SelectedEntry->IsPartitioned)
     {
         return FALSE;
     }
 
-    Error = PrimaryPartitionCreationChecks(List);
+    Error = PrimaryPartitionCreationChecks(SelectedEntry);
     if (Error != NOT_AN_ERROR)
     {
         DPRINT1("PrimaryPartitionCreationChecks() failed with error %lu\n",
Error);
         return FALSE;
     }
 
-    DiskEntry = List->CurrentDisk;
-    PartEntry = List->CurrentPartition;
-
     /* Convert the current entry, or insert and initialize a new partition entry */
-    PartEntry = InitializePartitionEntry(DiskEntry, PartEntry, SectorCount, AutoCreate);
+    PartEntry = InitializePartitionEntry(SelectedEntry->DiskEntry, SelectedEntry,
SectorCount, AutoCreate);
     if (PartEntry == NULL)
         return FALSE;
 
-    UpdateDiskLayout(DiskEntry);
+    UpdateDiskLayout(PartEntry->DiskEntry);
 
     AssignDriveLetters(List);
 
@@ -2681,34 +2691,31 @@ AddLogicalDiskSpace(
 BOOLEAN
 CreateExtendedPartition(
     IN PPARTLIST List,
+    IN PPARTENTRY SelectedEntry,
     IN ULONGLONG SectorCount)
 {
     ERROR_NUMBER Error;
-    PDISKENTRY DiskEntry;
     PPARTENTRY PartEntry;
 
     DPRINT1("CreateExtendedPartition(%I64u)\n", SectorCount);
 
     if (List == NULL ||
-        List->CurrentDisk == NULL ||
-        List->CurrentPartition == NULL ||
-        List->CurrentPartition->IsPartitioned)
+        SelectedEntry == NULL ||
+        SelectedEntry->DiskEntry == NULL ||
+        SelectedEntry->IsPartitioned)
     {
         return FALSE;
     }
 
-    Error = ExtendedPartitionCreationChecks(List);
+    Error = ExtendedPartitionCreationChecks(SelectedEntry);
     if (Error != NOT_AN_ERROR)
     {
         DPRINT1("ExtendedPartitionCreationChecks() failed with error %lu\n",
Error);
         return FALSE;
     }
 
-    DiskEntry = List->CurrentDisk;
-    PartEntry = List->CurrentPartition;
-
     /* Convert the current entry, or insert and initialize a new partition entry */
-    PartEntry = InitializePartitionEntry(DiskEntry, PartEntry, SectorCount, FALSE);
+    PartEntry = InitializePartitionEntry(SelectedEntry->DiskEntry, SelectedEntry,
SectorCount, FALSE);
     if (PartEntry == NULL)
         return FALSE;
 
@@ -2727,11 +2734,11 @@ CreateExtendedPartition(
     PartEntry->New = FALSE;
     PartEntry->FormatState = Formatted;
 
-    DiskEntry->ExtendedPartition = PartEntry;
+    PartEntry->DiskEntry->ExtendedPartition = PartEntry;
 
-    AddLogicalDiskSpace(DiskEntry);
+    AddLogicalDiskSpace(PartEntry->DiskEntry);
 
-    UpdateDiskLayout(DiskEntry);
+    UpdateDiskLayout(PartEntry->DiskEntry);
 
     AssignDriveLetters(List);
 
@@ -2741,41 +2748,38 @@ CreateExtendedPartition(
 BOOLEAN
 CreateLogicalPartition(
     IN PPARTLIST List,
+    IN PPARTENTRY SelectedEntry,
     IN ULONGLONG SectorCount,
     IN BOOLEAN AutoCreate)
 {
     ERROR_NUMBER Error;
-    PDISKENTRY DiskEntry;
     PPARTENTRY PartEntry;
 
     DPRINT1("CreateLogicalPartition(%I64u)\n", SectorCount);
 
     if (List == NULL ||
-        List->CurrentDisk == NULL ||
-        List->CurrentPartition == NULL ||
-        List->CurrentPartition->IsPartitioned)
+        SelectedEntry == NULL ||
+        SelectedEntry->DiskEntry == NULL ||
+        SelectedEntry->IsPartitioned)
     {
         return FALSE;
     }
 
-    Error = LogicalPartitionCreationChecks(List);
+    Error = LogicalPartitionCreationChecks(SelectedEntry);
     if (Error != NOT_AN_ERROR)
     {
         DPRINT1("LogicalPartitionCreationChecks() failed with error %lu\n",
Error);
         return FALSE;
     }
 
-    DiskEntry = List->CurrentDisk;
-    PartEntry = List->CurrentPartition;
-
     /* Convert the current entry, or insert and initialize a new partition entry */
-    PartEntry = InitializePartitionEntry(DiskEntry, PartEntry, SectorCount, AutoCreate);
+    PartEntry = InitializePartitionEntry(SelectedEntry->DiskEntry, SelectedEntry,
SectorCount, AutoCreate);
     if (PartEntry == NULL)
         return FALSE;
 
     PartEntry->LogicalPartition = TRUE;
 
-    UpdateDiskLayout(DiskEntry);
+    UpdateDiskLayout(PartEntry->DiskEntry);
 
     AssignDriveLetters(List);
 
@@ -2795,11 +2799,10 @@ DismountVolume(
     HANDLE PartitionHandle;
     WCHAR Buffer[MAX_PATH];
 
-    /* Check whether the partition is valid and may have been mounted in the system */
+    /* Check whether the partition is valid and was mounted by the system */
     if (!PartEntry->IsPartitioned ||
-        PartEntry->PartitionType == PARTITION_ENTRY_UNUSED ||
-        IsContainerPartition(PartEntry->PartitionType)     ||
-        !IsRecognizedPartition(PartEntry->PartitionType)   ||
+        IsContainerPartition(PartEntry->PartitionType)   ||
+        !IsRecognizedPartition(PartEntry->PartitionType) ||
         PartEntry->FormatState == Unformatted /* || PartEntry->FormatState ==
UnknownFormat */ ||
         !*PartEntry->FileSystem ||
         PartEntry->PartitionNumber == 0)
@@ -2808,6 +2811,8 @@ DismountVolume(
         return STATUS_SUCCESS;
     }
 
+    ASSERT(PartEntry->PartitionType != PARTITION_ENTRY_UNUSED);
+
     /* Open the volume */
     RtlStringCchPrintfW(Buffer, ARRAYSIZE(Buffer),
                         L"\\Device\\Harddisk%lu\\Partition%lu",
@@ -2888,32 +2893,37 @@ DismountVolume(
 }
 
 VOID
-DeleteCurrentPartition(
-    IN PPARTLIST List)
+DeletePartition(
+    IN PPARTLIST List,
+    IN PPARTENTRY PartEntry)
 {
     PDISKENTRY DiskEntry;
-    PPARTENTRY PartEntry;
     PPARTENTRY PrevPartEntry;
     PPARTENTRY NextPartEntry;
     PPARTENTRY LogicalPartEntry;
     PLIST_ENTRY Entry;
 
     if (List == NULL ||
-        List->CurrentDisk == NULL ||
-        List->CurrentPartition == NULL ||
-        List->CurrentPartition->IsPartitioned == FALSE)
+        PartEntry == NULL ||
+        PartEntry->DiskEntry == NULL ||
+        PartEntry->IsPartitioned == FALSE)
     {
         return;
     }
 
-    /* Clear the system disk and partition pointers if the system partition is being
deleted */
-    if (List->SystemPartition == List->CurrentPartition)
+    ASSERT(PartEntry->PartitionType != PARTITION_ENTRY_UNUSED);
+
+    /* Clear the system partition pointers if it is being deleted */
+    if (List->SystemPartition == PartEntry)
     {
+        ASSERT(List->SystemPartition);
+
+        if (List->SystemPartition == List->OriginalSystemPartition)
+            List->OriginalSystemPartition = NULL;
         List->SystemPartition = NULL;
     }
 
-    DiskEntry = List->CurrentDisk;
-    PartEntry = List->CurrentPartition;
+    DiskEntry = PartEntry->DiskEntry;
 
     /* Check which type of partition (primary/logical or extended) is being deleted */
     if (DiskEntry->ExtendedPartition == PartEntry)
@@ -2961,7 +2971,11 @@ DeleteCurrentPartition(
         RtlFreeHeap(ProcessHeap, 0, NextPartEntry);
 
         /* Update current partition */
-        List->CurrentPartition = PrevPartEntry;
+        if (List->CurrentPartition == PartEntry ||
+            List->CurrentPartition == NextPartEntry)
+        {
+            List->CurrentPartition = PrevPartEntry;
+        }
     }
     else if (PrevPartEntry != NULL && NextPartEntry == NULL)
     {
@@ -2975,7 +2989,10 @@ DeleteCurrentPartition(
         RtlFreeHeap(ProcessHeap, 0, PartEntry);
 
         /* Update current partition */
-        List->CurrentPartition = PrevPartEntry;
+        if (List->CurrentPartition == PartEntry)
+        {
+            List->CurrentPartition = PrevPartEntry;
+        }
     }
     else if (PrevPartEntry == NULL && NextPartEntry != NULL)
     {
@@ -2990,7 +3007,10 @@ DeleteCurrentPartition(
         RtlFreeHeap(ProcessHeap, 0, PartEntry);
 
         /* Update current partition */
-        List->CurrentPartition = NextPartEntry;
+        if (List->CurrentPartition == PartEntry)
+        {
+            List->CurrentPartition = NextPartEntry;
+        }
     }
     else
     {
@@ -3010,6 +3030,13 @@ DeleteCurrentPartition(
     AssignDriveLetters(List);
 }
 
+VOID
+DeleteCurrentPartition(
+    IN PPARTLIST List)
+{
+    DeletePartition(List, List->CurrentPartition);
+}
+
 static
 BOOLEAN
 IsSupportedActivePartition(
@@ -3292,10 +3319,8 @@ SetSystemPartition:
     List->SystemPartition->DiskEntry->Dirty = TRUE;
 }
 
-static
 NTSTATUS
 WritePartitions(
-    IN PPARTLIST List,
     IN PDISKENTRY DiskEntry)
 {
     NTSTATUS Status;
@@ -3312,6 +3337,10 @@ WritePartitions(
 
     DPRINT("WritePartitions() Disk: %lu\n", DiskEntry->DiskNumber);
 
+    /* If the disk is not dirty, there is nothing to do */
+    if (!DiskEntry->Dirty)
+        return STATUS_SUCCESS;
+
     RtlStringCchPrintfW(DstPath, ARRAYSIZE(DstPath),
                         L"\\Device\\Harddisk%lu\\Partition0",
                         DiskEntry->DiskNumber);
@@ -3394,6 +3423,7 @@ WritePartitions(
 
         if (PartEntry->IsPartitioned)
         {
+            ASSERT(PartEntry->PartitionType != PARTITION_ENTRY_UNUSED);
             PartitionInfo =
&DiskEntry->LayoutBuffer->PartitionEntry[PartEntry->PartitionIndex];
             PartEntry->PartitionNumber = PartitionInfo->PartitionNumber;
         }
@@ -3408,6 +3438,7 @@ WritePartitions(
 
         if (PartEntry->IsPartitioned)
         {
+            ASSERT(PartEntry->PartitionType != PARTITION_ENTRY_UNUSED);
             PartitionInfo =
&DiskEntry->LayoutBuffer->PartitionEntry[PartEntry->PartitionIndex];
             PartEntry->PartitionNumber = PartitionInfo->PartitionNumber;
         }
@@ -3423,7 +3454,8 @@ WritePartitions(
     // DiskEntry->NoMbr was TRUE (instead of NewDisk).
     //
 
-    // DiskEntry->Dirty = FALSE;
+    /* The layout has been successfully updated, the disk is not dirty anymore */
+    DiskEntry->Dirty = FALSE;
 
     return Status;
 }
@@ -3432,6 +3464,7 @@ BOOLEAN
 WritePartitionsToDisk(
     IN PPARTLIST List)
 {
+    NTSTATUS Status;
     PLIST_ENTRY Entry;
     PDISKENTRY DiskEntry;
 
@@ -3452,8 +3485,12 @@ WritePartitionsToDisk(
 
         if (DiskEntry->Dirty != FALSE)
         {
-            WritePartitions(List, DiskEntry);
-            DiskEntry->Dirty = FALSE;
+            Status = WritePartitions(DiskEntry);
+            if (!NT_SUCCESS(Status))
+            {
+                DPRINT1("WritePartitionsToDisk() failed to update disk %lu, Status
0x%08lx\n",
+                        DiskEntry->DiskNumber, Status);
+            }
         }
     }
 
@@ -3553,8 +3590,10 @@ SetMountedDeviceValues(
              Entry2 = Entry2->Flink)
         {
             PartEntry = CONTAINING_RECORD(Entry2, PARTENTRY, ListEntry);
-            if (PartEntry->IsPartitioned)
+            if (PartEntry->IsPartitioned) // &&
!IsContainerPartition(PartEntry->PartitionType)
             {
+                ASSERT(PartEntry->PartitionType != PARTITION_ENTRY_UNUSED);
+
                 /* Assign a "\DosDevices\#:" mount point to this partition */
                 if (PartEntry->DriveLetter)
                 {
@@ -3574,8 +3613,10 @@ SetMountedDeviceValues(
              Entry2 = Entry2->Flink)
         {
             PartEntry = CONTAINING_RECORD(Entry2, PARTENTRY, ListEntry);
-            if (PartEntry->IsPartitioned)
+            if (PartEntry->IsPartitioned) // &&
!IsContainerPartition(PartEntry->PartitionType)
             {
+                ASSERT(PartEntry->PartitionType != PARTITION_ENTRY_UNUSED);
+
                 /* Assign a "\DosDevices\#:" mount point to this partition */
                 if (PartEntry->DriveLetter)
                 {
@@ -3611,13 +3652,9 @@ SetPartitionType(
 
 ERROR_NUMBER
 PrimaryPartitionCreationChecks(
-    IN PPARTLIST List)
+    IN PPARTENTRY PartEntry)
 {
-    PDISKENTRY DiskEntry;
-    PPARTENTRY PartEntry;
-
-    DiskEntry = List->CurrentDisk;
-    PartEntry = List->CurrentPartition;
+    PDISKENTRY DiskEntry = PartEntry->DiskEntry;
 
     if (DiskEntry->DiskStyle == PARTITION_STYLE_GPT)
     {
@@ -3638,13 +3675,9 @@ PrimaryPartitionCreationChecks(
 
 ERROR_NUMBER
 ExtendedPartitionCreationChecks(
-    IN PPARTLIST List)
+    IN PPARTENTRY PartEntry)
 {
-    PDISKENTRY DiskEntry;
-    PPARTENTRY PartEntry;
-
-    DiskEntry = List->CurrentDisk;
-    PartEntry = List->CurrentPartition;
+    PDISKENTRY DiskEntry = PartEntry->DiskEntry;
 
     if (DiskEntry->DiskStyle == PARTITION_STYLE_GPT)
     {
@@ -3669,13 +3702,9 @@ ExtendedPartitionCreationChecks(
 
 ERROR_NUMBER
 LogicalPartitionCreationChecks(
-    IN PPARTLIST List)
+    IN PPARTENTRY PartEntry)
 {
-    PDISKENTRY DiskEntry;
-    PPARTENTRY PartEntry;
-
-    DiskEntry = List->CurrentDisk;
-    PartEntry = List->CurrentPartition;
+    PDISKENTRY DiskEntry = PartEntry->DiskEntry;
 
     if (DiskEntry->DiskStyle == PARTITION_STYLE_GPT)
     {
@@ -3778,7 +3807,7 @@ GetNextUncheckedPartition(
              Entry2 = Entry2->Flink)
         {
             PartEntry = CONTAINING_RECORD(Entry2, PARTENTRY, ListEntry);
-            if (PartEntry->NeedsCheck == TRUE)
+            if (PartEntry->IsPartitioned && PartEntry->NeedsCheck)
             {
                 ASSERT(DiskEntry == PartEntry->DiskEntry);
                 if (pDiskEntry) *pDiskEntry = DiskEntry;
@@ -3792,7 +3821,7 @@ GetNextUncheckedPartition(
              Entry2 = Entry2->Flink)
         {
             PartEntry = CONTAINING_RECORD(Entry2, PARTENTRY, ListEntry);
-            if (PartEntry->NeedsCheck == TRUE)
+            if (PartEntry->IsPartitioned && PartEntry->NeedsCheck)
             {
                 ASSERT(DiskEntry == PartEntry->DiskEntry);
                 if (pDiskEntry) *pDiskEntry = DiskEntry;
diff --git a/base/setup/lib/utils/partlist.h b/base/setup/lib/utils/partlist.h
index 57873a444e..ab4905a6d1 100644
--- a/base/setup/lib/utils/partlist.h
+++ b/base/setup/lib/utils/partlist.h
@@ -292,20 +292,28 @@ GetPrevPartition(
 BOOLEAN
 CreatePrimaryPartition(
     IN PPARTLIST List,
+    IN PPARTENTRY SelectedEntry,
     IN ULONGLONG SectorCount,
     IN BOOLEAN AutoCreate);
 
 BOOLEAN
 CreateExtendedPartition(
     IN PPARTLIST List,
+    IN PPARTENTRY SelectedEntry,
     IN ULONGLONG SectorCount);
 
 BOOLEAN
 CreateLogicalPartition(
     IN PPARTLIST List,
+    IN PPARTENTRY SelectedEntry,
     IN ULONGLONG SectorCount,
     IN BOOLEAN AutoCreate);
 
+VOID
+DeletePartition(
+    IN PPARTLIST List,
+    IN PPARTENTRY PartEntry);
+
 VOID
 DeleteCurrentPartition(
     IN PPARTLIST List);
@@ -314,6 +322,10 @@ VOID
 CheckActiveSystemPartition(
     IN PPARTLIST List);
 
+NTSTATUS
+WritePartitions(
+    IN PDISKENTRY DiskEntry);
+
 BOOLEAN
 WritePartitionsToDisk(
     IN PPARTLIST List);
@@ -335,15 +347,15 @@ SetPartitionType(
 
 ERROR_NUMBER
 PrimaryPartitionCreationChecks(
-    IN PPARTLIST List);
+    IN PPARTENTRY PartEntry);
 
 ERROR_NUMBER
 ExtendedPartitionCreationChecks(
-    IN PPARTLIST List);
+    IN PPARTENTRY PartEntry);
 
 ERROR_NUMBER
 LogicalPartitionCreationChecks(
-    IN PPARTLIST List);
+    IN PPARTENTRY PartEntry);
 
 BOOLEAN
 GetNextUnformattedPartition(
diff --git a/base/setup/usetup/usetup.c b/base/setup/usetup/usetup.c
index b8df7c9261..555d37e83a 100644
--- a/base/setup/usetup/usetup.c
+++ b/base/setup/usetup/usetup.c
@@ -1499,12 +1499,14 @@ SelectPartitionPage(PINPUT_RECORD Ir)
                 if (PartitionList->CurrentPartition->LogicalPartition)
                 {
                     CreateLogicalPartition(PartitionList,
+                                           PartitionList->CurrentPartition,
                                           
PartitionList->CurrentPartition->SectorCount.QuadPart,
                                            TRUE);
                 }
                 else
                 {
                     CreatePrimaryPartition(PartitionList,
+                                           PartitionList->CurrentPartition,
                                           
PartitionList->CurrentPartition->SectorCount.QuadPart,
                                            TRUE);
                 }
@@ -1607,7 +1609,7 @@ SelectPartitionPage(PINPUT_RECORD Ir)
             {
                 if (PartitionList->CurrentPartition->LogicalPartition)
                 {
-                    Error = LogicalPartitionCreationChecks(PartitionList);
+                    Error =
LogicalPartitionCreationChecks(PartitionList->CurrentPartition);
                     if (Error != NOT_AN_ERROR)
                     {
                         MUIDisplayError(Error, Ir, POPUP_WAIT_ANY_KEY);
@@ -1615,12 +1617,13 @@ SelectPartitionPage(PINPUT_RECORD Ir)
                     }
 
                     CreateLogicalPartition(PartitionList,
+                                           PartitionList->CurrentPartition,
                                            0ULL,
                                            TRUE);
                 }
                 else
                 {
-                    Error = PrimaryPartitionCreationChecks(PartitionList);
+                    Error =
PrimaryPartitionCreationChecks(PartitionList->CurrentPartition);
                     if (Error != NOT_AN_ERROR)
                     {
                         MUIDisplayError(Error, Ir, POPUP_WAIT_ANY_KEY);
@@ -1628,6 +1631,7 @@ SelectPartitionPage(PINPUT_RECORD Ir)
                     }
 
                     CreatePrimaryPartition(PartitionList,
+                                           PartitionList->CurrentPartition,
                                            0ULL,
                                            TRUE);
                 }
@@ -1646,7 +1650,7 @@ SelectPartitionPage(PINPUT_RECORD Ir)
         {
             if (PartitionList->CurrentPartition->LogicalPartition == FALSE)
             {
-                Error = PrimaryPartitionCreationChecks(PartitionList);
+                Error =
PrimaryPartitionCreationChecks(PartitionList->CurrentPartition);
                 if (Error != NOT_AN_ERROR)
                 {
                     MUIDisplayError(Error, Ir, POPUP_WAIT_ANY_KEY);
@@ -1660,7 +1664,7 @@ SelectPartitionPage(PINPUT_RECORD Ir)
         {
             if (PartitionList->CurrentPartition->LogicalPartition == FALSE)
             {
-                Error = ExtendedPartitionCreationChecks(PartitionList);
+                Error =
ExtendedPartitionCreationChecks(PartitionList->CurrentPartition);
                 if (Error != NOT_AN_ERROR)
                 {
                     MUIDisplayError(Error, Ir, POPUP_WAIT_ANY_KEY);
@@ -1674,7 +1678,7 @@ SelectPartitionPage(PINPUT_RECORD Ir)
         {
             if (PartitionList->CurrentPartition->LogicalPartition)
             {
-                Error = LogicalPartitionCreationChecks(PartitionList);
+                Error =
LogicalPartitionCreationChecks(PartitionList->CurrentPartition);
                 if (Error != NOT_AN_ERROR)
                 {
                     MUIDisplayError(Error, Ir, POPUP_WAIT_ANY_KEY);
@@ -2068,6 +2072,7 @@ CreatePrimaryPartitionPage(PINPUT_RECORD Ir)
             DPRINT ("Partition size: %I64u bytes\n", PartSize);
 
             CreatePrimaryPartition(PartitionList,
+                                   PartitionList->CurrentPartition,
                                    SectorCount,
                                    FALSE);
 
@@ -2231,6 +2236,7 @@ CreateExtendedPartitionPage(PINPUT_RECORD Ir)
             DPRINT ("Partition size: %I64u bytes\n", PartSize);
 
             CreateExtendedPartition(PartitionList,
+                                    PartitionList->CurrentPartition,
                                     SectorCount);
 
             return SELECT_PARTITION_PAGE;
@@ -2393,6 +2399,7 @@ CreateLogicalPartitionPage(PINPUT_RECORD Ir)
             DPRINT("Partition size: %I64u bytes\n", PartSize);
 
             CreateLogicalPartition(PartitionList,
+                                   PartitionList->CurrentPartition,
                                    SectorCount,
                                    FALSE);
 
@@ -2736,6 +2743,8 @@ SelectFileSystemPage(PINPUT_RECORD Ir)
     PartEntry = TempPartition;
     DiskEntry = PartEntry->DiskEntry;
 
+    ASSERT(PartEntry->IsPartitioned && PartEntry->PartitionNumber != 0);
+
     /* Adjust disk size */
     DiskSize = DiskEntry->SectorCount.QuadPart * DiskEntry->BytesPerSector;
     if (DiskSize >= 10 * GB) /* 10 GB */
@@ -2997,7 +3006,10 @@ FormatPartitionPage(PINPUT_RECORD Ir)
     PartEntry = TempPartition;
     DiskEntry = PartEntry->DiskEntry;
 
+    ASSERT(PartEntry->IsPartitioned && PartEntry->PartitionNumber != 0);
+
     SelectedFileSystem = FileSystemList->Selected;
+    ASSERT(SelectedFileSystem && SelectedFileSystem->FileSystem);
 
     while (TRUE)
     {
@@ -3163,6 +3175,8 @@ CheckFileSystemPage(PINPUT_RECORD Ir)
         return INSTALL_DIRECTORY_PAGE;
     }
 
+    ASSERT(PartEntry->IsPartitioned && PartEntry->PartitionNumber != 0);
+
     CONSOLE_SetTextXY(6, 8, MUIGetString(STRING_CHECKINGPART));
 
     CONSOLE_SetStatusText(MUIGetString(STRING_PLEASEWAIT));
@@ -3930,6 +3944,8 @@ BootLoaderPage(PINPUT_RECORD Ir)
 
     CONSOLE_SetStatusText(MUIGetString(STRING_PLEASEWAIT));
 
+    ASSERT(PartitionList->SystemPartition->IsPartitioned &&
PartitionList->SystemPartition->PartitionNumber != 0);
+
     RtlFreeUnicodeString(&USetupData.SystemRootPath);
     RtlStringCchPrintfW(PathBuffer, ARRAYSIZE(PathBuffer),
             L"\\Device\\Harddisk%lu\\Partition%lu\\",