[NTOSKRNL]

[reactos.git] / reactos / ntoskrnl / fstub / disksup.c

diff --git a/reactos/ntoskrnl/fstub/disksup.c b/reactos/ntoskrnl/fstub/disksup.c
index 725ac34..7c4f651 100644 (file)
--- a/reactos/ntoskrnl/fstub/disksup.c
+++ b/reactos/ntoskrnl/fstub/disksup.c
@@ -4,7 +4,7 @@
 * FILE:            ntoskrnl/fstub/disksup.c
 * PURPOSE:         I/O HAL Routines for Disk Access
 * PROGRAMMERS:     Alex Ionescu (alex.ionescu@reactos.org)
-*                  Eric Kohl (ekohl@rz-online.de)
+*                  Eric Kohl
 *                  Casper S. Hornstrup (chorns@users.sourceforge.net)
 */
 
@@ -20,10 +20,12 @@
 #if 1
 const WCHAR DiskMountString[] = L"\\DosDevices\\%C:";
 
-#define AUTO_DRIVE         ((ULONG)-1)
+#define AUTO_DRIVE         MAXULONG
 
 #define PARTITION_MAGIC    0xaa55
 
+#define EFI_PMBR_OSTYPE_EFI 0xEE
+
 #include <pshpack1.h>
 
 typedef struct _REG_DISK_MOUNT_INFO
@@ -47,7 +49,9 @@ HalpAssignDrive(IN PUNICODE_STRING PartitionName,
                 IN UCHAR DriveType,
                 IN ULONG Signature,
                 IN LARGE_INTEGER StartingOffset,
-                IN HANDLE hKey)
+                IN HANDLE hKey,
+                IN PUNICODE_STRING BootDevice,
+                OUT PUCHAR NtSystemPath)
 {
     WCHAR DriveNameBuffer[16];
     UNICODE_STRING DriveName;
@@ -128,6 +132,14 @@ HalpAssignDrive(IN PUNICODE_STRING PartitionName,
             DPRINT1("ZwCreateValueKey failed for %wZ, status=%x\n", &DriveName, Status);
         }
     }
+
+    /* Check if this is a boot partition */
+    if (RtlCompareUnicodeString(PartitionName, BootDevice, FALSE) == 0)
+    {
+        /* Set NtSystemPath to that partition's disk letter */
+        *NtSystemPath = (UCHAR)('A' + DriveNumber);
+    }
+
     return TRUE;
 }
 
@@ -420,6 +432,11 @@ xHalIoAssignDriveLetters(IN PLOADER_PARAMETER_BLOCK LoaderBlock,
     PKEY_VALUE_PARTIAL_INFORMATION PartialInformation;
     PREG_DISK_MOUNT_INFO DiskMountInfo;
     ULONG RDiskCount;
+    UNICODE_STRING BootDevice;
+
+    Status = RtlAnsiStringToUnicodeString(&BootDevice,
+                                          NtDeviceName,
+                                          TRUE);
 
     DPRINT("xHalIoAssignDriveLetters()\n");
 
@@ -437,6 +454,8 @@ xHalIoAssignDriveLetters(IN PLOADER_PARAMETER_BLOCK LoaderBlock,
     PartialInformation = (PKEY_VALUE_PARTIAL_INFORMATION)ExAllocatePool(PagedPool,
         sizeof(KEY_VALUE_PARTIAL_INFORMATION) + sizeof(REG_DISK_MOUNT_INFO));
 
+    if (!Buffer1 || !Buffer2 || !PartialInformation) return;
+
     DiskMountInfo = (PREG_DISK_MOUNT_INFO) PartialInformation->Data;
 
     /* Open or Create the 'MountedDevices' key */
@@ -507,8 +526,18 @@ xHalIoAssignDriveLetters(IN PLOADER_PARAMETER_BLOCK LoaderBlock,
     }
 
     /* Initialize layout array */
+    if (ConfigInfo->DiskCount == 0)
+        goto end_assign_disks;
     LayoutArray = ExAllocatePool(NonPagedPool,
         ConfigInfo->DiskCount * sizeof(PDRIVE_LAYOUT_INFORMATION));
+    if (!LayoutArray)
+    {
+        ExFreePool(PartialInformation);
+        ExFreePool(Buffer2);
+        ExFreePool(Buffer1);
+        if (hKey) ZwClose(hKey);
+    }
+
     RtlZeroMemory(LayoutArray,
         ConfigInfo->DiskCount * sizeof(PDRIVE_LAYOUT_INFORMATION));
     for (i = 0; i < ConfigInfo->DiskCount; i++)
@@ -612,7 +641,9 @@ xHalIoAssignDriveLetters(IN PLOADER_PARAMETER_BLOCK LoaderBlock,
                                             DOSDEVICE_DRIVE_FIXED,
                                             DiskMountInfo->Signature,
                                             DiskMountInfo->StartingOffset,
-                                            NULL);
+                                            NULL,
+                                            &BootDevice,
+                                            NtSystemPath);
                                         /* Mark the partition as assigned */
                                         LayoutArray[i]->PartitionEntry[j].RewritePartition = TRUE;
                                     }
@@ -662,7 +693,9 @@ xHalIoAssignDriveLetters(IN PLOADER_PARAMETER_BLOCK LoaderBlock,
                             DOSDEVICE_DRIVE_FIXED,
                             LayoutArray[DiskNumber]->Signature,
                             LayoutArray[DiskNumber]->PartitionEntry[j].StartingOffset,
-                            hKey);
+                            hKey,
+                            &BootDevice,
+                            NtSystemPath);
                         /* Mark the partition as assigned */
                         LayoutArray[DiskNumber]->PartitionEntry[j].RewritePartition = TRUE;
                     }
@@ -702,7 +735,9 @@ xHalIoAssignDriveLetters(IN PLOADER_PARAMETER_BLOCK LoaderBlock,
                         DOSDEVICE_DRIVE_FIXED,
                         LayoutArray[DiskNumber]->Signature,
                         LayoutArray[DiskNumber]->PartitionEntry[j].StartingOffset,
-                        hKey);
+                        hKey,
+                        &BootDevice,
+                        NtSystemPath);
                     /* Mark the partition as assigned */
                     LayoutArray[DiskNumber]->PartitionEntry[j].RewritePartition = TRUE;
                 }
@@ -741,7 +776,9 @@ xHalIoAssignDriveLetters(IN PLOADER_PARAMETER_BLOCK LoaderBlock,
                         DOSDEVICE_DRIVE_FIXED,
                         LayoutArray[DiskNumber]->Signature,
                         LayoutArray[DiskNumber]->PartitionEntry[j].StartingOffset,
-                        hKey);
+                        hKey,
+                        &BootDevice,
+                        NtSystemPath);
                     /* Mark the partition as assigned */
                     LayoutArray[DiskNumber]->PartitionEntry[j].RewritePartition = TRUE;
                 }
@@ -776,7 +813,9 @@ xHalIoAssignDriveLetters(IN PLOADER_PARAMETER_BLOCK LoaderBlock,
                         DOSDEVICE_DRIVE_FIXED,
                         LayoutArray[DiskNumber]->Signature,
                         LayoutArray[DiskNumber]->PartitionEntry[j].StartingOffset,
-                        hKey);
+                        hKey,
+                        &BootDevice,
+                        NtSystemPath);
                     /* Mark the partition as assigned */
                     LayoutArray[DiskNumber]->PartitionEntry[j].RewritePartition = TRUE;
                 }
@@ -812,7 +851,9 @@ xHalIoAssignDriveLetters(IN PLOADER_PARAMETER_BLOCK LoaderBlock,
                         DOSDEVICE_DRIVE_FIXED,
                         LayoutArray[DiskNumber]->Signature,
                         LayoutArray[DiskNumber]->PartitionEntry[j].StartingOffset,
-                        hKey);
+                        hKey,
+                        &BootDevice,
+                        NtSystemPath);
                     /* Mark the partition as assigned */
                     LayoutArray[DiskNumber]->PartitionEntry[j].RewritePartition = TRUE;
                 }
@@ -844,7 +885,9 @@ xHalIoAssignDriveLetters(IN PLOADER_PARAMETER_BLOCK LoaderBlock,
                     DOSDEVICE_DRIVE_REMOVABLE,
                     0,
                     RtlConvertLongToLargeInteger(0),
-                    hKey);
+                    hKey,
+                    &BootDevice,
+                    NtSystemPath);
             }
         }
     }
@@ -856,6 +899,7 @@ xHalIoAssignDriveLetters(IN PLOADER_PARAMETER_BLOCK LoaderBlock,
             ExFreePool(LayoutArray[i]);
     }
     ExFreePool(LayoutArray);
+end_assign_disks:
 
     /* Assign floppy drives */
     DPRINT("Floppy drives: %d\n", ConfigInfo->FloppyCount);
@@ -875,7 +919,9 @@ xHalIoAssignDriveLetters(IN PLOADER_PARAMETER_BLOCK LoaderBlock,
             DOSDEVICE_DRIVE_REMOVABLE,
             0,
             RtlConvertLongToLargeInteger(0),
-            hKey);
+            hKey,
+            &BootDevice,
+            NtSystemPath);
     }
 
     /* Assign cdrom drives */
@@ -895,7 +941,9 @@ xHalIoAssignDriveLetters(IN PLOADER_PARAMETER_BLOCK LoaderBlock,
             DOSDEVICE_DRIVE_CDROM,
             0,
             RtlConvertLongToLargeInteger(0),
-            hKey);
+            hKey,
+            &BootDevice,
+            NtSystemPath);
     }
 
     /* Anything else to do? */
@@ -1312,6 +1360,24 @@ xHalExamineMBR(IN PDEVICE_OBJECT DeviceObject,
     }
 }
 
+VOID
+NTAPI
+FstubFixupEfiPartition(IN PPARTITION_DESCRIPTOR PartitionDescriptor,
+                       IN ULONGLONG MaxOffset)
+{
+    ULONG PartitionLength;
+    PAGED_CODE();
+
+    /* Compute partition length (according to MBR entry) */
+    PartitionLength = PartitionDescriptor->StartingSectorLsb0 + PartitionDescriptor->PartitionLengthLsb0;
+    /* In case the partition length goes beyond disk size... */
+    if (PartitionLength > MaxOffset)
+    {
+        /* Resize partition to its maximum real length */
+        PartitionDescriptor->PartitionLengthLsb0 = MaxOffset - PartitionDescriptor->StartingSectorLsb0;
+    }
+}
+
 NTSTATUS
 FASTCALL
 xHalIoReadPartitionTable(IN PDEVICE_OBJECT DeviceObject,
@@ -1379,7 +1445,7 @@ xHalIoReadPartitionTable(IN PDEVICE_OBJECT DeviceObject,
             MaxOffset, MaxSector);
 
     /* Allocate our buffer */
-    Buffer = ExAllocatePoolWithTag(NonPagedPool, PAGE_SIZE, TAG_FILE_SYSTEM);
+    Buffer = ExAllocatePoolWithTag(NonPagedPool, InputSize, TAG_FILE_SYSTEM);
     if (!Buffer)
     {
         /* Fail, free the input buffer */
@@ -1457,9 +1523,6 @@ xHalIoReadPartitionTable(IN PDEVICE_OBJECT DeviceObject,
         PartitionDescriptor = (PPARTITION_DESCRIPTOR)
                                &(((PUSHORT)Buffer)[PARTITION_TABLE_OFFSET]);
 
-        /* Get the partition type */
-        PartitionType = PartitionDescriptor->PartitionType;
-
         /* Start looping partitions */
         j++;
         DPRINT("FSTUB: Partition Table %d:\n", j);
@@ -1478,6 +1541,14 @@ xHalIoReadPartitionTable(IN PDEVICE_OBJECT DeviceObject,
                     GET_STARTING_SECTOR(PartitionDescriptor),
                     GET_PARTITION_LENGTH(PartitionDescriptor));
 
+            /* Check whether we're facing a protective MBR */
+            if (PartitionType == EFI_PMBR_OSTYPE_EFI)
+            {
+                /* Partition length might be bigger than disk size */
+                FstubFixupEfiPartition(PartitionDescriptor,
+                                       MaxOffset);
+            }
+
             /* Make sure that the partition is valid, unless it's the first */
             if (!(HalpIsValidPartitionEntry(PartitionDescriptor,
                                             MaxOffset,
@@ -1640,7 +1711,7 @@ xHalIoReadPartitionTable(IN PDEVICE_OBJECT DeviceObject,
         for (Entry = 1; Entry <= 4; Entry++, PartitionDescriptor++)
         {
             /* Check if this is a container partition, since we skipped them */
-            if (IsContainerPartition(PartitionType))
+            if (IsContainerPartition(PartitionDescriptor->PartitionType))
             {
                 /* Get its offset */
                 Offset.QuadPart = VolumeOffset.QuadPart +