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