[PATCH 3/3] abootimg: Implement smart image load feature

[Roman Stratiienko]

Load only part of the boot partition that contains valuable
information, thus improving the boot time.

Signed-off-by: Roman Stratiienko <r.stratiienko at gmail.com>
---
 boot/image-android.c | 70 ++++++++++++++++++++++++++++++++++++++++++++
 cmd/abootimg.c       | 40 ++++++++++++++++++++++---
 include/image.h      | 12 ++++++++
 3 files changed, 118 insertions(+), 4 deletions(-)

diff --git a/boot/image-android.c b/boot/image-android.c
index da8003f370..d00a896a40 100644
--- a/boot/image-android.c
+++ b/boot/image-android.c
@@ -204,6 +204,76 @@ bool android_image_get_data(const void *boot_hdr, const void *vendor_boot_hdr,
 	return true;
 }
 
+/**
+ * android_image_get_valuable_size() - get the size of the android image
+ *
+ * This function checks if the image is Android boot image and returns the
+ * valuable size of the image.
+ *
+ * @hdr_addr: Boot image header address (boot or vendor_boot)
+ *
+ * @return size of the image on success, 0 on failure
+ */
+size_t android_image_get_valuable_size(const void *hdr_addr)
+{
+	int version, size;
+
+	if (is_android_boot_image_header(hdr_addr)) {
+		const struct andr_boot_img_hdr_v0 *hdr = hdr_addr;
+
+		version = ((struct andr_boot_img_hdr_v0 *)hdr_addr)->header_version;
+		if (version > 2) {
+			const struct andr_boot_img_hdr_v3 *hdr = hdr_addr;
+
+			size = ALIGN(hdr->header_size, ANDR_GKI_PAGE_SIZE);
+			size += ALIGN(hdr->kernel_size, ANDR_GKI_PAGE_SIZE);
+			size += ALIGN(hdr->ramdisk_size, ANDR_GKI_PAGE_SIZE);
+
+			if (version > 3)
+				size += ALIGN(hdr->signature_size, ANDR_GKI_PAGE_SIZE);
+
+			return size;
+		}
+
+		size = hdr->page_size;
+		size += ALIGN(hdr->kernel_size, hdr->page_size);
+		size += ALIGN(hdr->ramdisk_size, hdr->page_size);
+		size += ALIGN(hdr->second_size, hdr->page_size);
+
+		if (version > 0)
+			size += ALIGN(hdr->recovery_dtbo_size, hdr->page_size);
+
+		if (version > 1)
+			size += ALIGN(hdr->dtb_size, hdr->page_size);
+
+		return size;
+	}
+
+	if (is_android_vendor_boot_image_header(hdr_addr)) {
+		const struct andr_vnd_boot_img_hdr *hdr = hdr_addr;
+
+		version = ((struct andr_vnd_boot_img_hdr *)hdr_addr)->header_version;
+		if (version < 3) {
+			printf("Vendor boot image header version %d is not supported\n", version);
+
+			return 0;
+		}
+
+		size = ALIGN(hdr->header_size, hdr->page_size);
+		size += ALIGN(hdr->vendor_ramdisk_size, hdr->page_size);
+		size += ALIGN(hdr->dtb_size, hdr->page_size);
+
+		if (version > 3) {
+			size += ALIGN(hdr->vendor_ramdisk_table_size, hdr->page_size);
+			size += ALIGN(hdr->bootconfig_size, hdr->page_size);
+		}
+
+		return size;
+	}
+
+	return 0;
+}
+
 static ulong android_image_get_kernel_addr(struct andr_image_data *img_data)
 {
 	/*
diff --git a/cmd/abootimg.c b/cmd/abootimg.c
index 808c9c4941..fe7c5c5e2c 100644
--- a/cmd/abootimg.c
+++ b/cmd/abootimg.c
@@ -182,6 +182,35 @@ static int abootimg_get_dtb(int argc, char *const argv[])
 	return CMD_RET_USAGE;
 }
 
+static int abootimg_smart_load(struct blk_desc *desc, struct disk_partition *info, void *addr)
+{
+	int ret, size;
+
+	ret = blk_dread(desc, info->start, 4096 / info->blksz, addr);
+	if (ret < 0) {
+		printf("Error: Failed to read partition\n");
+		return CMD_RET_FAILURE;
+	}
+
+	size = android_image_get_valuable_size(addr);
+	if (size == 0)
+		return 0;
+
+	ret = blk_dread(desc, info->start, DIV_ROUND_UP(size, info->blksz), addr);
+	if (ret < 0) {
+		printf("Error: Failed to read partition\n");
+		return CMD_RET_FAILURE;
+	}
+
+	memset(addr + size, 0, info->size * info->blksz - size);
+
+	printf("Loaded Android boot image using smart load (%d/%d MB)\n",
+	       (int)DIV_ROUND_UP(size, 1024 * 1024),
+	       (int)DIV_ROUND_UP(info->size * info->blksz, 1024 * 1024));
+
+	return size;
+}
+
 static int do_abootimg_addr(struct cmd_tbl *cmdtp, int flag, int argc,
 			    char *const argv[])
 {
@@ -299,10 +328,13 @@ static int do_abootimg_load(struct cmd_tbl *cmdtp, int flag, int argc,
 		return CMD_RET_FAILURE;
 	}
 
-	ret = blk_dread(desc, info.start, info.size, addr);
-	if (ret < 0) {
-		printf("Error: Failed to read partition %s\n", buf);
-		goto fail;
+	ret = abootimg_smart_load(desc, &info, addr);
+	if (ret <= 0) {
+		ret = blk_dread(desc, info.start, info.size, addr);
+		if (ret < 0) {
+			printf("Error: Failed to read partition %s\n", buf);
+			goto fail;
+		}
 	}
 
 	sprintf(buf, "abootimg_%s_ptr", argv[3]);
diff --git a/include/image.h b/include/image.h
index c5b288f62b..7d8ff40c3f 100644
--- a/include/image.h
+++ b/include/image.h
@@ -1826,6 +1826,18 @@ struct andr_image_data;
 bool android_image_get_data(const void *boot_hdr, const void *vendor_boot_hdr,
 			    struct andr_image_data *data);
 
+/**
+ * android_image_get_valuable_size() - get the size of the android image
+ *
+ * This function checks if the image is Android boot image and returns the
+ * valuable size of the image.
+ *
+ * @hdr_addr: Boot image header address (boot or vendor_boot)
+ *
+ * @return size of the image on success, 0 on failure
+ */
+size_t android_image_get_valuable_size(const void *hdr_addr);
+
 struct andr_boot_img_hdr_v0;
 
 /**
-- 
2.40.1