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