[PATCH] disk: efi_loader: Improve disk image detection in efi_bootmgr

[Javier Tia]

Refines the disk image detection mechanism in the EFI boot manager.
Instead of relying on file extensions, the updated code uses the
partition driver lookup function to determine if a downloaded file is a
valid disk image.

The change enhances the robustness of the boot manager by providing a
more accurate disk image detection. It also improves error handling by
adding a cleanup function for temporary block devices.

The partition driver lookup function is also made public to enable its
use in the EFI boot manager. The function documentation is updated
accordingly.

Signed-off-by: Javier Tia <javier.tia at linaro.org>
---
 disk/part.c                  |  3 +-
 include/part.h               | 18 +++++++++++
 lib/efi_loader/efi_bootmgr.c | 59 ++++++++++++++++++++++++++++++------
 3 files changed, 70 insertions(+), 10 deletions(-)

diff --git a/disk/part.c b/disk/part.c
index 66e2b3a7219..adc24976624 100644
--- a/disk/part.c
+++ b/disk/part.c
@@ -63,7 +63,7 @@ static struct part_driver *part_driver_get_type(int part_type)
  * @dev_desc: Device descriptor
  * Return: Driver found, or NULL if none
  */
-static struct part_driver *part_driver_lookup_type(struct blk_desc *desc)
+struct part_driver *part_driver_lookup_type(struct blk_desc *desc)
 {
 	struct part_driver *drv =
 		ll_entry_start(struct part_driver, part_driver);
@@ -855,3 +855,4 @@ int part_get_bootable(struct blk_desc *desc)
 
 	return 0;
 }
+
diff --git a/include/part.h b/include/part.h
index b772fb34c8a..3fc3186f4c4 100644
--- a/include/part.h
+++ b/include/part.h
@@ -727,6 +727,24 @@ int part_get_type_by_name(const char *name);
  */
 int part_get_bootable(struct blk_desc *desc);
 
+/**
+ * part_driver_lookup_type() - Look up the partition driver for a blk device
+ *
+ * If @desc->part_type is PART_TYPE_UNKNOWN, this checks each partition driver
+ * against the blk device to see if there is a valid partition table acceptable
+ * to that driver.
+ *
+ * If @desc->part_type is already set, it just returns the driver for that
+ * type, without testing if the driver can find a valid partition on the
+ * descriptor.
+ *
+ * On success it updates @desc->part_type if set to PART_TYPE_UNKNOWN on entry
+ *
+ * @desc: Device descriptor
+ * Return: Driver found, or NULL if none
+ */
+struct part_driver *part_driver_lookup_type(struct blk_desc *desc);
+
 #else
 static inline int part_driver_get_count(void)
 { return 0; }
diff --git a/lib/efi_loader/efi_bootmgr.c b/lib/efi_loader/efi_bootmgr.c
index 662993fb809..7acb4f4e92b 100644
--- a/lib/efi_loader/efi_bootmgr.c
+++ b/lib/efi_loader/efi_bootmgr.c
@@ -17,6 +17,7 @@
 #include <log.h>
 #include <malloc.h>
 #include <net.h>
+#include <part.h>
 #include <efi_loader.h>
 #include <efi_variable.h>
 #include <asm/unaligned.h>
@@ -407,7 +408,7 @@ static efi_status_t efi_bootmgr_release_uridp(struct uridp_context *ctx)
 	if (!ctx)
 		return ret;
 
-	/* cleanup for iso or img image */
+	/* cleanup for disk image */
 	if (ctx->ramdisk_blk_dev) {
 		ret = efi_add_memory_map(ctx->image_addr, ctx->image_size,
 					 EFI_CONVENTIONAL_MEMORY);
@@ -452,6 +453,20 @@ static void EFIAPI efi_bootmgr_http_return(struct efi_event *event,
 	EFI_EXIT(ret);
 }
 
+/**
+ * cleanup_temp_blkdev() - Clean up temporary block device
+ *
+ * @temp_bdev: temporary block device to clean up
+ */
+static void cleanup_temp_blkdev(struct udevice *temp_bdev)
+{
+	if (!temp_bdev)
+		return;
+
+	if (blkmap_destroy(temp_bdev->parent))
+		log_err("Destroying temporary blkmap failed\n");
+}
+
 /**
  * try_load_from_uri_path() - Handle the URI device path
  *
@@ -466,7 +481,6 @@ static efi_status_t try_load_from_uri_path(struct efi_device_path_uri *uridp,
 {
 	char *s;
 	int err;
-	int uri_len;
 	efi_status_t ret;
 	void *source_buffer;
 	efi_uintn_t source_size;
@@ -516,13 +530,37 @@ static efi_status_t try_load_from_uri_path(struct efi_device_path_uri *uridp,
 	image_size = ALIGN(image_size, SZ_2M);
 
 	/*
-	 * If the file extension is ".iso" or ".img", mount it and try to load
-	 * the default file.
-	 * If the file is PE-COFF image, load the downloaded file.
+	 * Check if the downloaded file is a disk image or PE-COFF image.
+	 * Use partition driver lookup for disk image detection.
 	 */
-	uri_len = strlen(uridp->uri);
-	if (!strncmp(&uridp->uri[uri_len - 4], ".iso", 4) ||
-	    !strncmp(&uridp->uri[uri_len - 4], ".img", 4)) {
+	struct udevice *temp_bdev;
+	struct part_driver *part_drv = NULL;
+
+	/* Create temporary block device from the downloaded image for testing */
+	temp_bdev = mount_image(lo_label, image_addr, image_size);
+	if (!temp_bdev) {
+		log_debug("mount_image() failed, will attempt PE-COFF detection\n");
+	}
+	if (temp_bdev) {
+		struct blk_desc *desc = dev_get_uclass_plat(temp_bdev);
+		if (!desc) {
+			/* Clean up temporary device when dev_get_uclass_plat() returns NULL */
+			cleanup_temp_blkdev(temp_bdev);
+		} else {
+			/* Use part_driver_lookup_type for comprehensive partition detection */
+			part_drv = part_driver_lookup_type(desc);
+
+			/*
+			 * If it's a disk image with valid partition type,
+			 * prepare_loaded_image will create the final block device
+			 */
+			if (part_drv)
+				/* Clean up temporary device, prepare_loaded_image will create a new one */
+				cleanup_temp_blkdev(temp_bdev);
+		}
+	}
+
+	if (part_drv) {
 		ret = prepare_loaded_image(lo_label, image_addr, image_size,
 					   &loaded_dp, &blk);
 		if (ret != EFI_SUCCESS)
@@ -545,7 +583,10 @@ static efi_status_t try_load_from_uri_path(struct efi_device_path_uri *uridp,
 		source_buffer = (void *)image_addr;
 		source_size = image_size;
 	} else {
-		log_err("Error: file type is not supported\n");
+		/* Clean up temporary device if it was created but not a valid disk image */
+		cleanup_temp_blkdev(temp_bdev);
+
+		log_err("Error: disk image is not supported\n");
 		ret = EFI_UNSUPPORTED;
 		goto err;
 	}
-- 
2.50.1