From 723947eab127001405d07963f2046e147be6498f Mon Sep 17 00:00:00 2001 From: Eric Kohl Date: Sat, 28 May 2022 11:33:29 +0200 Subject: [PATCH] [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 + + 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= (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= (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=| */ - 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= */ - 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 + 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 +#include #define NDEBUG #include @@ -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 */