[PATCH] disk: efi_loader: Add function to detect disk image type

[Javier Tia]

This commit introduces a new function, detect_disk_image_type(), which
checks if a buffer contains a disk image. This function is used to
detect disk images based on their headers, rather than relying on file
extensions.

The function checks for ISO 9660, DOS MBR, and GPT protective MBR
signatures. It returns 1 if a disk image is detected, and 0 otherwise.

This function is then used in the efi_bootmgr.c file to determine if a
downloaded file is a PE-COFF image or a disk image. This change
improves the reliability of disk image detection, as it no longer
relies on file extensions, which can be easily changed or omitted.

Signed-off-by: Javier Tia <javier.tia at linaro.org>
---
 disk/part.c                  | 46 ++++++++++++++++++++++++++++++++++++
 include/part.h               |  9 +++++++
 lib/efi_loader/efi_bootmgr.c | 29 ++++++++++-------------
 3 files changed, 68 insertions(+), 16 deletions(-)

diff --git a/disk/part.c b/disk/part.c
index 66e2b3a7219..9d9483b216c 100644
--- a/disk/part.c
+++ b/disk/part.c
@@ -855,3 +855,49 @@ int part_get_bootable(struct blk_desc *desc)
 
 	return 0;
 }
+
+/**
+* detect_disk_image_type() - Detect if buffer contains a disk image
+*
+* @buffer:  Buffer containing potential disk image data
+* @size:    Size of the buffer
+* Return:   1 if disk image detected, 0 otherwise
+*/
+int detect_disk_image_type(void *buffer, size_t size)
+{
+	unsigned char *buf = (unsigned char *)buffer;
+
+	/* Need at least 512 bytes for basic checks */
+	if (!buffer || size < 512)
+		return 0;
+
+#if IS_ENABLED(CONFIG_ISO_PARTITION)
+	/* Check for ISO 9660 signature "CD001" at offset 0x8001 */
+	if (size > 0x8005) {
+		if (buf[0x8001] == 'C' && buf[0x8002] == 'D' &&
+		    buf[0x8003] == '0' && buf[0x8004] == '0' &&
+		    buf[0x8005] == '1') {
+			return 1;
+		}
+	}
+#endif
+
+#if IS_ENABLED(CONFIG_DOS_PARTITION)
+	/* Check for DOS MBR signature 0x55AA at offset 0x1FE */
+	if (buf[0x1FE] == 0x55 && buf[0x1FF] == 0xAA) {
+		return 1;
+	}
+#endif
+
+#if IS_ENABLED(CONFIG_EFI_PARTITION)
+	/* Check for GPT protective MBR */
+	if (buf[0x1FE] == 0x55 && buf[0x1FF] == 0xAA) {
+		/* Check if partition type is 0xEE (GPT protective) */
+		if (buf[0x1C2] == 0xEE) {
+			return 1;
+		}
+	}
+#endif
+
+	return 0;
+}
diff --git a/include/part.h b/include/part.h
index b772fb34c8a..39c04e0c8f7 100644
--- a/include/part.h
+++ b/include/part.h
@@ -727,6 +727,15 @@ int part_get_type_by_name(const char *name);
  */
 int part_get_bootable(struct blk_desc *desc);
 
+/**
+ * detect_disk_image_type() - Detect if buffer contains a disk image
+ *
+ * @buffer: Buffer containing potential disk image data
+ * @size:   Size of the buffer
+ * Return:  1 if disk image detected, 0 otherwise
+ */
+int detect_disk_image_type(void *buffer, size_t size);
+
 #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..da63ad41b40 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 image file */
 	if (ctx->ramdisk_blk_dev) {
 		ret = efi_add_memory_map(ctx->image_addr, ctx->image_size,
 					 EFI_CONVENTIONAL_MEMORY);
@@ -466,7 +467,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,21 +516,10 @@ 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 PE-COFF image or a disk image.
+	 * Use header-based detection instead of file extensions.
 	 */
-	uri_len = strlen(uridp->uri);
-	if (!strncmp(&uridp->uri[uri_len - 4], ".iso", 4) ||
-	    !strncmp(&uridp->uri[uri_len - 4], ".img", 4)) {
-		ret = prepare_loaded_image(lo_label, image_addr, image_size,
-					   &loaded_dp, &blk);
-		if (ret != EFI_SUCCESS)
-			goto err;
-
-		source_buffer = NULL;
-		source_size = 0;
-	} else if (efi_check_pe((void *)image_addr, image_size, NULL) == EFI_SUCCESS) {
+	if (efi_check_pe((void *)image_addr, image_size, NULL) == EFI_SUCCESS) {
 		/*
 		 * loaded_dp must exist until efi application returns,
 		 * will be freed in return_to_efibootmgr event callback.
@@ -544,6 +533,14 @@ 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 if (detect_disk_image_type((void *)image_addr, image_size)) {
+		ret = prepare_loaded_image(lo_label, image_addr, image_size,
+					   &loaded_dp, &blk);
+		if (ret != EFI_SUCCESS)
+			goto err;
+
+		source_buffer = NULL;
+		source_size = 0;
 	} else {
 		log_err("Error: file type is not supported\n");
 		ret = EFI_UNSUPPORTED;
-- 
2.50.1