[FREELDR] Improve disk detection

Author: iLauncherDev

boot/freeldr/freeldr/arch/uefi/uefidisk.c

@@ -58,7 +58,9 @@ static ULONG OffsetToBoot;
 static ULONG PublicBootArcDisk;
 static INTERNAL_UEFI_DISK* InternalUefiDisk = NULL;
 static EFI_GUID bioGuid = BLOCK_IO_PROTOCOL;
+static EFI_GUID loadedImageGuid = EFI_LOADED_IMAGE_PROTOCOL_GUID;
 static EFI_BLOCK_IO* bio;
+static EFI_LOADED_IMAGE_PROTOCOL *loadedImage;
 static EFI_HANDLE* handles = NULL;
 
 /* FUNCTIONS *****************************************************************/
@@ -434,28 +436,41 @@ UefiSetupBlockDevices(VOID)
             ULONG increment = 0;
             ULONG i;
 
-            /* 3) Grab the offset into the array of handles and decrement per volume (valid partition) */
-            for (increment = OffsetToBoot; increment > 0; increment--)
+            Status = GlobalSystemTable->BootServices->HandleProtocol(
+                GlobalImageHandle,
+                &loadedImageGuid,
+                (void**)&loadedImage
+            );
+            if (EFI_ERROR(Status))
+                continue;
+
+            for (increment = 0; increment < SystemHandleCount; ++increment) {
+                if (handles[increment] == loadedImage->DeviceHandle) {
+                    break;
+                }
+            }
+
+            if (SystemHandleCount == increment)
+                continue;
+
+            GlobalSystemTable->BootServices->HandleProtocol(handles[increment], &bioGuid, (void**)&bio);
+            if (bio->Media->LogicalPartition == FALSE)
             {
-                GlobalSystemTable->BootServices->HandleProtocol(handles[increment], &bioGuid, (void**)&bio);
-                if (bio->Media->LogicalPartition == FALSE)
-                {
-                    TRACE("Found root at increment %u\n", increment);
-                    UefiBootRootIdentifier = increment;
+                TRACE("Found root at increment %u\n", increment);
+                UefiBootRootIdentifier = increment;
 
-                    for (i = 0; i <= PcBiosDiskCount; ++i)
+                for (i = 0; i <= PcBiosDiskCount; ++i)
+                {
+                    /* Now only of the root drive number is equal to this drive we found above */
+                    if (InternalUefiDisk[i].UefiRootNumber == UefiBootRootIdentifier)
                     {
-                        /* Now only of the root drive number is equal to this drive we found above */
-                        if (InternalUefiDisk[i].UefiRootNumber == UefiBootRootIdentifier)
-                        {
-                            InternalUefiDisk[i].IsThisTheBootDrive = TRUE;
-                            PublicBootArcDisk = i;
-                            TRACE("Found Boot drive\n");
-                        }
+                        InternalUefiDisk[i].IsThisTheBootDrive = TRUE;
+                        PublicBootArcDisk = i;
+                        TRACE("Found Boot drive\n");
                     }
-
-                    break;
                 }
+
+                break;
             }
         }
     }
@@ -499,49 +514,51 @@ BOOLEAN
 UefiInitializeBootDevices(VOID)
 {
     ULONG i = 0;
+    PMASTER_BOOT_RECORD Mbr;
+    PULONG Buffer;
+    ULONG ChecksumSize = 512;
+    ULONG Checksum = 0;
+    ULONG Signature;
 
     DiskReadBufferSize = EFI_PAGE_SIZE;
     DiskReadBuffer = MmAllocateMemoryWithType(DiskReadBufferSize, LoaderFirmwareTemporary);
     UefiSetupBlockDevices();
     UefiSetBootpath();
 
-    /* Add it, if it's a cdrom */
+    /* Add it */
     GlobalSystemTable->BootServices->HandleProtocol(handles[UefiBootRootIdentifier], &bioGuid, (void**)&bio);
-    if (bio->Media->RemovableMedia == TRUE && bio->Media->BlockSize == 2048)
+
+    /* Read the MBR */
+    if (!MachDiskReadLogicalSectors(FrldrBootDrive, 16ULL, 1, DiskReadBuffer))
     {
-        PMASTER_BOOT_RECORD Mbr;
-        PULONG Buffer;
-        ULONG Checksum = 0;
-        ULONG Signature;
+        ERR("Reading MBR failed\n");
+        return FALSE;
+    }
 
-        /* Read the MBR */
-        if (!MachDiskReadLogicalSectors(FrldrBootDrive, 16ULL, 1, DiskReadBuffer))
-        {
-            ERR("Reading MBR failed\n");
-            return FALSE;
-        }
+    Buffer = (ULONG*)DiskReadBuffer;
+    Mbr = (PMASTER_BOOT_RECORD)DiskReadBuffer;
 
-        Buffer = (ULONG*)DiskReadBuffer;
-        Mbr = (PMASTER_BOOT_RECORD)DiskReadBuffer;
+    Signature = Mbr->Signature;
+    TRACE("Signature: %x\n", Signature);
 
-        Signature = Mbr->Signature;
-        TRACE("Signature: %x\n", Signature);
+    if (bio->Media->BlockSize == 2048)
+        ChecksumSize = bio->Media->BlockSize;
 
-        /* Calculate the MBR checksum */
-        for (i = 0; i < 2048 / sizeof(ULONG); i++)
-        {
-            Checksum += Buffer[i];
-        }
-        Checksum = ~Checksum + 1;
-        TRACE("Checksum: %x\n", Checksum);
+    /* Calculate the MBR checksum */
+    for (i = 0; i < ChecksumSize / sizeof(ULONG); i++)
+    {
+        Checksum += Buffer[i];
+    }
+    Checksum = ~Checksum + 1;
+    TRACE("Checksum: %x\n", Checksum);
+
+    /* Fill out the ARC disk block */
+    AddReactOSArcDiskInfo(FrLdrBootPath, Signature, Checksum, TRUE);
 
-        /* Fill out the ARC disk block */
-        AddReactOSArcDiskInfo(FrLdrBootPath, Signature, Checksum, TRUE);
+    FsRegisterDevice(FrLdrBootPath, &UefiDiskVtbl);
+    PcBiosDiskCount++; // This is not accounted for in the number of pre-enumerated BIOS drives!
+    TRACE("Additional boot drive detected: 0x%02X\n", (int)FrldrBootDrive);
 
-        FsRegisterDevice(FrLdrBootPath, &UefiDiskVtbl);
-        PcBiosDiskCount++; // This is not accounted for in the number of pre-enumerated BIOS drives!
-        TRACE("Additional boot drive detected: 0x%02X\n", (int)FrldrBootDrive);
-    }
     return TRUE;
 }