xref: /reactos/base/system/diskpart/partlist.c (revision 647d3512)
1 /*
2  * PROJECT:         ReactOS DiskPart
3  * LICENSE:         GPL - See COPYING in the top level directory
4  * FILE:            base/system/diskpart/partlist.c
5  * PURPOSE:         Manages all the partitions of the OS in an interactive way.
6  * PROGRAMMERS:     Eric Kohl
7  */
8 
9 /* INCLUDES *******************************************************************/
10 
11 #include "diskpart.h"
12 #include <ntddscsi.h>
13 
14 #define NDEBUG
15 #include <debug.h>
16 
17 #define InsertAscendingList(ListHead, NewEntry, Type, ListEntryField, SortField)\
18 {\
19   PLIST_ENTRY current;\
20 \
21   current = (ListHead)->Flink;\
22   while (current != (ListHead))\
23   {\
24     if (CONTAINING_RECORD(current, Type, ListEntryField)->SortField >=\
25         (NewEntry)->SortField)\
26     {\
27       break;\
28     }\
29     current = current->Flink;\
30   }\
31 \
32   InsertTailList(current, &((NewEntry)->ListEntryField));\
33 }
34 
35 /* We have to define it there, because it is not in the MS DDK */
36 #define PARTITION_LINUX 0x83
37 
38 #define PARTITION_TBL_SIZE 4
39 
40 #define MBR_MAGIC 0xAA55
41 
42 #include <pshpack1.h>
43 
44 typedef struct _PARTITION
45 {
46     unsigned char   BootFlags;        /* bootable?  0=no, 128=yes  */
47     unsigned char   StartingHead;     /* beginning head number */
48     unsigned char   StartingSector;   /* beginning sector number */
49     unsigned char   StartingCylinder; /* 10 bit nmbr, with high 2 bits put in begsect */
50     unsigned char   PartitionType;    /* Operating System type indicator code */
51     unsigned char   EndingHead;       /* ending head number */
52     unsigned char   EndingSector;     /* ending sector number */
53     unsigned char   EndingCylinder;   /* also a 10 bit nmbr, with same high 2 bit trick */
54     unsigned int  StartingBlock;      /* first sector relative to start of disk */
55     unsigned int  SectorCount;        /* number of sectors in partition */
56 } PARTITION, *PPARTITION;
57 
58 typedef struct _PARTITION_SECTOR
59 {
60     UCHAR BootCode[440];                     /* 0x000 */
61     ULONG Signature;                         /* 0x1B8 */
62     UCHAR Reserved[2];                       /* 0x1BC */
63     PARTITION Partition[PARTITION_TBL_SIZE]; /* 0x1BE */
64     USHORT Magic;                            /* 0x1FE */
65 } PARTITION_SECTOR, *PPARTITION_SECTOR;
66 
67 #include <poppack.h>
68 
69 
70 /* GLOBALS ********************************************************************/
71 
72 LIST_ENTRY DiskListHead;
73 LIST_ENTRY BiosDiskListHead;
74 LIST_ENTRY VolumeListHead;
75 
76 PDISKENTRY CurrentDisk = NULL;
77 PPARTENTRY CurrentPartition = NULL;
78 PVOLENTRY  CurrentVolume = NULL;
79 
80 
81 /* FUNCTIONS ******************************************************************/
82 
83 ULONGLONG
AlignDown(_In_ ULONGLONG Value,_In_ ULONG Alignment)84 AlignDown(
85     _In_ ULONGLONG Value,
86     _In_ ULONG Alignment)
87 {
88     ULONGLONG Temp;
89 
90     Temp = Value / Alignment;
91 
92     return Temp * Alignment;
93 }
94 
95 static
96 VOID
GetDriverName(PDISKENTRY DiskEntry)97 GetDriverName(
98     PDISKENTRY DiskEntry)
99 {
100     RTL_QUERY_REGISTRY_TABLE QueryTable[2];
101     WCHAR KeyName[32];
102     NTSTATUS Status;
103 
104     RtlInitUnicodeString(&DiskEntry->DriverName,
105                          NULL);
106 
107     StringCchPrintfW(KeyName, ARRAYSIZE(KeyName),
108                      L"\\Scsi\\Scsi Port %lu",
109                      DiskEntry->Port);
110 
111     RtlZeroMemory(&QueryTable,
112                   sizeof(QueryTable));
113 
114     QueryTable[0].Name = L"Driver";
115     QueryTable[0].Flags = RTL_QUERY_REGISTRY_DIRECT;
116     QueryTable[0].EntryContext = &DiskEntry->DriverName;
117 
118     Status = RtlQueryRegistryValues(RTL_REGISTRY_DEVICEMAP,
119                                     KeyName,
120                                     QueryTable,
121                                     NULL,
122                                     NULL);
123     if (!NT_SUCCESS(Status))
124     {
125         DPRINT1("RtlQueryRegistryValues() failed (Status %lx)\n", Status);
126     }
127 }
128 
129 static
130 NTSTATUS
131 NTAPI
DiskIdentifierQueryRoutine(PWSTR ValueName,ULONG ValueType,PVOID ValueData,ULONG ValueLength,PVOID Context,PVOID EntryContext)132 DiskIdentifierQueryRoutine(
133     PWSTR ValueName,
134     ULONG ValueType,
135     PVOID ValueData,
136     ULONG ValueLength,
137     PVOID Context,
138     PVOID EntryContext)
139 {
140     PBIOSDISKENTRY BiosDiskEntry = (PBIOSDISKENTRY)Context;
141     UNICODE_STRING NameU;
142 
143     if (ValueType == REG_SZ &&
144         ValueLength == 20 * sizeof(WCHAR))
145     {
146         NameU.Buffer = (PWCHAR)ValueData;
147         NameU.Length = NameU.MaximumLength = 8 * sizeof(WCHAR);
148         RtlUnicodeStringToInteger(&NameU, 16, &BiosDiskEntry->Checksum);
149 
150         NameU.Buffer = (PWCHAR)ValueData + 9;
151         RtlUnicodeStringToInteger(&NameU, 16, &BiosDiskEntry->Signature);
152 
153         return STATUS_SUCCESS;
154     }
155 
156     return STATUS_UNSUCCESSFUL;
157 }
158 
159 static
160 NTSTATUS
161 NTAPI
DiskConfigurationDataQueryRoutine(PWSTR ValueName,ULONG ValueType,PVOID ValueData,ULONG ValueLength,PVOID Context,PVOID EntryContext)162 DiskConfigurationDataQueryRoutine(
163     PWSTR ValueName,
164     ULONG ValueType,
165     PVOID ValueData,
166     ULONG ValueLength,
167     PVOID Context,
168     PVOID EntryContext)
169 {
170     PBIOSDISKENTRY BiosDiskEntry = (PBIOSDISKENTRY)Context;
171     PCM_FULL_RESOURCE_DESCRIPTOR FullResourceDescriptor;
172     PCM_DISK_GEOMETRY_DEVICE_DATA DiskGeometry;
173     ULONG i;
174 
175     if (ValueType != REG_FULL_RESOURCE_DESCRIPTOR ||
176         ValueLength < sizeof(CM_FULL_RESOURCE_DESCRIPTOR))
177         return STATUS_UNSUCCESSFUL;
178 
179     FullResourceDescriptor = (PCM_FULL_RESOURCE_DESCRIPTOR)ValueData;
180 
181     /* Hm. Version and Revision are not set on Microsoft Windows XP... */
182 #if 0
183     if (FullResourceDescriptor->PartialResourceList.Version != 1 ||
184         FullResourceDescriptor->PartialResourceList.Revision != 1)
185         return STATUS_UNSUCCESSFUL;
186 #endif
187 
188     for (i = 0; i < FullResourceDescriptor->PartialResourceList.Count; i++)
189     {
190         if (FullResourceDescriptor->PartialResourceList.PartialDescriptors[i].Type != CmResourceTypeDeviceSpecific ||
191             FullResourceDescriptor->PartialResourceList.PartialDescriptors[i].u.DeviceSpecificData.DataSize != sizeof(CM_DISK_GEOMETRY_DEVICE_DATA))
192             continue;
193 
194         DiskGeometry = (PCM_DISK_GEOMETRY_DEVICE_DATA)&FullResourceDescriptor->PartialResourceList.PartialDescriptors[i + 1];
195         BiosDiskEntry->DiskGeometry = *DiskGeometry;
196 
197         return STATUS_SUCCESS;
198     }
199 
200     return STATUS_UNSUCCESSFUL;
201 }
202 
203 static
204 NTSTATUS
205 NTAPI
SystemConfigurationDataQueryRoutine(PWSTR ValueName,ULONG ValueType,PVOID ValueData,ULONG ValueLength,PVOID Context,PVOID EntryContext)206 SystemConfigurationDataQueryRoutine(
207     PWSTR ValueName,
208     ULONG ValueType,
209     PVOID ValueData,
210     ULONG ValueLength,
211     PVOID Context,
212     PVOID EntryContext)
213 {
214     PCM_FULL_RESOURCE_DESCRIPTOR FullResourceDescriptor;
215     PCM_INT13_DRIVE_PARAMETER* Int13Drives = (PCM_INT13_DRIVE_PARAMETER*)Context;
216     ULONG i;
217 
218     if (ValueType != REG_FULL_RESOURCE_DESCRIPTOR ||
219         ValueLength < sizeof (CM_FULL_RESOURCE_DESCRIPTOR))
220         return STATUS_UNSUCCESSFUL;
221 
222     FullResourceDescriptor = (PCM_FULL_RESOURCE_DESCRIPTOR)ValueData;
223 
224     /* Hm. Version and Revision are not set on Microsoft Windows XP... */
225 #if 0
226     if (FullResourceDescriptor->PartialResourceList.Version != 1 ||
227         FullResourceDescriptor->PartialResourceList.Revision != 1)
228         return STATUS_UNSUCCESSFUL;
229 #endif
230 
231     for (i = 0; i < FullResourceDescriptor->PartialResourceList.Count; i++)
232     {
233         if (FullResourceDescriptor->PartialResourceList.PartialDescriptors[i].Type != CmResourceTypeDeviceSpecific ||
234             FullResourceDescriptor->PartialResourceList.PartialDescriptors[i].u.DeviceSpecificData.DataSize % sizeof(CM_INT13_DRIVE_PARAMETER) != 0)
235             continue;
236 
237         *Int13Drives = (CM_INT13_DRIVE_PARAMETER*)RtlAllocateHeap(RtlGetProcessHeap(), 0,
238                        FullResourceDescriptor->PartialResourceList.PartialDescriptors[i].u.DeviceSpecificData.DataSize);
239         if (*Int13Drives == NULL)
240             return STATUS_NO_MEMORY;
241 
242         memcpy(*Int13Drives,
243                &FullResourceDescriptor->PartialResourceList.PartialDescriptors[i + 1],
244                FullResourceDescriptor->PartialResourceList.PartialDescriptors[i].u.DeviceSpecificData.DataSize);
245         return STATUS_SUCCESS;
246     }
247 
248     return STATUS_UNSUCCESSFUL;
249 }
250 
251 
252 #define ROOT_NAME   L"\\Registry\\Machine\\HARDWARE\\DESCRIPTION\\System\\MultifunctionAdapter"
253 
254 static
255 VOID
EnumerateBiosDiskEntries(VOID)256 EnumerateBiosDiskEntries(VOID)
257 {
258     RTL_QUERY_REGISTRY_TABLE QueryTable[3];
259     WCHAR Name[120];
260     ULONG AdapterCount;
261     ULONG DiskCount;
262     NTSTATUS Status;
263     PCM_INT13_DRIVE_PARAMETER Int13Drives;
264     PBIOSDISKENTRY BiosDiskEntry;
265 
266     memset(QueryTable, 0, sizeof(QueryTable));
267 
268     QueryTable[1].Name = L"Configuration Data";
269     QueryTable[1].QueryRoutine = SystemConfigurationDataQueryRoutine;
270     Int13Drives = NULL;
271     Status = RtlQueryRegistryValues(RTL_REGISTRY_ABSOLUTE,
272                                     L"\\Registry\\Machine\\HARDWARE\\DESCRIPTION\\System",
273                                     &QueryTable[1],
274                                     (PVOID)&Int13Drives,
275                                     NULL);
276     if (!NT_SUCCESS(Status))
277     {
278         DPRINT1("Unable to query the 'Configuration Data' key in '\\Registry\\Machine\\HARDWARE\\DESCRIPTION\\System', status=%lx\n", Status);
279         return;
280     }
281 
282     AdapterCount = 0;
283     while (1)
284     {
285         StringCchPrintfW(Name, ARRAYSIZE(Name),
286                          L"%s\\%lu", ROOT_NAME, AdapterCount);
287         Status = RtlQueryRegistryValues(RTL_REGISTRY_ABSOLUTE,
288                                         Name,
289                                         &QueryTable[2],
290                                         NULL,
291                                         NULL);
292         if (!NT_SUCCESS(Status))
293         {
294             break;
295         }
296 
297         StringCchPrintfW(Name, ARRAYSIZE(Name),
298                          L"%s\\%lu\\DiskController", ROOT_NAME, AdapterCount);
299         Status = RtlQueryRegistryValues(RTL_REGISTRY_ABSOLUTE,
300                                         Name,
301                                         &QueryTable[2],
302                                         NULL,
303                                         NULL);
304         if (NT_SUCCESS(Status))
305         {
306             while (1)
307             {
308                 StringCchPrintfW(Name, ARRAYSIZE(Name),
309                                  L"%s\\%lu\\DiskController\\0", ROOT_NAME, AdapterCount);
310                 Status = RtlQueryRegistryValues(RTL_REGISTRY_ABSOLUTE,
311                                                 Name,
312                                                 &QueryTable[2],
313                                                 NULL,
314                                                 NULL);
315                 if (!NT_SUCCESS(Status))
316                 {
317                     RtlFreeHeap(RtlGetProcessHeap(), 0, Int13Drives);
318                     return;
319                 }
320 
321                 StringCchPrintfW(Name, ARRAYSIZE(Name),
322                                  L"%s\\%lu\\DiskController\\0\\DiskPeripheral", ROOT_NAME, AdapterCount);
323                 Status = RtlQueryRegistryValues(RTL_REGISTRY_ABSOLUTE,
324                                                 Name,
325                                                 &QueryTable[2],
326                                                 NULL,
327                                                 NULL);
328                 if (NT_SUCCESS(Status))
329                 {
330                     QueryTable[0].Name = L"Identifier";
331                     QueryTable[0].QueryRoutine = DiskIdentifierQueryRoutine;
332                     QueryTable[1].Name = L"Configuration Data";
333                     QueryTable[1].QueryRoutine = DiskConfigurationDataQueryRoutine;
334 
335                     DiskCount = 0;
336                     while (1)
337                     {
338                         BiosDiskEntry = (BIOSDISKENTRY*)RtlAllocateHeap(RtlGetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(BIOSDISKENTRY));
339                         if (BiosDiskEntry == NULL)
340                         {
341                             break;
342                         }
343 
344                         StringCchPrintfW(Name, ARRAYSIZE(Name),
345                                          L"%s\\%lu\\DiskController\\0\\DiskPeripheral\\%lu", ROOT_NAME, AdapterCount, DiskCount);
346                         Status = RtlQueryRegistryValues(RTL_REGISTRY_ABSOLUTE,
347                                                         Name,
348                                                         QueryTable,
349                                                         (PVOID)BiosDiskEntry,
350                                                         NULL);
351                         if (!NT_SUCCESS(Status))
352                         {
353                             RtlFreeHeap(RtlGetProcessHeap(), 0, BiosDiskEntry);
354                             break;
355                         }
356 
357                         BiosDiskEntry->DiskNumber = DiskCount;
358                         BiosDiskEntry->Recognized = FALSE;
359 
360                         if (DiskCount < Int13Drives[0].NumberDrives)
361                         {
362                             BiosDiskEntry->Int13DiskData = Int13Drives[DiskCount];
363                         }
364                         else
365                         {
366                             DPRINT1("Didn't find int13 drive datas for disk %u\n", DiskCount);
367                         }
368 
369                         InsertTailList(&BiosDiskListHead, &BiosDiskEntry->ListEntry);
370 
371                         DPRINT("DiskNumber:        %lu\n", BiosDiskEntry->DiskNumber);
372                         DPRINT("Signature:         %08lx\n", BiosDiskEntry->Signature);
373                         DPRINT("Checksum:          %08lx\n", BiosDiskEntry->Checksum);
374                         DPRINT("BytesPerSector:    %lu\n", BiosDiskEntry->DiskGeometry.BytesPerSector);
375                         DPRINT("NumberOfCylinders: %lu\n", BiosDiskEntry->DiskGeometry.NumberOfCylinders);
376                         DPRINT("NumberOfHeads:     %lu\n", BiosDiskEntry->DiskGeometry.NumberOfHeads);
377                         DPRINT("DriveSelect:       %02x\n", BiosDiskEntry->Int13DiskData.DriveSelect);
378                         DPRINT("MaxCylinders:      %lu\n", BiosDiskEntry->Int13DiskData.MaxCylinders);
379                         DPRINT("SectorsPerTrack:   %d\n", BiosDiskEntry->Int13DiskData.SectorsPerTrack);
380                         DPRINT("MaxHeads:          %d\n", BiosDiskEntry->Int13DiskData.MaxHeads);
381                         DPRINT("NumberDrives:      %d\n", BiosDiskEntry->Int13DiskData.NumberDrives);
382 
383                         DiskCount++;
384                     }
385                 }
386 
387                 RtlFreeHeap(RtlGetProcessHeap(), 0, Int13Drives);
388                 return;
389             }
390         }
391 
392         AdapterCount++;
393     }
394 
395     RtlFreeHeap(RtlGetProcessHeap(), 0, Int13Drives);
396 }
397 
398 
399 static
400 VOID
AddPartitionToDisk(ULONG DiskNumber,PDISKENTRY DiskEntry,ULONG PartitionIndex,BOOLEAN LogicalPartition)401 AddPartitionToDisk(
402     ULONG DiskNumber,
403     PDISKENTRY DiskEntry,
404     ULONG PartitionIndex,
405     BOOLEAN LogicalPartition)
406 {
407     PPARTITION_INFORMATION PartitionInfo;
408     PPARTENTRY PartEntry;
409 
410     PartitionInfo = &DiskEntry->LayoutBuffer->PartitionEntry[PartitionIndex];
411     if (PartitionInfo->PartitionType == 0 ||
412         (LogicalPartition == TRUE && IsContainerPartition(PartitionInfo->PartitionType)))
413         return;
414 
415     PartEntry = RtlAllocateHeap(RtlGetProcessHeap(),
416                                 HEAP_ZERO_MEMORY,
417                                 sizeof(PARTENTRY));
418     if (PartEntry == NULL)
419     {
420         return;
421     }
422 
423     PartEntry->DiskEntry = DiskEntry;
424 
425     PartEntry->StartSector.QuadPart = (ULONGLONG)PartitionInfo->StartingOffset.QuadPart / DiskEntry->BytesPerSector;
426     PartEntry->SectorCount.QuadPart = (ULONGLONG)PartitionInfo->PartitionLength.QuadPart / DiskEntry->BytesPerSector;
427 
428     PartEntry->BootIndicator = PartitionInfo->BootIndicator;
429     PartEntry->PartitionType = PartitionInfo->PartitionType;
430 
431     PartEntry->LogicalPartition = LogicalPartition;
432     PartEntry->IsPartitioned = TRUE;
433     PartEntry->PartitionNumber = PartitionInfo->PartitionNumber;
434     PartEntry->PartitionIndex = PartitionIndex;
435 
436     if (IsContainerPartition(PartEntry->PartitionType))
437     {
438         PartEntry->FormatState = Unformatted;
439 
440         if (LogicalPartition == FALSE && DiskEntry->ExtendedPartition == NULL)
441             DiskEntry->ExtendedPartition = PartEntry;
442     }
443     else if ((PartEntry->PartitionType == PARTITION_FAT_12) ||
444              (PartEntry->PartitionType == PARTITION_FAT_16) ||
445              (PartEntry->PartitionType == PARTITION_HUGE) ||
446              (PartEntry->PartitionType == PARTITION_XINT13) ||
447              (PartEntry->PartitionType == PARTITION_FAT32) ||
448              (PartEntry->PartitionType == PARTITION_FAT32_XINT13))
449     {
450 #if 0
451         if (CheckFatFormat())
452         {
453             PartEntry->FormatState = Preformatted;
454         }
455         else
456         {
457             PartEntry->FormatState = Unformatted;
458         }
459 #endif
460         PartEntry->FormatState = Preformatted;
461     }
462     else if (PartEntry->PartitionType == PARTITION_LINUX)
463     {
464 #if 0
465         if (CheckExt2Format())
466         {
467             PartEntry->FormatState = Preformatted;
468         }
469         else
470         {
471             PartEntry->FormatState = Unformatted;
472         }
473 #endif
474         PartEntry->FormatState = Preformatted;
475     }
476     else if (PartEntry->PartitionType == PARTITION_IFS)
477     {
478 #if 0
479         if (CheckNtfsFormat())
480         {
481             PartEntry->FormatState = Preformatted;
482         }
483         else if (CheckHpfsFormat())
484         {
485             PartEntry->FormatState = Preformatted;
486         }
487         else
488         {
489             PartEntry->FormatState = Unformatted;
490         }
491 #endif
492         PartEntry->FormatState = Preformatted;
493     }
494     else
495     {
496         PartEntry->FormatState = UnknownFormat;
497     }
498 
499     if (LogicalPartition)
500         InsertTailList(&DiskEntry->LogicalPartListHead,
501                        &PartEntry->ListEntry);
502     else
503         InsertTailList(&DiskEntry->PrimaryPartListHead,
504                        &PartEntry->ListEntry);
505 }
506 
507 
508 static
509 VOID
ScanForUnpartitionedDiskSpace(PDISKENTRY DiskEntry)510 ScanForUnpartitionedDiskSpace(
511     PDISKENTRY DiskEntry)
512 {
513     ULONGLONG LastStartSector;
514     ULONGLONG LastSectorCount;
515     ULONGLONG LastUnusedSectorCount;
516     PPARTENTRY PartEntry;
517     PPARTENTRY NewPartEntry;
518     PLIST_ENTRY Entry;
519 
520     DPRINT("ScanForUnpartitionedDiskSpace()\n");
521 
522     if (IsListEmpty(&DiskEntry->PrimaryPartListHead))
523     {
524         DPRINT1("No primary partition!\n");
525 
526         /* Create a partition table that represents the empty disk */
527         NewPartEntry = RtlAllocateHeap(RtlGetProcessHeap(),
528                                        HEAP_ZERO_MEMORY,
529                                        sizeof(PARTENTRY));
530         if (NewPartEntry == NULL)
531             return;
532 
533         NewPartEntry->DiskEntry = DiskEntry;
534 
535         NewPartEntry->IsPartitioned = FALSE;
536         NewPartEntry->StartSector.QuadPart = (ULONGLONG)DiskEntry->SectorAlignment;
537         NewPartEntry->SectorCount.QuadPart = AlignDown(DiskEntry->SectorCount.QuadPart, DiskEntry->SectorAlignment) -
538                                              NewPartEntry->StartSector.QuadPart;
539 
540         DPRINT1("First Sector: %I64u\n", NewPartEntry->StartSector.QuadPart);
541         DPRINT1("Last Sector: %I64u\n", NewPartEntry->StartSector.QuadPart + NewPartEntry->SectorCount.QuadPart - 1);
542         DPRINT1("Total Sectors: %I64u\n", NewPartEntry->SectorCount.QuadPart);
543 
544         NewPartEntry->FormatState = Unformatted;
545 
546         InsertTailList(&DiskEntry->PrimaryPartListHead,
547                        &NewPartEntry->ListEntry);
548 
549         return;
550     }
551 
552     /* Start partition at head 1, cylinder 0 */
553     LastStartSector = DiskEntry->SectorAlignment;
554     LastSectorCount = 0ULL;
555     LastUnusedSectorCount = 0ULL;
556 
557     Entry = DiskEntry->PrimaryPartListHead.Flink;
558     while (Entry != &DiskEntry->PrimaryPartListHead)
559     {
560         PartEntry = CONTAINING_RECORD(Entry, PARTENTRY, ListEntry);
561 
562         if (PartEntry->PartitionType != PARTITION_ENTRY_UNUSED ||
563             PartEntry->SectorCount.QuadPart != 0ULL)
564         {
565             LastUnusedSectorCount =
566                 PartEntry->StartSector.QuadPart - (LastStartSector + LastSectorCount);
567 
568             if (PartEntry->StartSector.QuadPart > (LastStartSector + LastSectorCount) &&
569                 LastUnusedSectorCount >= (ULONGLONG)DiskEntry->SectorAlignment)
570             {
571                 DPRINT("Unpartitioned disk space %I64u sectors\n", LastUnusedSectorCount);
572 
573                 NewPartEntry = RtlAllocateHeap(RtlGetProcessHeap(),
574                                                HEAP_ZERO_MEMORY,
575                                                sizeof(PARTENTRY));
576                 if (NewPartEntry == NULL)
577                     return;
578 
579                 NewPartEntry->DiskEntry = DiskEntry;
580 
581                 NewPartEntry->IsPartitioned = FALSE;
582                 NewPartEntry->StartSector.QuadPart = LastStartSector + LastSectorCount;
583                 NewPartEntry->SectorCount.QuadPart = AlignDown(NewPartEntry->StartSector.QuadPart + LastUnusedSectorCount, DiskEntry->SectorAlignment) -
584                                                      NewPartEntry->StartSector.QuadPart;
585 
586                 DPRINT1("First Sector: %I64u\n", NewPartEntry->StartSector.QuadPart);
587                 DPRINT1("Last Sector: %I64u\n", NewPartEntry->StartSector.QuadPart + NewPartEntry->SectorCount.QuadPart - 1);
588                 DPRINT1("Total Sectors: %I64u\n", NewPartEntry->SectorCount.QuadPart);
589 
590                 NewPartEntry->FormatState = Unformatted;
591 
592                 /* Insert the table into the list */
593                 InsertTailList(&PartEntry->ListEntry,
594                                &NewPartEntry->ListEntry);
595             }
596 
597             LastStartSector = PartEntry->StartSector.QuadPart;
598             LastSectorCount = PartEntry->SectorCount.QuadPart;
599         }
600 
601         Entry = Entry->Flink;
602     }
603 
604     /* Check for trailing unpartitioned disk space */
605     if ((LastStartSector + LastSectorCount) < DiskEntry->SectorCount.QuadPart)
606     {
607         LastUnusedSectorCount = AlignDown(DiskEntry->SectorCount.QuadPart - (LastStartSector + LastSectorCount), DiskEntry->SectorAlignment);
608 
609         if (LastUnusedSectorCount >= (ULONGLONG)DiskEntry->SectorAlignment)
610         {
611             DPRINT1("Unpartitioned disk space: %I64u sectors\n", LastUnusedSectorCount);
612 
613             NewPartEntry = RtlAllocateHeap(RtlGetProcessHeap(),
614                                            HEAP_ZERO_MEMORY,
615                                            sizeof(PARTENTRY));
616             if (NewPartEntry == NULL)
617                 return;
618 
619             NewPartEntry->DiskEntry = DiskEntry;
620 
621             NewPartEntry->IsPartitioned = FALSE;
622             NewPartEntry->StartSector.QuadPart = LastStartSector + LastSectorCount;
623             NewPartEntry->SectorCount.QuadPart = AlignDown(NewPartEntry->StartSector.QuadPart + LastUnusedSectorCount, DiskEntry->SectorAlignment) -
624                                                  NewPartEntry->StartSector.QuadPart;
625 
626             DPRINT1("First Sector: %I64u\n", NewPartEntry->StartSector.QuadPart);
627             DPRINT1("Last Sector: %I64u\n", NewPartEntry->StartSector.QuadPart + NewPartEntry->SectorCount.QuadPart - 1);
628             DPRINT1("Total Sectors: %I64u\n", NewPartEntry->SectorCount.QuadPart);
629 
630             NewPartEntry->FormatState = Unformatted;
631 
632             /* Append the table to the list */
633             InsertTailList(&DiskEntry->PrimaryPartListHead,
634                            &NewPartEntry->ListEntry);
635         }
636     }
637 
638     if (DiskEntry->ExtendedPartition != NULL)
639     {
640         if (IsListEmpty(&DiskEntry->LogicalPartListHead))
641         {
642             DPRINT1("No logical partition!\n");
643 
644             /* Create a partition table entry that represents the empty extended partition */
645             NewPartEntry = RtlAllocateHeap(RtlGetProcessHeap(),
646                                            HEAP_ZERO_MEMORY,
647                                            sizeof(PARTENTRY));
648             if (NewPartEntry == NULL)
649                 return;
650 
651             NewPartEntry->DiskEntry = DiskEntry;
652             NewPartEntry->LogicalPartition = TRUE;
653 
654             NewPartEntry->IsPartitioned = FALSE;
655             NewPartEntry->StartSector.QuadPart = DiskEntry->ExtendedPartition->StartSector.QuadPart + (ULONGLONG)DiskEntry->SectorAlignment;
656             NewPartEntry->SectorCount.QuadPart = DiskEntry->ExtendedPartition->SectorCount.QuadPart - (ULONGLONG)DiskEntry->SectorAlignment;
657 
658             DPRINT1("First Sector: %I64u\n", NewPartEntry->StartSector.QuadPart);
659             DPRINT1("Last Sector: %I64u\n", NewPartEntry->StartSector.QuadPart + NewPartEntry->SectorCount.QuadPart - 1);
660             DPRINT1("Total Sectors: %I64u\n", NewPartEntry->SectorCount.QuadPart);
661 
662             NewPartEntry->FormatState = Unformatted;
663 
664             InsertTailList(&DiskEntry->LogicalPartListHead,
665                            &NewPartEntry->ListEntry);
666 
667             return;
668         }
669 
670         /* Start partition at head 1, cylinder 0 */
671         LastStartSector = DiskEntry->ExtendedPartition->StartSector.QuadPart + (ULONGLONG)DiskEntry->SectorAlignment;
672         LastSectorCount = 0ULL;
673         LastUnusedSectorCount = 0ULL;
674 
675         Entry = DiskEntry->LogicalPartListHead.Flink;
676         while (Entry != &DiskEntry->LogicalPartListHead)
677         {
678             PartEntry = CONTAINING_RECORD(Entry, PARTENTRY, ListEntry);
679 
680             if (PartEntry->PartitionType != PARTITION_ENTRY_UNUSED ||
681                 PartEntry->SectorCount.QuadPart != 0ULL)
682             {
683                 LastUnusedSectorCount =
684                     PartEntry->StartSector.QuadPart - (ULONGLONG)DiskEntry->SectorAlignment - (LastStartSector + LastSectorCount);
685 
686                 if ((PartEntry->StartSector.QuadPart - (ULONGLONG)DiskEntry->SectorAlignment) > (LastStartSector + LastSectorCount) &&
687                     LastUnusedSectorCount >= (ULONGLONG)DiskEntry->SectorAlignment)
688                 {
689                     DPRINT("Unpartitioned disk space %I64u sectors\n", LastUnusedSectorCount);
690 
691                     NewPartEntry = RtlAllocateHeap(RtlGetProcessHeap(),
692                                                    HEAP_ZERO_MEMORY,
693                                                    sizeof(PARTENTRY));
694                     if (NewPartEntry == NULL)
695                         return;
696 
697                     NewPartEntry->DiskEntry = DiskEntry;
698                     NewPartEntry->LogicalPartition = TRUE;
699 
700                     NewPartEntry->IsPartitioned = FALSE;
701                     NewPartEntry->StartSector.QuadPart = LastStartSector + LastSectorCount;
702                     NewPartEntry->SectorCount.QuadPart = AlignDown(NewPartEntry->StartSector.QuadPart + LastUnusedSectorCount, DiskEntry->SectorAlignment) -
703                                                          NewPartEntry->StartSector.QuadPart;
704 
705                     DPRINT("First Sector: %I64u\n", NewPartEntry->StartSector.QuadPart);
706                     DPRINT("Last Sector: %I64u\n", NewPartEntry->StartSector.QuadPart + NewPartEntry->SectorCount.QuadPart - 1);
707                     DPRINT("Total Sectors: %I64u\n", NewPartEntry->SectorCount.QuadPart);
708 
709                     NewPartEntry->FormatState = Unformatted;
710 
711                     /* Insert the table into the list */
712                     InsertTailList(&PartEntry->ListEntry,
713                                    &NewPartEntry->ListEntry);
714                 }
715 
716                 LastStartSector = PartEntry->StartSector.QuadPart;
717                 LastSectorCount = PartEntry->SectorCount.QuadPart;
718             }
719 
720             Entry = Entry->Flink;
721         }
722 
723         /* Check for trailing unpartitioned disk space */
724         if ((LastStartSector + LastSectorCount) < DiskEntry->ExtendedPartition->StartSector.QuadPart + DiskEntry->ExtendedPartition->SectorCount.QuadPart)
725         {
726             LastUnusedSectorCount = AlignDown(DiskEntry->ExtendedPartition->StartSector.QuadPart + DiskEntry->ExtendedPartition->SectorCount.QuadPart - (LastStartSector + LastSectorCount),
727                                               DiskEntry->SectorAlignment);
728 
729             if (LastUnusedSectorCount >= (ULONGLONG)DiskEntry->SectorAlignment)
730             {
731                 DPRINT("Unpartitioned disk space: %I64u sectors\n", LastUnusedSectorCount);
732 
733                 NewPartEntry = RtlAllocateHeap(RtlGetProcessHeap(),
734                                                HEAP_ZERO_MEMORY,
735                                                sizeof(PARTENTRY));
736                 if (NewPartEntry == NULL)
737                     return;
738 
739                 NewPartEntry->DiskEntry = DiskEntry;
740                 NewPartEntry->LogicalPartition = TRUE;
741 
742                 NewPartEntry->IsPartitioned = FALSE;
743                 NewPartEntry->StartSector.QuadPart = LastStartSector + LastSectorCount;
744                 NewPartEntry->SectorCount.QuadPart = AlignDown(NewPartEntry->StartSector.QuadPart + LastUnusedSectorCount, DiskEntry->SectorAlignment) -
745                                                      NewPartEntry->StartSector.QuadPart;
746 
747                 DPRINT("First Sector: %I64u\n", NewPartEntry->StartSector.QuadPart);
748                 DPRINT("Last Sector: %I64u\n", NewPartEntry->StartSector.QuadPart + NewPartEntry->SectorCount.QuadPart - 1);
749                 DPRINT("Total Sectors: %I64u\n", NewPartEntry->SectorCount.QuadPart);
750 
751                 NewPartEntry->FormatState = Unformatted;
752 
753                 /* Append the table to the list */
754                 InsertTailList(&DiskEntry->LogicalPartListHead,
755                                &NewPartEntry->ListEntry);
756             }
757         }
758     }
759 
760     DPRINT("ScanForUnpartitionedDiskSpace() done\n");
761 }
762 
763 
764 static
765 VOID
AddDiskToList(HANDLE FileHandle,ULONG DiskNumber)766 AddDiskToList(
767     HANDLE FileHandle,
768     ULONG DiskNumber)
769 {
770     DISK_GEOMETRY DiskGeometry;
771     SCSI_ADDRESS ScsiAddress;
772     PDISKENTRY DiskEntry;
773     IO_STATUS_BLOCK Iosb;
774     NTSTATUS Status;
775     PPARTITION_SECTOR Mbr;
776     PULONG Buffer;
777     LARGE_INTEGER FileOffset;
778     WCHAR Identifier[20];
779     ULONG Checksum;
780     ULONG Signature;
781     ULONG i;
782     PLIST_ENTRY ListEntry;
783     PBIOSDISKENTRY BiosDiskEntry;
784     ULONG LayoutBufferSize;
785     PDRIVE_LAYOUT_INFORMATION NewLayoutBuffer;
786 
787     Status = NtDeviceIoControlFile(FileHandle,
788                                    NULL,
789                                    NULL,
790                                    NULL,
791                                    &Iosb,
792                                    IOCTL_DISK_GET_DRIVE_GEOMETRY,
793                                    NULL,
794                                    0,
795                                    &DiskGeometry,
796                                    sizeof(DISK_GEOMETRY));
797     if (!NT_SUCCESS(Status))
798     {
799         return;
800     }
801 
802     if (DiskGeometry.MediaType != FixedMedia &&
803         DiskGeometry.MediaType != RemovableMedia)
804     {
805         return;
806     }
807 
808     Status = NtDeviceIoControlFile(FileHandle,
809                                    NULL,
810                                    NULL,
811                                    NULL,
812                                    &Iosb,
813                                    IOCTL_SCSI_GET_ADDRESS,
814                                    NULL,
815                                    0,
816                                    &ScsiAddress,
817                                    sizeof(SCSI_ADDRESS));
818     if (!NT_SUCCESS(Status))
819     {
820         return;
821     }
822 
823     Mbr = (PARTITION_SECTOR*)RtlAllocateHeap(RtlGetProcessHeap(),
824                                              0,
825                                              DiskGeometry.BytesPerSector);
826     if (Mbr == NULL)
827     {
828         return;
829     }
830 
831     FileOffset.QuadPart = 0;
832     Status = NtReadFile(FileHandle,
833                         NULL,
834                         NULL,
835                         NULL,
836                         &Iosb,
837                         (PVOID)Mbr,
838                         DiskGeometry.BytesPerSector,
839                         &FileOffset,
840                         NULL);
841     if (!NT_SUCCESS(Status))
842     {
843         RtlFreeHeap(RtlGetProcessHeap(), 0, Mbr);
844         DPRINT1("NtReadFile failed, status=%x\n", Status);
845         return;
846     }
847     Signature = Mbr->Signature;
848 
849     /* Calculate the MBR checksum */
850     Checksum = 0;
851     Buffer = (PULONG)Mbr;
852     for (i = 0; i < 128; i++)
853     {
854         Checksum += Buffer[i];
855     }
856     Checksum = ~Checksum + 1;
857 
858     StringCchPrintfW(Identifier, ARRAYSIZE(Identifier),
859                      L"%08x-%08x-A", Checksum, Signature);
860     DPRINT("Identifier: %S\n", Identifier);
861 
862     DiskEntry = RtlAllocateHeap(RtlGetProcessHeap(),
863                                 HEAP_ZERO_MEMORY,
864                                 sizeof(DISKENTRY));
865     if (DiskEntry == NULL)
866     {
867         return;
868     }
869 
870 //    DiskEntry->Checksum = Checksum;
871 //    DiskEntry->Signature = Signature;
872     DiskEntry->BiosFound = FALSE;
873 
874     /* Check if this disk has a valid MBR */
875     if (Mbr->Magic != MBR_MAGIC)
876         DiskEntry->NoMbr = TRUE;
877     else
878         DiskEntry->NoMbr = FALSE;
879 
880     /* Free Mbr sector buffer */
881     RtlFreeHeap(RtlGetProcessHeap(), 0, Mbr);
882 
883     ListEntry = BiosDiskListHead.Flink;
884     while (ListEntry != &BiosDiskListHead)
885     {
886         BiosDiskEntry = CONTAINING_RECORD(ListEntry, BIOSDISKENTRY, ListEntry);
887         /* FIXME:
888          *   Compare the size from bios and the reported size from driver.
889          *   If we have more than one disk with a zero or with the same signatur
890          *   we must create new signatures and reboot. After the reboot,
891          *   it is possible to identify the disks.
892          */
893         if (BiosDiskEntry->Signature == Signature &&
894             BiosDiskEntry->Checksum == Checksum &&
895             !BiosDiskEntry->Recognized)
896         {
897             if (!DiskEntry->BiosFound)
898             {
899                 DiskEntry->BiosDiskNumber = BiosDiskEntry->DiskNumber;
900                 DiskEntry->BiosFound = TRUE;
901                 BiosDiskEntry->Recognized = TRUE;
902             }
903             else
904             {
905             }
906         }
907         ListEntry = ListEntry->Flink;
908     }
909 
910     if (!DiskEntry->BiosFound)
911     {
912 #if 0
913         RtlFreeHeap(ProcessHeap, 0, DiskEntry);
914         return;
915 #else
916         DPRINT1("WARNING: Setup could not find a matching BIOS disk entry. Disk %d is not be bootable by the BIOS!\n", DiskNumber);
917 #endif
918     }
919 
920     InitializeListHead(&DiskEntry->PrimaryPartListHead);
921     InitializeListHead(&DiskEntry->LogicalPartListHead);
922 
923     DiskEntry->Cylinders = DiskGeometry.Cylinders.QuadPart;
924     DiskEntry->TracksPerCylinder = DiskGeometry.TracksPerCylinder;
925     DiskEntry->SectorsPerTrack = DiskGeometry.SectorsPerTrack;
926     DiskEntry->BytesPerSector = DiskGeometry.BytesPerSector;
927 
928     DPRINT("Cylinders %I64u\n", DiskEntry->Cylinders);
929     DPRINT("TracksPerCylinder %I64u\n", DiskEntry->TracksPerCylinder);
930     DPRINT("SectorsPerTrack %I64u\n", DiskEntry->SectorsPerTrack);
931     DPRINT("BytesPerSector %I64u\n", DiskEntry->BytesPerSector);
932 
933     DiskEntry->SectorCount.QuadPart = DiskGeometry.Cylinders.QuadPart *
934                                       (ULONGLONG)DiskGeometry.TracksPerCylinder *
935                                       (ULONGLONG)DiskGeometry.SectorsPerTrack;
936 
937 //    DiskEntry->SectorAlignment = DiskGeometry.SectorsPerTrack;
938 //    DiskEntry->CylinderAlignment = DiskGeometry.SectorsPerTrack * DiskGeometry.TracksPerCylinder;
939     DiskEntry->SectorAlignment = (1024 * 1024) / DiskGeometry.BytesPerSector;
940     DiskEntry->CylinderAlignment = (1024 * 1024) / DiskGeometry.BytesPerSector;
941 
942     DPRINT1("SectorCount: %I64u\n", DiskEntry->SectorCount);
943     DPRINT1("SectorAlignment: %lu\n", DiskEntry->SectorAlignment);
944     DPRINT1("CylinderAlignment: %lu\n", DiskEntry->CylinderAlignment);
945 
946     DiskEntry->DiskNumber = DiskNumber;
947     DiskEntry->Port = ScsiAddress.PortNumber;
948     DiskEntry->PathId = ScsiAddress.PathId;
949     DiskEntry->TargetId = ScsiAddress.TargetId;
950     DiskEntry->Lun = ScsiAddress.Lun;
951 
952     GetDriverName(DiskEntry);
953 
954     InsertAscendingList(&DiskListHead, DiskEntry, DISKENTRY, ListEntry, DiskNumber);
955 
956     /* Allocate a layout buffer with 4 partition entries first */
957     LayoutBufferSize = sizeof(DRIVE_LAYOUT_INFORMATION) +
958                        ((4 - ANYSIZE_ARRAY) * sizeof(PARTITION_INFORMATION));
959     DiskEntry->LayoutBuffer = RtlAllocateHeap(RtlGetProcessHeap(),
960                                               HEAP_ZERO_MEMORY,
961                                               LayoutBufferSize);
962     if (DiskEntry->LayoutBuffer == NULL)
963     {
964         DPRINT1("Failed to allocate the disk layout buffer!\n");
965         return;
966     }
967 
968     for (;;)
969     {
970         DPRINT1("Buffer size: %lu\n", LayoutBufferSize);
971         Status = NtDeviceIoControlFile(FileHandle,
972                                        NULL,
973                                        NULL,
974                                        NULL,
975                                        &Iosb,
976                                        IOCTL_DISK_GET_DRIVE_LAYOUT,
977                                        NULL,
978                                        0,
979                                        DiskEntry->LayoutBuffer,
980                                        LayoutBufferSize);
981         if (NT_SUCCESS(Status))
982             break;
983 
984         if (Status != STATUS_BUFFER_TOO_SMALL)
985         {
986             DPRINT1("NtDeviceIoControlFile() failed (Status: 0x%08lx)\n", Status);
987             return;
988         }
989 
990         LayoutBufferSize += 4 * sizeof(PARTITION_INFORMATION);
991         NewLayoutBuffer = RtlReAllocateHeap(RtlGetProcessHeap(),
992                                             HEAP_ZERO_MEMORY,
993                                             DiskEntry->LayoutBuffer,
994                                             LayoutBufferSize);
995         if (NewLayoutBuffer == NULL)
996         {
997             DPRINT1("Failed to reallocate the disk layout buffer!\n");
998             return;
999         }
1000 
1001         DiskEntry->LayoutBuffer = NewLayoutBuffer;
1002     }
1003 
1004     DPRINT1("PartitionCount: %lu\n", DiskEntry->LayoutBuffer->PartitionCount);
1005 
1006 #ifdef DUMP_PARTITION_TABLE
1007     DumpPartitionTable(DiskEntry);
1008 #endif
1009 
1010     if (DiskEntry->LayoutBuffer->PartitionEntry[0].StartingOffset.QuadPart != 0 &&
1011         DiskEntry->LayoutBuffer->PartitionEntry[0].PartitionLength.QuadPart != 0 &&
1012         DiskEntry->LayoutBuffer->PartitionEntry[0].PartitionType != 0)
1013     {
1014         if ((DiskEntry->LayoutBuffer->PartitionEntry[0].StartingOffset.QuadPart / DiskEntry->BytesPerSector) % DiskEntry->SectorsPerTrack == 0)
1015         {
1016             DPRINT("Use %lu Sector alignment!\n", DiskEntry->SectorsPerTrack);
1017         }
1018         else if (DiskEntry->LayoutBuffer->PartitionEntry[0].StartingOffset.QuadPart % (1024 * 1024) == 0)
1019         {
1020             DPRINT1("Use megabyte (%lu Sectors) alignment!\n", (1024 * 1024) / DiskEntry->BytesPerSector);
1021         }
1022         else
1023         {
1024             DPRINT1("No matching alignment found! Partition 1 starts at %I64u\n", DiskEntry->LayoutBuffer->PartitionEntry[0].StartingOffset.QuadPart);
1025         }
1026     }
1027     else
1028     {
1029         DPRINT1("No valid partition table found! Use megabyte (%lu Sectors) alignment!\n", (1024 * 1024) / DiskEntry->BytesPerSector);
1030     }
1031 
1032 
1033     if (DiskEntry->LayoutBuffer->PartitionCount == 0)
1034     {
1035         DiskEntry->NewDisk = TRUE;
1036         DiskEntry->LayoutBuffer->PartitionCount = 4;
1037 
1038         for (i = 0; i < 4; i++)
1039             DiskEntry->LayoutBuffer->PartitionEntry[i].RewritePartition = TRUE;
1040     }
1041     else
1042     {
1043         for (i = 0; i < 4; i++)
1044         {
1045             AddPartitionToDisk(DiskNumber, DiskEntry, i, FALSE);
1046         }
1047 
1048         for (i = 4; i < DiskEntry->LayoutBuffer->PartitionCount; i += 4)
1049         {
1050             AddPartitionToDisk(DiskNumber, DiskEntry, i, TRUE);
1051         }
1052     }
1053 
1054     ScanForUnpartitionedDiskSpace(DiskEntry);
1055 }
1056 
1057 
1058 NTSTATUS
CreatePartitionList(VOID)1059 CreatePartitionList(VOID)
1060 {
1061     OBJECT_ATTRIBUTES ObjectAttributes;
1062     SYSTEM_DEVICE_INFORMATION Sdi;
1063     IO_STATUS_BLOCK Iosb;
1064     ULONG ReturnSize;
1065     NTSTATUS Status;
1066     ULONG DiskNumber;
1067     WCHAR Buffer[MAX_PATH];
1068     UNICODE_STRING Name;
1069     HANDLE FileHandle;
1070 
1071     CurrentDisk = NULL;
1072     CurrentPartition = NULL;
1073 
1074 //    BootDisk = NULL;
1075 //    BootPartition = NULL;
1076 
1077 //    TempDisk = NULL;
1078 //    TempPartition = NULL;
1079 //    FormatState = Start;
1080 
1081     InitializeListHead(&DiskListHead);
1082     InitializeListHead(&BiosDiskListHead);
1083 
1084     EnumerateBiosDiskEntries();
1085 
1086     Status = NtQuerySystemInformation(SystemDeviceInformation,
1087                                       &Sdi,
1088                                       sizeof(SYSTEM_DEVICE_INFORMATION),
1089                                       &ReturnSize);
1090     if (!NT_SUCCESS(Status))
1091     {
1092         return Status;
1093     }
1094 
1095     for (DiskNumber = 0; DiskNumber < Sdi.NumberOfDisks; DiskNumber++)
1096     {
1097         StringCchPrintfW(Buffer, ARRAYSIZE(Buffer),
1098                          L"\\Device\\Harddisk%d\\Partition0",
1099                          DiskNumber);
1100 
1101         RtlInitUnicodeString(&Name,
1102                              Buffer);
1103 
1104         InitializeObjectAttributes(&ObjectAttributes,
1105                                    &Name,
1106                                    0,
1107                                    NULL,
1108                                    NULL);
1109 
1110         Status = NtOpenFile(&FileHandle,
1111                             FILE_READ_DATA | FILE_READ_ATTRIBUTES | SYNCHRONIZE,
1112                             &ObjectAttributes,
1113                             &Iosb,
1114                             FILE_SHARE_READ,
1115                             FILE_SYNCHRONOUS_IO_NONALERT);
1116         if (NT_SUCCESS(Status))
1117         {
1118             AddDiskToList(FileHandle, DiskNumber);
1119 
1120             NtClose(FileHandle);
1121         }
1122     }
1123 
1124 //    UpdateDiskSignatures(List);
1125 
1126 //    AssignDriveLetters(List);
1127 
1128     return STATUS_SUCCESS;
1129 }
1130 
1131 
1132 VOID
DestroyPartitionList(VOID)1133 DestroyPartitionList(VOID)
1134 {
1135     PDISKENTRY DiskEntry;
1136     PBIOSDISKENTRY BiosDiskEntry;
1137     PPARTENTRY PartEntry;
1138     PLIST_ENTRY Entry;
1139 
1140     CurrentDisk = NULL;
1141     CurrentPartition = NULL;
1142 
1143     /* Release disk and partition info */
1144     while (!IsListEmpty(&DiskListHead))
1145     {
1146         Entry = RemoveHeadList(&DiskListHead);
1147         DiskEntry = CONTAINING_RECORD(Entry, DISKENTRY, ListEntry);
1148 
1149         /* Release driver name */
1150         RtlFreeUnicodeString(&DiskEntry->DriverName);
1151 
1152         /* Release primary partition list */
1153         while (!IsListEmpty(&DiskEntry->PrimaryPartListHead))
1154         {
1155             Entry = RemoveHeadList(&DiskEntry->PrimaryPartListHead);
1156             PartEntry = CONTAINING_RECORD(Entry, PARTENTRY, ListEntry);
1157 
1158             RtlFreeHeap(RtlGetProcessHeap(), 0, PartEntry);
1159         }
1160 
1161         /* Release logical partition list */
1162         while (!IsListEmpty(&DiskEntry->LogicalPartListHead))
1163         {
1164             Entry = RemoveHeadList(&DiskEntry->LogicalPartListHead);
1165             PartEntry = CONTAINING_RECORD(Entry, PARTENTRY, ListEntry);
1166 
1167             RtlFreeHeap(RtlGetProcessHeap(), 0, PartEntry);
1168         }
1169 
1170         /* Release layout buffer */
1171         if (DiskEntry->LayoutBuffer != NULL)
1172             RtlFreeHeap(RtlGetProcessHeap(), 0, DiskEntry->LayoutBuffer);
1173 
1174 
1175         /* Release disk entry */
1176         RtlFreeHeap(RtlGetProcessHeap(), 0, DiskEntry);
1177     }
1178 
1179     /* Release the bios disk info */
1180     while (!IsListEmpty(&BiosDiskListHead))
1181     {
1182         Entry = RemoveHeadList(&BiosDiskListHead);
1183         BiosDiskEntry = CONTAINING_RECORD(Entry, BIOSDISKENTRY, ListEntry);
1184 
1185         RtlFreeHeap(RtlGetProcessHeap(), 0, BiosDiskEntry);
1186     }
1187 }
1188 
1189 
1190 static
1191 VOID
GetVolumeExtents(_In_ HANDLE VolumeHandle,_In_ PVOLENTRY VolumeEntry)1192 GetVolumeExtents(
1193     _In_ HANDLE VolumeHandle,
1194     _In_ PVOLENTRY VolumeEntry)
1195 {
1196     DWORD dwBytesReturned = 0, dwLength, i;
1197     PVOLUME_DISK_EXTENTS pExtents;
1198     BOOL bResult;
1199     DWORD dwError;
1200 
1201     dwLength = sizeof(VOLUME_DISK_EXTENTS);
1202     pExtents = RtlAllocateHeap(RtlGetProcessHeap(), HEAP_ZERO_MEMORY, dwLength);
1203     if (pExtents == NULL)
1204         return;
1205 
1206     bResult = DeviceIoControl(VolumeHandle,
1207                               IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS,
1208                               NULL,
1209                               0,
1210                               pExtents,
1211                               dwLength,
1212                               &dwBytesReturned,
1213                               NULL);
1214     if (!bResult)
1215     {
1216         dwError = GetLastError();
1217 
1218         if (dwError != ERROR_MORE_DATA)
1219         {
1220             RtlFreeHeap(RtlGetProcessHeap(), 0, pExtents);
1221             return;
1222         }
1223         else
1224         {
1225             dwLength = sizeof(VOLUME_DISK_EXTENTS) + ((pExtents->NumberOfDiskExtents - 1) * sizeof(DISK_EXTENT));
1226             RtlFreeHeap(RtlGetProcessHeap(), 0, pExtents);
1227             pExtents = RtlAllocateHeap(RtlGetProcessHeap(), HEAP_ZERO_MEMORY, dwLength);
1228             if (pExtents == NULL)
1229             {
1230                 return;
1231             }
1232 
1233             bResult = DeviceIoControl(VolumeHandle,
1234                                       IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS,
1235                                       NULL,
1236                                       0,
1237                                       pExtents,
1238                                       dwLength,
1239                                       &dwBytesReturned,
1240                                       NULL);
1241             if (!bResult)
1242             {
1243                 RtlFreeHeap(RtlGetProcessHeap(), 0, pExtents);
1244                 return;
1245             }
1246         }
1247     }
1248 
1249     for (i = 0; i < pExtents->NumberOfDiskExtents; i++)
1250         VolumeEntry->Size.QuadPart += pExtents->Extents[i].ExtentLength.QuadPart;
1251 
1252     VolumeEntry->pExtents = pExtents;
1253 }
1254 
1255 
1256 static
1257 VOID
GetVolumeType(_In_ HANDLE VolumeHandle,_In_ PVOLENTRY VolumeEntry)1258 GetVolumeType(
1259     _In_ HANDLE VolumeHandle,
1260     _In_ PVOLENTRY VolumeEntry)
1261 {
1262     FILE_FS_DEVICE_INFORMATION DeviceInfo;
1263     IO_STATUS_BLOCK IoStatusBlock;
1264     NTSTATUS Status;
1265 
1266     Status = NtQueryVolumeInformationFile(VolumeHandle,
1267                                           &IoStatusBlock,
1268                                           &DeviceInfo,
1269                                           sizeof(FILE_FS_DEVICE_INFORMATION),
1270                                           FileFsDeviceInformation);
1271     if (!NT_SUCCESS(Status))
1272         return;
1273 
1274     switch (DeviceInfo.DeviceType)
1275     {
1276         case FILE_DEVICE_CD_ROM:
1277         case FILE_DEVICE_CD_ROM_FILE_SYSTEM:
1278             VolumeEntry->VolumeType = VOLUME_TYPE_CDROM;
1279             break;
1280 
1281         case FILE_DEVICE_DISK:
1282         case FILE_DEVICE_DISK_FILE_SYSTEM:
1283             if (DeviceInfo.Characteristics & FILE_REMOVABLE_MEDIA)
1284                 VolumeEntry->VolumeType = VOLUME_TYPE_REMOVABLE;
1285             else
1286                 VolumeEntry->VolumeType = VOLUME_TYPE_PARTITION;
1287             break;
1288 
1289         default:
1290             VolumeEntry->VolumeType = VOLUME_TYPE_UNKNOWN;
1291             break;
1292     }
1293 }
1294 
1295 
1296 static
1297 VOID
AddVolumeToList(ULONG ulVolumeNumber,PWSTR pszVolumeName)1298 AddVolumeToList(
1299     ULONG ulVolumeNumber,
1300     PWSTR pszVolumeName)
1301 {
1302     PVOLENTRY VolumeEntry;
1303     HANDLE VolumeHandle;
1304 
1305     DWORD dwError, dwLength;
1306     WCHAR szPathNames[MAX_PATH + 1];
1307     WCHAR szVolumeName[MAX_PATH + 1];
1308     WCHAR szFilesystem[MAX_PATH + 1];
1309 
1310     DWORD  CharCount            = 0;
1311     size_t Index                = 0;
1312 
1313     OBJECT_ATTRIBUTES ObjectAttributes;
1314     UNICODE_STRING Name;
1315     IO_STATUS_BLOCK Iosb;
1316     NTSTATUS Status;
1317 
1318     DPRINT("AddVolumeToList(%S)\n", pszVolumeName);
1319 
1320     VolumeEntry = RtlAllocateHeap(RtlGetProcessHeap(),
1321                                   HEAP_ZERO_MEMORY,
1322                                   sizeof(VOLENTRY));
1323     if (VolumeEntry == NULL)
1324         return;
1325 
1326     VolumeEntry->VolumeNumber = ulVolumeNumber;
1327     wcscpy(VolumeEntry->VolumeName, pszVolumeName);
1328 
1329     Index = wcslen(pszVolumeName) - 1;
1330 
1331     pszVolumeName[Index] = L'\0';
1332 
1333     CharCount = QueryDosDeviceW(&pszVolumeName[4], VolumeEntry->DeviceName, ARRAYSIZE(VolumeEntry->DeviceName));
1334 
1335     pszVolumeName[Index] = L'\\';
1336 
1337     if (CharCount == 0)
1338     {
1339         RtlFreeHeap(RtlGetProcessHeap(), 0, VolumeEntry);
1340         return;
1341     }
1342 
1343     DPRINT("DeviceName: %S\n", VolumeEntry->DeviceName);
1344 
1345     RtlInitUnicodeString(&Name, VolumeEntry->DeviceName);
1346 
1347     InitializeObjectAttributes(&ObjectAttributes,
1348                                &Name,
1349                                0,
1350                                NULL,
1351                                NULL);
1352 
1353     Status = NtOpenFile(&VolumeHandle,
1354                         SYNCHRONIZE,
1355                         &ObjectAttributes,
1356                         &Iosb,
1357                         0,
1358                         FILE_DIRECTORY_FILE | FILE_SYNCHRONOUS_IO_NONALERT | FILE_OPEN_FOR_BACKUP_INTENT);
1359     if (NT_SUCCESS(Status))
1360     {
1361         GetVolumeType(VolumeHandle, VolumeEntry);
1362         GetVolumeExtents(VolumeHandle, VolumeEntry);
1363         NtClose(VolumeHandle);
1364     }
1365 
1366     if (GetVolumeInformationW(pszVolumeName,
1367                               szVolumeName,
1368                               MAX_PATH + 1,
1369                               NULL, //  [out, optional] LPDWORD lpVolumeSerialNumber,
1370                               NULL, //  [out, optional] LPDWORD lpMaximumComponentLength,
1371                               NULL, //  [out, optional] LPDWORD lpFileSystemFlags,
1372                               szFilesystem,
1373                               MAX_PATH + 1))
1374     {
1375         VolumeEntry->pszLabel = RtlAllocateHeap(RtlGetProcessHeap(),
1376                                                 0,
1377                                                 (wcslen(szVolumeName) + 1) * sizeof(WCHAR));
1378         if (VolumeEntry->pszLabel)
1379             wcscpy(VolumeEntry->pszLabel, szVolumeName);
1380 
1381         VolumeEntry->pszFilesystem = RtlAllocateHeap(RtlGetProcessHeap(),
1382                                                      0,
1383                                                      (wcslen(szFilesystem) + 1) * sizeof(WCHAR));
1384         if (VolumeEntry->pszFilesystem)
1385             wcscpy(VolumeEntry->pszFilesystem, szFilesystem);
1386     }
1387     else
1388     {
1389         dwError = GetLastError();
1390         if (dwError == ERROR_UNRECOGNIZED_VOLUME)
1391         {
1392             VolumeEntry->pszFilesystem = RtlAllocateHeap(RtlGetProcessHeap(),
1393                                                          0,
1394                                                          (3 + 1) * sizeof(WCHAR));
1395             if (VolumeEntry->pszFilesystem)
1396                 wcscpy(VolumeEntry->pszFilesystem, L"RAW");
1397         }
1398     }
1399 
1400     if (GetVolumePathNamesForVolumeNameW(pszVolumeName,
1401                                          szPathNames,
1402                                          ARRAYSIZE(szPathNames),
1403                                          &dwLength))
1404     {
1405         VolumeEntry->DriveLetter = szPathNames[0];
1406     }
1407 
1408     InsertTailList(&VolumeListHead,
1409                    &VolumeEntry->ListEntry);
1410 }
1411 
1412 
1413 NTSTATUS
CreateVolumeList(VOID)1414 CreateVolumeList(VOID)
1415 {
1416     HANDLE hVolume = INVALID_HANDLE_VALUE;
1417     WCHAR szVolumeName[MAX_PATH];
1418     ULONG ulVolumeNumber = 0;
1419     BOOL Success;
1420 
1421     CurrentVolume = NULL;
1422 
1423     InitializeListHead(&VolumeListHead);
1424 
1425     hVolume = FindFirstVolumeW(szVolumeName, ARRAYSIZE(szVolumeName));
1426     if (hVolume == INVALID_HANDLE_VALUE)
1427     {
1428 
1429         return STATUS_UNSUCCESSFUL;
1430     }
1431 
1432     AddVolumeToList(ulVolumeNumber++, szVolumeName);
1433 
1434     for (;;)
1435     {
1436         Success = FindNextVolumeW(hVolume, szVolumeName, ARRAYSIZE(szVolumeName));
1437         if (!Success)
1438         {
1439             break;
1440         }
1441 
1442         AddVolumeToList(ulVolumeNumber++, szVolumeName);
1443     }
1444 
1445     FindVolumeClose(hVolume);
1446 
1447     return STATUS_SUCCESS;
1448 }
1449 
1450 
1451 VOID
DestroyVolumeList(VOID)1452 DestroyVolumeList(VOID)
1453 {
1454     PLIST_ENTRY Entry;
1455     PVOLENTRY VolumeEntry;
1456 
1457     CurrentVolume = NULL;
1458 
1459     /* Release disk and partition info */
1460     while (!IsListEmpty(&VolumeListHead))
1461     {
1462         Entry = RemoveHeadList(&VolumeListHead);
1463         VolumeEntry = CONTAINING_RECORD(Entry, VOLENTRY, ListEntry);
1464 
1465         if (VolumeEntry->pszLabel)
1466             RtlFreeHeap(RtlGetProcessHeap(), 0, VolumeEntry->pszLabel);
1467 
1468         if (VolumeEntry->pszFilesystem)
1469             RtlFreeHeap(RtlGetProcessHeap(), 0, VolumeEntry->pszFilesystem);
1470 
1471         if (VolumeEntry->pExtents)
1472             RtlFreeHeap(RtlGetProcessHeap(), 0, VolumeEntry->pExtents);
1473 
1474         /* Release disk entry */
1475         RtlFreeHeap(RtlGetProcessHeap(), 0, VolumeEntry);
1476     }
1477 }
1478 
1479 
1480 NTSTATUS
WritePartitions(_In_ PDISKENTRY DiskEntry)1481 WritePartitions(
1482     _In_ PDISKENTRY DiskEntry)
1483 {
1484     NTSTATUS Status;
1485     OBJECT_ATTRIBUTES ObjectAttributes;
1486     UNICODE_STRING Name;
1487     HANDLE FileHandle;
1488     IO_STATUS_BLOCK Iosb;
1489     ULONG BufferSize;
1490     PPARTITION_INFORMATION PartitionInfo;
1491     ULONG PartitionCount;
1492     PLIST_ENTRY ListEntry;
1493     PPARTENTRY PartEntry;
1494     WCHAR DstPath[MAX_PATH];
1495 
1496     DPRINT("WritePartitions() Disk: %lu\n", DiskEntry->DiskNumber);
1497 
1498     /* If the disk is not dirty, there is nothing to do */
1499     if (!DiskEntry->Dirty)
1500         return STATUS_SUCCESS;
1501 
1502     StringCchPrintfW(DstPath, ARRAYSIZE(DstPath),
1503                      L"\\Device\\Harddisk%lu\\Partition0",
1504                      DiskEntry->DiskNumber);
1505     RtlInitUnicodeString(&Name, DstPath);
1506 
1507     InitializeObjectAttributes(&ObjectAttributes,
1508                                &Name,
1509                                OBJ_CASE_INSENSITIVE,
1510                                NULL,
1511                                NULL);
1512 
1513     Status = NtOpenFile(&FileHandle,
1514                         GENERIC_READ | GENERIC_WRITE | SYNCHRONIZE,
1515                         &ObjectAttributes,
1516                         &Iosb,
1517                         0,
1518                         FILE_SYNCHRONOUS_IO_NONALERT);
1519     if (!NT_SUCCESS(Status))
1520     {
1521         DPRINT1("NtOpenFile() failed (Status %lx)\n", Status);
1522         return Status;
1523     }
1524 
1525     //
1526     // FIXME: We first *MUST* use IOCTL_DISK_CREATE_DISK to initialize
1527     // the disk in MBR or GPT format in case the disk was not initialized!!
1528     // For this we must ask the user which format to use.
1529     //
1530 
1531     /* Save the original partition count to be restored later (see comment below) */
1532     PartitionCount = DiskEntry->LayoutBuffer->PartitionCount;
1533 
1534     /* Set the new disk layout and retrieve its updated version with possibly modified partition numbers */
1535     BufferSize = sizeof(DRIVE_LAYOUT_INFORMATION) +
1536                  ((PartitionCount - 1) * sizeof(PARTITION_INFORMATION));
1537     Status = NtDeviceIoControlFile(FileHandle,
1538                                    NULL,
1539                                    NULL,
1540                                    NULL,
1541                                    &Iosb,
1542                                    IOCTL_DISK_SET_DRIVE_LAYOUT,
1543                                    DiskEntry->LayoutBuffer,
1544                                    BufferSize,
1545                                    DiskEntry->LayoutBuffer,
1546                                    BufferSize);
1547     NtClose(FileHandle);
1548 
1549     /*
1550      * IOCTL_DISK_SET_DRIVE_LAYOUT calls IoWritePartitionTable(), which converts
1551      * DiskEntry->LayoutBuffer->PartitionCount into a partition *table* count,
1552      * where such a table is expected to enumerate up to 4 partitions:
1553      * partition *table* count == ROUND_UP(PartitionCount, 4) / 4 .
1554      * Due to this we need to restore the original PartitionCount number.
1555      */
1556     DiskEntry->LayoutBuffer->PartitionCount = PartitionCount;
1557 
1558     /* Check whether the IOCTL_DISK_SET_DRIVE_LAYOUT call succeeded */
1559     if (!NT_SUCCESS(Status))
1560     {
1561         DPRINT1("IOCTL_DISK_SET_DRIVE_LAYOUT failed (Status 0x%08lx)\n", Status);
1562         return Status;
1563     }
1564 
1565     /* Update the partition numbers */
1566 
1567     /* Update the primary partition table */
1568     for (ListEntry = DiskEntry->PrimaryPartListHead.Flink;
1569          ListEntry != &DiskEntry->PrimaryPartListHead;
1570          ListEntry = ListEntry->Flink)
1571     {
1572         PartEntry = CONTAINING_RECORD(ListEntry, PARTENTRY, ListEntry);
1573 
1574         if (PartEntry->IsPartitioned)
1575         {
1576             ASSERT(PartEntry->PartitionType != PARTITION_ENTRY_UNUSED);
1577             PartitionInfo = &DiskEntry->LayoutBuffer->PartitionEntry[PartEntry->PartitionIndex];
1578             PartEntry->PartitionNumber = PartitionInfo->PartitionNumber;
1579         }
1580     }
1581 
1582     /* Update the logical partition table */
1583     for (ListEntry = DiskEntry->LogicalPartListHead.Flink;
1584          ListEntry != &DiskEntry->LogicalPartListHead;
1585          ListEntry = ListEntry->Flink)
1586     {
1587         PartEntry = CONTAINING_RECORD(ListEntry, PARTENTRY, ListEntry);
1588 
1589         if (PartEntry->IsPartitioned)
1590         {
1591             ASSERT(PartEntry->PartitionType != PARTITION_ENTRY_UNUSED);
1592             PartitionInfo = &DiskEntry->LayoutBuffer->PartitionEntry[PartEntry->PartitionIndex];
1593             PartEntry->PartitionNumber = PartitionInfo->PartitionNumber;
1594         }
1595     }
1596 
1597     /* The layout has been successfully updated, the disk is not dirty anymore */
1598     DiskEntry->Dirty = FALSE;
1599 
1600     return Status;
1601 }
1602 
1603 
1604 static
1605 BOOLEAN
IsEmptyLayoutEntry(IN PPARTITION_INFORMATION PartitionInfo)1606 IsEmptyLayoutEntry(
1607     IN PPARTITION_INFORMATION PartitionInfo)
1608 {
1609     if (PartitionInfo->StartingOffset.QuadPart == 0 &&
1610         PartitionInfo->PartitionLength.QuadPart == 0)
1611     {
1612         return TRUE;
1613     }
1614 
1615     return FALSE;
1616 }
1617 
1618 
1619 static
1620 BOOLEAN
IsSamePrimaryLayoutEntry(IN PPARTITION_INFORMATION PartitionInfo,IN PDISKENTRY DiskEntry,IN PPARTENTRY PartEntry)1621 IsSamePrimaryLayoutEntry(
1622     IN PPARTITION_INFORMATION PartitionInfo,
1623     IN PDISKENTRY DiskEntry,
1624     IN PPARTENTRY PartEntry)
1625 {
1626     if ((PartitionInfo->StartingOffset.QuadPart == PartEntry->StartSector.QuadPart * DiskEntry->BytesPerSector) &&
1627         (PartitionInfo->PartitionLength.QuadPart == PartEntry->SectorCount.QuadPart * DiskEntry->BytesPerSector))
1628     {
1629         return TRUE;
1630     }
1631 
1632     return FALSE;
1633 }
1634 
1635 
1636 ULONG
GetPrimaryPartitionCount(_In_ PDISKENTRY DiskEntry)1637 GetPrimaryPartitionCount(
1638     _In_ PDISKENTRY DiskEntry)
1639 {
1640     PLIST_ENTRY Entry;
1641     PPARTENTRY PartEntry;
1642     ULONG Count = 0;
1643 
1644     for (Entry = DiskEntry->PrimaryPartListHead.Flink;
1645          Entry != &DiskEntry->PrimaryPartListHead;
1646          Entry = Entry->Flink)
1647     {
1648         PartEntry = CONTAINING_RECORD(Entry, PARTENTRY, ListEntry);
1649         if (PartEntry->IsPartitioned)
1650             Count++;
1651     }
1652 
1653     return Count;
1654 }
1655 
1656 
1657 static
1658 ULONG
GetLogicalPartitionCount(_In_ PDISKENTRY DiskEntry)1659 GetLogicalPartitionCount(
1660     _In_ PDISKENTRY DiskEntry)
1661 {
1662     PLIST_ENTRY ListEntry;
1663     PPARTENTRY PartEntry;
1664     ULONG Count = 0;
1665 
1666     for (ListEntry = DiskEntry->LogicalPartListHead.Flink;
1667          ListEntry != &DiskEntry->LogicalPartListHead;
1668          ListEntry = ListEntry->Flink)
1669     {
1670         PartEntry = CONTAINING_RECORD(ListEntry, PARTENTRY, ListEntry);
1671         if (PartEntry->IsPartitioned)
1672             Count++;
1673     }
1674 
1675     return Count;
1676 }
1677 
1678 
1679 static
1680 BOOLEAN
ReAllocateLayoutBuffer(_In_ PDISKENTRY DiskEntry)1681 ReAllocateLayoutBuffer(
1682     _In_ PDISKENTRY DiskEntry)
1683 {
1684     PDRIVE_LAYOUT_INFORMATION NewLayoutBuffer;
1685     ULONG NewPartitionCount;
1686     ULONG CurrentPartitionCount = 0;
1687     ULONG LayoutBufferSize;
1688     ULONG i;
1689 
1690     DPRINT1("ReAllocateLayoutBuffer()\n");
1691 
1692     NewPartitionCount = 4 + GetLogicalPartitionCount(DiskEntry) * 4;
1693 
1694     if (DiskEntry->LayoutBuffer)
1695         CurrentPartitionCount = DiskEntry->LayoutBuffer->PartitionCount;
1696 
1697     DPRINT1("CurrentPartitionCount: %lu ; NewPartitionCount: %lu\n",
1698             CurrentPartitionCount, NewPartitionCount);
1699 
1700     if (CurrentPartitionCount == NewPartitionCount)
1701         return TRUE;
1702 
1703     LayoutBufferSize = sizeof(DRIVE_LAYOUT_INFORMATION) +
1704                        ((NewPartitionCount - ANYSIZE_ARRAY) * sizeof(PARTITION_INFORMATION));
1705     NewLayoutBuffer = RtlReAllocateHeap(RtlGetProcessHeap(),
1706                                         HEAP_ZERO_MEMORY,
1707                                         DiskEntry->LayoutBuffer,
1708                                         LayoutBufferSize);
1709     if (NewLayoutBuffer == NULL)
1710     {
1711         DPRINT1("Failed to allocate the new layout buffer (size: %lu)\n", LayoutBufferSize);
1712         return FALSE;
1713     }
1714 
1715     NewLayoutBuffer->PartitionCount = NewPartitionCount;
1716 
1717     /* If the layout buffer grows, make sure the new (empty) entries are written to the disk */
1718     if (NewPartitionCount > CurrentPartitionCount)
1719     {
1720         for (i = CurrentPartitionCount; i < NewPartitionCount; i++)
1721         {
1722             NewLayoutBuffer->PartitionEntry[i].RewritePartition = TRUE;
1723         }
1724     }
1725 
1726     DiskEntry->LayoutBuffer = NewLayoutBuffer;
1727 
1728     return TRUE;
1729 }
1730 
1731 
1732 VOID
UpdateDiskLayout(_In_ PDISKENTRY DiskEntry)1733 UpdateDiskLayout(
1734     _In_ PDISKENTRY DiskEntry)
1735 {
1736     PPARTITION_INFORMATION PartitionInfo;
1737     PPARTITION_INFORMATION LinkInfo = NULL;
1738     PLIST_ENTRY ListEntry;
1739     PPARTENTRY PartEntry;
1740     LARGE_INTEGER HiddenSectors64;
1741     ULONG Index;
1742     ULONG PartitionNumber = 1;
1743 
1744     DPRINT1("UpdateDiskLayout()\n");
1745 
1746     /* Resize the layout buffer if necessary */
1747     if (ReAllocateLayoutBuffer(DiskEntry) == FALSE)
1748     {
1749         DPRINT("ReAllocateLayoutBuffer() failed.\n");
1750         return;
1751     }
1752 
1753     /* Update the primary partition table */
1754     Index = 0;
1755     for (ListEntry = DiskEntry->PrimaryPartListHead.Flink;
1756          ListEntry != &DiskEntry->PrimaryPartListHead;
1757          ListEntry = ListEntry->Flink)
1758     {
1759         PartEntry = CONTAINING_RECORD(ListEntry, PARTENTRY, ListEntry);
1760 
1761         if (PartEntry->IsPartitioned)
1762         {
1763             ASSERT(PartEntry->PartitionType != PARTITION_ENTRY_UNUSED);
1764 
1765             PartitionInfo = &DiskEntry->LayoutBuffer->PartitionEntry[Index];
1766             PartEntry->PartitionIndex = Index;
1767 
1768             /* Reset the current partition number only for newly-created (unmounted) partitions */
1769             if (PartEntry->New)
1770                 PartEntry->PartitionNumber = 0;
1771 
1772             PartEntry->OnDiskPartitionNumber = (!IsContainerPartition(PartEntry->PartitionType) ? PartitionNumber : 0);
1773 
1774             if (!IsSamePrimaryLayoutEntry(PartitionInfo, DiskEntry, PartEntry))
1775             {
1776                 DPRINT1("Updating primary partition entry %lu\n", Index);
1777 
1778                 PartitionInfo->StartingOffset.QuadPart = PartEntry->StartSector.QuadPart * DiskEntry->BytesPerSector;
1779                 PartitionInfo->PartitionLength.QuadPart = PartEntry->SectorCount.QuadPart * DiskEntry->BytesPerSector;
1780                 PartitionInfo->HiddenSectors = PartEntry->StartSector.LowPart;
1781                 PartitionInfo->PartitionNumber = PartEntry->PartitionNumber;
1782                 PartitionInfo->PartitionType = PartEntry->PartitionType;
1783                 PartitionInfo->BootIndicator = PartEntry->BootIndicator;
1784                 PartitionInfo->RecognizedPartition = IsRecognizedPartition(PartEntry->PartitionType);
1785                 PartitionInfo->RewritePartition = TRUE;
1786             }
1787 
1788             if (!IsContainerPartition(PartEntry->PartitionType))
1789                 PartitionNumber++;
1790 
1791             Index++;
1792         }
1793     }
1794 
1795     ASSERT(Index <= 4);
1796 
1797     /* Update the logical partition table */
1798     Index = 4;
1799     for (ListEntry = DiskEntry->LogicalPartListHead.Flink;
1800          ListEntry != &DiskEntry->LogicalPartListHead;
1801          ListEntry = ListEntry->Flink)
1802     {
1803         PartEntry = CONTAINING_RECORD(ListEntry, PARTENTRY, ListEntry);
1804 
1805         if (PartEntry->IsPartitioned)
1806         {
1807             ASSERT(PartEntry->PartitionType != PARTITION_ENTRY_UNUSED);
1808 
1809             PartitionInfo = &DiskEntry->LayoutBuffer->PartitionEntry[Index];
1810             PartEntry->PartitionIndex = Index;
1811 
1812             /* Reset the current partition number only for newly-created (unmounted) partitions */
1813             if (PartEntry->New)
1814                 PartEntry->PartitionNumber = 0;
1815 
1816             PartEntry->OnDiskPartitionNumber = PartitionNumber;
1817 
1818             DPRINT1("Updating logical partition entry %lu\n", Index);
1819 
1820             PartitionInfo->StartingOffset.QuadPart = PartEntry->StartSector.QuadPart * DiskEntry->BytesPerSector;
1821             PartitionInfo->PartitionLength.QuadPart = PartEntry->SectorCount.QuadPart * DiskEntry->BytesPerSector;
1822             PartitionInfo->HiddenSectors = DiskEntry->SectorAlignment;
1823             PartitionInfo->PartitionNumber = PartEntry->PartitionNumber;
1824             PartitionInfo->PartitionType = PartEntry->PartitionType;
1825             PartitionInfo->BootIndicator = FALSE;
1826             PartitionInfo->RecognizedPartition = IsRecognizedPartition(PartEntry->PartitionType);
1827             PartitionInfo->RewritePartition = TRUE;
1828 
1829             /* Fill the link entry of the previous partition entry */
1830             if (LinkInfo != NULL)
1831             {
1832                 LinkInfo->StartingOffset.QuadPart = (PartEntry->StartSector.QuadPart - DiskEntry->SectorAlignment) * DiskEntry->BytesPerSector;
1833                 LinkInfo->PartitionLength.QuadPart = (PartEntry->StartSector.QuadPart + DiskEntry->SectorAlignment) * DiskEntry->BytesPerSector;
1834                 HiddenSectors64.QuadPart = PartEntry->StartSector.QuadPart - DiskEntry->SectorAlignment - DiskEntry->ExtendedPartition->StartSector.QuadPart;
1835                 LinkInfo->HiddenSectors = HiddenSectors64.LowPart;
1836                 LinkInfo->PartitionNumber = 0;
1837                 LinkInfo->PartitionType = PARTITION_EXTENDED;
1838                 LinkInfo->BootIndicator = FALSE;
1839                 LinkInfo->RecognizedPartition = FALSE;
1840                 LinkInfo->RewritePartition = TRUE;
1841             }
1842 
1843             /* Save a pointer to the link entry of the current partition entry */
1844             LinkInfo = &DiskEntry->LayoutBuffer->PartitionEntry[Index + 1];
1845 
1846             PartitionNumber++;
1847             Index += 4;
1848         }
1849     }
1850 
1851     /* Wipe unused primary partition entries */
1852     for (Index = GetPrimaryPartitionCount(DiskEntry); Index < 4; Index++)
1853     {
1854         DPRINT1("Primary partition entry %lu\n", Index);
1855 
1856         PartitionInfo = &DiskEntry->LayoutBuffer->PartitionEntry[Index];
1857 
1858         if (!IsEmptyLayoutEntry(PartitionInfo))
1859         {
1860             DPRINT1("Wiping primary partition entry %lu\n", Index);
1861 
1862             PartitionInfo->StartingOffset.QuadPart = 0;
1863             PartitionInfo->PartitionLength.QuadPart = 0;
1864             PartitionInfo->HiddenSectors = 0;
1865             PartitionInfo->PartitionNumber = 0;
1866             PartitionInfo->PartitionType = PARTITION_ENTRY_UNUSED;
1867             PartitionInfo->BootIndicator = FALSE;
1868             PartitionInfo->RecognizedPartition = FALSE;
1869             PartitionInfo->RewritePartition = TRUE;
1870         }
1871     }
1872 
1873     /* Wipe unused logical partition entries */
1874     for (Index = 4; Index < DiskEntry->LayoutBuffer->PartitionCount; Index++)
1875     {
1876         if (Index % 4 >= 2)
1877         {
1878             DPRINT1("Logical partition entry %lu\n", Index);
1879 
1880             PartitionInfo = &DiskEntry->LayoutBuffer->PartitionEntry[Index];
1881 
1882             if (!IsEmptyLayoutEntry(PartitionInfo))
1883             {
1884                 DPRINT1("Wiping partition entry %lu\n", Index);
1885 
1886                 PartitionInfo->StartingOffset.QuadPart = 0;
1887                 PartitionInfo->PartitionLength.QuadPart = 0;
1888                 PartitionInfo->HiddenSectors = 0;
1889                 PartitionInfo->PartitionNumber = 0;
1890                 PartitionInfo->PartitionType = PARTITION_ENTRY_UNUSED;
1891                 PartitionInfo->BootIndicator = FALSE;
1892                 PartitionInfo->RecognizedPartition = FALSE;
1893                 PartitionInfo->RewritePartition = TRUE;
1894             }
1895         }
1896     }
1897 
1898     DiskEntry->Dirty = TRUE;
1899 }
1900 
1901 
1902 PPARTENTRY
GetPrevUnpartitionedEntry(_In_ PPARTENTRY PartEntry)1903 GetPrevUnpartitionedEntry(
1904     _In_ PPARTENTRY PartEntry)
1905 {
1906     PDISKENTRY DiskEntry = PartEntry->DiskEntry;
1907     PPARTENTRY PrevPartEntry;
1908     PLIST_ENTRY ListHead;
1909 
1910     if (PartEntry->LogicalPartition)
1911         ListHead = &DiskEntry->LogicalPartListHead;
1912     else
1913         ListHead = &DiskEntry->PrimaryPartListHead;
1914 
1915     if (PartEntry->ListEntry.Blink != ListHead)
1916     {
1917         PrevPartEntry = CONTAINING_RECORD(PartEntry->ListEntry.Blink,
1918                                           PARTENTRY,
1919                                           ListEntry);
1920         if (!PrevPartEntry->IsPartitioned)
1921         {
1922             ASSERT(PrevPartEntry->PartitionType == PARTITION_ENTRY_UNUSED);
1923             return PrevPartEntry;
1924         }
1925     }
1926 
1927     return NULL;
1928 }
1929 
1930 
1931 PPARTENTRY
GetNextUnpartitionedEntry(_In_ PPARTENTRY PartEntry)1932 GetNextUnpartitionedEntry(
1933     _In_ PPARTENTRY PartEntry)
1934 {
1935     PDISKENTRY DiskEntry = PartEntry->DiskEntry;
1936     PPARTENTRY NextPartEntry;
1937     PLIST_ENTRY ListHead;
1938 
1939     if (PartEntry->LogicalPartition)
1940         ListHead = &DiskEntry->LogicalPartListHead;
1941     else
1942         ListHead = &DiskEntry->PrimaryPartListHead;
1943 
1944     if (PartEntry->ListEntry.Flink != ListHead)
1945     {
1946         NextPartEntry = CONTAINING_RECORD(PartEntry->ListEntry.Flink,
1947                                           PARTENTRY,
1948                                           ListEntry);
1949         if (!NextPartEntry->IsPartitioned)
1950         {
1951             ASSERT(NextPartEntry->PartitionType == PARTITION_ENTRY_UNUSED);
1952             return NextPartEntry;
1953         }
1954     }
1955 
1956     return NULL;
1957 }
1958 
1959 NTSTATUS
DismountVolume(_In_ PPARTENTRY PartEntry)1960 DismountVolume(
1961     _In_ PPARTENTRY PartEntry)
1962 {
1963     NTSTATUS Status;
1964     NTSTATUS LockStatus;
1965     UNICODE_STRING Name;
1966     OBJECT_ATTRIBUTES ObjectAttributes;
1967     IO_STATUS_BLOCK IoStatusBlock;
1968     HANDLE PartitionHandle;
1969     WCHAR Buffer[MAX_PATH];
1970 
1971     /* Check whether the partition is valid and was mounted by the system */
1972     if (!PartEntry->IsPartitioned ||
1973         IsContainerPartition(PartEntry->PartitionType)   ||
1974         !IsRecognizedPartition(PartEntry->PartitionType) ||
1975         PartEntry->FormatState == UnknownFormat ||
1976         // NOTE: If FormatState == Unformatted but *FileSystem != 0 this means
1977         // it has been usually mounted with RawFS and thus needs to be dismounted.
1978 /*        !*PartEntry->FileSystem || */
1979         PartEntry->PartitionNumber == 0)
1980     {
1981         /* The partition is not mounted, so just return success */
1982         return STATUS_SUCCESS;
1983     }
1984 
1985     ASSERT(PartEntry->PartitionType != PARTITION_ENTRY_UNUSED);
1986 
1987     /* Open the volume */
1988     StringCchPrintfW(Buffer, ARRAYSIZE(Buffer),
1989                      L"\\Device\\Harddisk%lu\\Partition%lu",
1990                      PartEntry->DiskEntry->DiskNumber,
1991                      PartEntry->PartitionNumber);
1992     RtlInitUnicodeString(&Name, Buffer);
1993 
1994     InitializeObjectAttributes(&ObjectAttributes,
1995                                &Name,
1996                                OBJ_CASE_INSENSITIVE,
1997                                NULL,
1998                                NULL);
1999 
2000     Status = NtOpenFile(&PartitionHandle,
2001                         GENERIC_READ | GENERIC_WRITE | SYNCHRONIZE,
2002                         &ObjectAttributes,
2003                         &IoStatusBlock,
2004                         FILE_SHARE_READ | FILE_SHARE_WRITE,
2005                         FILE_SYNCHRONOUS_IO_NONALERT);
2006     if (!NT_SUCCESS(Status))
2007     {
2008         DPRINT1("ERROR: Cannot open volume %wZ for dismounting! (Status 0x%lx)\n", &Name, Status);
2009         return Status;
2010     }
2011 
2012     /* Lock the volume */
2013     LockStatus = NtFsControlFile(PartitionHandle,
2014                                  NULL,
2015                                  NULL,
2016                                  NULL,
2017                                  &IoStatusBlock,
2018                                  FSCTL_LOCK_VOLUME,
2019                                  NULL,
2020                                  0,
2021                                  NULL,
2022                                  0);
2023     if (!NT_SUCCESS(LockStatus))
2024     {
2025         DPRINT1("WARNING: Failed to lock volume! Operations may fail! (Status 0x%lx)\n", LockStatus);
2026     }
2027 
2028     /* Dismount the volume */
2029     Status = NtFsControlFile(PartitionHandle,
2030                              NULL,
2031                              NULL,
2032                              NULL,
2033                              &IoStatusBlock,
2034                              FSCTL_DISMOUNT_VOLUME,
2035                              NULL,
2036                              0,
2037                              NULL,
2038                              0);
2039     if (!NT_SUCCESS(Status))
2040     {
2041         DPRINT1("Failed to unmount volume (Status 0x%lx)\n", Status);
2042     }
2043 
2044     /* Unlock the volume */
2045     LockStatus = NtFsControlFile(PartitionHandle,
2046                                  NULL,
2047                                  NULL,
2048                                  NULL,
2049                                  &IoStatusBlock,
2050                                  FSCTL_UNLOCK_VOLUME,
2051                                  NULL,
2052                                  0,
2053                                  NULL,
2054                                  0);
2055     if (!NT_SUCCESS(LockStatus))
2056     {
2057         DPRINT1("Failed to unlock volume (Status 0x%lx)\n", LockStatus);
2058     }
2059 
2060     /* Close the volume */
2061     NtClose(PartitionHandle);
2062 
2063     return Status;
2064 }
2065 
2066 
2067 PVOLENTRY
GetVolumeFromPartition(_In_ PPARTENTRY PartEntry)2068 GetVolumeFromPartition(
2069     _In_ PPARTENTRY PartEntry)
2070 {
2071     PLIST_ENTRY Entry;
2072     PVOLENTRY VolumeEntry;
2073     ULONG i;
2074 
2075     if ((PartEntry == NULL) ||
2076         (PartEntry->DiskEntry == NULL))
2077         return NULL;
2078 
2079     Entry = VolumeListHead.Flink;
2080     while (Entry != &VolumeListHead)
2081     {
2082         VolumeEntry = CONTAINING_RECORD(Entry, VOLENTRY, ListEntry);
2083 
2084         if (VolumeEntry->pExtents == NULL)
2085             return NULL;
2086 
2087         for (i = 0; i < VolumeEntry->pExtents->NumberOfDiskExtents; i++)
2088         {
2089             if (VolumeEntry->pExtents->Extents[i].DiskNumber == PartEntry->DiskEntry->DiskNumber)
2090             {
2091                 if ((VolumeEntry->pExtents->Extents[i].StartingOffset.QuadPart == PartEntry->StartSector.QuadPart * PartEntry->DiskEntry->BytesPerSector) &&
2092                     (VolumeEntry->pExtents->Extents[i].ExtentLength.QuadPart == PartEntry->SectorCount.QuadPart * PartEntry->DiskEntry->BytesPerSector))
2093                     return VolumeEntry;
2094             }
2095         }
2096 
2097         Entry = Entry->Flink;
2098     }
2099 
2100     return NULL;
2101 }
2102 
2103 
2104 VOID
RemoveVolume(_In_ PVOLENTRY VolumeEntry)2105 RemoveVolume(
2106     _In_ PVOLENTRY VolumeEntry)
2107 {
2108     if (VolumeEntry == NULL)
2109         return;
2110 
2111     if (VolumeEntry == CurrentVolume)
2112         CurrentVolume = NULL;
2113 
2114     RemoveEntryList(&VolumeEntry->ListEntry);
2115 
2116     if (VolumeEntry->pszLabel)
2117         RtlFreeHeap(RtlGetProcessHeap(), 0, VolumeEntry->pszLabel);
2118 
2119     if (VolumeEntry->pszFilesystem)
2120         RtlFreeHeap(RtlGetProcessHeap(), 0, VolumeEntry->pszFilesystem);
2121 
2122     if (VolumeEntry->pExtents)
2123         RtlFreeHeap(RtlGetProcessHeap(), 0, VolumeEntry->pExtents);
2124 
2125     /* Release disk entry */
2126     RtlFreeHeap(RtlGetProcessHeap(), 0, VolumeEntry);
2127 }
2128 
2129 /* EOF */
2130