From 2d1ef29a25ee9006800ae6242ffba4b2e5c7a909 Mon Sep 17 00:00:00 2001
From: Tianling Shen <cnsztl@immortalwrt.org>
Date: Wed, 13 Aug 2025 02:56:46 +0800
Subject: [PATCH] uboot-mtk-20250711: add fit_get_totalsize from OpenWrt

Signed-off-by: Tianling Shen <cnsztl@immortalwrt.org>
---
 uboot-mtk-20250711/boot/image-fit.c | 41 +++++++++++++++++
 uboot-mtk-20250711/cmd/bootm.c      | 70 +++++++++++++++++++++++++++++
 uboot-mtk-20250711/include/image.h  |  1 +
 3 files changed, 112 insertions(+)

diff --git a/uboot-mtk-20250711/boot/image-fit.c b/uboot-mtk-20250711/boot/image-fit.c
index d8ab4d963..d292cfe03 100644
--- a/uboot-mtk-20250711/boot/image-fit.c
+++ b/uboot-mtk-20250711/boot/image-fit.c
@@ -2071,6 +2071,47 @@ static const char *fit_get_image_type_property(int ph_type)
 	return "unknown";
 }
 
+size_t fit_get_totalsize(const void *fit)
+{
+	int ret, ndepth, noffset, images_noffset;
+	size_t data_size, hdrsize, img_total, max_size = 0;
+	const void *data;
+
+	ret = fdt_check_header(fit);
+	if (ret) {
+		debug("Wrong FIT format: not a flattened device tree (err=%d)\n",
+			  ret);
+		return 0;
+	}
+
+	hdrsize = fdt_totalsize(fit);
+
+	/* take care of simple FIT with internal images */
+	max_size = hdrsize;
+
+	images_noffset = fdt_path_offset(fit, FIT_IMAGES_PATH);
+	if (images_noffset < 0)
+		goto out;
+
+	for (ndepth = 0,
+	     noffset = fdt_next_node(fit, images_noffset, &ndepth);
+	     (noffset >= 0) && (ndepth > 0);
+	     noffset = fdt_next_node(fit, noffset, &ndepth)) {
+		if (ndepth == 1) {
+			ret = fit_image_get_data(fit, noffset, &data, &data_size);
+			if (ret)
+				goto out;
+
+			img_total = data_size + (data - fit);
+
+			max_size = (max_size > img_total) ? max_size : img_total;
+		}
+	}
+
+out:
+	return max_size;
+}
+
 int fit_image_load(struct bootm_headers *images, ulong addr,
 		   const char **fit_unamep, const char **fit_uname_configp,
 		   int arch, int ph_type, int bootstage_id,
diff --git a/uboot-mtk-20250711/cmd/bootm.c b/uboot-mtk-20250711/cmd/bootm.c
index bee683d05..b8100bfa9 100644
--- a/uboot-mtk-20250711/cmd/bootm.c
+++ b/uboot-mtk-20250711/cmd/bootm.c
@@ -260,6 +260,76 @@ U_BOOT_CMD(
 /* iminfo - print header info for a requested image */
 /*******************************************************************/
 #if defined(CONFIG_CMD_IMI)
+#define SECTOR_SHIFT 9
+static int image_totalsize(struct cmd_tbl *cmdtp, int flag, int argc,
+			   char *const argv[], short int in_blocks)
+{
+	ulong addr;
+	void *hdr;
+	uint32_t bsize, tsize = 0;
+	char buf[16];
+
+	if (argc >= 2)
+		addr = simple_strtoul(argv[1], NULL, 16);
+	else
+		addr = image_load_addr;
+
+	hdr = (void *)map_sysmem(addr, 0);
+
+	switch (genimg_get_format(hdr)) {
+	case IMAGE_FORMAT_LEGACY:
+		if(CONFIG_IS_ENABLED(LEGACY_IMAGE_FORMAT))
+			tsize = image_get_image_size(hdr);
+		break;
+	case IMAGE_FORMAT_FIT:
+		if(CONFIG_IS_ENABLED(FIT))
+			tsize = fit_get_totalsize(hdr);
+		break;
+	}
+
+	unmap_sysmem(hdr);
+	if (tsize == 0)
+		return 1;
+
+	bsize = (tsize >> SECTOR_SHIFT) + ((tsize & ((1 << SECTOR_SHIFT) - 1))?1:0);
+
+	if (!in_blocks)
+		snprintf(buf, sizeof(buf), "%x", tsize);
+	else
+		snprintf(buf, sizeof(buf), "%x", bsize);
+
+	if (argc >= 3)
+		return env_set(argv[2], buf);
+	else
+		printf("%s\n", buf);
+
+	return 0;
+}
+
+static int do_imsz(struct cmd_tbl *cmdtp, int flag, int argc,
+		     char *const argv[])
+{
+	return image_totalsize(cmdtp, flag, argc, argv, 0);
+}
+
+static int do_imszb(struct cmd_tbl *cmdtp, int flag, int argc,
+		     char *const argv[])
+{
+	return image_totalsize(cmdtp, flag, argc, argv, 1);
+}
+
+U_BOOT_CMD(
+	imsz,	CONFIG_SYS_MAXARGS,	1,	do_imsz,
+	"get image total size (in bytes)",
+	"addr [maxhdrlen] [varname]\n"
+);
+
+U_BOOT_CMD(
+	imszb,	CONFIG_SYS_MAXARGS,	1,	do_imszb,
+	"get image total size (in blocks)",
+	"addr [maxhdrlen] [varname]\n"
+);
+
 static int do_iminfo(struct cmd_tbl *cmdtp, int flag, int argc,
 		     char *const argv[])
 {
diff --git a/uboot-mtk-20250711/include/image.h b/uboot-mtk-20250711/include/image.h
index b26382237..1d9542b06 100644
--- a/uboot-mtk-20250711/include/image.h
+++ b/uboot-mtk-20250711/include/image.h
@@ -1117,6 +1117,7 @@ int fit_parse_subimage(const char *spec, ulong addr_curr,
 		ulong *addr, const char **image_name);
 
 int fit_get_subimage_count(const void *fit, int images_noffset);
+size_t fit_get_totalsize(const void *fit);
 void fit_print_contents(const void *fit);
 void fit_image_print(const void *fit, int noffset, const char *p);