abootimg: Add init_boot image support

Quote from [1]:

"For devices launching with Android 13, the generic ramdisk is removed
from the boot image and placed in a separate init_boot image.
This change leaves the boot image with only the GKI kernel."

While at it, update wrong error handling message when vendor_boot
cannot be loaded.

[1]: https://source.android.com/docs/core/architecture/partitions/generic-boot

Signed-off-by: Roman Stratiienko <r.stratiienko@gmail.com>
Reviewed-by: Mattijs Korpershoek <mkorpershoek@baylibre.com>
diff --git a/cmd/abootimg.c b/cmd/abootimg.c
index 88c77d9..327712a 100644
--- a/cmd/abootimg.c
+++ b/cmd/abootimg.c
@@ -14,6 +14,7 @@
 
 /* Please use abootimg_addr() macro to obtain the boot image address */
 static ulong _abootimg_addr = -1;
+static ulong _ainit_bootimg_addr = -1;
 static ulong _avendor_bootimg_addr = -1;
 
 ulong get_abootimg_addr(void)
@@ -21,6 +22,11 @@
 	return (_abootimg_addr == -1 ? image_load_addr : _abootimg_addr);
 }
 
+ulong get_ainit_bootimg_addr(void)
+{
+	return _ainit_bootimg_addr;
+}
+
 ulong get_avendor_bootimg_addr(void)
 {
 	return _avendor_bootimg_addr;
@@ -179,7 +185,7 @@
 	char *endp;
 	ulong img_addr;
 
-	if (argc < 2 || argc > 3)
+	if (argc < 2 || argc > 4)
 		return CMD_RET_USAGE;
 
 	img_addr = hextoul(argv[1], &endp);
@@ -190,16 +196,26 @@
 
 	_abootimg_addr = img_addr;
 
-	if (argc == 3) {
+	if (argc > 2) {
 		img_addr = simple_strtoul(argv[2], &endp, 16);
 		if (*endp != '\0') {
-			printf("Error: Wrong vendor image address\n");
+			printf("Error: Wrong vendor_boot image address\n");
 			return CMD_RET_FAILURE;
 		}
 
 		_avendor_bootimg_addr = img_addr;
 	}
 
+	if (argc == 4) {
+		img_addr = simple_strtoul(argv[3], &endp, 16);
+		if (*endp != '\0') {
+			printf("Error: Wrong init_boot image address\n");
+			return CMD_RET_FAILURE;
+		}
+
+		_ainit_bootimg_addr = img_addr;
+	}
+
 	return CMD_RET_SUCCESS;
 }
 
@@ -243,7 +259,7 @@
 }
 
 static struct cmd_tbl cmd_abootimg_sub[] = {
-	U_BOOT_CMD_MKENT(addr, 3, 1, do_abootimg_addr, "", ""),
+	U_BOOT_CMD_MKENT(addr, 4, 1, do_abootimg_addr, "", ""),
 	U_BOOT_CMD_MKENT(dump, 2, 1, do_abootimg_dump, "", ""),
 	U_BOOT_CMD_MKENT(get, 5, 1, do_abootimg_get, "", ""),
 };
@@ -271,7 +287,7 @@
 U_BOOT_CMD(
 	abootimg, CONFIG_SYS_MAXARGS, 0, do_abootimg,
 	"manipulate Android Boot Image",
-	"addr <boot_img_addr> [<vendor_boot_img_addr>]>\n"
+	"addr <boot_img_addr> [<vendor_boot_img_addr> [<init_boot_img_addr>]]\n"
 	"    - set the address in RAM where boot image is located\n"
 	"      ($loadaddr is used by default)\n"
 	"abootimg dump dtb\n"