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