efi/multifs: correctly enumerate partitions per diskmultifs

The initial multifs support for EFI was saving all device handles which
are logical partitions regardless which disk they belong to, so indexing
disks other than 0 (or even partitions outside disk 0) would not work.

This patch fixes enumeration of all logical partitions per disk thus
making possible to index both disk and partition correctly.

Signed-off-by: Paulo Alcantara <pcacjr@zytor.com>

1 files changed, 159 insertions, 54 deletions

diff --git a/efi/multifs.c b/efi/multifs.c
index 2f217d9d..9c4b7510 100644
--- a/efi/multifs.c
+++ b/efi/multifs.c
@@ -21,80 +21,186 @@
 
 #include "efi.h"
 
-static EFI_HANDLE *logical_parts = NULL;
-static unsigned int logical_parts_no = 0;
+#define DISKS_MAX 0xff
+#define PARTS_MAX 0xff
 
-/* Find all device handles which support EFI_BLOCK_IO_PROTOCOL and are logical
- * partitions */
-static EFI_STATUS find_all_logical_parts(void)
+struct efi_disk_entry {
+    uint8_t sig_type;
+    union {
+        uint32_t id;
+        EFI_GUID guid;
+    };
+    EFI_HANDLE parts[PARTS_MAX];
+};
+
+static uint8_t disks_no = 0;
+static struct efi_disk_entry efi_disks[DISKS_MAX] = { 0, };
+
+/* Find all device handles that support EFI_BLOCK_IO_PROTOCOL */
+static EFI_STATUS find_all_block_devs(EFI_HANDLE **bdevs, unsigned long *len)
 {
     EFI_STATUS status;
-    unsigned long len = 0;
-    EFI_HANDLE *handles = NULL;
-    unsigned long i;
-    EFI_BLOCK_IO *bio;
-
-    if (logical_parts) {
-        status = EFI_SUCCESS;
-        goto out;
-    }
 
+    *len = 0;
     status = uefi_call_wrapper(BS->LocateHandle, 5, ByProtocol,
-                               &BlockIoProtocol, NULL, &len, NULL);
-    if (EFI_ERROR(status) && status != EFI_BUFFER_TOO_SMALL)
+                               &BlockIoProtocol, NULL, len, NULL);
+    if (EFI_ERROR(status) && status != EFI_BUFFER_TOO_SMALL) {
+        Print(L"%a: failed to locate BlockIo handles\n", __func__);
         goto out;
+    }
 
-    handles = malloc(len);
-    if (!handles) {
+    *bdevs = malloc(*len);
+    if (!*bdevs) {
+        Print(L"%a: malloc failed\n", __func__);
         status = EFI_OUT_OF_RESOURCES;
         goto out;
     }
 
-    logical_parts = malloc(len);
-    if (!logical_parts) {
-        status = EFI_OUT_OF_RESOURCES;
+    status = uefi_call_wrapper(BS->LocateHandle, 5, ByProtocol,
+                               &BlockIoProtocol, NULL, len,
+                               (void **)*bdevs);
+    if (EFI_ERROR(status)) {
+        Print(L"%a: failed to locate BlockIo handles\n", __func__);
         goto out_free;
     }
+    return status;
 
-    status = uefi_call_wrapper(BS->LocateHandle, 5, ByProtocol,
-                               &BlockIoProtocol, NULL, &len,
-                               (void **)handles);
-    if (EFI_ERROR(status))
-        goto out_free;
+out_free:
+    free(*bdevs);
+out:
+    return status;
+}
 
-    for (i = 0; i < len / sizeof(EFI_HANDLE); i++) {
-        status = uefi_call_wrapper(BS->HandleProtocol, 3, handles[i],
-                                   &BlockIoProtocol, (void **)&bio);
-        if (EFI_ERROR(status))
-            goto out_free;
-        if (bio->Media->LogicalPartition) {
-            logical_parts[logical_parts_no++] = handles[i];
+static inline PCI_DEVICE_PATH *get_pci_dev_path(EFI_HANDLE dev)
+{
+    EFI_DEVICE_PATH *path;
+
+    for (path = DevicePathFromHandle(dev); !IsDevicePathEndType(path);
+         path = NextDevicePathNode(path)) {
+        if (DevicePathType(path) == HARDWARE_DEVICE_PATH &&
+            DevicePathSubType(path) == HW_PCI_DP) {
+            return (PCI_DEVICE_PATH *)path;
         }
     }
+    return NULL;
+}
 
-    free(handles);
-    return status;
+static inline HARDDRIVE_DEVICE_PATH *get_hd_dev_path(EFI_HANDLE dev)
+{
+    EFI_DEVICE_PATH *path;
 
-out_free:
-    if (handles)
-        free(handles);
-    if (logical_parts)
-        free(logical_parts);
-out:
+    for (path = DevicePathFromHandle(dev); !IsDevicePathEndType(path);
+         path = NextDevicePathNode(path)) {
+        if (DevicePathType(path) == MEDIA_DEVICE_PATH &&
+            DevicePathSubType(path) == MEDIA_HARDDRIVE_DP) {
+            return (HARDDRIVE_DEVICE_PATH *)path;
+        }
+    }
+    return NULL;
+}
+
+static int set_efi_disk_info(uint8_t di, EFI_HANDLE dev)
+{
+    HARDDRIVE_DEVICE_PATH *hd;
+
+    hd = get_hd_dev_path(dev);
+    if (!hd)
+        return -1;
+
+    switch (hd->SignatureType) {
+    case SIGNATURE_TYPE_MBR:
+        efi_disks[di].id = *((uint32_t *)&hd->Signature[0]);
+        break;
+    case SIGNATURE_TYPE_GUID:
+        memcpy(&efi_disks[di].guid, &hd->Signature[0],
+               sizeof(efi_disks[di].guid));
+        break;
+    }
+    efi_disks[di].sig_type = hd->SignatureType;
+    return 0;
+}
+
+static EFI_STATUS find_rem_parts(uint8_t di, EFI_HANDLE *bdevs,
+                                 unsigned long *bi, unsigned long bdevsno)
+{
+    unsigned long i = *bi + 1;
+    EFI_STATUS status = EFI_SUCCESS;
+    EFI_BLOCK_IO *bio;
+    PCI_DEVICE_PATH *dev_pci = get_pci_dev_path(bdevs[*bi]);
+    PCI_DEVICE_PATH *pci;
+    HARDDRIVE_DEVICE_PATH *hd;
+    uint8_t pi = 1;
+
+    while (i < bdevsno) {
+        pci = get_pci_dev_path(bdevs[i]);
+        if (dev_pci->Function != pci->Function &&
+            dev_pci->Device != pci->Device) {
+            break;
+        }
+        status = uefi_call_wrapper(BS->HandleProtocol, 3, bdevs[i],
+                                   &BlockIoProtocol, (void **)&bio);
+        if (EFI_ERROR(status)) {
+            Print(L"%a: failed to find BlockIO protocol: %r\n", __func__,
+                  status);
+            break;
+        }
+        if (!bio->Media->LogicalPartition)
+            break;
+        hd = get_hd_dev_path(bdevs[i]);
+        if (hd->SignatureType != efi_disks[di].sig_type)
+            break;
+
+        efi_disks[di].parts[pi++] = bdevs[i++];
+    }
+    *bi = i - 1;
     return status;
 }
 
-static inline EFI_HANDLE get_logical_part(unsigned int partno)
+static EFI_STATUS find_all_partitions(void)
 {
-    if (!logical_parts || partno > logical_parts_no)
-        return NULL;
-    return logical_parts[partno - 1];
+    EFI_STATUS status;
+    EFI_HANDLE *bdevs;
+    unsigned long len;
+    unsigned long i;
+    EFI_BLOCK_IO *bio;
+    int ret;
+
+    status = find_all_block_devs(&bdevs, &len);
+    if (EFI_ERROR(status)) {
+        Print(L"%a: failed to find block devices: %r\n", __func__, status);
+        return status;
+    }
+
+    for (i = 0; i < len / sizeof(EFI_HANDLE); i++) {
+        status = uefi_call_wrapper(BS->HandleProtocol, 3, bdevs[i],
+                                   &BlockIoProtocol, (void **)&bio);
+        if (EFI_ERROR(status))
+            break;
+        if (!bio->Media->LogicalPartition)
+            continue;
+        ret = set_efi_disk_info(disks_no, bdevs[i]);
+        if (ret) {
+            Print(L"%a: failed to set EFI disk info\n", __func__);
+            status = EFI_INVALID_PARAMETER;
+            break;
+        }
+        efi_disks[disks_no].parts[0] = bdevs[i]; /* first partition */
+        status = find_rem_parts(disks_no, bdevs, &i, len / sizeof(EFI_HANDLE));
+        if (EFI_ERROR(status)) {
+            Print(L"%a: failed to find partitions: %r\n", __func__, status);
+            break;
+        }
+        disks_no++;
+    }
+    free(bdevs);
+    return status;
 }
 
-static inline EFI_HANDLE find_device_handle(unsigned int diskno,
-                                            unsigned int partno)
+static inline EFI_HANDLE find_device_handle(uint8_t di, uint8_t pi)
 {
-    return get_logical_part(partno);
+    if (di >= disks_no || pi > PARTS_MAX)
+        return NULL;
+    return efi_disks[di].parts[pi - 1];
 }
 
 static inline void *get_dev_info_priv(EFI_HANDLE handle)
@@ -131,7 +237,7 @@ __export struct fs_info *efi_multifs_get_fs_info(const char **path)
     handle = find_device_handle(diskno, partno);
     if (!handle)
         goto free_fsp;
-    dprintf("%s: found partition %d\n", __func__, partno);
+    dprintf("%s: found disk %u part %u\n", __func__, diskno, partno);
 
     priv = get_dev_info_priv(handle);
     if (!priv)
@@ -139,9 +245,10 @@ __export struct fs_info *efi_multifs_get_fs_info(const char **path)
 
     ret = multifs_setup_fs_info(fsp, diskno, partno, priv);
     if (ret) {
-        dprintf("%s: failed to set up fs info\n", __func__);
+        Print(L"%a: failed to set up fs info\n", __func__);
         goto free_priv;
     }
+
     return fsp;
 
 free_priv:
@@ -155,11 +262,9 @@ __export void efi_multifs_init(void *addr __attribute__((unused)))
 {
     EFI_STATUS status;
 
-    status = find_all_logical_parts();
+    status = find_all_partitions();
     if (EFI_ERROR(status)) {
-        printf("%s: failed to locate device handles of logical partitions\n",
-               __func__);
-        printf("%s: EFI status code: 0x%08X\n", __func__, status);
+        Print(L"%a: failed to find disk partitions: %r\n", __func__, status);
         return;
     }
     dprintf("%s: initialised multifs support\n", __func__);