https://git.reactos.org/?p=reactos.git;a=commitdiff;h=723947eab127001405d07963f2046e147be6498f

commit 723947eab127001405d07963f2046e147be6498f
Author:     Eric Kohl <[email protected]>
AuthorDate: Sat May 28 11:33:29 2022 +0200
Commit:     Eric Kohl <[email protected]>
CommitDate: Sat May 28 11:33:29 2022 +0200

    [DISKPART] Implement CREATE PARTITION PRIMARY and DELETE PARTITION
---
 base/system/diskpart/create.c   | 146 ++++++++--
 base/system/diskpart/delete.c   | 121 ++++++++
 base/system/diskpart/diskpart.h |  34 ++-
 base/system/diskpart/partlist.c | 630 ++++++++++++++++++++++++++++++++++++++--
 4 files changed, 891 insertions(+), 40 deletions(-)

diff --git a/base/system/diskpart/create.c b/base/system/diskpart/create.c
index a933da2c5ae..095bf4e3d10 100644
--- a/base/system/diskpart/create.c
+++ b/base/system/diskpart/create.c
@@ -8,6 +8,10 @@
 
 #include "diskpart.h"
 
+#define NDEBUG
+#include <debug.h>
+
+
 BOOL
 CreateExtendedPartition(
     INT argc,
@@ -47,32 +51,30 @@ CreatePrimaryPartition(
     INT argc,
     PWSTR *argv)
 {
-    LARGE_INTEGER liSize, liOffset;
-    INT i;
+    PPARTENTRY PartEntry, NewPartEntry;
+    PLIST_ENTRY ListEntry;
+    ULONGLONG ullSize = 0ULL, ullOffset = 0ULL;
+    ULONGLONG ullSectorCount;
+    UCHAR PartitionType = 6;
+    INT i, length;
 //    BOOL bNoErr = FALSE;
     PWSTR pszSuffix = NULL;
 
-    liSize.QuadPart = -1;
-    liOffset.QuadPart = -1;
-
-/*
     if (CurrentDisk == NULL)
     {
         ConResPuts(StdOut, IDS_SELECT_NO_DISK);
         return TRUE;
     }
-*/
 
     for (i = 3; i < argc; i++)
     {
         if (HasPrefix(argv[i], L"size=", &pszSuffix))
         {
             /* size=<N> (MB) */
-            ConPrintf(StdOut, L"Size : %s\n", pszSuffix);
+            DPRINT("Size : %s\n", pszSuffix);
 
-            liSize.QuadPart = _wcstoui64(pszSuffix, NULL, 10);
-            if (((liSize.QuadPart == 0) && (errno == ERANGE)) ||
-                (liSize.QuadPart < 0))
+            ullSize = _wcstoui64(pszSuffix, NULL, 10);
+            if ((ullSize == 0) && (errno == ERANGE))
             {
                 ConResPuts(StdErr, IDS_ERROR_INVALID_ARGS);
                 return TRUE;
@@ -81,11 +83,10 @@ CreatePrimaryPartition(
         else if (HasPrefix(argv[i], L"offset=", &pszSuffix))
         {
             /* offset=<N> (KB) */
-            ConPrintf(StdOut, L"Offset : %s\n", pszSuffix);
+            DPRINT("Offset : %s\n", pszSuffix);
 
-            liOffset.QuadPart = _wcstoui64(pszSuffix, NULL, 10);
-            if (((liOffset.QuadPart == 0) && (errno == ERANGE)) ||
-                (liOffset.QuadPart < 0))
+            ullOffset = _wcstoui64(pszSuffix, NULL, 10);
+            if ((ullOffset == 0) && (errno == ERANGE))
             {
                 ConResPuts(StdErr, IDS_ERROR_INVALID_ARGS);
                 return TRUE;
@@ -94,17 +95,41 @@ CreatePrimaryPartition(
         else if (HasPrefix(argv[i], L"id=", &pszSuffix))
         {
             /* id=<Byte>|<GUID> */
-            ConPrintf(StdOut, L"Id : %s\n", pszSuffix);
+            DPRINT("Id : %s\n", pszSuffix);
+
+            length = wcslen(pszSuffix);
+            if ((length == 1) || (length == 2))
+            {
+                /* Byte */
+                PartitionType = (UCHAR)wcstoul(pszSuffix, NULL, 16);
+                if ((PartitionType == 0) && (errno == ERANGE))
+                {
+                    ConResPuts(StdErr, IDS_ERROR_INVALID_ARGS);
+                    return TRUE;
+                }
+            }
+#if 0
+            else if ()
+            {
+                /* GUID */
+            }
+#endif
+            else
+            {
+                ConResPuts(StdErr, IDS_ERROR_INVALID_ARGS);
+                return TRUE; 
+            }
         }
         else if (HasPrefix(argv[i], L"align=", &pszSuffix))
         {
             /* align=<N> */
-            ConPrintf(StdOut, L"Align : %s\n", pszSuffix);
+            DPRINT("Align : %s\n", pszSuffix);
+//            bAlign = TRUE;
         }
         else if (_wcsicmp(argv[i], L"noerr") == 0)
         {
             /* noerr */
-            ConPrintf(StdOut, L"NoErr\n", pszSuffix);
+            DPRINT("NoErr\n", pszSuffix);
 //            bNoErr = TRUE;
         }
         else
@@ -114,8 +139,89 @@ CreatePrimaryPartition(
         }
     }
 
-    ConPrintf(StdOut, L"Size: %I64d\n", liSize.QuadPart);
-    ConPrintf(StdOut, L"Offset: %I64d\n", liOffset.QuadPart);
+    DPRINT1("Size: %I64u\n", ullSize);
+    DPRINT1("Offset: %I64u\n", ullOffset);
+    DPRINT1("Partition Type: %hx\n", PartitionType);
+
+    if (GetPrimaryPartitionCount(CurrentDisk) >= 4)
+    {
+        ConPuts(StdOut, L"No space left for another primary partition!\n");
+        return TRUE;
+    }
+
+    if (ullSize != 0)
+        ullSectorCount = (ullSize * 1024 * 1024) / CurrentDisk->BytesPerSector;
+    else
+        ullSectorCount = 0;
+
+    DPRINT1("SectorCount: %I64u\n", ullSectorCount);
+
+    for (ListEntry = CurrentDisk->PrimaryPartListHead.Flink;
+         ListEntry != &CurrentDisk->PrimaryPartListHead;
+         ListEntry = ListEntry->Flink)
+    {
+        PartEntry = CONTAINING_RECORD(ListEntry, PARTENTRY, ListEntry);
+        if (PartEntry->IsPartitioned)
+            continue;
+
+        if (ullSectorCount == 0)
+        {
+            PartEntry->IsPartitioned = TRUE;
+            PartEntry->New = TRUE;
+            PartEntry->PartitionType = PartitionType;
+            PartEntry->FormatState = Unformatted;
+            PartEntry->FileSystemName[0] = L'\0';
+
+            CurrentDisk->Dirty = TRUE;
+            break;
+        }
+        else
+        {
+            if (PartEntry->SectorCount.QuadPart == ullSectorCount)
+            {
+                PartEntry->IsPartitioned = TRUE;
+                PartEntry->New = TRUE;
+                PartEntry->PartitionType = PartitionType;
+                PartEntry->FormatState = Unformatted;
+                PartEntry->FileSystemName[0] = L'\0';
+
+                CurrentDisk->Dirty = TRUE;
+                break;
+            }
+            else if (PartEntry->SectorCount.QuadPart > ullSectorCount)
+            {
+                NewPartEntry = RtlAllocateHeap(RtlGetProcessHeap(), 
HEAP_ZERO_MEMORY, sizeof(PPARTENTRY));
+                if (NewPartEntry == NULL)
+                {
+                    ConPuts(StdOut, L"Memory allocation failed!\n");
+                    return TRUE;
+                }
+
+                NewPartEntry->DiskEntry = PartEntry->DiskEntry;
+
+                NewPartEntry->StartSector.QuadPart = 
PartEntry->StartSector.QuadPart;
+                NewPartEntry->SectorCount.QuadPart = ullSectorCount;
+
+                NewPartEntry->LogicalPartition = FALSE;
+                NewPartEntry->IsPartitioned = TRUE;
+                NewPartEntry->New = TRUE;
+                NewPartEntry->PartitionType = PartitionType;
+                NewPartEntry->FormatState = Unformatted;
+                NewPartEntry->FileSystemName[0] = L'\0';
+
+                PartEntry->StartSector.QuadPart += ullSectorCount;
+                PartEntry->SectorCount.QuadPart -= ullSectorCount;
+
+                InsertTailList(ListEntry, &NewPartEntry->ListEntry);
+
+                CurrentDisk->Dirty = TRUE;
+                break;
+            }
+        }
+    }
+
+    UpdateDiskLayout(CurrentDisk);
+    WritePartitions(CurrentDisk);
 
     return TRUE;
 }
diff --git a/base/system/diskpart/delete.c b/base/system/diskpart/delete.c
index de712944dcc..c17498bd863 100644
--- a/base/system/diskpart/delete.c
+++ b/base/system/diskpart/delete.c
@@ -8,6 +8,9 @@
 
 #include "diskpart.h"
 
+#define NDEBUG
+#include <debug.h>
+
 BOOL
 DeleteDisk(
     _In_ INT argc,
@@ -16,14 +19,132 @@ DeleteDisk(
     return TRUE;
 }
 
+
 BOOL
 DeletePartition(
     _In_ INT argc,
     _In_ PWSTR *argv)
 {
+    PPARTENTRY PrevPartEntry;
+    PPARTENTRY NextPartEntry;
+    PPARTENTRY LogicalPartEntry;
+    PLIST_ENTRY Entry;
+
+    DPRINT("DeletePartition()\n");
+
+    if (CurrentDisk == NULL)
+    {
+        ConResPuts(StdOut, IDS_SELECT_NO_DISK);
+        return TRUE;
+    }
+
+    if (CurrentPartition == NULL)
+    {
+        ConResPuts(StdOut, IDS_SELECT_NO_PARTITION);
+        return TRUE;
+    }
+
+    ASSERT(CurrentPartition->PartitionType != PARTITION_ENTRY_UNUSED);
+
+    /* Clear the system partition if it is being deleted */
+#if 0
+    if (List->SystemPartition == PartEntry)
+    {
+        ASSERT(List->SystemPartition);
+        List->SystemPartition = NULL;
+    }
+#endif
+
+    /* Check which type of partition (primary/logical or extended) is being 
deleted */
+    if (CurrentDisk->ExtendedPartition == CurrentPartition)
+    {
+        /* An extended partition is being deleted: delete all logical 
partition entries */
+        while (!IsListEmpty(&CurrentDisk->LogicalPartListHead))
+        {
+            Entry = RemoveHeadList(&CurrentDisk->LogicalPartListHead);
+            LogicalPartEntry = CONTAINING_RECORD(Entry, PARTENTRY, ListEntry);
+
+            /* Dismount the logical partition */
+            DismountVolume(LogicalPartEntry);
+
+            /* Delete it */
+            RtlFreeHeap(RtlGetProcessHeap(), 0, LogicalPartEntry);
+        }
+
+        CurrentDisk->ExtendedPartition = NULL;
+    }
+    else
+    {
+        /* A primary partition is being deleted: dismount it */
+        DismountVolume(CurrentPartition);
+    }
+
+    /* Adjust the unpartitioned disk space entries */
+
+    /* Get pointer to previous and next unpartitioned entries */
+    PrevPartEntry = GetPrevUnpartitionedEntry(CurrentPartition);
+    NextPartEntry = GetNextUnpartitionedEntry(CurrentPartition);
+
+    if (PrevPartEntry != NULL && NextPartEntry != NULL)
+    {
+        /* Merge the previous, current and next unpartitioned entries */
+
+        /* Adjust the previous entry length */
+        PrevPartEntry->SectorCount.QuadPart += 
(CurrentPartition->SectorCount.QuadPart + NextPartEntry->SectorCount.QuadPart);
+
+        /* Remove the current and next entries */
+        RemoveEntryList(&CurrentPartition->ListEntry);
+        RtlFreeHeap(RtlGetProcessHeap(), 0, CurrentPartition);
+        RemoveEntryList(&NextPartEntry->ListEntry);
+        RtlFreeHeap(RtlGetProcessHeap(), 0, NextPartEntry);
+    }
+    else if (PrevPartEntry != NULL && NextPartEntry == NULL)
+    {
+        /* Merge the current and the previous unpartitioned entries */
+
+        /* Adjust the previous entry length */
+        PrevPartEntry->SectorCount.QuadPart += 
CurrentPartition->SectorCount.QuadPart;
+
+        /* Remove the current entry */
+        RemoveEntryList(&CurrentPartition->ListEntry);
+        RtlFreeHeap(RtlGetProcessHeap(), 0, CurrentPartition);
+    }
+    else if (PrevPartEntry == NULL && NextPartEntry != NULL)
+    {
+        /* Merge the current and the next unpartitioned entries */
+
+        /* Adjust the next entry offset and length */
+        NextPartEntry->StartSector.QuadPart = 
CurrentPartition->StartSector.QuadPart;
+        NextPartEntry->SectorCount.QuadPart += 
CurrentPartition->SectorCount.QuadPart;
+
+        /* Remove the current entry */
+        RemoveEntryList(&CurrentPartition->ListEntry);
+        RtlFreeHeap(RtlGetProcessHeap(), 0, CurrentPartition);
+    }
+    else
+    {
+        /* Nothing to merge but change the current entry */
+        CurrentPartition->IsPartitioned = FALSE;
+        CurrentPartition->OnDiskPartitionNumber = 0;
+        CurrentPartition->PartitionNumber = 0;
+        // CurrentPartition->PartitionIndex = 0;
+        CurrentPartition->BootIndicator = FALSE;
+        CurrentPartition->PartitionType = PARTITION_ENTRY_UNUSED;
+        CurrentPartition->FormatState = Unformatted;
+        CurrentPartition->FileSystemName[0] = L'\0';
+        CurrentPartition->DriveLetter = 0;
+        RtlZeroMemory(CurrentPartition->VolumeLabel, 
sizeof(CurrentPartition->VolumeLabel));
+    }
+
+    CurrentPartition = NULL;
+
+    UpdateDiskLayout(CurrentDisk);
+    WritePartitions(CurrentDisk);
+
     return TRUE;
 }
 
+
 BOOL
 DeleteVolume(
     _In_ INT argc,
diff --git a/base/system/diskpart/diskpart.h b/base/system/diskpart/diskpart.h
index f13e6724262..218b6030a17 100644
--- a/base/system/diskpart/diskpart.h
+++ b/base/system/diskpart/diskpart.h
@@ -94,13 +94,14 @@ typedef struct _PARTENTRY
 
     BOOLEAN BootIndicator;
     UCHAR PartitionType;
-    ULONG HiddenSectors;
+    ULONG OnDiskPartitionNumber;
     ULONG PartitionNumber;
     ULONG PartitionIndex;
 
     CHAR DriveLetter;
     CHAR VolumeLabel[17];
     CHAR FileSystemName[9];
+    FORMATSTATE FormatState;
 
     BOOLEAN LogicalPartition;
 
@@ -113,8 +114,6 @@ typedef struct _PARTENTRY
     /* Partition was created automatically. */
     BOOLEAN AutoCreate;
 
-    FORMATSTATE FormatState;
-
     /* Partition must be checked */
     BOOLEAN NeedsCheck;
 
@@ -383,6 +382,11 @@ BOOL offline_main(INT argc, LPWSTR *argv);
 BOOL online_main(INT argc, LPWSTR *argv);
 
 /* partlist.c */
+ULONGLONG
+AlignDown(
+    _In_ ULONGLONG Value,
+    _In_ ULONG Alignment);
+
 NTSTATUS
 CreatePartitionList(VOID);
 
@@ -395,6 +399,30 @@ CreateVolumeList(VOID);
 VOID
 DestroyVolumeList(VOID);
 
+NTSTATUS
+WritePartitions(
+    _In_ PDISKENTRY DiskEntry);
+
+VOID
+UpdateDiskLayout(
+    _In_ PDISKENTRY DiskEntry);
+
+PPARTENTRY
+GetPrevUnpartitionedEntry(
+    _In_ PPARTENTRY PartEntry);
+
+PPARTENTRY
+GetNextUnpartitionedEntry(
+    _In_ PPARTENTRY PartEntry);
+
+ULONG
+GetPrimaryPartitionCount(
+    _In_ PDISKENTRY DiskEntry);
+
+NTSTATUS
+DismountVolume(
+    IN PPARTENTRY PartEntry);
+
 /* recover.c */
 BOOL recover_main(INT argc, LPWSTR *argv);
 
diff --git a/base/system/diskpart/partlist.c b/base/system/diskpart/partlist.c
index 6d91e385571..568f902ea38 100644
--- a/base/system/diskpart/partlist.c
+++ b/base/system/diskpart/partlist.c
@@ -10,6 +10,7 @@
 
 #include "diskpart.h"
 #include <ntddscsi.h>
+#include <strsafe.h>
 
 #define NDEBUG
 #include <debug.h>
@@ -80,8 +81,8 @@ PVOLENTRY  CurrentVolume = NULL;
 
 ULONGLONG
 AlignDown(
-    IN ULONGLONG Value,
-    IN ULONG Alignment)
+    _In_ ULONGLONG Value,
+    _In_ ULONG Alignment)
 {
     ULONGLONG Temp;
 
@@ -102,9 +103,9 @@ GetDriverName(
     RtlInitUnicodeString(&DiskEntry->DriverName,
                          NULL);
 
-    swprintf(KeyName,
-             L"\\Scsi\\Scsi Port %lu",
-             DiskEntry->Port);
+    StringCchPrintfW(KeyName, ARRAYSIZE(KeyName),
+                     L"\\Scsi\\Scsi Port %lu",
+                     DiskEntry->Port);
 
     RtlZeroMemory(&QueryTable,
                   sizeof(QueryTable));
@@ -280,7 +281,8 @@ EnumerateBiosDiskEntries(VOID)
     AdapterCount = 0;
     while (1)
     {
-        swprintf(Name, L"%s\\%lu", ROOT_NAME, AdapterCount);
+        StringCchPrintfW(Name, ARRAYSIZE(Name),
+                         L"%s\\%lu", ROOT_NAME, AdapterCount);
         Status = RtlQueryRegistryValues(RTL_REGISTRY_ABSOLUTE,
                                         Name,
                                         &QueryTable[2],
@@ -291,7 +293,8 @@ EnumerateBiosDiskEntries(VOID)
             break;
         }
 
-        swprintf(Name, L"%s\\%lu\\DiskController", ROOT_NAME, AdapterCount);
+        StringCchPrintfW(Name, ARRAYSIZE(Name),
+                         L"%s\\%lu\\DiskController", ROOT_NAME, AdapterCount);
         Status = RtlQueryRegistryValues(RTL_REGISTRY_ABSOLUTE,
                                         Name,
                                         &QueryTable[2],
@@ -301,7 +304,8 @@ EnumerateBiosDiskEntries(VOID)
         {
             while (1)
             {
-                swprintf(Name, L"%s\\%lu\\DiskController\\0", ROOT_NAME, 
AdapterCount);
+                StringCchPrintfW(Name, ARRAYSIZE(Name),
+                                 L"%s\\%lu\\DiskController\\0", ROOT_NAME, 
AdapterCount);
                 Status = RtlQueryRegistryValues(RTL_REGISTRY_ABSOLUTE,
                                                 Name,
                                                 &QueryTable[2],
@@ -313,7 +317,8 @@ EnumerateBiosDiskEntries(VOID)
                     return;
                 }
 
-                swprintf(Name, L"%s\\%lu\\DiskController\\0\\DiskPeripheral", 
ROOT_NAME, AdapterCount);
+                StringCchPrintfW(Name, ARRAYSIZE(Name),
+                                 
L"%s\\%lu\\DiskController\\0\\DiskPeripheral", ROOT_NAME, AdapterCount);
                 Status = RtlQueryRegistryValues(RTL_REGISTRY_ABSOLUTE,
                                                 Name,
                                                 &QueryTable[2],
@@ -335,7 +340,8 @@ EnumerateBiosDiskEntries(VOID)
                             break;
                         }
 
-                        swprintf(Name, 
L"%s\\%lu\\DiskController\\0\\DiskPeripheral\\%lu", ROOT_NAME, AdapterCount, 
DiskCount);
+                        StringCchPrintfW(Name, ARRAYSIZE(Name),
+                                         
L"%s\\%lu\\DiskController\\0\\DiskPeripheral\\%lu", ROOT_NAME, AdapterCount, 
DiskCount);
                         Status = RtlQueryRegistryValues(RTL_REGISTRY_ABSOLUTE,
                                                         Name,
                                                         QueryTable,
@@ -420,7 +426,6 @@ AddPartitionToDisk(
 
     PartEntry->BootIndicator = PartitionInfo->BootIndicator;
     PartEntry->PartitionType = PartitionInfo->PartitionType;
-    PartEntry->HiddenSectors = PartitionInfo->HiddenSectors;
 
     PartEntry->LogicalPartition = LogicalPartition;
     PartEntry->IsPartitioned = TRUE;
@@ -849,7 +854,8 @@ AddDiskToList(
     }
     Checksum = ~Checksum + 1;
 
-    swprintf(Identifier, L"%08x-%08x-A", Checksum, Signature);
+    StringCchPrintfW(Identifier, ARRAYSIZE(Identifier),
+                     L"%08x-%08x-A", Checksum, Signature);
     DPRINT("Identifier: %S\n", Identifier);
 
     DiskEntry = RtlAllocateHeap(RtlGetProcessHeap(),
@@ -927,8 +933,10 @@ AddDiskToList(
                                       
(ULONGLONG)DiskGeometry.TracksPerCylinder *
                                       (ULONGLONG)DiskGeometry.SectorsPerTrack;
 
-    DiskEntry->SectorAlignment = DiskGeometry.SectorsPerTrack;
-    DiskEntry->CylinderAlignment = DiskGeometry.SectorsPerTrack * 
DiskGeometry.TracksPerCylinder;
+//    DiskEntry->SectorAlignment = DiskGeometry.SectorsPerTrack;
+//    DiskEntry->CylinderAlignment = DiskGeometry.SectorsPerTrack * 
DiskGeometry.TracksPerCylinder;
+    DiskEntry->SectorAlignment = (1024 * 1024) / DiskGeometry.BytesPerSector;
+    DiskEntry->CylinderAlignment = (1024 * 1024) / DiskGeometry.BytesPerSector;
 
     DPRINT1("SectorCount: %I64u\n", DiskEntry->SectorCount);
     DPRINT1("SectorAlignment: %lu\n", DiskEntry->SectorAlignment);
@@ -1085,9 +1093,10 @@ CreatePartitionList(VOID)
 
     for (DiskNumber = 0; DiskNumber < Sdi.NumberOfDisks; DiskNumber++)
     {
-        swprintf(Buffer,
-                 L"\\Device\\Harddisk%d\\Partition0",
-                 DiskNumber);
+        StringCchPrintfW(Buffer, ARRAYSIZE(Buffer),
+                         L"\\Device\\Harddisk%d\\Partition0",
+                         DiskNumber);
+
         RtlInitUnicodeString(&Name,
                              Buffer);
 
@@ -1309,4 +1318,591 @@ DestroyVolumeList(VOID)
     }
 }
 
+
+NTSTATUS
+WritePartitions(
+    _In_ PDISKENTRY DiskEntry)
+{
+    NTSTATUS Status;
+    OBJECT_ATTRIBUTES ObjectAttributes;
+    UNICODE_STRING Name;
+    HANDLE FileHandle;
+    IO_STATUS_BLOCK Iosb;
+    ULONG BufferSize;
+    PPARTITION_INFORMATION PartitionInfo;
+    ULONG PartitionCount;
+    PLIST_ENTRY ListEntry;
+    PPARTENTRY PartEntry;
+    WCHAR DstPath[MAX_PATH];
+
+    DPRINT("WritePartitions() Disk: %lu\n", DiskEntry->DiskNumber);
+
+    /* If the disk is not dirty, there is nothing to do */
+    if (!DiskEntry->Dirty)
+        return STATUS_SUCCESS;
+
+    StringCchPrintfW(DstPath, ARRAYSIZE(DstPath),
+                     L"\\Device\\Harddisk%lu\\Partition0",
+                     DiskEntry->DiskNumber);
+    RtlInitUnicodeString(&Name, DstPath);
+
+    InitializeObjectAttributes(&ObjectAttributes,
+                               &Name,
+                               OBJ_CASE_INSENSITIVE,
+                               NULL,
+                               NULL);
+
+    Status = NtOpenFile(&FileHandle,
+                        GENERIC_READ | GENERIC_WRITE | SYNCHRONIZE,
+                        &ObjectAttributes,
+                        &Iosb,
+                        0,
+                        FILE_SYNCHRONOUS_IO_NONALERT);
+    if (!NT_SUCCESS(Status))
+    {
+        DPRINT1("NtOpenFile() failed (Status %lx)\n", Status);
+        return Status;
+    }
+
+    //
+    // FIXME: We first *MUST* use IOCTL_DISK_CREATE_DISK to initialize
+    // the disk in MBR or GPT format in case the disk was not initialized!!
+    // For this we must ask the user which format to use.
+    //
+
+    /* Save the original partition count to be restored later (see comment 
below) */
+    PartitionCount = DiskEntry->LayoutBuffer->PartitionCount;
+
+    /* Set the new disk layout and retrieve its updated version with possibly 
modified partition numbers */
+    BufferSize = sizeof(DRIVE_LAYOUT_INFORMATION) +
+                 ((PartitionCount - 1) * sizeof(PARTITION_INFORMATION));
+    Status = NtDeviceIoControlFile(FileHandle,
+                                   NULL,
+                                   NULL,
+                                   NULL,
+                                   &Iosb,
+                                   IOCTL_DISK_SET_DRIVE_LAYOUT,
+                                   DiskEntry->LayoutBuffer,
+                                   BufferSize,
+                                   DiskEntry->LayoutBuffer,
+                                   BufferSize);
+    NtClose(FileHandle);
+
+    /*
+     * IOCTL_DISK_SET_DRIVE_LAYOUT calls IoWritePartitionTable(), which 
converts
+     * DiskEntry->LayoutBuffer->PartitionCount into a partition *table* count,
+     * where such a table is expected to enumerate up to 4 partitions:
+     * partition *table* count == ROUND_UP(PartitionCount, 4) / 4 .
+     * Due to this we need to restore the original PartitionCount number.
+     */
+    DiskEntry->LayoutBuffer->PartitionCount = PartitionCount;
+
+    /* Check whether the IOCTL_DISK_SET_DRIVE_LAYOUT call succeeded */
+    if (!NT_SUCCESS(Status))
+    {
+        DPRINT1("IOCTL_DISK_SET_DRIVE_LAYOUT failed (Status 0x%08lx)\n", 
Status);
+        return Status;
+    }
+
+    /* Update the partition numbers */
+
+    /* Update the primary partition table */
+    for (ListEntry = DiskEntry->PrimaryPartListHead.Flink;
+         ListEntry != &DiskEntry->PrimaryPartListHead;
+         ListEntry = ListEntry->Flink)
+    {
+        PartEntry = CONTAINING_RECORD(ListEntry, PARTENTRY, ListEntry);
+
+        if (PartEntry->IsPartitioned)
+        {
+            ASSERT(PartEntry->PartitionType != PARTITION_ENTRY_UNUSED);
+            PartitionInfo = 
&DiskEntry->LayoutBuffer->PartitionEntry[PartEntry->PartitionIndex];
+            PartEntry->PartitionNumber = PartitionInfo->PartitionNumber;
+        }
+    }
+
+    /* Update the logical partition table */
+    for (ListEntry = DiskEntry->LogicalPartListHead.Flink;
+         ListEntry != &DiskEntry->LogicalPartListHead;
+         ListEntry = ListEntry->Flink)
+    {
+        PartEntry = CONTAINING_RECORD(ListEntry, PARTENTRY, ListEntry);
+
+        if (PartEntry->IsPartitioned)
+        {
+            ASSERT(PartEntry->PartitionType != PARTITION_ENTRY_UNUSED);
+            PartitionInfo = 
&DiskEntry->LayoutBuffer->PartitionEntry[PartEntry->PartitionIndex];
+            PartEntry->PartitionNumber = PartitionInfo->PartitionNumber;
+        }
+    }
+
+    /* The layout has been successfully updated, the disk is not dirty anymore 
*/
+    DiskEntry->Dirty = FALSE;
+
+    return Status;
+}
+
+
+static
+BOOLEAN
+IsEmptyLayoutEntry(
+    IN PPARTITION_INFORMATION PartitionInfo)
+{
+    if (PartitionInfo->StartingOffset.QuadPart == 0 &&
+        PartitionInfo->PartitionLength.QuadPart == 0)
+    {
+        return TRUE;
+    }
+
+    return FALSE;
+}
+
+
+static
+BOOLEAN
+IsSamePrimaryLayoutEntry(
+    IN PPARTITION_INFORMATION PartitionInfo,
+    IN PDISKENTRY DiskEntry,
+    IN PPARTENTRY PartEntry)
+{
+    if ((PartitionInfo->StartingOffset.QuadPart == 
PartEntry->StartSector.QuadPart * DiskEntry->BytesPerSector) &&
+        (PartitionInfo->PartitionLength.QuadPart == 
PartEntry->SectorCount.QuadPart * DiskEntry->BytesPerSector))
+    {
+        return TRUE;
+    }
+
+    return FALSE;
+}
+
+
+ULONG
+GetPrimaryPartitionCount(
+    IN PDISKENTRY DiskEntry)
+{
+    PLIST_ENTRY Entry;
+    PPARTENTRY PartEntry;
+    ULONG Count = 0;
+
+    for (Entry = DiskEntry->PrimaryPartListHead.Flink;
+         Entry != &DiskEntry->PrimaryPartListHead;
+         Entry = Entry->Flink)
+    {
+        PartEntry = CONTAINING_RECORD(Entry, PARTENTRY, ListEntry);
+        if (PartEntry->IsPartitioned)
+            Count++;
+    }
+
+    return Count;
+}
+
+
+static
+ULONG
+GetLogicalPartitionCount(
+    IN PDISKENTRY DiskEntry)
+{
+    PLIST_ENTRY ListEntry;
+    PPARTENTRY PartEntry;
+    ULONG Count = 0;
+
+    for (ListEntry = DiskEntry->LogicalPartListHead.Flink;
+         ListEntry != &DiskEntry->LogicalPartListHead;
+         ListEntry = ListEntry->Flink)
+    {
+        PartEntry = CONTAINING_RECORD(ListEntry, PARTENTRY, ListEntry);
+        if (PartEntry->IsPartitioned)
+            Count++;
+    }
+
+    return Count;
+}
+
+
+static
+BOOLEAN
+ReAllocateLayoutBuffer(
+    IN PDISKENTRY DiskEntry)
+{
+    PDRIVE_LAYOUT_INFORMATION NewLayoutBuffer;
+    ULONG NewPartitionCount;
+    ULONG CurrentPartitionCount = 0;
+    ULONG LayoutBufferSize;
+    ULONG i;
+
+    DPRINT1("ReAllocateLayoutBuffer()\n");
+
+    NewPartitionCount = 4 + GetLogicalPartitionCount(DiskEntry) * 4;
+
+    if (DiskEntry->LayoutBuffer)
+        CurrentPartitionCount = DiskEntry->LayoutBuffer->PartitionCount;
+
+    DPRINT1("CurrentPartitionCount: %lu ; NewPartitionCount: %lu\n",
+            CurrentPartitionCount, NewPartitionCount);
+
+    if (CurrentPartitionCount == NewPartitionCount)
+        return TRUE;
+
+    LayoutBufferSize = sizeof(DRIVE_LAYOUT_INFORMATION) +
+                       ((NewPartitionCount - ANYSIZE_ARRAY) * 
sizeof(PARTITION_INFORMATION));
+    NewLayoutBuffer = RtlReAllocateHeap(RtlGetProcessHeap(),
+                                        HEAP_ZERO_MEMORY,
+                                        DiskEntry->LayoutBuffer,
+                                        LayoutBufferSize);
+    if (NewLayoutBuffer == NULL)
+    {
+        DPRINT1("Failed to allocate the new layout buffer (size: %lu)\n", 
LayoutBufferSize);
+        return FALSE;
+    }
+
+    NewLayoutBuffer->PartitionCount = NewPartitionCount;
+
+    /* If the layout buffer grows, make sure the new (empty) entries are 
written to the disk */
+    if (NewPartitionCount > CurrentPartitionCount)
+    {
+        for (i = CurrentPartitionCount; i < NewPartitionCount; i++)
+        {
+            NewLayoutBuffer->PartitionEntry[i].RewritePartition = TRUE;
+        }
+    }
+
+    DiskEntry->LayoutBuffer = NewLayoutBuffer;
+
+    return TRUE;
+}
+
+
+VOID
+UpdateDiskLayout(
+    IN PDISKENTRY DiskEntry)
+{
+    PPARTITION_INFORMATION PartitionInfo;
+    PPARTITION_INFORMATION LinkInfo = NULL;
+    PLIST_ENTRY ListEntry;
+    PPARTENTRY PartEntry;
+    LARGE_INTEGER HiddenSectors64;
+    ULONG Index;
+    ULONG PartitionNumber = 1;
+
+    DPRINT1("UpdateDiskLayout()\n");
+
+    /* Resize the layout buffer if necessary */
+    if (ReAllocateLayoutBuffer(DiskEntry) == FALSE)
+    {
+        DPRINT("ReAllocateLayoutBuffer() failed.\n");
+        return;
+    }
+
+    /* Update the primary partition table */
+    Index = 0;
+    for (ListEntry = DiskEntry->PrimaryPartListHead.Flink;
+         ListEntry != &DiskEntry->PrimaryPartListHead;
+         ListEntry = ListEntry->Flink)
+    {
+        PartEntry = CONTAINING_RECORD(ListEntry, PARTENTRY, ListEntry);
+
+        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 
(unmounted) partitions */
+            if (PartEntry->New)
+                PartEntry->PartitionNumber = 0;
+
+            PartEntry->OnDiskPartitionNumber = 
(!IsContainerPartition(PartEntry->PartitionType) ? PartitionNumber : 0);
+
+            if (!IsSamePrimaryLayoutEntry(PartitionInfo, DiskEntry, PartEntry))
+            {
+                DPRINT1("Updating primary partition entry %lu\n", Index);
+
+                PartitionInfo->StartingOffset.QuadPart = 
PartEntry->StartSector.QuadPart * DiskEntry->BytesPerSector;
+                PartitionInfo->PartitionLength.QuadPart = 
PartEntry->SectorCount.QuadPart * DiskEntry->BytesPerSector;
+                PartitionInfo->HiddenSectors = PartEntry->StartSector.LowPart;
+                PartitionInfo->PartitionNumber = PartEntry->PartitionNumber;
+                PartitionInfo->PartitionType = PartEntry->PartitionType;
+                PartitionInfo->BootIndicator = PartEntry->BootIndicator;
+                PartitionInfo->RecognizedPartition = 
IsRecognizedPartition(PartEntry->PartitionType);
+                PartitionInfo->RewritePartition = TRUE;
+            }
+
+            if (!IsContainerPartition(PartEntry->PartitionType))
+                PartitionNumber++;
+
+            Index++;
+        }
+    }
+
+    ASSERT(Index <= 4);
+
+    /* Update the logical partition table */
+    Index = 4;
+    for (ListEntry = DiskEntry->LogicalPartListHead.Flink;
+         ListEntry != &DiskEntry->LogicalPartListHead;
+         ListEntry = ListEntry->Flink)
+    {
+        PartEntry = CONTAINING_RECORD(ListEntry, PARTENTRY, ListEntry);
+
+        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 
(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;
+            PartitionInfo->PartitionNumber = PartEntry->PartitionNumber;
+            PartitionInfo->PartitionType = PartEntry->PartitionType;
+            PartitionInfo->BootIndicator = FALSE;
+            PartitionInfo->RecognizedPartition = 
IsRecognizedPartition(PartEntry->PartitionType);
+            PartitionInfo->RewritePartition = TRUE;
+
+            /* Fill the link entry of the previous partition entry */
+            if (LinkInfo != NULL)
+            {
+                LinkInfo->StartingOffset.QuadPart = 
(PartEntry->StartSector.QuadPart - DiskEntry->SectorAlignment) * 
DiskEntry->BytesPerSector;
+                LinkInfo->PartitionLength.QuadPart = 
(PartEntry->StartSector.QuadPart + DiskEntry->SectorAlignment) * 
DiskEntry->BytesPerSector;
+                HiddenSectors64.QuadPart = PartEntry->StartSector.QuadPart - 
DiskEntry->SectorAlignment - DiskEntry->ExtendedPartition->StartSector.QuadPart;
+                LinkInfo->HiddenSectors = HiddenSectors64.LowPart;
+                LinkInfo->PartitionNumber = 0;
+                LinkInfo->PartitionType = PARTITION_EXTENDED;
+                LinkInfo->BootIndicator = FALSE;
+                LinkInfo->RecognizedPartition = FALSE;
+                LinkInfo->RewritePartition = TRUE;
+            }
+
+            /* Save a pointer to the link entry of the current partition entry 
*/
+            LinkInfo = &DiskEntry->LayoutBuffer->PartitionEntry[Index + 1];
+
+            PartitionNumber++;
+            Index += 4;
+        }
+    }
+
+    /* Wipe unused primary partition entries */
+    for (Index = GetPrimaryPartitionCount(DiskEntry); Index < 4; Index++)
+    {
+        DPRINT1("Primary partition entry %lu\n", Index);
+
+        PartitionInfo = &DiskEntry->LayoutBuffer->PartitionEntry[Index];
+
+        if (!IsEmptyLayoutEntry(PartitionInfo))
+        {
+            DPRINT1("Wiping primary partition entry %lu\n", Index);
+
+            PartitionInfo->StartingOffset.QuadPart = 0;
+            PartitionInfo->PartitionLength.QuadPart = 0;
+            PartitionInfo->HiddenSectors = 0;
+            PartitionInfo->PartitionNumber = 0;
+            PartitionInfo->PartitionType = PARTITION_ENTRY_UNUSED;
+            PartitionInfo->BootIndicator = FALSE;
+            PartitionInfo->RecognizedPartition = FALSE;
+            PartitionInfo->RewritePartition = TRUE;
+        }
+    }
+
+    /* Wipe unused logical partition entries */
+    for (Index = 4; Index < DiskEntry->LayoutBuffer->PartitionCount; Index++)
+    {
+        if (Index % 4 >= 2)
+        {
+            DPRINT1("Logical partition entry %lu\n", Index);
+
+            PartitionInfo = &DiskEntry->LayoutBuffer->PartitionEntry[Index];
+
+            if (!IsEmptyLayoutEntry(PartitionInfo))
+            {
+                DPRINT1("Wiping partition entry %lu\n", Index);
+
+                PartitionInfo->StartingOffset.QuadPart = 0;
+                PartitionInfo->PartitionLength.QuadPart = 0;
+                PartitionInfo->HiddenSectors = 0;
+                PartitionInfo->PartitionNumber = 0;
+                PartitionInfo->PartitionType = PARTITION_ENTRY_UNUSED;
+                PartitionInfo->BootIndicator = FALSE;
+                PartitionInfo->RecognizedPartition = FALSE;
+                PartitionInfo->RewritePartition = TRUE;
+            }
+        }
+    }
+
+    DiskEntry->Dirty = TRUE;
+}
+
+
+PPARTENTRY
+GetPrevUnpartitionedEntry(
+    IN PPARTENTRY PartEntry)
+{
+    PDISKENTRY DiskEntry = PartEntry->DiskEntry;
+    PPARTENTRY PrevPartEntry;
+    PLIST_ENTRY ListHead;
+
+    if (PartEntry->LogicalPartition)
+        ListHead = &DiskEntry->LogicalPartListHead;
+    else
+        ListHead = &DiskEntry->PrimaryPartListHead;
+
+    if (PartEntry->ListEntry.Blink != ListHead)
+    {
+        PrevPartEntry = CONTAINING_RECORD(PartEntry->ListEntry.Blink,
+                                          PARTENTRY,
+                                          ListEntry);
+        if (!PrevPartEntry->IsPartitioned)
+        {
+            ASSERT(PrevPartEntry->PartitionType == PARTITION_ENTRY_UNUSED);
+            return PrevPartEntry;
+        }
+    }
+
+    return NULL;
+}
+
+
+PPARTENTRY
+GetNextUnpartitionedEntry(
+    IN PPARTENTRY PartEntry)
+{
+    PDISKENTRY DiskEntry = PartEntry->DiskEntry;
+    PPARTENTRY NextPartEntry;
+    PLIST_ENTRY ListHead;
+
+    if (PartEntry->LogicalPartition)
+        ListHead = &DiskEntry->LogicalPartListHead;
+    else
+        ListHead = &DiskEntry->PrimaryPartListHead;
+
+    if (PartEntry->ListEntry.Flink != ListHead)
+    {
+        NextPartEntry = CONTAINING_RECORD(PartEntry->ListEntry.Flink,
+                                          PARTENTRY,
+                                          ListEntry);
+        if (!NextPartEntry->IsPartitioned)
+        {
+            ASSERT(NextPartEntry->PartitionType == PARTITION_ENTRY_UNUSED);
+            return NextPartEntry;
+        }
+    }
+
+    return NULL;
+}
+
+NTSTATUS
+DismountVolume(
+    IN PPARTENTRY PartEntry)
+{
+    NTSTATUS Status;
+    NTSTATUS LockStatus;
+    UNICODE_STRING Name;
+    OBJECT_ATTRIBUTES ObjectAttributes;
+    IO_STATUS_BLOCK IoStatusBlock;
+    HANDLE PartitionHandle;
+    WCHAR Buffer[MAX_PATH];
+
+    /* Check whether the partition is valid and was mounted by the system */
+    if (!PartEntry->IsPartitioned ||
+        IsContainerPartition(PartEntry->PartitionType)   ||
+        !IsRecognizedPartition(PartEntry->PartitionType) ||
+        PartEntry->FormatState == UnknownFormat ||
+        // NOTE: If FormatState == Unformatted but *FileSystem != 0 this means
+        // it has been usually mounted with RawFS and thus needs to be 
dismounted.
+/*        !*PartEntry->FileSystem || */
+        PartEntry->PartitionNumber == 0)
+    {
+        /* The partition is not mounted, so just return success */
+        return STATUS_SUCCESS;
+    }
+
+    ASSERT(PartEntry->PartitionType != PARTITION_ENTRY_UNUSED);
+
+    /* Open the volume */
+    StringCchPrintfW(Buffer, ARRAYSIZE(Buffer),
+                     L"\\Device\\Harddisk%lu\\Partition%lu",
+                     PartEntry->DiskEntry->DiskNumber,
+                     PartEntry->PartitionNumber);
+    RtlInitUnicodeString(&Name, Buffer);
+
+    InitializeObjectAttributes(&ObjectAttributes,
+                               &Name,
+                               OBJ_CASE_INSENSITIVE,
+                               NULL,
+                               NULL);
+
+    Status = NtOpenFile(&PartitionHandle,
+                        GENERIC_READ | GENERIC_WRITE | SYNCHRONIZE,
+                        &ObjectAttributes,
+                        &IoStatusBlock,
+                        FILE_SHARE_READ | FILE_SHARE_WRITE,
+                        FILE_SYNCHRONOUS_IO_NONALERT);
+    if (!NT_SUCCESS(Status))
+    {
+        DPRINT1("ERROR: Cannot open volume %wZ for dismounting! (Status 
0x%lx)\n", &Name, Status);
+        return Status;
+    }
+
+    /* Lock the volume */
+    LockStatus = NtFsControlFile(PartitionHandle,
+                                 NULL,
+                                 NULL,
+                                 NULL,
+                                 &IoStatusBlock,
+                                 FSCTL_LOCK_VOLUME,
+                                 NULL,
+                                 0,
+                                 NULL,
+                                 0);
+    if (!NT_SUCCESS(LockStatus))
+    {
+        DPRINT1("WARNING: Failed to lock volume! Operations may fail! (Status 
0x%lx)\n", LockStatus);
+    }
+
+    /* Dismount the volume */
+    Status = NtFsControlFile(PartitionHandle,
+                             NULL,
+                             NULL,
+                             NULL,
+                             &IoStatusBlock,
+                             FSCTL_DISMOUNT_VOLUME,
+                             NULL,
+                             0,
+                             NULL,
+                             0);
+    if (!NT_SUCCESS(Status))
+    {
+        DPRINT1("Failed to unmount volume (Status 0x%lx)\n", Status);
+    }
+
+    /* Unlock the volume */
+    LockStatus = NtFsControlFile(PartitionHandle,
+                                 NULL,
+                                 NULL,
+                                 NULL,
+                                 &IoStatusBlock,
+                                 FSCTL_UNLOCK_VOLUME,
+                                 NULL,
+                                 0,
+                                 NULL,
+                                 0);
+    if (!NT_SUCCESS(LockStatus))
+    {
+        DPRINT1("Failed to unlock volume (Status 0x%lx)\n", LockStatus);
+    }
+
+    /* Close the volume */
+    NtClose(PartitionHandle);
+
+    return Status;
+}
+
 /* EOF */

Reply via email to