xref: /reactos/ntoskrnl/fstub/fstubex.c (revision 1f28f715)
1 /*
2 * PROJECT:         ReactOS Kernel
3 * LICENSE:         GPL - See COPYING in the top level directory
4 * FILE:            ntoskrnl/fstub/fstubex.c
5 * PURPOSE:         Extended FSTUB Routines (not linked to HAL)
6 * PROGRAMMERS:     Pierre Schweitzer (pierre.schweitzer@reactos.org)
7 */
8 
9 /* INCLUDES ******************************************************************/
10 
11 #include <ntoskrnl.h>
12 #define NDEBUG
13 #include <debug.h>
14 
15 /* PRIVATE FUNCTIONS *********************************************************/
16 
17 typedef struct _DISK_INFORMATION
18 {
19     PDEVICE_OBJECT DeviceObject;
20     ULONG SectorSize;
21     DISK_GEOMETRY_EX DiskGeometry;
22     PUCHAR Buffer;
23     ULONGLONG SectorCount;
24 } DISK_INFORMATION, *PDISK_INFORMATION;
25 
26 #include <pshpack1.h>
27 typedef struct _EFI_PARTITION_HEADER
28 {
29     ULONGLONG Signature;         // 0
30     ULONG Revision;              // 8
31     ULONG HeaderSize;            // 12
32     ULONG HeaderCRC32;           // 16
33     ULONG Reserved;              // 20
34     ULONGLONG MyLBA;             // 24
35     ULONGLONG AlternateLBA;      // 32
36     ULONGLONG FirstUsableLBA;    // 40
37     ULONGLONG LastUsableLBA;     // 48
38     GUID DiskGUID;               // 56
39     ULONGLONG PartitionEntryLBA; // 72
40     ULONG NumberOfEntries;       // 80
41     ULONG SizeOfPartitionEntry;  // 84
42     ULONG PartitionEntryCRC32;   // 88
43 } EFI_PARTITION_HEADER, *PEFI_PARTITION_HEADER;
44 C_ASSERT(sizeof(EFI_PARTITION_HEADER) == 92);
45 #include <poppack.h>
46 
47 typedef struct _EFI_PARTITION_ENTRY
48 {
49     GUID PartitionType;    // 0
50     GUID UniquePartition;  // 16
51     ULONGLONG StartingLBA; // 32
52     ULONGLONG EndingLBA;   // 40
53     ULONGLONG Attributes;  // 48
54     WCHAR Name[0x24];      // 56
55 } EFI_PARTITION_ENTRY, *PEFI_PARTITION_ENTRY;
56 C_ASSERT(sizeof(EFI_PARTITION_ENTRY) == 128);
57 
58 typedef struct _PARTITION_TABLE_ENTRY
59 {
60     UCHAR BootIndicator;
61     UCHAR StartHead;
62     UCHAR StartSector;
63     UCHAR StartCylinder;
64     UCHAR SystemIndicator;
65     UCHAR EndHead;
66     UCHAR EndSector;
67     UCHAR EndCylinder;
68     ULONG SectorCountBeforePartition;
69     ULONG PartitionSectorCount;
70 } PARTITION_TABLE_ENTRY, *PPARTITION_TABLE_ENTRY;
71 C_ASSERT(sizeof(PARTITION_TABLE_ENTRY) == 16);
72 
73 #include <pshpack1.h>
74 typedef struct _MASTER_BOOT_RECORD
75 {
76     UCHAR MasterBootRecordCodeAndData[0x1B8]; // 0
77     ULONG Signature;                          // 440
78     USHORT Reserved;                          // 444
79     PARTITION_TABLE_ENTRY PartitionTable[4];  // 446
80     USHORT MasterBootRecordMagic;             // 510
81 } MASTER_BOOT_RECORD, *PMASTER_BOOT_RECORD;
82 C_ASSERT(sizeof(MASTER_BOOT_RECORD) == 512);
83 #include <poppack.h>
84 
85 /* Partition entry size (bytes) - FIXME: It's hardcoded as Microsoft does, but according to specs, it shouldn't be */
86 #define PARTITION_ENTRY_SIZE 128
87 /* Defines "EFI PART" */
88 #define EFI_HEADER_SIGNATURE  0x5452415020494645ULL
89 /* Defines version 1.0 */
90 #define EFI_HEADER_REVISION_1 0x00010000
91 /* Defines system type for MBR showing that a GPT is following */
92 #define EFI_PMBR_OSTYPE_EFI 0xEE
93 /* Defines size to store a complete GUID + null char */
94 #define EFI_GUID_STRING_SIZE 0x27
95 
96 #define IS_VALID_DISK_INFO(Disk) \
97   (Disk)               &&        \
98   (Disk->DeviceObject) &&        \
99   (Disk->SectorSize)   &&        \
100   (Disk->Buffer)       &&        \
101   (Disk->SectorCount)
102 
103 VOID
104 NTAPI
105 FstubDbgPrintPartitionEx(IN PPARTITION_INFORMATION_EX PartitionEntry,
106                          IN ULONG PartitionNumber
107 );
108 
109 NTSTATUS
110 NTAPI
111 FstubDetectPartitionStyle(IN PDISK_INFORMATION Disk,
112                           IN PARTITION_STYLE * PartitionStyle
113 );
114 
115 VOID
116 NTAPI
117 FstubFreeDiskInformation(IN PDISK_INFORMATION DiskBuffer
118 );
119 
120 NTSTATUS
121 NTAPI
122 FstubGetDiskGeometry(IN PDEVICE_OBJECT DeviceObject,
123                      OUT PDISK_GEOMETRY_EX Geometry
124 );
125 
126 NTSTATUS
127 NTAPI
128 FstubReadSector(IN PDEVICE_OBJECT DeviceObject,
129                 IN ULONG SectorSize,
130                 IN ULONGLONG StartingSector OPTIONAL,
131                 OUT PVOID Buffer
132 );
133 
134 NTSTATUS
135 NTAPI
136 FstubWriteBootSectorEFI(IN PDISK_INFORMATION Disk
137 );
138 
139 NTSTATUS
140 NTAPI
141 FstubWriteHeaderEFI(IN PDISK_INFORMATION Disk,
142                     IN ULONG PartitionsSizeSector,
143                     IN GUID DiskGUID,
144                     IN ULONG NumberOfEntries,
145                     IN ULONGLONG FirstUsableLBA,
146                     IN ULONGLONG LastUsableLBA,
147                     IN ULONG PartitionEntryCRC32,
148                     IN BOOLEAN WriteBackupTable);
149 
150 NTSTATUS
151 NTAPI
152 FstubWritePartitionTableEFI(IN PDISK_INFORMATION Disk,
153                             IN GUID DiskGUID,
154                             IN ULONG MaxPartitionCount,
155                             IN ULONGLONG FirstUsableLBA,
156                             IN ULONGLONG LastUsableLBA,
157                             IN BOOLEAN WriteBackupTable,
158                             IN ULONG PartitionCount,
159                             IN PPARTITION_INFORMATION_EX PartitionEntries OPTIONAL
160 );
161 
162 NTSTATUS
163 NTAPI
164 FstubWriteSector(IN PDEVICE_OBJECT DeviceObject,
165                  IN ULONG SectorSize,
166                  IN ULONGLONG StartingSector OPTIONAL,
167                  IN PVOID Buffer
168 );
169 
170 VOID
171 NTAPI
FstubAdjustPartitionCount(IN ULONG SectorSize,IN OUT PULONG PartitionCount)172 FstubAdjustPartitionCount(IN ULONG SectorSize,
173                           IN OUT PULONG PartitionCount)
174 {
175     ULONG Count;
176 
177     PAGED_CODE();
178 
179     ASSERT(SectorSize);
180     ASSERT(PartitionCount);
181 
182     /* Get partition count */
183     Count = *PartitionCount;
184     /* We need at least 128 entries */
185     if (Count < 128)
186     {
187         Count = 128;
188     }
189 
190     /* Then, ensure that we will have a round value,
191      * ie, all sectors will be full of entries
192      * There won't be lonely entries
193      */
194     Count = (Count * PARTITION_ENTRY_SIZE) / SectorSize;
195     Count = (Count * SectorSize) / PARTITION_ENTRY_SIZE;
196     ASSERT(*PartitionCount <= Count);
197     /* Return result */
198     *PartitionCount = Count;
199 
200     /* One more sanity check */
201     if (SectorSize == 512)
202     {
203         ASSERT(Count % 4 == 0);
204     }
205 }
206 
207 NTSTATUS
208 NTAPI
FstubAllocateDiskInformation(IN PDEVICE_OBJECT DeviceObject,OUT PDISK_INFORMATION * DiskBuffer,IN PDISK_GEOMETRY_EX DiskGeometry OPTIONAL)209 FstubAllocateDiskInformation(IN PDEVICE_OBJECT DeviceObject,
210                              OUT PDISK_INFORMATION * DiskBuffer,
211                              IN PDISK_GEOMETRY_EX DiskGeometry OPTIONAL)
212 {
213     NTSTATUS Status;
214     PDISK_INFORMATION DiskInformation;
215 
216     PAGED_CODE();
217 
218     ASSERT(DeviceObject);
219     ASSERT(DiskBuffer);
220 
221     /* Allocate internal structure */
222     DiskInformation = ExAllocatePoolWithTag(NonPagedPool, sizeof(DISK_INFORMATION), TAG_FSTUB);
223     if (!DiskInformation)
224     {
225         return STATUS_INSUFFICIENT_RESOURCES;
226     }
227 
228     /* If caller don't pass needed information, let's get them */
229     if (!DiskGeometry)
230     {
231         Status = FstubGetDiskGeometry(DeviceObject, &(DiskInformation->DiskGeometry));
232         if (!NT_SUCCESS(Status))
233         {
234             ExFreePoolWithTag(DiskInformation, TAG_FSTUB);
235             return Status;
236         }
237     }
238     else
239     {
240         RtlCopyMemory(&DiskInformation->DiskGeometry, DiskGeometry, sizeof(DISK_GEOMETRY_EX));
241     }
242 
243     /* Ensure read/received information are correct */
244     if (DiskInformation->DiskGeometry.Geometry.BytesPerSector == 0 ||
245         DiskInformation->DiskGeometry.DiskSize.QuadPart == 0)
246     {
247         ExFreePoolWithTag(DiskInformation, TAG_FSTUB);
248         return STATUS_DEVICE_NOT_READY;
249     }
250 
251     /* Store vital information as well */
252     DiskInformation->DeviceObject = DeviceObject;
253     DiskInformation->SectorSize = DiskInformation->DiskGeometry.Geometry.BytesPerSector;
254     DiskInformation->SectorCount = DiskInformation->DiskGeometry.DiskSize.QuadPart / DiskInformation->SectorSize;
255 
256     /* Finally, allocate the buffer that will be used for different read */
257     DiskInformation->Buffer = ExAllocatePoolWithTag(NonPagedPoolCacheAligned, DiskInformation->SectorSize, TAG_FSTUB);
258     if (!DiskInformation->Buffer)
259     {
260         ExFreePoolWithTag(DiskInformation, TAG_FSTUB);
261         return STATUS_INSUFFICIENT_RESOURCES;
262     }
263 
264     /* Return allocated internal structure */
265     *DiskBuffer = DiskInformation;
266 
267     return STATUS_SUCCESS;
268 }
269 
270 PDRIVE_LAYOUT_INFORMATION
271 NTAPI
FstubConvertExtendedToLayout(IN PDRIVE_LAYOUT_INFORMATION_EX LayoutEx)272 FstubConvertExtendedToLayout(IN PDRIVE_LAYOUT_INFORMATION_EX LayoutEx)
273 {
274     ULONG i;
275     PDRIVE_LAYOUT_INFORMATION DriveLayout;
276 
277     PAGED_CODE();
278 
279     ASSERT(LayoutEx);
280 
281     /* Check whether we're dealing with MBR partition table */
282     if (LayoutEx->PartitionStyle != PARTITION_STYLE_MBR)
283     {
284         ASSERT(FALSE);
285         return NULL;
286     }
287 
288     /* Allocate needed buffer */
289     DriveLayout = ExAllocatePoolWithTag(NonPagedPool,
290                                         FIELD_OFFSET(DRIVE_LAYOUT_INFORMATION, PartitionEntry) +
291                                         LayoutEx->PartitionCount * sizeof(PARTITION_INFORMATION),
292                                         TAG_FSTUB);
293     if (!DriveLayout)
294     {
295         return NULL;
296     }
297 
298     /* Convert information about partition table */
299     DriveLayout->PartitionCount = LayoutEx->PartitionCount;
300     DriveLayout->Signature = LayoutEx->Mbr.Signature;
301 
302     /* Convert each partition */
303     for (i = 0; i < LayoutEx->PartitionCount; i++)
304     {
305         DriveLayout->PartitionEntry[i].StartingOffset = LayoutEx->PartitionEntry[i].StartingOffset;
306         DriveLayout->PartitionEntry[i].PartitionLength = LayoutEx->PartitionEntry[i].PartitionLength;
307         DriveLayout->PartitionEntry[i].HiddenSectors = LayoutEx->PartitionEntry[i].Mbr.HiddenSectors;
308         DriveLayout->PartitionEntry[i].PartitionNumber = LayoutEx->PartitionEntry[i].PartitionNumber;
309         DriveLayout->PartitionEntry[i].PartitionType = LayoutEx->PartitionEntry[i].Mbr.PartitionType;
310         DriveLayout->PartitionEntry[i].BootIndicator = LayoutEx->PartitionEntry[i].Mbr.BootIndicator;
311         DriveLayout->PartitionEntry[i].RecognizedPartition = LayoutEx->PartitionEntry[i].Mbr.RecognizedPartition;
312         DriveLayout->PartitionEntry[i].RewritePartition = LayoutEx->PartitionEntry[i].RewritePartition;
313     }
314 
315     return DriveLayout;
316 }
317 
318 VOID
319 NTAPI
FstubCopyEntryEFI(OUT PEFI_PARTITION_ENTRY Entry,IN PPARTITION_INFORMATION_EX Partition,ULONG SectorSize)320 FstubCopyEntryEFI(OUT PEFI_PARTITION_ENTRY Entry,
321                   IN PPARTITION_INFORMATION_EX Partition,
322                   ULONG SectorSize)
323 {
324     PAGED_CODE();
325 
326     ASSERT(Entry);
327     ASSERT(Partition);
328     ASSERT(SectorSize);
329 
330     /* Just convert data to EFI partition entry type */
331     Entry->PartitionType = Partition->Gpt.PartitionType;
332     Entry->UniquePartition = Partition->Gpt.PartitionId;
333     Entry->StartingLBA = Partition->StartingOffset.QuadPart / SectorSize;
334     Entry->EndingLBA = (Partition->StartingOffset.QuadPart + Partition->PartitionLength.QuadPart - 1) / SectorSize;
335     Entry->Attributes = Partition->Gpt.Attributes;
336     RtlCopyMemory(Entry->Name, Partition->Gpt.Name, sizeof(Entry->Name));
337 }
338 
339 NTSTATUS
340 NTAPI
FstubCreateDiskMBR(IN PDEVICE_OBJECT DeviceObject,IN PCREATE_DISK_MBR DiskInfo)341 FstubCreateDiskMBR(IN PDEVICE_OBJECT DeviceObject,
342                    IN PCREATE_DISK_MBR DiskInfo)
343 {
344     NTSTATUS Status;
345     PDISK_INFORMATION Disk = NULL;
346     PMASTER_BOOT_RECORD MasterBootRecord;
347 
348     PAGED_CODE();
349 
350     ASSERT(DeviceObject);
351 
352     /* Allocate internal structure */
353     Status = FstubAllocateDiskInformation(DeviceObject, &Disk, NULL);
354     if (!NT_SUCCESS(Status))
355     {
356         return Status;
357     }
358 
359     /* Read previous MBR, if any */
360     Status = FstubReadSector(Disk->DeviceObject,
361                              Disk->SectorSize,
362                              0ULL,
363                              Disk->Buffer);
364     if (!NT_SUCCESS(Status))
365     {
366         FstubFreeDiskInformation(Disk);
367         return Status;
368     }
369     /* Fill the buffer with needed information, we won't overwrite boot code */
370     MasterBootRecord = (PMASTER_BOOT_RECORD)Disk->Buffer;
371     MasterBootRecord->Signature = DiskInfo->Signature;
372     RtlZeroMemory(MasterBootRecord->PartitionTable, sizeof(PARTITION_TABLE_ENTRY) * NUM_PARTITION_TABLE_ENTRIES);
373     MasterBootRecord->MasterBootRecordMagic = BOOT_RECORD_SIGNATURE;
374 
375     /* Finally, write MBR */
376     Status = FstubWriteSector(Disk->DeviceObject,
377                               Disk->SectorSize,
378                               0ULL,
379                               Disk->Buffer);
380 
381     /* Release internal structure and return */
382     FstubFreeDiskInformation(Disk);
383     return Status;
384 }
385 
386 NTSTATUS
387 NTAPI
FstubCreateDiskEFI(IN PDEVICE_OBJECT DeviceObject,IN PCREATE_DISK_GPT DiskInfo)388 FstubCreateDiskEFI(IN PDEVICE_OBJECT DeviceObject,
389                    IN PCREATE_DISK_GPT DiskInfo)
390 {
391     NTSTATUS Status;
392     PDISK_INFORMATION Disk = NULL;
393     ULONGLONG FirstUsableLBA, LastUsableLBA;
394     ULONG MaxPartitionCount, SectorsForPartitions;
395 
396     PAGED_CODE();
397 
398     ASSERT(DeviceObject);
399     ASSERT(DiskInfo);
400 
401     /* Allocate internal structure */
402     Status = FstubAllocateDiskInformation(DeviceObject, &Disk, NULL);
403     if (!NT_SUCCESS(Status))
404     {
405         return Status;
406     }
407     ASSERT(Disk);
408 
409     /* Write legacy MBR */
410     Status = FstubWriteBootSectorEFI(Disk);
411     if (!NT_SUCCESS(Status))
412     {
413         FstubFreeDiskInformation(Disk);
414         return Status;
415     }
416 
417     /* Get max entries and adjust its number */
418     MaxPartitionCount = DiskInfo->MaxPartitionCount;
419     FstubAdjustPartitionCount(Disk->SectorSize, &MaxPartitionCount);
420 
421     /* Count number of sectors needed to store partitions */
422     SectorsForPartitions = (MaxPartitionCount * PARTITION_ENTRY_SIZE) / Disk->SectorSize;
423     /* Set first usable LBA: Legacy MBR + GPT header + Partitions entries */
424     FirstUsableLBA = SectorsForPartitions + 2;
425     /* Set last usable LBA: Last sector - GPT header - Partitions entries */
426     LastUsableLBA = Disk->SectorCount - SectorsForPartitions - 1;
427 
428     /* First, write primary table */
429     Status = FstubWritePartitionTableEFI(Disk,
430                                          DiskInfo->DiskId,
431                                          MaxPartitionCount,
432                                          FirstUsableLBA,
433                                          LastUsableLBA,
434                                          FALSE,
435                                          0,
436                                          NULL);
437     /* Then, write backup table */
438     if (NT_SUCCESS(Status))
439     {
440         Status = FstubWritePartitionTableEFI(Disk,
441                                              DiskInfo->DiskId,
442                                              MaxPartitionCount,
443                                              FirstUsableLBA,
444                                              LastUsableLBA,
445                                              TRUE,
446                                              0,
447                                              NULL);
448     }
449 
450     /* Release internal structure and return */
451     FstubFreeDiskInformation(Disk);
452     return Status;
453 }
454 
455 NTSTATUS
456 NTAPI
FstubCreateDiskRaw(IN PDEVICE_OBJECT DeviceObject)457 FstubCreateDiskRaw(IN PDEVICE_OBJECT DeviceObject)
458 {
459     NTSTATUS Status;
460     PDISK_INFORMATION Disk = NULL;
461     PARTITION_STYLE PartitionStyle;
462     PMASTER_BOOT_RECORD MasterBootRecord;
463 
464     PAGED_CODE();
465 
466     ASSERT(DeviceObject);
467 
468     /* Allocate internal structure */
469     Status = FstubAllocateDiskInformation(DeviceObject, &Disk, NULL);
470     if (!NT_SUCCESS(Status))
471     {
472         return Status;
473     }
474 
475     /* Detect current disk partition style */
476     Status = FstubDetectPartitionStyle(Disk, &PartitionStyle);
477     if (!NT_SUCCESS(Status))
478     {
479         FstubFreeDiskInformation(Disk);
480         return Status;
481     }
482 
483     /* Read MBR, if any */
484     Status = FstubReadSector(Disk->DeviceObject,
485                              Disk->SectorSize,
486                              0ULL,
487                              Disk->Buffer);
488     if (!NT_SUCCESS(Status))
489     {
490         FstubFreeDiskInformation(Disk);
491         return Status;
492     }
493 
494     /* Only zero useful stuff */
495     MasterBootRecord = (PMASTER_BOOT_RECORD)Disk->Buffer;
496     MasterBootRecord->Signature = 0;
497     RtlZeroMemory(MasterBootRecord->PartitionTable, sizeof(PARTITION_TABLE_ENTRY));
498     MasterBootRecord->MasterBootRecordMagic = 0;
499 
500     /* Write back that destroyed MBR */
501     Status = FstubWriteSector(Disk->DeviceObject,
502                               Disk->SectorSize,
503                               0ULL,
504                               Disk->Buffer);
505     /* If previous style wasn't GPT, we're done here */
506     if (PartitionStyle != PARTITION_STYLE_GPT)
507     {
508         FstubFreeDiskInformation(Disk);
509         return Status;
510     }
511 
512     /* Otherwise, we've to zero the two GPT headers */
513     RtlZeroMemory(Disk->Buffer, Disk->SectorSize);
514     /* Erase primary header */
515     Status = FstubWriteSector(Disk->DeviceObject,
516                               Disk->SectorSize,
517                               1ULL,
518                               Disk->Buffer);
519     /* In case of success, erase backup header */
520     if (NT_SUCCESS(Status))
521     {
522         Status = FstubWriteSector(Disk->DeviceObject,
523                                   Disk->SectorSize,
524                                   Disk->SectorCount - 1ULL,
525                                   Disk->Buffer);
526     }
527 
528     /* Release internal structure and return */
529     FstubFreeDiskInformation(Disk);
530     return Status;
531 }
532 
533 PCHAR
534 NTAPI
FstubDbgGuidToString(IN PGUID Guid,OUT PCHAR String)535 FstubDbgGuidToString(IN PGUID Guid,
536                      OUT PCHAR String)
537 {
538     sprintf(String,
539             "{%08x-%04x-%04x-%02x%02x-%02x%02x%02x%02x%02x%02x}",
540             Guid->Data1,
541             Guid->Data2,
542             Guid->Data3,
543             Guid->Data4[0],
544             Guid->Data4[1],
545             Guid->Data4[2],
546             Guid->Data4[3],
547             Guid->Data4[4],
548             Guid->Data4[5],
549             Guid->Data4[6],
550             Guid->Data4[7]);
551 
552     return String;
553 }
554 
555 VOID
556 NTAPI
FstubDbgPrintDriveLayoutEx(IN PDRIVE_LAYOUT_INFORMATION_EX DriveLayout)557 FstubDbgPrintDriveLayoutEx(IN PDRIVE_LAYOUT_INFORMATION_EX DriveLayout)
558 {
559     ULONG i;
560     CHAR Guid[EFI_GUID_STRING_SIZE];
561 
562     PAGED_CODE();
563 
564     DPRINT("FSTUB: DRIVE_LAYOUT_INFORMATION_EX: %p\n", DriveLayout);
565     switch (DriveLayout->PartitionStyle)
566     {
567         case PARTITION_STYLE_MBR:
568             if (DriveLayout->PartitionCount % 4 != 0)
569             {
570                 DPRINT("Warning: Partition count isn't a 4-factor: %lu!\n", DriveLayout->PartitionCount);
571             }
572 
573             DPRINT("Signature: %8.8x\n", DriveLayout->Mbr.Signature);
574             for (i = 0; i < DriveLayout->PartitionCount; i++)
575             {
576                 FstubDbgPrintPartitionEx(DriveLayout->PartitionEntry, i);
577             }
578 
579             break;
580         case PARTITION_STYLE_GPT:
581             FstubDbgGuidToString(&(DriveLayout->Gpt.DiskId), Guid);
582             DPRINT("DiskId: %s\n", Guid);
583             DPRINT("StartingUsableOffset: %I64x\n", DriveLayout->Gpt.StartingUsableOffset.QuadPart);
584             DPRINT("UsableLength: %I64x\n", DriveLayout->Gpt.UsableLength.QuadPart);
585             DPRINT("MaxPartitionCount: %lu\n", DriveLayout->Gpt.MaxPartitionCount);
586             for (i = 0; i < DriveLayout->PartitionCount; i++)
587             {
588                 FstubDbgPrintPartitionEx(DriveLayout->PartitionEntry, i);
589             }
590 
591             break;
592         default:
593             DPRINT("Unsupported partition style: %lu\n", DriveLayout->PartitionStyle);
594     }
595 }
596 
597 VOID
598 NTAPI
FstubDbgPrintPartitionEx(IN PPARTITION_INFORMATION_EX PartitionEntry,IN ULONG PartitionNumber)599 FstubDbgPrintPartitionEx(IN PPARTITION_INFORMATION_EX PartitionEntry,
600                          IN ULONG PartitionNumber)
601 {
602     CHAR Guid[EFI_GUID_STRING_SIZE];
603 
604     PAGED_CODE();
605 
606     DPRINT("Printing partition %lu\n", PartitionNumber);
607 
608     switch (PartitionEntry[PartitionNumber].PartitionStyle)
609     {
610         case PARTITION_STYLE_MBR:
611             DPRINT("  StartingOffset: %I64x\n", PartitionEntry[PartitionNumber].StartingOffset.QuadPart);
612             DPRINT("  PartitionLength: %I64x\n", PartitionEntry[PartitionNumber].PartitionLength.QuadPart);
613             DPRINT("  RewritePartition: %u\n", PartitionEntry[PartitionNumber].RewritePartition);
614             DPRINT("  PartitionType: %02x\n", PartitionEntry[PartitionNumber].Mbr.PartitionType);
615             DPRINT("  BootIndicator: %u\n", PartitionEntry[PartitionNumber].Mbr.BootIndicator);
616             DPRINT("  RecognizedPartition: %u\n", PartitionEntry[PartitionNumber].Mbr.RecognizedPartition);
617             DPRINT("  HiddenSectors: %lu\n", PartitionEntry[PartitionNumber].Mbr.HiddenSectors);
618 
619             break;
620         case PARTITION_STYLE_GPT:
621             DPRINT("  StartingOffset: %I64x\n", PartitionEntry[PartitionNumber].StartingOffset.QuadPart);
622             DPRINT("  PartitionLength: %I64x\n", PartitionEntry[PartitionNumber].PartitionLength.QuadPart);
623             DPRINT("  RewritePartition: %u\n", PartitionEntry[PartitionNumber].RewritePartition);
624             FstubDbgGuidToString(&(PartitionEntry[PartitionNumber].Gpt.PartitionType), Guid);
625             DPRINT("  PartitionType: %s\n", Guid);
626             FstubDbgGuidToString(&(PartitionEntry[PartitionNumber].Gpt.PartitionId), Guid);
627             DPRINT("  PartitionId: %s\n", Guid);
628             DPRINT("  Attributes: %I64x\n", PartitionEntry[PartitionNumber].Gpt.Attributes);
629             DPRINT("  Name: %ws\n", PartitionEntry[PartitionNumber].Gpt.Name);
630 
631             break;
632         default:
633             DPRINT("  Unsupported partition style: %ld\n", PartitionEntry[PartitionNumber].PartitionStyle);
634     }
635 }
636 
637 VOID
638 NTAPI
FstubDbgPrintSetPartitionEx(IN PSET_PARTITION_INFORMATION_EX PartitionEntry,IN ULONG PartitionNumber)639 FstubDbgPrintSetPartitionEx(IN PSET_PARTITION_INFORMATION_EX PartitionEntry,
640                             IN ULONG PartitionNumber)
641 {
642     CHAR Guid[EFI_GUID_STRING_SIZE];
643 
644     PAGED_CODE();
645 
646     DPRINT("FSTUB: SET_PARTITION_INFORMATION_EX: %p\n", PartitionEntry);
647     DPRINT("Modifying partition %lu\n", PartitionNumber);
648     switch (PartitionEntry->PartitionStyle)
649     {
650         case PARTITION_STYLE_MBR:
651             DPRINT("  PartitionType: %02x\n", PartitionEntry->Mbr.PartitionType);
652 
653             break;
654         case PARTITION_STYLE_GPT:
655             FstubDbgGuidToString(&(PartitionEntry->Gpt.PartitionType), Guid);
656             DPRINT("  PartitionType: %s\n", Guid);
657             FstubDbgGuidToString(&(PartitionEntry->Gpt.PartitionId), Guid);
658             DPRINT("  PartitionId: %s\n", Guid);
659             DPRINT("  Attributes: %I64x\n", PartitionEntry->Gpt.Attributes);
660             DPRINT("  Name: %ws\n", PartitionEntry->Gpt.Name);
661 
662             break;
663         default:
664             DPRINT("  Unsupported partition style: %ld\n", PartitionEntry[PartitionNumber].PartitionStyle);
665     }
666 }
667 
668 NTSTATUS
669 NTAPI
FstubDetectPartitionStyle(IN PDISK_INFORMATION Disk,IN PARTITION_STYLE * PartitionStyle)670 FstubDetectPartitionStyle(IN PDISK_INFORMATION Disk,
671                           IN PARTITION_STYLE * PartitionStyle)
672 {
673     NTSTATUS Status;
674     PPARTITION_DESCRIPTOR PartitionDescriptor;
675 
676     PAGED_CODE();
677 
678     ASSERT(IS_VALID_DISK_INFO(Disk));
679     ASSERT(PartitionStyle);
680 
681     /* Read disk first sector */
682     Status = FstubReadSector(Disk->DeviceObject,
683                              Disk->SectorSize,
684                              0,
685                              Disk->Buffer);
686     if (!NT_SUCCESS(Status))
687     {
688         return Status;
689     }
690 
691     /* Get the partition descriptor array */
692     PartitionDescriptor = (PPARTITION_DESCRIPTOR)&Disk->Buffer[PARTITION_TABLE_OFFSET];
693     /* If we have not the 0xAA55 then it's raw partition */
694     if (*(PUINT16)&Disk->Buffer[BOOT_SIGNATURE_OFFSET] != BOOT_RECORD_SIGNATURE)
695     {
696         *PartitionStyle = PARTITION_STYLE_RAW;
697     }
698     /* Check partitions types: if first is 0xEE and all the others 0, we have GPT */
699     else if (PartitionDescriptor[0].PartitionType == EFI_PMBR_OSTYPE_EFI &&
700              PartitionDescriptor[1].PartitionType == 0 &&
701              PartitionDescriptor[2].PartitionType == 0 &&
702              PartitionDescriptor[3].PartitionType == 0)
703     {
704         *PartitionStyle = PARTITION_STYLE_GPT;
705     }
706     /* Otherwise, partition table is in MBR */
707     else
708     {
709         *PartitionStyle = PARTITION_STYLE_MBR;
710     }
711 
712     return STATUS_SUCCESS;
713 }
714 
715 VOID
716 NTAPI
FstubFreeDiskInformation(IN PDISK_INFORMATION DiskBuffer)717 FstubFreeDiskInformation(IN PDISK_INFORMATION DiskBuffer)
718 {
719     if (DiskBuffer)
720     {
721         if (DiskBuffer->Buffer)
722         {
723             ExFreePoolWithTag(DiskBuffer->Buffer, TAG_FSTUB);
724         }
725         ExFreePoolWithTag(DiskBuffer, TAG_FSTUB);
726     }
727 }
728 
729 NTSTATUS
730 NTAPI
FstubGetDiskGeometry(IN PDEVICE_OBJECT DeviceObject,OUT PDISK_GEOMETRY_EX Geometry)731 FstubGetDiskGeometry(IN PDEVICE_OBJECT DeviceObject,
732                      OUT PDISK_GEOMETRY_EX Geometry)
733 {
734     NTSTATUS Status;
735     PIRP Irp;
736     PKEVENT Event = NULL;
737     PDISK_GEOMETRY_EX DiskGeometry = NULL;
738     PIO_STATUS_BLOCK IoStatusBlock = NULL;
739 
740     PAGED_CODE();
741 
742     ASSERT(DeviceObject);
743     ASSERT(Geometry);
744 
745     /* Allocate needed components */
746     DiskGeometry = ExAllocatePoolWithTag(NonPagedPool, sizeof(DISK_GEOMETRY_EX), TAG_FSTUB);
747     if (!DiskGeometry)
748     {
749         Status = STATUS_INSUFFICIENT_RESOURCES;
750         goto Cleanup;
751     }
752 
753     IoStatusBlock = ExAllocatePoolWithTag(NonPagedPool, sizeof(IO_STATUS_BLOCK), TAG_FSTUB);
754     if (!IoStatusBlock)
755     {
756         Status = STATUS_INSUFFICIENT_RESOURCES;
757         goto Cleanup;
758     }
759 
760     Event = ExAllocatePoolWithTag(NonPagedPool, sizeof(KEVENT), TAG_FSTUB);
761     if (!Event)
762     {
763         Status = STATUS_INSUFFICIENT_RESOURCES;
764         goto Cleanup;
765     }
766     /* Initialize the waiting event */
767     KeInitializeEvent(Event, NotificationEvent, FALSE);
768 
769     /* Build the request to get disk geometry */
770     Irp = IoBuildDeviceIoControlRequest(IOCTL_DISK_GET_DRIVE_GEOMETRY_EX,
771                                         DeviceObject,
772                                         0,
773                                         0,
774                                         DiskGeometry,
775                                         sizeof(DISK_GEOMETRY_EX),
776                                         FALSE,
777                                         Event,
778                                         IoStatusBlock);
779     if (!Irp)
780     {
781         Status = STATUS_INSUFFICIENT_RESOURCES;
782         goto Cleanup;
783     }
784 
785     /* Call the driver and wait for completion if needed */
786     Status = IoCallDriver(DeviceObject, Irp);
787     if (Status == STATUS_PENDING)
788     {
789         KeWaitForSingleObject(Event, Executive, KernelMode, FALSE, NULL);
790         Status = IoStatusBlock->Status;
791     }
792 
793     /* In case of a success, return read data */
794     if (NT_SUCCESS(Status))
795     {
796         *Geometry = *DiskGeometry;
797     }
798 
799 Cleanup:
800     if (DiskGeometry)
801     {
802         ExFreePoolWithTag(DiskGeometry, TAG_FSTUB);
803 
804         if (NT_SUCCESS(Status))
805         {
806             ASSERT(Geometry->Geometry.BytesPerSector % PARTITION_ENTRY_SIZE == 0);
807         }
808     }
809 
810     if (IoStatusBlock)
811     {
812         ExFreePoolWithTag(IoStatusBlock, TAG_FSTUB);
813     }
814 
815     if (Event)
816     {
817         ExFreePoolWithTag(Event, TAG_FSTUB);
818     }
819 
820     return Status;
821 }
822 
823 NTSTATUS
824 NTAPI
FstubReadHeaderEFI(IN PDISK_INFORMATION Disk,IN BOOLEAN ReadBackupTable,PEFI_PARTITION_HEADER * HeaderBuffer)825 FstubReadHeaderEFI(IN PDISK_INFORMATION Disk,
826                    IN BOOLEAN ReadBackupTable,
827                    PEFI_PARTITION_HEADER * HeaderBuffer)
828 {
829     NTSTATUS Status;
830     PUCHAR Sector = NULL;
831     ULONGLONG StartingSector;
832     PEFI_PARTITION_HEADER EFIHeader;
833     ULONG i, HeaderCRC32, PreviousCRC32, SectoredPartitionEntriesSize, LonelyPartitions;
834 
835     PAGED_CODE();
836 
837     ASSERT(Disk);
838     ASSERT(IS_VALID_DISK_INFO(Disk));
839     ASSERT(HeaderBuffer);
840 
841     /* In case we want to read backup table, we read last disk sector */
842     if (ReadBackupTable)
843     {
844         StartingSector = Disk->SectorCount - 1ULL;
845     }
846     else
847     {
848         /* Otherwise we start at first sector (as sector 0 is the MBR) */
849         StartingSector = 1ULL;
850     }
851 
852     Status = FstubReadSector(Disk->DeviceObject,
853                              Disk->SectorSize,
854                              StartingSector,
855                              Disk->Buffer);
856     if (!NT_SUCCESS(Status))
857     {
858         DPRINT("EFI::Failed reading header!\n");
859         return Status;
860     }
861     /* Let's use read buffer as EFI_PARTITION_HEADER */
862     EFIHeader = (PEFI_PARTITION_HEADER)Disk->Buffer;
863 
864 
865     /* First check signature
866      * Then, check version (we only support v1)
867      * Finally check header size
868      */
869     if (EFIHeader->Signature != EFI_HEADER_SIGNATURE ||
870         EFIHeader->Revision != EFI_HEADER_REVISION_1 ||
871         EFIHeader->HeaderSize != sizeof(EFI_PARTITION_HEADER))
872     {
873         DPRINT("EFI::Wrong signature/version/header size!\n");
874         DPRINT("%I64x (expected: %I64x)\n", EFIHeader->Signature, EFI_HEADER_SIGNATURE);
875         DPRINT("%03x (expected: %03x)\n", EFIHeader->Revision, EFI_HEADER_REVISION_1);
876         DPRINT("%02x (expected: %02x)\n", EFIHeader->HeaderSize, sizeof(EFI_PARTITION_HEADER));
877         return STATUS_DISK_CORRUPT_ERROR;
878     }
879 
880     /* Save current checksum */
881     HeaderCRC32 = EFIHeader->HeaderCRC32;
882     /* Then zero the one in EFI header. This is needed to compute header checksum */
883     EFIHeader->HeaderCRC32 = 0;
884     /* Compute header checksum and compare with the one present in partition table */
885     if (RtlComputeCrc32(0, Disk->Buffer, sizeof(EFI_PARTITION_HEADER)) != HeaderCRC32)
886     {
887         DPRINT("EFI::Not matching header checksum!\n");
888         return STATUS_DISK_CORRUPT_ERROR;
889     }
890     /* Put back removed checksum in header */
891     EFIHeader->HeaderCRC32 = HeaderCRC32;
892 
893     /* Check if current LBA is matching with ours */
894     if (EFIHeader->MyLBA != StartingSector)
895     {
896         DPRINT("EFI::Not matching starting sector!\n");
897         return STATUS_DISK_CORRUPT_ERROR;
898     }
899 
900     /* Allocate a buffer to read a sector on the disk */
901     Sector = ExAllocatePoolWithTag(NonPagedPool,
902                                    Disk->SectorSize,
903                                    TAG_FSTUB);
904     if (!Sector)
905     {
906         DPRINT("EFI::Lacking resources!\n");
907         return STATUS_INSUFFICIENT_RESOURCES;
908     }
909 
910     /* Count how much sectors we'll have to read to read the whole partition table */
911     SectoredPartitionEntriesSize = (EFIHeader->NumberOfEntries * PARTITION_ENTRY_SIZE) / Disk->SectorSize;
912     /* Compute partition table checksum */
913     for (i = 0, PreviousCRC32 = 0; i < SectoredPartitionEntriesSize; i++)
914     {
915         Status = FstubReadSector(Disk->DeviceObject,
916                                  Disk->SectorSize,
917                                  EFIHeader->PartitionEntryLBA + i,
918                                  Sector);
919         if (!NT_SUCCESS(Status))
920         {
921             ExFreePoolWithTag(Sector, TAG_FSTUB);
922             DPRINT("EFI::Failed reading sector for partition entry!\n");
923             return Status;
924         }
925 
926         PreviousCRC32 = RtlComputeCrc32(PreviousCRC32, Sector, Disk->SectorSize);
927     }
928 
929     /* Check whether we have a last sector not full of partitions */
930     LonelyPartitions = (EFIHeader->NumberOfEntries * PARTITION_ENTRY_SIZE) % Disk->SectorSize;
931     /* In such case, we have to complete checksum computation */
932     if (LonelyPartitions != 0)
933     {
934         /* Read the sector that contains those partitions */
935         Status = FstubReadSector(Disk->DeviceObject,
936                                  Disk->SectorSize,
937                                  EFIHeader->PartitionEntryLBA + i,
938                                  Sector);
939         if (!NT_SUCCESS(Status))
940         {
941             ExFreePoolWithTag(Sector, TAG_FSTUB);
942             DPRINT("EFI::Failed reading sector for partition entry!\n");
943             return Status;
944         }
945 
946         /* Then complete checksum by computing on each partition */
947         for (i = 0; i < LonelyPartitions; i++)
948         {
949             PreviousCRC32 = RtlComputeCrc32(PreviousCRC32, Sector + i * PARTITION_ENTRY_SIZE, PARTITION_ENTRY_SIZE);
950         }
951     }
952 
953     /* Finally, release memory */
954     ExFreePoolWithTag(Sector, TAG_FSTUB);
955 
956     /* Compare checksums */
957     if (PreviousCRC32 == EFIHeader->PartitionEntryCRC32)
958     {
959         /* In case of a success, return read header */
960         *HeaderBuffer = EFIHeader;
961         return STATUS_SUCCESS;
962     }
963     else
964     {
965         DPRINT("EFI::Not matching partition table checksum!\n");
966         DPRINT("EFI::Expected: %x, received: %x\n", EFIHeader->PartitionEntryCRC32, PreviousCRC32);
967         return STATUS_DISK_CORRUPT_ERROR;
968     }
969 }
970 
971 NTSTATUS
972 NTAPI
FstubReadPartitionTableEFI(IN PDISK_INFORMATION Disk,IN BOOLEAN ReadBackupTable,OUT PDRIVE_LAYOUT_INFORMATION_EX * DriveLayout)973 FstubReadPartitionTableEFI(IN PDISK_INFORMATION Disk,
974                            IN BOOLEAN ReadBackupTable,
975                            OUT PDRIVE_LAYOUT_INFORMATION_EX* DriveLayout)
976 {
977     NTSTATUS Status;
978     ULONG NumberOfEntries;
979     PEFI_PARTITION_HEADER EfiHeader;
980     EFI_PARTITION_ENTRY PartitionEntry;
981 #if 0
982     BOOLEAN UpdatedPartitionTable = FALSE;
983     ULONGLONG SectorsForPartitions, PartitionEntryLBA;
984 #else
985     ULONGLONG PartitionEntryLBA;
986 #endif
987     PDRIVE_LAYOUT_INFORMATION_EX DriveLayoutEx = NULL;
988     ULONG i, PartitionCount, PartitionIndex, PartitionsPerSector;
989 
990     PAGED_CODE();
991 
992     ASSERT(Disk);
993 
994     /* Zero output */
995     *DriveLayout = NULL;
996 
997     /* Read EFI header */
998     Status = FstubReadHeaderEFI(Disk,
999                                 ReadBackupTable,
1000                                 &EfiHeader);
1001     if (!NT_SUCCESS(Status))
1002     {
1003         return Status;
1004     }
1005 
1006     /* Backup the number of entries, will be used later on */
1007     NumberOfEntries = EfiHeader->NumberOfEntries;
1008 
1009     /* Allocate a DRIVE_LAYOUT_INFORMATION_EX struct big enough */
1010     DriveLayoutEx = ExAllocatePoolWithTag(NonPagedPool,
1011                                           FIELD_OFFSET(DRIVE_LAYOUT_INFORMATION_EX, PartitionEntry) +
1012                                           EfiHeader->NumberOfEntries * sizeof(PARTITION_INFORMATION_EX),
1013                                           TAG_FSTUB);
1014     if (!DriveLayoutEx)
1015     {
1016         return STATUS_INSUFFICIENT_RESOURCES;
1017     }
1018 
1019 #if 0
1020     if (!ReadBackupTable)
1021     {
1022         /* If we weren't ask to read backup table,
1023          * check the status of the backup table.
1024          * In case it's not where we're expecting it, move it and ask
1025          * for a partition table rewrite.
1026          */
1027         if ((Disk->SectorCount - 1ULL) != EfiHeader->AlternateLBA)
1028         {
1029             /* We'll update it. First, count number of sectors needed to store partitions */
1030             SectorsForPartitions = ((ULONGLONG)EfiHeader->NumberOfEntries * PARTITION_ENTRY_SIZE) / Disk->SectorSize;
1031             /* Then set first usable LBA: Legacy MBR + GPT header + Partitions entries */
1032             EfiHeader->FirstUsableLBA = SectorsForPartitions + 2;
1033             /* Then set last usable LBA: Last sector - GPT header - Partitions entries */
1034             EfiHeader->LastUsableLBA = Disk->SectorCount - SectorsForPartitions - 1;
1035             /* Inform that we'll rewrite partition table */
1036             UpdatedPartitionTable = TRUE;
1037         }
1038     }
1039 #endif
1040 
1041     DriveLayoutEx->PartitionStyle = PARTITION_STYLE_GPT;
1042     /* Translate LBA -> Offset */
1043     DriveLayoutEx->Gpt.StartingUsableOffset.QuadPart = EfiHeader->FirstUsableLBA * Disk->SectorSize;
1044     DriveLayoutEx->Gpt.UsableLength.QuadPart = EfiHeader->LastUsableLBA - EfiHeader->FirstUsableLBA * Disk->SectorSize;
1045     DriveLayoutEx->Gpt.MaxPartitionCount = EfiHeader->NumberOfEntries;
1046     DriveLayoutEx->Gpt.DiskId = EfiHeader->DiskGUID;
1047 
1048     /* Backup partition entry position */
1049     PartitionEntryLBA = EfiHeader->PartitionEntryLBA;
1050     /* Count number of partitions per sector */
1051     PartitionsPerSector = (Disk->SectorSize / PARTITION_ENTRY_SIZE);
1052     /* Read all partitions and fill in structure
1053      * BEWARE! Past that point EfiHeader IS NOT VALID ANYMORE
1054      * It will be erased by the reading of the partition entry
1055      */
1056     for (i = 0, PartitionCount = 0, PartitionIndex = PartitionsPerSector;
1057          i < NumberOfEntries;
1058          i++)
1059     {
1060         /* Only read following sector if we finished with previous sector */
1061         if (PartitionIndex == PartitionsPerSector)
1062         {
1063             Status = FstubReadSector(Disk->DeviceObject,
1064                                      Disk->SectorSize,
1065                                      PartitionEntryLBA + (i / PartitionsPerSector),
1066                                      Disk->Buffer);
1067             if (!NT_SUCCESS(Status))
1068             {
1069                 ExFreePoolWithTag(DriveLayoutEx, TAG_FSTUB);
1070                 return Status;
1071             }
1072 
1073             PartitionIndex = 0;
1074         }
1075         /* Read following partition */
1076         PartitionEntry = ((PEFI_PARTITION_ENTRY)Disk->Buffer)[PartitionIndex];
1077         PartitionIndex++;
1078 
1079         /* If partition GUID is 00000000-0000-0000-0000-000000000000, then it's unused, skip it */
1080         if (PartitionEntry.PartitionType.Data1 == 0 &&
1081             PartitionEntry.PartitionType.Data2 == 0 &&
1082             PartitionEntry.PartitionType.Data3 == 0 &&
1083             ((PULONGLONG)PartitionEntry.PartitionType.Data4)[0] == 0)
1084         {
1085             continue;
1086         }
1087 
1088         /* Write data to structure. Don't forget GPT is using sectors, Windows offsets */
1089         DriveLayoutEx->PartitionEntry[PartitionCount].StartingOffset.QuadPart = PartitionEntry.StartingLBA * Disk->SectorSize;
1090         DriveLayoutEx->PartitionEntry[PartitionCount].PartitionLength.QuadPart = (PartitionEntry.EndingLBA -
1091                                                                                   PartitionEntry.StartingLBA + 1) *
1092                                                                                  Disk->SectorSize;
1093         /* This number starts from 1 */
1094         DriveLayoutEx->PartitionEntry[PartitionCount].PartitionNumber = PartitionCount + 1;
1095         DriveLayoutEx->PartitionEntry[PartitionCount].RewritePartition = FALSE;
1096         DriveLayoutEx->PartitionEntry[PartitionCount].PartitionStyle = PARTITION_STYLE_GPT;
1097         DriveLayoutEx->PartitionEntry[PartitionCount].Gpt.PartitionType = PartitionEntry.PartitionType;
1098         DriveLayoutEx->PartitionEntry[PartitionCount].Gpt.PartitionId = PartitionEntry.UniquePartition;
1099         DriveLayoutEx->PartitionEntry[PartitionCount].Gpt.Attributes = PartitionEntry.Attributes;
1100         RtlCopyMemory(DriveLayoutEx->PartitionEntry[PartitionCount].Gpt.Name,
1101                       PartitionEntry.Name, sizeof(PartitionEntry.Name));
1102 
1103         /* Update partition count */
1104         PartitionCount++;
1105     }
1106     DriveLayoutEx->PartitionCount = PartitionCount;
1107 
1108 #if 0
1109     /* If we updated partition table using backup table, rewrite partition table */
1110     if (UpdatedPartitionTable)
1111     {
1112         IoWritePartitionTableEx(Disk->DeviceObject,
1113                                 DriveLayoutEx);
1114     }
1115 #endif
1116 
1117     /* Finally, return read data */
1118     *DriveLayout = DriveLayoutEx;
1119 
1120     return Status;
1121 }
1122 
1123 NTSTATUS
1124 NTAPI
FstubReadPartitionTableMBR(IN PDISK_INFORMATION Disk,IN BOOLEAN ReturnRecognizedPartitions,OUT PDRIVE_LAYOUT_INFORMATION_EX * ReturnedDriveLayout)1125 FstubReadPartitionTableMBR(IN PDISK_INFORMATION Disk,
1126                            IN BOOLEAN ReturnRecognizedPartitions,
1127                            OUT PDRIVE_LAYOUT_INFORMATION_EX* ReturnedDriveLayout)
1128 {
1129     NTSTATUS Status;
1130     ULONG i;
1131     PDRIVE_LAYOUT_INFORMATION DriveLayout = NULL;
1132     PDRIVE_LAYOUT_INFORMATION_EX DriveLayoutEx = NULL;
1133 
1134     PAGED_CODE();
1135 
1136     ASSERT(IS_VALID_DISK_INFO(Disk));
1137     ASSERT(ReturnedDriveLayout);
1138 
1139     /* Zero output */
1140     *ReturnedDriveLayout = NULL;
1141 
1142     /* Read partition table the old way */
1143     Status = IoReadPartitionTable(Disk->DeviceObject,
1144                                   Disk->SectorSize,
1145                                   ReturnRecognizedPartitions,
1146                                   &DriveLayout);
1147     if (!NT_SUCCESS(Status))
1148     {
1149         return Status;
1150     }
1151 
1152     /* Allocate a DRIVE_LAYOUT_INFORMATION_EX struct big enough */
1153     DriveLayoutEx = ExAllocatePoolWithTag(NonPagedPool,
1154                                           FIELD_OFFSET(DRIVE_LAYOUT_INFORMATION_EX, PartitionEntry) +
1155                                           DriveLayout->PartitionCount * sizeof(PARTITION_INFORMATION_EX),
1156                                           TAG_FSTUB);
1157     if (!DriveLayoutEx)
1158     {
1159         /* Let's not leak memory as in Windows 2003 */
1160         ExFreePool(DriveLayout);
1161         return STATUS_INSUFFICIENT_RESOURCES;
1162     }
1163 
1164     /* Start converting the DRIVE_LAYOUT_INFORMATION structure */
1165     DriveLayoutEx->PartitionStyle = PARTITION_STYLE_MBR;
1166     DriveLayoutEx->PartitionCount = DriveLayout->PartitionCount;
1167     DriveLayoutEx->Mbr.Signature = DriveLayout->Signature;
1168 
1169     /* Convert each found partition */
1170     for (i = 0; i < DriveLayout->PartitionCount; i++)
1171     {
1172         DriveLayoutEx->PartitionEntry[i].PartitionStyle = PARTITION_STYLE_MBR;
1173         DriveLayoutEx->PartitionEntry[i].StartingOffset = DriveLayout->PartitionEntry[i].StartingOffset;
1174         DriveLayoutEx->PartitionEntry[i].PartitionLength = DriveLayout->PartitionEntry[i].PartitionLength;
1175         DriveLayoutEx->PartitionEntry[i].PartitionNumber = DriveLayout->PartitionEntry[i].PartitionNumber;
1176         DriveLayoutEx->PartitionEntry[i].RewritePartition = DriveLayout->PartitionEntry[i].RewritePartition;
1177         DriveLayoutEx->PartitionEntry[i].Mbr.PartitionType = DriveLayout->PartitionEntry[i].PartitionType;
1178         DriveLayoutEx->PartitionEntry[i].Mbr.BootIndicator = DriveLayout->PartitionEntry[i].BootIndicator;
1179         DriveLayoutEx->PartitionEntry[i].Mbr.RecognizedPartition = DriveLayout->PartitionEntry[i].RecognizedPartition;
1180         DriveLayoutEx->PartitionEntry[i].Mbr.HiddenSectors = DriveLayout->PartitionEntry[i].HiddenSectors;
1181     }
1182 
1183     /* Finally, return data and free old structure */
1184     *ReturnedDriveLayout = DriveLayoutEx;
1185     ExFreePool(DriveLayout);
1186 
1187     return STATUS_SUCCESS;
1188 }
1189 
1190 NTSTATUS
1191 NTAPI
FstubReadSector(IN PDEVICE_OBJECT DeviceObject,IN ULONG SectorSize,IN ULONGLONG StartingSector OPTIONAL,OUT PVOID Buffer)1192 FstubReadSector(IN PDEVICE_OBJECT DeviceObject,
1193                 IN ULONG SectorSize,
1194                 IN ULONGLONG StartingSector OPTIONAL,
1195                 OUT PVOID Buffer)
1196 {
1197     NTSTATUS Status;
1198     PIRP Irp;
1199     KEVENT Event;
1200     LARGE_INTEGER StartingOffset;
1201     IO_STATUS_BLOCK IoStatusBlock;
1202     PIO_STACK_LOCATION IoStackLocation;
1203 
1204     PAGED_CODE();
1205 
1206     ASSERT(DeviceObject);
1207     ASSERT(Buffer);
1208     ASSERT(SectorSize);
1209 
1210     /* Compute starting offset */
1211     StartingOffset.QuadPart = StartingSector * SectorSize;
1212 
1213     /* Initialize waiting event */
1214     KeInitializeEvent(&Event, NotificationEvent, FALSE);
1215 
1216     /* Prepare IRP */
1217     Irp = IoBuildSynchronousFsdRequest(IRP_MJ_READ,
1218                                        DeviceObject,
1219                                        Buffer,
1220                                        SectorSize,
1221                                        &StartingOffset,
1222                                        &Event,
1223                                        &IoStatusBlock);
1224     if (!Irp)
1225     {
1226         return STATUS_INSUFFICIENT_RESOURCES;
1227     }
1228 
1229     /* Override volume verify */
1230     IoStackLocation = IoGetNextIrpStackLocation(Irp);
1231     IoStackLocation->Flags |= SL_OVERRIDE_VERIFY_VOLUME;
1232 
1233     /* Then call driver, and wait for completion if needed */
1234     Status = IoCallDriver(DeviceObject, Irp);
1235     if (Status == STATUS_PENDING)
1236     {
1237         KeWaitForSingleObject(&Event, Executive, KernelMode, FALSE, NULL);
1238         Status = IoStatusBlock.Status;
1239     }
1240 
1241     return Status;
1242 }
1243 
1244 NTSTATUS
1245 NTAPI
FstubSetPartitionInformationEFI(IN PDISK_INFORMATION Disk,IN ULONG PartitionNumber,IN SET_PARTITION_INFORMATION_GPT * PartitionInfo)1246 FstubSetPartitionInformationEFI(IN PDISK_INFORMATION Disk,
1247                                 IN ULONG PartitionNumber,
1248                                 IN SET_PARTITION_INFORMATION_GPT * PartitionInfo)
1249 {
1250     NTSTATUS Status;
1251     PDRIVE_LAYOUT_INFORMATION_EX Layout = NULL;
1252 
1253     PAGED_CODE();
1254 
1255     ASSERT(Disk);
1256     ASSERT(PartitionInfo);
1257 
1258     /* Partition 0 isn't correct (should start at 1) */
1259     if (PartitionNumber == 0)
1260     {
1261         return STATUS_INVALID_PARAMETER;
1262     }
1263 
1264     /* Read partition table */
1265     Status = IoReadPartitionTableEx(Disk->DeviceObject, &Layout);
1266     if (!NT_SUCCESS(Status))
1267     {
1268         return Status;
1269     }
1270     ASSERT(Layout);
1271 
1272     /* If our partition (started at 0 now) is higher than partition count, then, there's an issue */
1273     if (Layout->PartitionCount <= --PartitionNumber)
1274     {
1275         ExFreePool(Layout);
1276         return STATUS_INVALID_PARAMETER;
1277     }
1278 
1279     /* Erase actual partition entry data with provided ones */
1280     Layout->PartitionEntry[PartitionNumber].Gpt.PartitionType = PartitionInfo->PartitionType;
1281     Layout->PartitionEntry[PartitionNumber].Gpt.PartitionId = PartitionInfo->PartitionId;
1282     Layout->PartitionEntry[PartitionNumber].Gpt.Attributes = PartitionInfo->Attributes;
1283     RtlCopyMemory(Layout->PartitionEntry[PartitionNumber].Gpt.Name, PartitionInfo->Name, sizeof(PartitionInfo->Name));
1284 
1285     /* Rewrite the whole partition table to update the modified entry */
1286     Status = IoWritePartitionTableEx(Disk->DeviceObject, Layout);
1287 
1288     /* Free partition table and return */
1289     ExFreePool(Layout);
1290     return Status;
1291 }
1292 
1293 NTSTATUS
1294 NTAPI
FstubVerifyPartitionTableEFI(IN PDISK_INFORMATION Disk,IN BOOLEAN FixErrors)1295 FstubVerifyPartitionTableEFI(IN PDISK_INFORMATION Disk,
1296                              IN BOOLEAN FixErrors)
1297 {
1298     NTSTATUS Status;
1299     PEFI_PARTITION_HEADER EFIHeader, ReadEFIHeader;
1300     BOOLEAN PrimaryValid = FALSE, BackupValid = FALSE, WriteBackup;
1301     ULONGLONG ReadPosition, WritePosition, SectorsForPartitions, PartitionIndex;
1302 
1303     PAGED_CODE();
1304 
1305     EFIHeader = ExAllocatePoolWithTag(NonPagedPool, sizeof(EFI_PARTITION_HEADER), TAG_FSTUB);
1306     if (!EFIHeader)
1307     {
1308         return STATUS_INSUFFICIENT_RESOURCES;
1309     }
1310 
1311     Status = FstubReadHeaderEFI(Disk, FALSE, &ReadEFIHeader);
1312     if (NT_SUCCESS(Status))
1313     {
1314         PrimaryValid = TRUE;
1315         ASSERT(ReadEFIHeader);
1316         RtlCopyMemory(EFIHeader, ReadEFIHeader, sizeof(EFI_PARTITION_HEADER));
1317     }
1318 
1319     Status = FstubReadHeaderEFI(Disk, TRUE, &ReadEFIHeader);
1320     if (NT_SUCCESS(Status))
1321     {
1322         BackupValid = TRUE;
1323         ASSERT(ReadEFIHeader);
1324         RtlCopyMemory(EFIHeader, ReadEFIHeader, sizeof(EFI_PARTITION_HEADER));
1325     }
1326 
1327     /* If both are sane, just return */
1328     if (PrimaryValid && BackupValid)
1329     {
1330         ExFreePoolWithTag(EFIHeader, TAG_FSTUB);
1331         return STATUS_SUCCESS;
1332     }
1333 
1334     /* If both are damaged OR if we have not been ordered to fix
1335      * Then, quit and warn about disk corruption
1336      */
1337     if ((!PrimaryValid && !BackupValid) || !FixErrors)
1338     {
1339         ExFreePoolWithTag(EFIHeader, TAG_FSTUB);
1340         return STATUS_DISK_CORRUPT_ERROR;
1341     }
1342 
1343     /* Compute sectors taken by partitions */
1344     SectorsForPartitions = (((ULONGLONG)EFIHeader->NumberOfEntries * PARTITION_ENTRY_SIZE) + Disk->SectorSize - 1) / Disk->SectorSize;
1345     if (PrimaryValid)
1346     {
1347         WriteBackup = TRUE;
1348         /* Take position at backup table for writing */
1349         WritePosition = Disk->SectorCount - SectorsForPartitions;
1350         /* And read from primary table */
1351         ReadPosition = 2ULL;
1352 
1353         DPRINT("EFI::Will repair backup table from primary\n");
1354     }
1355     else
1356     {
1357         ASSERT(BackupValid);
1358         WriteBackup = FALSE;
1359         /* Take position at primary table for writing */
1360         WritePosition = 2ULL;
1361         /* And read from backup table */
1362         ReadPosition = Disk->SectorCount - SectorsForPartitions;
1363 
1364         DPRINT("EFI::Will repair primary table from backup\n");
1365     }
1366 
1367     PartitionIndex = 0ULL;
1368 
1369     /* If no partitions are to be copied, just restore header */
1370     if (SectorsForPartitions <= 0)
1371     {
1372         Status = FstubWriteHeaderEFI(Disk,
1373                                      SectorsForPartitions,
1374                                      EFIHeader->DiskGUID,
1375                                      EFIHeader->NumberOfEntries,
1376                                      EFIHeader->FirstUsableLBA,
1377                                      EFIHeader->LastUsableLBA,
1378                                      EFIHeader->PartitionEntryCRC32,
1379                                      WriteBackup);
1380 
1381         goto Cleanup;
1382     }
1383 
1384     /* Copy all the partitions */
1385     for (; PartitionIndex < SectorsForPartitions; ++PartitionIndex)
1386     {
1387         /* First, read the partition from the first table */
1388         Status = FstubReadSector(Disk->DeviceObject,
1389                                  Disk->SectorSize,
1390                                  ReadPosition + PartitionIndex,
1391                                  Disk->Buffer);
1392         if (!NT_SUCCESS(Status))
1393         {
1394             goto Cleanup;
1395         }
1396 
1397         /* Then, write it in the other table */
1398         Status = FstubWriteSector(Disk->DeviceObject,
1399                                   Disk->SectorSize,
1400                                   WritePosition + PartitionIndex,
1401                                   Disk->Buffer);
1402         if (!NT_SUCCESS(Status))
1403         {
1404             goto Cleanup;
1405         }
1406     }
1407 
1408     /* Now we're done, write the header */
1409     Status = FstubWriteHeaderEFI(Disk,
1410                                  SectorsForPartitions,
1411                                  EFIHeader->DiskGUID,
1412                                  EFIHeader->NumberOfEntries,
1413                                  EFIHeader->FirstUsableLBA,
1414                                  EFIHeader->LastUsableLBA,
1415                                  EFIHeader->PartitionEntryCRC32,
1416                                  WriteBackup);
1417 
1418 Cleanup:
1419     ExFreePoolWithTag(EFIHeader, TAG_FSTUB);
1420     return Status;
1421 }
1422 
1423 NTSTATUS
1424 NTAPI
FstubWriteBootSectorEFI(IN PDISK_INFORMATION Disk)1425 FstubWriteBootSectorEFI(IN PDISK_INFORMATION Disk)
1426 {
1427     NTSTATUS Status;
1428     ULONG Signature = 0;
1429     PMASTER_BOOT_RECORD MasterBootRecord;
1430 
1431     PAGED_CODE();
1432 
1433     ASSERT(Disk);
1434     ASSERT(IS_VALID_DISK_INFO(Disk));
1435 
1436     /* Read if a MBR is already present */
1437     Status = FstubReadSector(Disk->DeviceObject,
1438                              Disk->SectorSize,
1439                              0ULL,
1440                              Disk->Buffer);
1441     MasterBootRecord = (PMASTER_BOOT_RECORD)Disk->Buffer;
1442     /* If one has been found */
1443     if (NT_SUCCESS(Status) && MasterBootRecord->MasterBootRecordMagic == BOOT_RECORD_SIGNATURE)
1444     {
1445         /* Save its signature */
1446         Signature = MasterBootRecord->Signature;
1447     }
1448 
1449     /* Reset the MBR */
1450     RtlZeroMemory(MasterBootRecord, Disk->SectorSize);
1451     /* Then create a fake MBR matching those purposes:
1452      * It must have only partition. Type of this partition
1453      * has to be 0xEE to signal a GPT is following.
1454      * This partition has to cover the whole disk. To prevent
1455      * any disk modification by a program that wouldn't
1456      * understand anything to GPT.
1457      */
1458     MasterBootRecord->Signature = Signature;
1459     MasterBootRecord->PartitionTable[0].StartSector = 2;
1460     MasterBootRecord->PartitionTable[0].SystemIndicator = EFI_PMBR_OSTYPE_EFI;
1461     MasterBootRecord->PartitionTable[0].EndHead = 0xFF;
1462     MasterBootRecord->PartitionTable[0].EndSector = 0xFF;
1463     MasterBootRecord->PartitionTable[0].EndCylinder = 0xFF;
1464     MasterBootRecord->PartitionTable[0].SectorCountBeforePartition = 1;
1465     MasterBootRecord->PartitionTable[0].PartitionSectorCount = 0xFFFFFFFF;
1466     MasterBootRecord->MasterBootRecordMagic = BOOT_RECORD_SIGNATURE;
1467 
1468     /* Finally, write that MBR */
1469     return FstubWriteSector(Disk->DeviceObject,
1470                             Disk->SectorSize,
1471                             0,
1472                             Disk->Buffer);
1473 }
1474 
1475 NTSTATUS
1476 NTAPI
FstubWriteEntryEFI(IN PDISK_INFORMATION Disk,IN ULONG PartitionsSizeSector,IN ULONG PartitionEntryNumber,IN PEFI_PARTITION_ENTRY PartitionEntry,IN BOOLEAN WriteBackupTable,IN BOOLEAN ForceWrite,OUT PULONG PartitionEntryCRC32 OPTIONAL)1477 FstubWriteEntryEFI(IN PDISK_INFORMATION Disk,
1478                    IN ULONG PartitionsSizeSector,
1479                    IN ULONG PartitionEntryNumber,
1480                    IN PEFI_PARTITION_ENTRY PartitionEntry,
1481                    IN BOOLEAN WriteBackupTable,
1482                    IN BOOLEAN ForceWrite,
1483                    OUT PULONG PartitionEntryCRC32 OPTIONAL)
1484 {
1485     NTSTATUS Status = STATUS_SUCCESS;
1486     ULONG Offset;
1487     ULONGLONG FirstEntryLBA;
1488 
1489     PAGED_CODE();
1490 
1491     ASSERT(Disk);
1492     ASSERT(IS_VALID_DISK_INFO(Disk));
1493 
1494     /* Get the first LBA where the partition table is:
1495      * On primary table, it's sector 2 (skip MBR & Header)
1496      * On backup table, it's ante last sector (Header) minus partition table size
1497      */
1498     if (!WriteBackupTable)
1499     {
1500         FirstEntryLBA = 2ULL;
1501     }
1502     else
1503     {
1504         FirstEntryLBA = Disk->SectorCount - PartitionsSizeSector - 1;
1505     }
1506 
1507     /* Copy the entry at the proper place into the buffer
1508      * That way, we don't erase previous entries
1509      */
1510     RtlCopyMemory((PVOID)((ULONG_PTR)Disk->Buffer + ((PartitionEntryNumber * PARTITION_ENTRY_SIZE) % Disk->SectorSize)),
1511                   PartitionEntry,
1512                   sizeof(EFI_PARTITION_ENTRY));
1513     /* Compute size of buffer */
1514     Offset = (PartitionEntryNumber * PARTITION_ENTRY_SIZE) % Disk->SectorSize + PARTITION_ENTRY_SIZE;
1515     ASSERT(Offset <= Disk->SectorSize);
1516 
1517     /* If it's full of partition entries, or if call ask for it, write down the data */
1518     if (Offset == Disk->SectorSize || ForceWrite)
1519     {
1520         /* We will write at first entry LBA + a shift made by already present/written entries */
1521         Status = FstubWriteSector(Disk->DeviceObject,
1522                                   Disk->SectorSize,
1523                                   FirstEntryLBA + ((PartitionEntryNumber * PARTITION_ENTRY_SIZE) / Disk->SectorSize),
1524                                   Disk->Buffer);
1525         if (!NT_SUCCESS(Status))
1526         {
1527             return Status;
1528         }
1529 
1530         /* We clean buffer */
1531         RtlZeroMemory(Disk->Buffer, Disk->SectorSize);
1532     }
1533 
1534     /* If we have a buffer for CRC32, then compute it */
1535     if (PartitionEntryCRC32)
1536     {
1537         *PartitionEntryCRC32 = RtlComputeCrc32(*PartitionEntryCRC32, (PUCHAR)PartitionEntry, PARTITION_ENTRY_SIZE);
1538     }
1539 
1540     return Status;
1541 }
1542 
1543 NTSTATUS
1544 NTAPI
FstubWriteHeaderEFI(IN PDISK_INFORMATION Disk,IN ULONG PartitionsSizeSector,IN GUID DiskGUID,IN ULONG NumberOfEntries,IN ULONGLONG FirstUsableLBA,IN ULONGLONG LastUsableLBA,IN ULONG PartitionEntryCRC32,IN BOOLEAN WriteBackupTable)1545 FstubWriteHeaderEFI(IN PDISK_INFORMATION Disk,
1546                     IN ULONG PartitionsSizeSector,
1547                     IN GUID DiskGUID,
1548                     IN ULONG NumberOfEntries,
1549                     IN ULONGLONG FirstUsableLBA,
1550                     IN ULONGLONG LastUsableLBA,
1551                     IN ULONG PartitionEntryCRC32,
1552                     IN BOOLEAN WriteBackupTable)
1553 {
1554     PEFI_PARTITION_HEADER EFIHeader;
1555 
1556     PAGED_CODE();
1557 
1558     ASSERT(Disk);
1559     ASSERT(IS_VALID_DISK_INFO(Disk));
1560 
1561     /* Let's use read buffer as EFI_PARTITION_HEADER */
1562     EFIHeader = (PEFI_PARTITION_HEADER)Disk->Buffer;
1563 
1564     /* Complete standard header information */
1565     EFIHeader->Signature = EFI_HEADER_SIGNATURE;
1566     EFIHeader->Revision = EFI_HEADER_REVISION_1;
1567     EFIHeader->HeaderSize = sizeof(EFI_PARTITION_HEADER);
1568     /* Set no CRC32 checksum at the moment */
1569     EFIHeader->HeaderCRC32 = 0;
1570     EFIHeader->Reserved = 0;
1571     /* Check whether we're writing primary or backup
1572      * That way, we can ajust LBA setting:
1573      * Primary is on first sector
1574      * Backup is on last sector
1575      */
1576     if (!WriteBackupTable)
1577     {
1578         EFIHeader->MyLBA = 1ULL;
1579         EFIHeader->AlternateLBA = Disk->SectorCount - 1ULL;
1580     }
1581     else
1582     {
1583         EFIHeader->MyLBA = Disk->SectorCount - 1ULL;
1584         EFIHeader->AlternateLBA = 1ULL;
1585     }
1586     /* Fill in with received data */
1587     EFIHeader->FirstUsableLBA = FirstUsableLBA;
1588     EFIHeader->LastUsableLBA = LastUsableLBA;
1589     EFIHeader->DiskGUID = DiskGUID;
1590     /* Check whether we're writing primary or backup
1591      * That way, we can ajust LBA setting:
1592      * On primary, partition entries are just after header, so sector 2
1593      * On backup, partition entries are just before header, so, last sector minus partition table size
1594      */
1595     if (!WriteBackupTable)
1596     {
1597         EFIHeader->PartitionEntryLBA = EFIHeader->MyLBA + 1ULL;
1598     }
1599     else
1600     {
1601         EFIHeader->PartitionEntryLBA = EFIHeader->MyLBA - PartitionsSizeSector;
1602     }
1603     /* Complete filling in */
1604     EFIHeader->NumberOfEntries = NumberOfEntries;
1605     EFIHeader->SizeOfPartitionEntry = PARTITION_ENTRY_SIZE;
1606     EFIHeader->PartitionEntryCRC32 = PartitionEntryCRC32;
1607     /* Finally, compute header checksum */
1608     EFIHeader->HeaderCRC32 = RtlComputeCrc32(0, (PUCHAR)EFIHeader, sizeof(EFI_PARTITION_HEADER));
1609 
1610     /* Debug the way we'll break disk, to let user pray */
1611     DPRINT("FSTUB: About to write the following header for %s table\n", (WriteBackupTable ? "backup" : "primary"));
1612     DPRINT(" Signature: %I64x\n", EFIHeader->Signature);
1613     DPRINT(" Revision: %x\n", EFIHeader->Revision);
1614     DPRINT(" HeaderSize: %x\n", EFIHeader->HeaderSize);
1615     DPRINT(" HeaderCRC32: %x\n", EFIHeader->HeaderCRC32);
1616     DPRINT(" MyLBA: %I64x\n", EFIHeader->MyLBA);
1617     DPRINT(" AlternateLBA: %I64x\n", EFIHeader->AlternateLBA);
1618     DPRINT(" FirstUsableLBA: %I64x\n", EFIHeader->FirstUsableLBA);
1619     DPRINT(" LastUsableLBA: %I64x\n", EFIHeader->LastUsableLBA);
1620     DPRINT(" PartitionEntryLBA: %I64x\n", EFIHeader->PartitionEntryLBA);
1621     DPRINT(" NumberOfEntries: %x\n", EFIHeader->NumberOfEntries);
1622     DPRINT(" SizeOfPartitionEntry: %x\n", EFIHeader->SizeOfPartitionEntry);
1623     DPRINT(" PartitionEntryCRC32: %x\n", EFIHeader->PartitionEntryCRC32);
1624 
1625     /* Write header to disk */
1626     return FstubWriteSector(Disk->DeviceObject,
1627                             Disk->SectorSize,
1628                             EFIHeader->MyLBA,
1629                             Disk->Buffer);
1630 }
1631 
1632 NTSTATUS
1633 NTAPI
FstubWritePartitionTableEFI(IN PDISK_INFORMATION Disk,IN GUID DiskGUID,IN ULONG MaxPartitionCount,IN ULONGLONG FirstUsableLBA,IN ULONGLONG LastUsableLBA,IN BOOLEAN WriteBackupTable,IN ULONG PartitionCount,IN PPARTITION_INFORMATION_EX PartitionEntries OPTIONAL)1634 FstubWritePartitionTableEFI(IN PDISK_INFORMATION Disk,
1635                             IN GUID DiskGUID,
1636                             IN ULONG MaxPartitionCount,
1637                             IN ULONGLONG FirstUsableLBA,
1638                             IN ULONGLONG LastUsableLBA,
1639                             IN BOOLEAN WriteBackupTable,
1640                             IN ULONG PartitionCount,
1641                             IN PPARTITION_INFORMATION_EX PartitionEntries OPTIONAL)
1642 {
1643     NTSTATUS Status;
1644     EFI_PARTITION_ENTRY Entry;
1645     ULONG i, WrittenPartitions, SectoredPartitionEntriesSize, PartitionEntryCRC32;
1646 
1647     PAGED_CODE();
1648 
1649     ASSERT(Disk);
1650     ASSERT(MaxPartitionCount >= 128);
1651     ASSERT(PartitionCount <= MaxPartitionCount);
1652 
1653     PartitionEntryCRC32 = 0;
1654     /* Count how much sectors we'll have to read to read the whole partition table */
1655     SectoredPartitionEntriesSize = (MaxPartitionCount * PARTITION_ENTRY_SIZE) / Disk->SectorSize;
1656 
1657     for (i = 0, WrittenPartitions = 0; i < PartitionCount; i++)
1658     {
1659         /* If partition GUID is 00000000-0000-0000-0000-000000000000, then it's unused, skip it */
1660         if (PartitionEntries[i].Gpt.PartitionType.Data1 == 0 &&
1661             PartitionEntries[i].Gpt.PartitionType.Data2 == 0 &&
1662             PartitionEntries[i].Gpt.PartitionType.Data3 == 0 &&
1663             ((PULONGLONG)PartitionEntries[i].Gpt.PartitionType.Data4)[0] == 0)
1664         {
1665             continue;
1666         }
1667 
1668         /* Copy the entry in the partition entry format */
1669         FstubCopyEntryEFI(&Entry, &PartitionEntries[i], Disk->SectorSize);
1670         /* Then write the entry to the disk */
1671         Status = FstubWriteEntryEFI(Disk,
1672                                     SectoredPartitionEntriesSize,
1673                                     WrittenPartitions,
1674                                     &Entry,
1675                                     WriteBackupTable,
1676                                     FALSE,
1677                                     &PartitionEntryCRC32);
1678         if (!NT_SUCCESS(Status))
1679         {
1680             return Status;
1681         }
1682         WrittenPartitions++;
1683     }
1684 
1685     /* Zero the buffer to write zeros to the disk */
1686     RtlZeroMemory(&Entry, sizeof(EFI_PARTITION_ENTRY));
1687     /* Write the disks with zeros for every unused remaining partition entry */
1688     for (i = WrittenPartitions; i < MaxPartitionCount; i++)
1689     {
1690         Status = FstubWriteEntryEFI(Disk,
1691                                     SectoredPartitionEntriesSize,
1692                                     i,
1693                                     &Entry,
1694                                     WriteBackupTable,
1695                                     FALSE,
1696                                     &PartitionEntryCRC32);
1697         if (!NT_SUCCESS(Status))
1698         {
1699             return Status;
1700         }
1701     }
1702 
1703     /* Once we're done, write the GPT header */
1704     return FstubWriteHeaderEFI(Disk,
1705                                SectoredPartitionEntriesSize,
1706                                DiskGUID,
1707                                MaxPartitionCount,
1708                                FirstUsableLBA,
1709                                LastUsableLBA,
1710                                PartitionEntryCRC32,
1711                                WriteBackupTable);
1712 }
1713 
1714 NTSTATUS
1715 NTAPI
FstubWritePartitionTableMBR(IN PDISK_INFORMATION Disk,IN PDRIVE_LAYOUT_INFORMATION_EX LayoutEx)1716 FstubWritePartitionTableMBR(IN PDISK_INFORMATION Disk,
1717                             IN PDRIVE_LAYOUT_INFORMATION_EX LayoutEx)
1718 {
1719     NTSTATUS Status;
1720     PDRIVE_LAYOUT_INFORMATION DriveLayout;
1721 
1722     PAGED_CODE();
1723 
1724     ASSERT(IS_VALID_DISK_INFO(Disk));
1725     ASSERT(LayoutEx);
1726 
1727     /* Convert data to the correct format */
1728     DriveLayout = FstubConvertExtendedToLayout(LayoutEx);
1729     if (!DriveLayout)
1730     {
1731         return STATUS_INSUFFICIENT_RESOURCES;
1732     }
1733 
1734     /* Really write information */
1735     Status = IoWritePartitionTable(Disk->DeviceObject,
1736                                    Disk->SectorSize,
1737                                    Disk->DiskGeometry.Geometry.SectorsPerTrack,
1738                                    Disk->DiskGeometry.Geometry.TracksPerCylinder,
1739                                    DriveLayout);
1740 
1741     /* Free allocated structure and return */
1742     ExFreePoolWithTag(DriveLayout, TAG_FSTUB);
1743     return Status;
1744 }
1745 
1746 NTSTATUS
1747 NTAPI
FstubWriteSector(IN PDEVICE_OBJECT DeviceObject,IN ULONG SectorSize,IN ULONGLONG StartingSector OPTIONAL,IN PVOID Buffer)1748 FstubWriteSector(IN PDEVICE_OBJECT DeviceObject,
1749                  IN ULONG SectorSize,
1750                  IN ULONGLONG StartingSector OPTIONAL,
1751                  IN PVOID Buffer)
1752 {
1753     NTSTATUS Status;
1754     PIRP Irp;
1755     KEVENT Event;
1756     LARGE_INTEGER StartingOffset;
1757     IO_STATUS_BLOCK IoStatusBlock;
1758     PIO_STACK_LOCATION IoStackLocation;
1759 
1760     PAGED_CODE();
1761 
1762     ASSERT(DeviceObject);
1763     ASSERT(Buffer);
1764     ASSERT(SectorSize);
1765 
1766     /* Compute starting offset */
1767     StartingOffset.QuadPart = StartingSector * SectorSize;
1768 
1769     /* Initialize waiting event */
1770     KeInitializeEvent(&Event, NotificationEvent, FALSE);
1771 
1772     /* Prepare IRP */
1773     Irp = IoBuildSynchronousFsdRequest(IRP_MJ_WRITE,
1774                                        DeviceObject,
1775                                        Buffer,
1776                                        SectorSize,
1777                                        &StartingOffset,
1778                                        &Event,
1779                                        &IoStatusBlock);
1780     if (!Irp)
1781     {
1782         return STATUS_INSUFFICIENT_RESOURCES;
1783     }
1784 
1785     /* Override volume verify */
1786     IoStackLocation = IoGetNextIrpStackLocation(Irp);
1787     IoStackLocation->Flags |= SL_OVERRIDE_VERIFY_VOLUME;
1788 
1789     /* Then call driver, and wait for completion if needed */
1790     Status = IoCallDriver(DeviceObject, Irp);
1791     if (Status == STATUS_PENDING)
1792     {
1793         KeWaitForSingleObject(&Event, Executive, KernelMode, FALSE, NULL);
1794         Status = IoStatusBlock.Status;
1795     }
1796 
1797     return Status;
1798 }
1799 
1800 /* FUNCTIONS *****************************************************************/
1801 
1802 /*
1803  * @implemented
1804  */
1805 NTSTATUS
1806 NTAPI
IoCreateDisk(IN PDEVICE_OBJECT DeviceObject,IN PCREATE_DISK Disk)1807 IoCreateDisk(IN PDEVICE_OBJECT DeviceObject,
1808              IN PCREATE_DISK Disk)
1809 {
1810     PARTITION_STYLE PartitionStyle;
1811 
1812     PAGED_CODE();
1813 
1814     ASSERT(DeviceObject);
1815 
1816     /* Get partition style. If caller didn't provided data, assume it's raw */
1817     PartitionStyle = ((Disk) ? Disk->PartitionStyle : PARTITION_STYLE_RAW);
1818     /* Then, call appropriate internal function */
1819     switch (PartitionStyle)
1820     {
1821         case PARTITION_STYLE_MBR:
1822             return FstubCreateDiskMBR(DeviceObject, &(Disk->Mbr));
1823         case PARTITION_STYLE_GPT:
1824             return FstubCreateDiskEFI(DeviceObject, &(Disk->Gpt));
1825         case PARTITION_STYLE_RAW:
1826             return FstubCreateDiskRaw(DeviceObject);
1827         default:
1828             return STATUS_NOT_SUPPORTED;
1829     }
1830 }
1831 
1832 /*
1833  * @implemented
1834  */
1835 NTSTATUS
1836 NTAPI
IoGetBootDiskInformation(IN OUT PBOOTDISK_INFORMATION BootDiskInformation,IN ULONG Size)1837 IoGetBootDiskInformation(IN OUT PBOOTDISK_INFORMATION BootDiskInformation,
1838                          IN ULONG Size)
1839 {
1840     NTSTATUS Status = STATUS_SUCCESS;
1841     PIRP Irp;
1842     KEVENT Event;
1843     PLIST_ENTRY NextEntry;
1844     PFILE_OBJECT FileObject;
1845     DISK_GEOMETRY DiskGeometry;
1846     PDEVICE_OBJECT DeviceObject;
1847     UNICODE_STRING DeviceStringW;
1848     IO_STATUS_BLOCK IoStatusBlock;
1849     CHAR Buffer[128], ArcBuffer[128];
1850     BOOLEAN SingleDisk, IsBootDiskInfoEx;
1851     PARC_DISK_SIGNATURE ArcDiskSignature;
1852     PARC_DISK_INFORMATION ArcDiskInformation;
1853     PARTITION_INFORMATION_EX PartitionInformation;
1854     PDRIVE_LAYOUT_INFORMATION_EX DriveLayout = NULL;
1855     ULONG DiskCount, DiskNumber, Signature, PartitionNumber;
1856     ANSI_STRING ArcBootString, ArcSystemString, DeviceStringA, ArcNameStringA;
1857     extern PLOADER_PARAMETER_BLOCK IopLoaderBlock;
1858 
1859     PAGED_CODE();
1860 
1861     /* Get loader block. If it's null, we come too late */
1862     if (!IopLoaderBlock)
1863     {
1864         return STATUS_TOO_LATE;
1865     }
1866 
1867     /* Check buffer size */
1868     if (Size < sizeof(BOOTDISK_INFORMATION))
1869     {
1870         return STATUS_INVALID_PARAMETER;
1871     }
1872 
1873     /* Init some useful stuff:
1874      * Get ARC disks information
1875      * Check whether we have a single disk on the machine
1876      * Check received structure size (extended or not?)
1877      * Init boot strings (system/boot)
1878      * Finaly, get disk count
1879      */
1880     ArcDiskInformation = IopLoaderBlock->ArcDiskInformation;
1881     SingleDisk = (ArcDiskInformation->DiskSignatureListHead.Flink->Flink ==
1882                  &ArcDiskInformation->DiskSignatureListHead);
1883 
1884     IsBootDiskInfoEx = (Size >= sizeof(BOOTDISK_INFORMATION_EX));
1885     RtlInitAnsiString(&ArcBootString, IopLoaderBlock->ArcBootDeviceName);
1886     RtlInitAnsiString(&ArcSystemString, IopLoaderBlock->ArcHalDeviceName);
1887     DiskCount = IoGetConfigurationInformation()->DiskCount;
1888 
1889     /* If no disk, return success */
1890     if (DiskCount == 0)
1891     {
1892         return STATUS_SUCCESS;
1893     }
1894 
1895     /* Now, browse all disks */
1896     for (DiskNumber = 0; DiskNumber < DiskCount; DiskNumber++)
1897     {
1898         /* Create the device name */
1899         sprintf(Buffer, "\\Device\\Harddisk%lu\\Partition0", DiskNumber);
1900         RtlInitAnsiString(&DeviceStringA, Buffer);
1901         Status = RtlAnsiStringToUnicodeString(&DeviceStringW, &DeviceStringA, TRUE);
1902         if (!NT_SUCCESS(Status))
1903         {
1904             continue;
1905         }
1906 
1907         /* Get its device object */
1908         Status = IoGetDeviceObjectPointer(&DeviceStringW,
1909                                           FILE_READ_ATTRIBUTES,
1910                                           &FileObject,
1911                                           &DeviceObject);
1912         RtlFreeUnicodeString(&DeviceStringW);
1913         if (!NT_SUCCESS(Status))
1914         {
1915             continue;
1916         }
1917 
1918         /* Prepare for getting disk geometry */
1919         KeInitializeEvent(&Event, NotificationEvent, FALSE);
1920         Irp = IoBuildDeviceIoControlRequest(IOCTL_DISK_GET_DRIVE_GEOMETRY,
1921                                             DeviceObject,
1922                                             NULL,
1923                                             0,
1924                                             &DiskGeometry,
1925                                             sizeof(DiskGeometry),
1926                                             FALSE,
1927                                             &Event,
1928                                             &IoStatusBlock);
1929         if (!Irp)
1930         {
1931             ObDereferenceObject(FileObject);
1932             continue;
1933         }
1934 
1935         /* Then, call the drive, and wait for it if needed */
1936         Status = IoCallDriver(DeviceObject, Irp);
1937         if (Status == STATUS_PENDING)
1938         {
1939             KeWaitForSingleObject(&Event, Executive, KernelMode, FALSE, NULL);
1940             Status = IoStatusBlock.Status;
1941         }
1942         if (!NT_SUCCESS(Status))
1943         {
1944             ObDereferenceObject(FileObject);
1945             continue;
1946         }
1947 
1948         /* Read partition table */
1949         Status = IoReadPartitionTableEx(DeviceObject,
1950                                         &DriveLayout);
1951 
1952         /* FileObject, you can go! */
1953         ObDereferenceObject(FileObject);
1954 
1955         if (!NT_SUCCESS(Status))
1956         {
1957             continue;
1958         }
1959 
1960         /* Ensure we have at least 512 bytes per sector */
1961         if (DiskGeometry.BytesPerSector < 512)
1962         {
1963             DiskGeometry.BytesPerSector = 512;
1964         }
1965 
1966         /* Now, for each ARC disk, try to find the matching */
1967         for (NextEntry = ArcDiskInformation->DiskSignatureListHead.Flink;
1968              NextEntry != &ArcDiskInformation->DiskSignatureListHead;
1969              NextEntry = NextEntry->Flink)
1970         {
1971             ArcDiskSignature = CONTAINING_RECORD(NextEntry,
1972                                                  ARC_DISK_SIGNATURE,
1973                                                  ListEntry);
1974             /* If they match, i.e.
1975              * - There's only one disk for both BIOS and detected
1976              * - Signatures are matching
1977              * - This is MBR
1978              * (We don't check checksums here)
1979              */
1980             if (((SingleDisk && DiskCount == 1) ||
1981                 (IopVerifyDiskSignature(DriveLayout, ArcDiskSignature, &Signature))) &&
1982                 (DriveLayout->PartitionStyle == PARTITION_STYLE_MBR))
1983             {
1984                 /* Create ARC name */
1985                 sprintf(ArcBuffer, "\\ArcName\\%s", ArcDiskSignature->ArcName);
1986                 RtlInitAnsiString(&ArcNameStringA, ArcBuffer);
1987 
1988                 /* Browse all partitions */
1989                 for (PartitionNumber = 1; PartitionNumber <= DriveLayout->PartitionCount; PartitionNumber++)
1990                 {
1991                     /* Create its device name */
1992                     sprintf(Buffer, "\\Device\\Harddisk%lu\\Partition%lu", DiskNumber, PartitionNumber);
1993                     RtlInitAnsiString(&DeviceStringA, Buffer);
1994                     Status = RtlAnsiStringToUnicodeString(&DeviceStringW, &DeviceStringA, TRUE);
1995                     if (!NT_SUCCESS(Status))
1996                     {
1997                         continue;
1998                     }
1999 
2000                     /* If IopVerifyDiskSignature returned no signature, take the one from DriveLayout */
2001                     if (!Signature)
2002                     {
2003                         Signature = DriveLayout->Mbr.Signature;
2004                     }
2005 
2006                     /* Create partial ARC name */
2007                     sprintf(ArcBuffer, "%spartition(%lu)", ArcDiskSignature->ArcName, PartitionNumber);
2008                     RtlInitAnsiString(&ArcNameStringA, ArcBuffer);
2009 
2010                     /* If it's matching boot string */
2011                     if (RtlEqualString(&ArcNameStringA, &ArcBootString, TRUE))
2012                     {
2013                         /*  Then, fill in information about boot device */
2014                         BootDiskInformation->BootDeviceSignature = Signature;
2015 
2016                         /* Get its device object */
2017                         Status = IoGetDeviceObjectPointer(&DeviceStringW,
2018                                                           FILE_READ_ATTRIBUTES,
2019                                                           &FileObject,
2020                                                           &DeviceObject);
2021                         if (!NT_SUCCESS(Status))
2022                         {
2023                             RtlFreeUnicodeString(&DeviceStringW);
2024                             continue;
2025                         }
2026 
2027                         /* And call the drive to get information about partition */
2028                         KeInitializeEvent(&Event, NotificationEvent, FALSE);
2029                         Irp = IoBuildDeviceIoControlRequest(IOCTL_DISK_GET_PARTITION_INFO_EX,
2030                                                             DeviceObject,
2031                                                             NULL,
2032                                                             0,
2033                                                             &PartitionInformation,
2034                                                             sizeof(PartitionInformation),
2035                                                             FALSE,
2036                                                             &Event,
2037                                                             &IoStatusBlock);
2038                         if (!Irp)
2039                         {
2040                             ObDereferenceObject(FileObject);
2041                             RtlFreeUnicodeString(&DeviceStringW);
2042                             continue;
2043                         }
2044 
2045                         /* Call & wait if needed */
2046                         Status = IoCallDriver(DeviceObject, Irp);
2047                         if (Status == STATUS_PENDING)
2048                         {
2049                             KeWaitForSingleObject(&Event, Executive, KernelMode, FALSE, NULL);
2050                             Status = IoStatusBlock.Status;
2051                         }
2052                         if (!NT_SUCCESS(Status))
2053                         {
2054                             ObDereferenceObject(FileObject);
2055                             RtlFreeUnicodeString(&DeviceStringW);
2056                             continue;
2057                         }
2058 
2059                         /* We get partition offset as demanded and return it */
2060                         BootDiskInformation->BootPartitionOffset = PartitionInformation.StartingOffset.QuadPart;
2061 
2062                         /* If called passed a BOOTDISK_INFORMATION_EX structure, give more intel */
2063                         if (IsBootDiskInfoEx)
2064                         {
2065                             /* Is partition style MBR or GPT? */
2066                             if (DriveLayout->PartitionStyle == PARTITION_STYLE_GPT)
2067                             {
2068                                 ((PBOOTDISK_INFORMATION_EX)BootDiskInformation)->BootDeviceGuid = DriveLayout->Gpt.DiskId;
2069                                 ((PBOOTDISK_INFORMATION_EX)BootDiskInformation)->BootDeviceIsGpt = TRUE;
2070                             }
2071                             else
2072                             {
2073                                 ((PBOOTDISK_INFORMATION_EX)BootDiskInformation)->BootDeviceIsGpt = FALSE;
2074                             }
2075                         }
2076 
2077                         /* Dereference FileObject */
2078                         ObDereferenceObject(FileObject);
2079                     }
2080 
2081                     /* If it's matching system string */
2082                     if (RtlEqualString(&ArcNameStringA, &ArcSystemString, TRUE))
2083                     {
2084                         /* Then, fill in information about the system device */
2085                         BootDiskInformation->SystemDeviceSignature = Signature;
2086 
2087                         /* Get its device object */
2088                         Status = IoGetDeviceObjectPointer(&DeviceStringW,
2089                                                           FILE_READ_ATTRIBUTES,
2090                                                           &FileObject,
2091                                                           &DeviceObject);
2092                         if (!NT_SUCCESS(Status))
2093                         {
2094                             RtlFreeUnicodeString(&DeviceStringW);
2095                             continue;
2096                         }
2097 
2098                         /* And call the drive to get information about partition */
2099                         KeInitializeEvent(&Event, NotificationEvent, FALSE);
2100                         Irp = IoBuildDeviceIoControlRequest(IOCTL_DISK_GET_PARTITION_INFO_EX,
2101                                                             DeviceObject,
2102                                                             NULL,
2103                                                             0,
2104                                                             &PartitionInformation,
2105                                                             sizeof(PartitionInformation),
2106                                                             FALSE,
2107                                                             &Event,
2108                                                             &IoStatusBlock);
2109                         if (!Irp)
2110                         {
2111                             ObDereferenceObject(FileObject);
2112                             RtlFreeUnicodeString(&DeviceStringW);
2113                             continue;
2114                         }
2115 
2116                         /* Call & wait if needed */
2117                         Status = IoCallDriver(DeviceObject, Irp);
2118                         if (Status == STATUS_PENDING)
2119                         {
2120                             KeWaitForSingleObject(&Event, Executive, KernelMode, FALSE, NULL);
2121                             Status = IoStatusBlock.Status;
2122                         }
2123                         if (!NT_SUCCESS(Status))
2124                         {
2125                             ObDereferenceObject(FileObject);
2126                             RtlFreeUnicodeString(&DeviceStringW);
2127                             continue;
2128                         }
2129 
2130                         /* We get partition offset as demanded and return it */
2131                         BootDiskInformation->SystemPartitionOffset = PartitionInformation.StartingOffset.QuadPart;
2132 
2133                         /* If called passed a BOOTDISK_INFORMATION_EX structure, give more intel */
2134                         if (IsBootDiskInfoEx)
2135                         {
2136                             /* Is partition style MBR or GPT? */
2137                             if (DriveLayout->PartitionStyle == PARTITION_STYLE_GPT)
2138                             {
2139                                 ((PBOOTDISK_INFORMATION_EX)BootDiskInformation)->SystemDeviceGuid = DriveLayout->Gpt.DiskId;
2140                                 ((PBOOTDISK_INFORMATION_EX)BootDiskInformation)->SystemDeviceIsGpt = TRUE;
2141                             }
2142                             else
2143                             {
2144                                 ((PBOOTDISK_INFORMATION_EX)BootDiskInformation)->SystemDeviceIsGpt = FALSE;
2145                             }
2146                         }
2147 
2148                         /* Dereference FileObject */
2149                         ObDereferenceObject(FileObject);
2150                     }
2151 
2152                     /* Release device string */
2153                     RtlFreeUnicodeString(&DeviceStringW);
2154                 }
2155             }
2156         }
2157 
2158         /* Finally, release drive layout */
2159         ExFreePool(DriveLayout);
2160     }
2161 
2162     /* And return */
2163     return Status;
2164 }
2165 
2166 /*
2167  * @implemented
2168  */
2169 NTSTATUS
2170 NTAPI
IoReadDiskSignature(IN PDEVICE_OBJECT DeviceObject,IN ULONG BytesPerSector,OUT PDISK_SIGNATURE Signature)2171 IoReadDiskSignature(IN PDEVICE_OBJECT DeviceObject,
2172                     IN ULONG BytesPerSector,
2173                     OUT PDISK_SIGNATURE Signature)
2174 {
2175     NTSTATUS Status;
2176     PUCHAR Buffer;
2177     ULONG HeaderCRC32, i, CheckSum;
2178     PEFI_PARTITION_HEADER EFIHeader;
2179     PPARTITION_DESCRIPTOR PartitionDescriptor;
2180 
2181     PAGED_CODE();
2182 
2183     /* Ensure we'll read at least 512 bytes */
2184     if (BytesPerSector < 512)
2185     {
2186         BytesPerSector = 512;
2187     }
2188 
2189     /* Allocate a buffer for reading operations */
2190     Buffer = ExAllocatePoolWithTag(NonPagedPoolCacheAligned, BytesPerSector, TAG_FSTUB);
2191     if (!Buffer)
2192     {
2193         return STATUS_NO_MEMORY;
2194     }
2195 
2196     /* Read first sector (sector 0) for MBR */
2197     Status = FstubReadSector(DeviceObject,
2198                              BytesPerSector,
2199                              0ULL,
2200                              Buffer);
2201     if (!NT_SUCCESS(Status))
2202     {
2203         goto Cleanup;
2204     }
2205 
2206     /* Get the partition descriptor array */
2207     PartitionDescriptor = (PPARTITION_DESCRIPTOR)
2208                           &(Buffer[PARTITION_TABLE_OFFSET]);
2209     /* Check partitions types: if first is 0xEE and all the others 0, we have GPT */
2210     if (PartitionDescriptor[0].PartitionType == EFI_PMBR_OSTYPE_EFI &&
2211         PartitionDescriptor[1].PartitionType == 0 &&
2212         PartitionDescriptor[2].PartitionType == 0 &&
2213         PartitionDescriptor[3].PartitionType == 0)
2214     {
2215         /* If we have GPT, read second sector (sector 1) for GPT header */
2216         Status = FstubReadSector(DeviceObject,
2217                                  BytesPerSector,
2218                                  1ULL,
2219                                  Buffer);
2220         if (!NT_SUCCESS(Status))
2221         {
2222             goto Cleanup;
2223         }
2224         EFIHeader = (PEFI_PARTITION_HEADER)Buffer;
2225 
2226         /* First check signature
2227          * Then, check version (we only support v1
2228          * Finally check header size
2229          */
2230         if (EFIHeader->Signature != EFI_HEADER_SIGNATURE ||
2231             EFIHeader->Revision != EFI_HEADER_REVISION_1 ||
2232             EFIHeader->HeaderSize != sizeof(EFI_PARTITION_HEADER))
2233         {
2234             Status = STATUS_DISK_CORRUPT_ERROR;
2235             goto Cleanup;
2236         }
2237 
2238         /* Save current checksum */
2239         HeaderCRC32 = EFIHeader->HeaderCRC32;
2240         /* Then zero the one in EFI header. This is needed to compute header checksum */
2241         EFIHeader->HeaderCRC32 = 0;
2242         /* Compute header checksum and compare with the one present in partition table */
2243         if (RtlComputeCrc32(0, Buffer, sizeof(EFI_PARTITION_HEADER)) != HeaderCRC32)
2244         {
2245             Status = STATUS_DISK_CORRUPT_ERROR;
2246             goto Cleanup;
2247         }
2248 
2249         /* Set partition table style to GPT and return disk GUID */
2250         Signature->PartitionStyle = PARTITION_STYLE_GPT;
2251         Signature->Gpt.DiskId = EFIHeader->DiskGUID;
2252     }
2253     else
2254     {
2255         /* Compute MBR checksum */
2256         for (i = 0, CheckSum = 0; i < 512; i += sizeof(UINT32))
2257         {
2258             CheckSum += *(PUINT32)&Buffer[i];
2259         }
2260         CheckSum = ~CheckSum + 1;
2261 
2262         /* Set partition table style to MBR and return signature (offset 440) and checksum */
2263         Signature->PartitionStyle = PARTITION_STYLE_MBR;
2264         Signature->Mbr.Signature = *(PUINT32)&Buffer[DISK_SIGNATURE_OFFSET];
2265         Signature->Mbr.CheckSum = CheckSum;
2266     }
2267 
2268 Cleanup:
2269     /* Free buffer and return */
2270     ExFreePoolWithTag(Buffer, TAG_FSTUB);
2271     return Status;
2272 }
2273 
2274 /*
2275  * @implemented
2276  */
2277 NTSTATUS
2278 NTAPI
IoReadPartitionTableEx(IN PDEVICE_OBJECT DeviceObject,IN PDRIVE_LAYOUT_INFORMATION_EX * DriveLayout)2279 IoReadPartitionTableEx(IN PDEVICE_OBJECT DeviceObject,
2280                        IN PDRIVE_LAYOUT_INFORMATION_EX* DriveLayout)
2281 {
2282     NTSTATUS Status;
2283     PDISK_INFORMATION Disk;
2284     PARTITION_STYLE PartitionStyle;
2285 
2286     PAGED_CODE();
2287 
2288     ASSERT(DeviceObject);
2289     ASSERT(DriveLayout);
2290 
2291     /* First of all, allocate internal structure */
2292     Status = FstubAllocateDiskInformation(DeviceObject, &Disk, NULL);
2293     if (!NT_SUCCESS(Status))
2294     {
2295         return Status;
2296     }
2297     ASSERT(Disk);
2298 
2299     /* Then, detect partition style (MBR? GPT/EFI? RAW?) */
2300     Status = FstubDetectPartitionStyle(Disk, &PartitionStyle);
2301     if (!NT_SUCCESS(Status))
2302     {
2303         FstubFreeDiskInformation(Disk);
2304         return Status;
2305     }
2306 
2307     /* Here partition table is really read, depending on its style */
2308     switch (PartitionStyle)
2309     {
2310         case PARTITION_STYLE_MBR:
2311         case PARTITION_STYLE_RAW:
2312             Status = FstubReadPartitionTableMBR(Disk, FALSE, DriveLayout);
2313             break;
2314 
2315         case PARTITION_STYLE_GPT:
2316             /* Read primary table */
2317             Status = FstubReadPartitionTableEFI(Disk, FALSE, DriveLayout);
2318             /* If it failed, try reading backup table */
2319             if (!NT_SUCCESS(Status))
2320             {
2321                 Status = FstubReadPartitionTableEFI(Disk, TRUE, DriveLayout);
2322             }
2323             break;
2324 
2325         default:
2326             DPRINT("Unknown partition type\n");
2327             Status = STATUS_UNSUCCESSFUL;
2328     }
2329 
2330     /* It's over, internal structure not needed anymore */
2331     FstubFreeDiskInformation(Disk);
2332 
2333     /* In case of success, print data */
2334     if (NT_SUCCESS(Status))
2335     {
2336         FstubDbgPrintDriveLayoutEx(*DriveLayout);
2337     }
2338 
2339     return Status;
2340 }
2341 
2342 /*
2343  * @implemented
2344  */
2345 NTSTATUS
2346 NTAPI
IoSetPartitionInformationEx(IN PDEVICE_OBJECT DeviceObject,IN ULONG PartitionNumber,IN PSET_PARTITION_INFORMATION_EX PartitionInfo)2347 IoSetPartitionInformationEx(IN PDEVICE_OBJECT DeviceObject,
2348                             IN ULONG PartitionNumber,
2349                             IN PSET_PARTITION_INFORMATION_EX PartitionInfo)
2350 {
2351     NTSTATUS Status;
2352     PDISK_INFORMATION Disk;
2353     PARTITION_STYLE PartitionStyle;
2354 
2355     PAGED_CODE();
2356 
2357     ASSERT(DeviceObject);
2358     ASSERT(PartitionInfo);
2359 
2360     /* Debug given modifications */
2361     FstubDbgPrintSetPartitionEx(PartitionInfo, PartitionNumber);
2362 
2363     /* Allocate internal structure */
2364     Status = FstubAllocateDiskInformation(DeviceObject, &Disk, NULL);
2365     if (!NT_SUCCESS(Status))
2366     {
2367         return Status;
2368     }
2369 
2370     /* Get partition table style on disk */
2371     Status = FstubDetectPartitionStyle(Disk, &PartitionStyle);
2372     if (!NT_SUCCESS(Status))
2373     {
2374         FstubFreeDiskInformation(Disk);
2375         return Status;
2376     }
2377 
2378     /* If it's not matching partition style given in modifications, give up */
2379     if (PartitionInfo->PartitionStyle != PartitionStyle)
2380     {
2381         FstubFreeDiskInformation(Disk);
2382         return STATUS_INVALID_PARAMETER;
2383     }
2384 
2385     /* Finally, handle modifications using proper function */
2386     switch (PartitionStyle)
2387     {
2388         case PARTITION_STYLE_MBR:
2389             Status = IoSetPartitionInformation(DeviceObject,
2390                                                Disk->SectorSize,
2391                                                PartitionNumber,
2392                                                PartitionInfo->Mbr.PartitionType);
2393             break;
2394         case PARTITION_STYLE_GPT:
2395             Status = FstubSetPartitionInformationEFI(Disk,
2396                                                      PartitionNumber,
2397                                                      &(PartitionInfo->Gpt));
2398             break;
2399         default:
2400             Status = STATUS_NOT_SUPPORTED;
2401     }
2402 
2403     /* Release internal structure and return */
2404     FstubFreeDiskInformation(Disk);
2405     return Status;
2406 }
2407 
2408 /*
2409  * @implemented
2410  */
2411 NTSTATUS
2412 NTAPI
IoVerifyPartitionTable(IN PDEVICE_OBJECT DeviceObject,IN BOOLEAN FixErrors)2413 IoVerifyPartitionTable(IN PDEVICE_OBJECT DeviceObject,
2414                        IN BOOLEAN FixErrors)
2415 {
2416     NTSTATUS Status;
2417     PDISK_INFORMATION Disk;
2418     PARTITION_STYLE PartitionStyle;
2419 
2420     PAGED_CODE();
2421 
2422     ASSERT(DeviceObject);
2423 
2424     /* Allocate internal structure */
2425     Status = FstubAllocateDiskInformation(DeviceObject, &Disk, NULL);
2426     if (!NT_SUCCESS(Status))
2427     {
2428         return Status;
2429     }
2430     ASSERT(Disk);
2431 
2432     /* Get partition table style on disk */
2433     Status = FstubDetectPartitionStyle(Disk, &PartitionStyle);
2434     if (!NT_SUCCESS(Status))
2435     {
2436         FstubFreeDiskInformation(Disk);
2437         return Status;
2438     }
2439 
2440     /* Action will depend on partition style */
2441     switch (PartitionStyle)
2442     {
2443         /* For MBR, assume it's always OK */
2444         case PARTITION_STYLE_MBR:
2445             Status = STATUS_SUCCESS;
2446             break;
2447         /* For GPT, call internal function */
2448         case PARTITION_STYLE_GPT:
2449             Status = FstubVerifyPartitionTableEFI(Disk, FixErrors);
2450             break;
2451         /* Otherwise, signal we can't work */
2452         default:
2453             Status = STATUS_NOT_SUPPORTED;
2454     }
2455 
2456     /* Release internal structure and return */
2457     FstubFreeDiskInformation(Disk);
2458     return Status;
2459 }
2460 
2461 /*
2462  * @implemented
2463  */
2464 NTSTATUS
2465 NTAPI
IoWritePartitionTableEx(IN PDEVICE_OBJECT DeviceObject,IN PDRIVE_LAYOUT_INFORMATION_EX DriveLayout)2466 IoWritePartitionTableEx(IN PDEVICE_OBJECT DeviceObject,
2467                         IN PDRIVE_LAYOUT_INFORMATION_EX DriveLayout)
2468 {
2469     NTSTATUS Status;
2470     GUID DiskGuid;
2471     ULONG NumberOfEntries;
2472     PDISK_INFORMATION Disk;
2473     PEFI_PARTITION_HEADER EfiHeader;
2474     ULONGLONG SectorsForPartitions, FirstUsableLBA, LastUsableLBA;
2475 
2476     PAGED_CODE();
2477 
2478     ASSERT(DeviceObject);
2479     ASSERT(DriveLayout);
2480 
2481     /* Debug partition table that must be written */
2482     FstubDbgPrintDriveLayoutEx(DriveLayout);
2483 
2484     /* Allocate internal structure */
2485     Status = FstubAllocateDiskInformation(DeviceObject, &Disk, NULL);
2486     if (!NT_SUCCESS(Status))
2487     {
2488         return Status;
2489     }
2490     ASSERT(Disk);
2491 
2492     switch (DriveLayout->PartitionStyle)
2493     {
2494         case PARTITION_STYLE_MBR:
2495             Status = FstubWritePartitionTableMBR(Disk, DriveLayout);
2496             break;
2497 
2498         case PARTITION_STYLE_GPT:
2499             /* Read primary table header */
2500             Status = FstubReadHeaderEFI(Disk,
2501                                         FALSE,
2502                                         &EfiHeader);
2503             /* If it failed, try reading back table header */
2504             if (!NT_SUCCESS(Status))
2505             {
2506                 Status = FstubReadHeaderEFI(Disk,
2507                                             TRUE,
2508                                             &EfiHeader);
2509             }
2510 
2511             /* We have a header! */
2512             if (NT_SUCCESS(Status))
2513             {
2514                 /* Check if there are enough places for the partitions to be written */
2515                 if (DriveLayout->PartitionCount <= EfiHeader->NumberOfEntries)
2516                 {
2517                     /* Backup data */
2518                     NumberOfEntries = EfiHeader->NumberOfEntries;
2519                     RtlCopyMemory(&DiskGuid, &EfiHeader->DiskGUID, sizeof(GUID));
2520                     /* Count number of sectors needed to store partitions */
2521                     SectorsForPartitions = ((ULONGLONG)NumberOfEntries * PARTITION_ENTRY_SIZE) / Disk->SectorSize;
2522                     /* Set first usable LBA: Legacy MBR + GPT header + Partitions entries */
2523                     FirstUsableLBA = SectorsForPartitions + 2;
2524                     /* Set last usable LBA: Last sector - GPT header - Partitions entries */
2525                     LastUsableLBA = Disk->SectorCount - SectorsForPartitions - 1;
2526                     /* Write primary table */
2527                     Status = FstubWritePartitionTableEFI(Disk,
2528                                                          DiskGuid,
2529                                                          NumberOfEntries,
2530                                                          FirstUsableLBA,
2531                                                          LastUsableLBA,
2532                                                          FALSE,
2533                                                          DriveLayout->PartitionCount,
2534                                                          DriveLayout->PartitionEntry);
2535                     /* If it succeeded, also update backup table */
2536                     if (NT_SUCCESS(Status))
2537                     {
2538                         Status = FstubWritePartitionTableEFI(Disk,
2539                                                              DiskGuid,
2540                                                              NumberOfEntries,
2541                                                              FirstUsableLBA,
2542                                                              LastUsableLBA,
2543                                                              TRUE,
2544                                                              DriveLayout->PartitionCount,
2545                                                              DriveLayout->PartitionEntry);
2546                     }
2547                 }
2548                 else
2549                 {
2550                     Status = STATUS_INVALID_PARAMETER;
2551                 }
2552             }
2553             break;
2554 
2555         default:
2556             DPRINT("Unsupported partition style: %lu\n", DriveLayout->PartitionStyle);
2557             Status = STATUS_NOT_SUPPORTED;
2558     }
2559 
2560     /* It's over, internal structure not needed anymore */
2561     FstubFreeDiskInformation(Disk);
2562 
2563     return Status;
2564 }
2565 
2566 /* EOF */
2567