[DISKPART] Implement the clean command

base/system/diskpart/clean.c

@@ -8,7 +8,241 @@
 
 #include "diskpart.h"
 
-BOOL clean_main(INT argc, LPWSTR *argv)
+#define NDEBUG
+#include <debug.h>
+
+
+BOOL
+clean_main(
+    _In_ INT argc,
+    _In_ PWSTR *argv)
 {
+    PLIST_ENTRY Entry;
+    PPARTENTRY PartEntry;
+    PVOLENTRY VolumeEntry;
+    BOOL bAll = FALSE;
+    PUCHAR SectorsBuffer = NULL;
+    ULONG LayoutBufferSize, Size;
+    INT i;
+    WCHAR Buffer[MAX_PATH];
+    UNICODE_STRING Name;
+    OBJECT_ATTRIBUTES ObjectAttributes;
+    IO_STATUS_BLOCK IoStatusBlock;
+    HANDLE FileHandle = NULL;
+    LARGE_INTEGER Offset, Count, MaxCount;
+    NTSTATUS Status;
+
+    DPRINT("Clean()\n");
+
+    if (CurrentDisk == NULL)
+    {
+        ConResPuts(StdOut, IDS_SELECT_NO_DISK);
+        return TRUE;
+    }
+
+    /* Do not allow to clean the boot disk */
+    if ((CurrentDisk->BiosFound == TRUE) &&
+        (CurrentDisk->BiosDiskNumber == 0))
+    {
+        ConResPuts(StdOut, IDS_CLEAN_SYSTEM);
+        return TRUE;
+    }
+
+    for (i = 1; i < argc; i++)
+    {
+        if (_wcsicmp(argv[1], L"all") == 0)
+        {
+            bAll = TRUE;
+        }
+    }
+
+    /* Dismount and remove all logical partitions */
+    while (!IsListEmpty(&CurrentDisk->LogicalPartListHead))
+    {
+        Entry = RemoveHeadList(&CurrentDisk->LogicalPartListHead);
+        PartEntry = CONTAINING_RECORD(Entry, PARTENTRY, ListEntry);
+
+        /* Dismount the logical partition */
+        if (PartEntry->PartitionType != 0)
+        {
+            DismountVolume(PartEntry);
+            VolumeEntry = GetVolumeFromPartition(PartEntry);
+            if (VolumeEntry)
+                RemoveVolume(VolumeEntry);
+        }
+
+        /* Delete it */
+        RtlFreeHeap(RtlGetProcessHeap(), 0, PartEntry);
+    }
+
+    /* Dismount and remove all primary partitions */
+    while (!IsListEmpty(&CurrentDisk->PrimaryPartListHead))
+    {
+        Entry = RemoveHeadList(&CurrentDisk->PrimaryPartListHead);
+        PartEntry = CONTAINING_RECORD(Entry, PARTENTRY, ListEntry);
+
+        /* Dismount the primary partition */
+        if ((PartEntry->PartitionType != 0) &&
+            (IsContainerPartition(PartEntry->PartitionType) == FALSE))
+        {
+            DismountVolume(PartEntry);
+            VolumeEntry = GetVolumeFromPartition(PartEntry);
+            if (VolumeEntry)
+                RemoveVolume(VolumeEntry);
+        }
+
+        /* Delete it */
+        RtlFreeHeap(RtlGetProcessHeap(), 0, PartEntry);
+    }
+
+    /* Initialize the disk entry */
+    CurrentDisk->ExtendedPartition = NULL;
+    CurrentDisk->Dirty = FALSE;
+    CurrentDisk->NewDisk = TRUE;
+    CurrentDisk->NoMbr = TRUE;
+
+    /* Wipe the layout buffer */
+    RtlFreeHeap(RtlGetProcessHeap(), 0, CurrentDisk->LayoutBuffer);
+
+    LayoutBufferSize = sizeof(DRIVE_LAYOUT_INFORMATION) +
+                       ((4 - ANYSIZE_ARRAY) * sizeof(PARTITION_INFORMATION));
+    CurrentDisk->LayoutBuffer = RtlAllocateHeap(RtlGetProcessHeap(),
+                                                HEAP_ZERO_MEMORY,
+                                                LayoutBufferSize);
+    if (CurrentDisk->LayoutBuffer == NULL)
+    {
+        DPRINT1("Failed to allocate the disk layout buffer!\n");
+        return TRUE;
+    }
+
+    /* Allocate a 1MB sectors buffer */
+    SectorsBuffer = RtlAllocateHeap(RtlGetProcessHeap(),
+                                    HEAP_ZERO_MEMORY,
+                                    1024 * 1024);
+    if (SectorsBuffer == NULL)
+    {
+        DPRINT1("Failed to allocate the sectors buffer!\n");
+        goto done;
+    }
+
+    /* Open the disk for writing */
+    StringCchPrintfW(Buffer, ARRAYSIZE(Buffer),
+                     L"\\Device\\Harddisk%d\\Partition0",
+                     CurrentDisk->DiskNumber);
+
+    RtlInitUnicodeString(&Name, Buffer);
+
+    InitializeObjectAttributes(&ObjectAttributes,
+                               &Name,
+                               OBJ_CASE_INSENSITIVE,
+                               NULL,
+                               NULL);
+
+    Status = NtOpenFile(&FileHandle,
+                        GENERIC_READ | GENERIC_WRITE | SYNCHRONIZE,
+                        &ObjectAttributes,
+                        &IoStatusBlock,
+                        0,
+                        FILE_SYNCHRONOUS_IO_NONALERT);
+    if (!NT_SUCCESS(Status))
+    {
+        DPRINT1("Failed to open the disk! (Status 0x%08lx)\n", Status);
+        ConResPuts(StdOut, IDS_CLEAN_FAIL);
+        goto done;
+    }
+
+    /* Clean sectors */
+    if (bAll)
+    {
+        MaxCount.QuadPart = (CurrentDisk->SectorCount.QuadPart * CurrentDisk->BytesPerSector) / (1024 * 1024);
+        for (Count.QuadPart = 0; Count.QuadPart < MaxCount.QuadPart; Count.QuadPart++)
+        {
+            Offset.QuadPart = Count.QuadPart * (1024 * 1024);
+            Status = NtWriteFile(FileHandle,
+                                 NULL,
+                                 NULL,
+                                 NULL,
+                                 &IoStatusBlock,
+                                 SectorsBuffer,
+                                 1024 * 1024,
+                                 &Offset,
+                                 NULL);
+            if (!NT_SUCCESS(Status))
+            {
+                DPRINT1("Failed to write MB! (Status 0x%08lx)\n", Status);
+                ConResPuts(StdOut, IDS_CLEAN_FAIL);
+                goto done;
+            }
+        }
+
+        Size = (ULONG)(CurrentDisk->SectorCount.QuadPart * CurrentDisk->BytesPerSector) % (1024 * 1024);
+        if (Size != 0)
+        {
+            Offset.QuadPart += (1024 * 1024);
+            Status = NtWriteFile(FileHandle,
+                                 NULL,
+                                 NULL,
+                                 NULL,
+                                 &IoStatusBlock,
+                                 SectorsBuffer,
+                                 Size,
+                                 &Offset,
+                                 NULL);
+            if (!NT_SUCCESS(Status))
+            {
+                DPRINT1("Failed to write the last part! (Status 0x%08lx)\n", Status);
+                ConResPuts(StdOut, IDS_CLEAN_FAIL);
+                goto done;
+            }
+        }
+    }
+    else
+    {
+        /* Clean the first MB */
+        Offset.QuadPart = 0;
+        Status = NtWriteFile(FileHandle,
+                             NULL,
+                             NULL,
+                             NULL,
+                             &IoStatusBlock,
+                             SectorsBuffer,
+                             1024 * 1024,
+                             &Offset,
+                             NULL);
+        if (!NT_SUCCESS(Status))
+        {
+            DPRINT1("Failed to write the first MB! (Status 0x%08lx)\n", Status);
+            ConResPuts(StdOut, IDS_CLEAN_FAIL);
+            goto done;
+        }
+
+        /* Clean the last MB */
+        Offset.QuadPart = (CurrentDisk->SectorCount.QuadPart * CurrentDisk->BytesPerSector) - (1024 * 1024);
+        Status = NtWriteFile(FileHandle,
+                             NULL,
+                             NULL,
+                             NULL,
+                             &IoStatusBlock,
+                             SectorsBuffer,
+                             1024 * 1024,
+                             &Offset,
+                             NULL);
+        if (!NT_SUCCESS(Status))
+        {
+            DPRINT1("Failed to write the last MB! (Status 0x%08lx)\n", Status);
+            ConResPuts(StdOut, IDS_CLEAN_FAIL);
+            goto done;
+        }
+    }
+
+    ConResPuts(StdOut, IDS_CLEAN_SUCCESS);
+
+done:
+    if (FileHandle != NULL)
+        NtClose(FileHandle);
+
+    if (SectorsBuffer != NULL)
+        RtlFreeHeap(RtlGetProcessHeap(), 0, SectorsBuffer);
+
     return TRUE;
 }

base/system/diskpart/diskpart.h

@@ -22,6 +22,7 @@
 #include <winioctl.h>
 
 #include <errno.h>
+#include <strsafe.h>
 
 #include <conutils.h>
 
@@ -439,7 +440,16 @@ GetPrimaryPartitionCount(
 
 NTSTATUS
 DismountVolume(
-    IN PPARTENTRY PartEntry);
+    _In_ PPARTENTRY PartEntry);
+
+PVOLENTRY
+GetVolumeFromPartition(
+    _In_ PPARTENTRY PartEntry);
+
+VOID
+RemoveVolume(
+    _In_ PVOLENTRY VolumeEntry);
+
 
 /* recover.c */
 BOOL recover_main(INT argc, LPWSTR *argv);

base/system/diskpart/lang/de-DE.rc

@@ -22,6 +22,13 @@ BEGIN
     IDS_ACTIVE_ALREADY "\nDie aktuelle Partition wurde bereits als aktiv markiert.\n"
 END
 
+STRINGTABLE
+BEGIN
+    IDS_CLEAN_FAIL "\nDer Datenträger konnte nicht bereinigt werden.\nMöglicherweise können die Daten auf diesem Datenträger nicht wiederhergestellt werden.\n"
+    IDS_CLEAN_SUCCESS "\nDer Datenträger wurde bereinigt.\n"
+    IDS_CLEAN_SYSTEM "\nDer gewählte Datenträger ist zum Ausführen des Computers erforderlich und kann nicht bereinigt werden.\n"
+END
+
 STRINGTABLE
 BEGIN
     IDS_CREATE_PARTITION_FAIL "\nDie angegebene Partition konnte nicht erstellt werden.\n"

base/system/diskpart/lang/en-US.rc

@@ -22,6 +22,13 @@ BEGIN
     IDS_ACTIVE_ALREADY "\nThe current partition is already marked as active.\n"
 END
 
+STRINGTABLE
+BEGIN
+    IDS_CLEAN_FAIL "\nDiskPart was unable to clean the disk.\nThe data on this disk may be unrecoverable.\n"
+    IDS_CLEAN_SUCCESS "\nDiskPart succeeded in cleaning the disk.\n"
+    IDS_CLEAN_SYSTEM "\nThe selected disk is neccessary to the operation of your computer, and may not be cleaned.\n"
+END
+
 STRINGTABLE
 BEGIN
     IDS_CREATE_PARTITION_FAIL "\nDiskPart was unable to create the specified partition.\n"

base/system/diskpart/lang/pl-PL.rc

@@ -22,6 +22,13 @@ BEGIN
     IDS_ACTIVE_ALREADY "\nThe current partition is already marked as active.\n"
 END
 
+STRINGTABLE
+BEGIN
+    IDS_CLEAN_FAIL "\nDiskPart was unable to clean the disk.\nThe data on this disk may be unrecoverable.\n"
+    IDS_CLEAN_SUCCESS "\nDiskPart succeeded in cleaning the disk.\n"
+    IDS_CLEAN_SYSTEM "\nThe selected disk is neccessary to the operation of your computer, and may not be cleaned.\n"
+END
+
 STRINGTABLE
 BEGIN
     IDS_CREATE_PARTITION_FAIL "\nDiskPart was unable to create the specified partition.\n"

base/system/diskpart/lang/pt-PT.rc

@@ -24,6 +24,13 @@ BEGIN
     IDS_ACTIVE_ALREADY "\nThe current partition is already marked as active.\n"
 END
 
+STRINGTABLE
+BEGIN
+    IDS_CLEAN_FAIL "\nDiskPart was unable to clean the disk.\nThe data on this disk may be unrecoverable.\n"
+    IDS_CLEAN_SUCCESS "\nDiskPart succeeded in cleaning the disk.\n"
+    IDS_CLEAN_SYSTEM "\nThe selected disk is neccessary to the operation of your computer, and may not be cleaned.\n"
+END
+
 STRINGTABLE
 BEGIN
     IDS_CREATE_PARTITION_FAIL "\nDiskPart was unable to create the specified partition.\n"

base/system/diskpart/lang/ro-RO.rc

@@ -24,6 +24,13 @@ BEGIN
     IDS_ACTIVE_ALREADY "\nThe current partition is already marked as active.\n"
 END
 
+STRINGTABLE
+BEGIN
+    IDS_CLEAN_FAIL "\nDiskPart was unable to clean the disk.\nThe data on this disk may be unrecoverable.\n"
+    IDS_CLEAN_SUCCESS "\nDiskPart succeeded in cleaning the disk.\n"
+    IDS_CLEAN_SYSTEM "\nThe selected disk is neccessary to the operation of your computer, and may not be cleaned.\n"
+END
+
 STRINGTABLE
 BEGIN
     IDS_CREATE_PARTITION_FAIL "\nDiskPart was unable to create the specified partition.\n"

base/system/diskpart/lang/ru-RU.rc

@@ -24,6 +24,13 @@ BEGIN
     IDS_ACTIVE_ALREADY "\nThe current partition is already marked as active.\n"
 END
 
+STRINGTABLE
+BEGIN
+    IDS_CLEAN_FAIL "\nDiskPart was unable to clean the disk.\nThe data on this disk may be unrecoverable.\n"
+    IDS_CLEAN_SUCCESS "\nDiskPart succeeded in cleaning the disk.\n"
+    IDS_CLEAN_SYSTEM "\nThe selected disk is neccessary to the operation of your computer, and may not be cleaned.\n"
+END
+
 STRINGTABLE
 BEGIN
     IDS_CREATE_PARTITION_FAIL "\nDiskPart was unable to create the specified partition.\n"

base/system/diskpart/lang/sq-AL.rc

@@ -26,6 +26,13 @@ BEGIN
     IDS_ACTIVE_ALREADY "\nThe current partition is already marked as active.\n"
 END
 
+STRINGTABLE
+BEGIN
+    IDS_CLEAN_FAIL "\nDiskPart was unable to clean the disk.\nThe data on this disk may be unrecoverable.\n"
+    IDS_CLEAN_SUCCESS "\nDiskPart succeeded in cleaning the disk.\n"
+    IDS_CLEAN_SYSTEM "\nThe selected disk is neccessary to the operation of your computer, and may not be cleaned.\n"
+END
+
 STRINGTABLE
 BEGIN
     IDS_CREATE_PARTITION_FAIL "\nDiskPart was unable to create the specified partition.\n"

base/system/diskpart/lang/tr-TR.rc

@@ -24,6 +24,13 @@ BEGIN
     IDS_ACTIVE_ALREADY "\nThe current partition is already marked as active.\n"
 END
 
+STRINGTABLE
+BEGIN
+    IDS_CLEAN_FAIL "\nDiskPart was unable to clean the disk.\nThe data on this disk may be unrecoverable.\n"
+    IDS_CLEAN_SUCCESS "\nDiskPart succeeded in cleaning the disk.\n"
+    IDS_CLEAN_SYSTEM "\nThe selected disk is neccessary to the operation of your computer, and may not be cleaned.\n"
+END
+
 STRINGTABLE
 BEGIN
     IDS_CREATE_PARTITION_FAIL "\nDiskPart was unable to create the specified partition.\n"

base/system/diskpart/lang/zh-CN.rc

@@ -31,6 +31,13 @@ BEGIN
     IDS_ACTIVE_ALREADY "\nThe current partition is already marked as active.\n"
 END
 
+STRINGTABLE
+BEGIN
+    IDS_CLEAN_FAIL "\nDiskPart was unable to clean the disk.\nThe data on this disk may be unrecoverable.\n"
+    IDS_CLEAN_SUCCESS "\nDiskPart succeeded in cleaning the disk.\n"
+    IDS_CLEAN_SYSTEM "\nThe selected disk is neccessary to the operation of your computer, and may not be cleaned.\n"
+END
+
 STRINGTABLE
 BEGIN
     IDS_CREATE_PARTITION_FAIL "\nDiskPart was unable to create the specified partition.\n"

base/system/diskpart/lang/zh-TW.rc

@@ -25,6 +25,13 @@ BEGIN
     IDS_ACTIVE_ALREADY "\nThe current partition is already marked as active.\n"
 END
 
+STRINGTABLE
+BEGIN
+    IDS_CLEAN_FAIL "\nDiskPart was unable to clean the disk.\nThe data on this disk may be unrecoverable.\n"
+    IDS_CLEAN_SUCCESS "\nDiskPart succeeded in cleaning the disk.\n"
+    IDS_CLEAN_SYSTEM "\nThe selected disk is neccessary to the operation of your computer, and may not be cleaned.\n"
+END
+
 STRINGTABLE
 BEGIN
     IDS_CREATE_PARTITION_FAIL "\nDiskPart was unable to create the specified partition.\n"

base/system/diskpart/partlist.c

@@ -10,7 +10,6 @@
 
 #include "diskpart.h"
 #include <ntddscsi.h>
-#include <strsafe.h>
 
 #define NDEBUG
 #include <debug.h>
@@ -1635,7 +1634,7 @@ IsSamePrimaryLayoutEntry(
 
 ULONG
 GetPrimaryPartitionCount(
-    IN PDISKENTRY DiskEntry)
+    _In_ PDISKENTRY DiskEntry)
 {
     PLIST_ENTRY Entry;
     PPARTENTRY PartEntry;
@@ -1657,7 +1656,7 @@ GetPrimaryPartitionCount(
 static
 ULONG
 GetLogicalPartitionCount(
-    IN PDISKENTRY DiskEntry)
+    _In_ PDISKENTRY DiskEntry)
 {
     PLIST_ENTRY ListEntry;
     PPARTENTRY PartEntry;
@@ -1679,7 +1678,7 @@ GetLogicalPartitionCount(
 static
 BOOLEAN
 ReAllocateLayoutBuffer(
-    IN PDISKENTRY DiskEntry)
+    _In_ PDISKENTRY DiskEntry)
 {
     PDRIVE_LAYOUT_INFORMATION NewLayoutBuffer;
     ULONG NewPartitionCount;
@@ -1731,7 +1730,7 @@ ReAllocateLayoutBuffer(
 
 VOID
 UpdateDiskLayout(
-    IN PDISKENTRY DiskEntry)
+    _In_ PDISKENTRY DiskEntry)
 {
     PPARTITION_INFORMATION PartitionInfo;
     PPARTITION_INFORMATION LinkInfo = NULL;
@@ -1901,7 +1900,7 @@ UpdateDiskLayout(
 
 PPARTENTRY
 GetPrevUnpartitionedEntry(
-    IN PPARTENTRY PartEntry)
+    _In_ PPARTENTRY PartEntry)
 {
     PDISKENTRY DiskEntry = PartEntry->DiskEntry;
     PPARTENTRY PrevPartEntry;
@@ -1930,7 +1929,7 @@ GetPrevUnpartitionedEntry(
 
 PPARTENTRY
 GetNextUnpartitionedEntry(
-    IN PPARTENTRY PartEntry)
+    _In_ PPARTENTRY PartEntry)
 {
     PDISKENTRY DiskEntry = PartEntry->DiskEntry;
     PPARTENTRY NextPartEntry;
@@ -1958,7 +1957,7 @@ GetNextUnpartitionedEntry(
 
 NTSTATUS
 DismountVolume(
-    IN PPARTENTRY PartEntry)
+    _In_ PPARTENTRY PartEntry)
 {
     NTSTATUS Status;
     NTSTATUS LockStatus;
@@ -2063,4 +2062,67 @@ DismountVolume(
     return Status;
 }
 
+
+PVOLENTRY
+GetVolumeFromPartition(
+    _In_ PPARTENTRY PartEntry)
+{
+    PLIST_ENTRY Entry;
+    PVOLENTRY VolumeEntry;
+    ULONG i;
+
+    if ((PartEntry == NULL) ||
+        (PartEntry->DiskEntry == NULL))
+        return NULL;
+
+    Entry = VolumeListHead.Flink;
+    while (Entry != &VolumeListHead)
+    {
+        VolumeEntry = CONTAINING_RECORD(Entry, VOLENTRY, ListEntry);
+
+        if (VolumeEntry->pExtents == NULL)
+            return NULL;
+
+        for (i = 0; i < VolumeEntry->pExtents->NumberOfDiskExtents; i++)
+        {
+            if (VolumeEntry->pExtents->Extents[i].DiskNumber == PartEntry->DiskEntry->DiskNumber)
+            {
+                if ((VolumeEntry->pExtents->Extents[i].StartingOffset.QuadPart == PartEntry->StartSector.QuadPart * PartEntry->DiskEntry->BytesPerSector) &&
+                    (VolumeEntry->pExtents->Extents[i].ExtentLength.QuadPart == PartEntry->SectorCount.QuadPart * PartEntry->DiskEntry->BytesPerSector))
+                    return VolumeEntry;
+            }
+        }
+
+        Entry = Entry->Flink;
+    }
+
+    return NULL;
+}
+
+
+VOID
+RemoveVolume(
+    _In_ PVOLENTRY VolumeEntry)
+{
+    if (VolumeEntry == NULL)
+        return;
+
+    if (VolumeEntry == CurrentVolume)
+        CurrentVolume = NULL;
+
+    RemoveEntryList(&VolumeEntry->ListEntry);
+
+    if (VolumeEntry->pszLabel)
+        RtlFreeHeap(RtlGetProcessHeap(), 0, VolumeEntry->pszLabel);
+
+    if (VolumeEntry->pszFilesystem)
+        RtlFreeHeap(RtlGetProcessHeap(), 0, VolumeEntry->pszFilesystem);
+
+    if (VolumeEntry->pExtents)
+        RtlFreeHeap(RtlGetProcessHeap(), 0, VolumeEntry->pExtents);
+
+    /* Release disk entry */
+    RtlFreeHeap(RtlGetProcessHeap(), 0, VolumeEntry);
+}
+
 /* EOF */

base/system/diskpart/resource.h

@@ -22,6 +22,10 @@
 #define IDS_ACTIVE_SUCCESS             1001
 #define IDS_ACTIVE_ALREADY             1002
 
+#define IDS_CLEAN_FAIL                 1020
+#define IDS_CLEAN_SUCCESS              1021
+#define IDS_CLEAN_SYSTEM               1022
+
 #define IDS_CREATE_PARTITION_FAIL      1050
 #define IDS_CREATE_PARTITION_SUCCESS   1051