From ea91031b223223667b8dbbaec7fa7bbb51f96a0b Mon Sep 17 00:00:00 2001
From: Parthiban Nallathambi <pn@denx.de>
Date: Wed, 21 Nov 2018 14:50:40 +0100
Subject: [PATCH] imx: hab: extend hab_auth_img to calculate ivt_offset

Current implementation of hab_auth_img command needs ivt_offset to
authenticate the image. But ivt header is placed at the end of image
date after padding.

This leaves the usage of hab_auth_img command to fixed size or static
offset for ivt header. New function "get_image_ivt_offset" is introduced
to find the ivt offset during runtime. The case conditional check in this
function is same as boot_get_kernel in common/bootm.c

With this variable length image e.g. FIT image with any random size can
have IVT at the end and ivt_offset option can be left optional

Can be used as "hab_auth_img $loadaddr $filesize" from u-boot script

Signed-off-by: Parthiban Nallathambi <pn@denx.de>
Reviewed-by: Breno Lima <breno.lima@nxp.com>
---
 arch/arm/mach-imx/hab.c | 29 +++++++++++++++++++++++++++--
 1 file changed, 27 insertions(+), 2 deletions(-)

diff --git a/arch/arm/mach-imx/hab.c b/arch/arm/mach-imx/hab.c
index b88acd13da..dbfd692fa3 100644
--- a/arch/arm/mach-imx/hab.c
+++ b/arch/arm/mach-imx/hab.c
@@ -6,6 +6,8 @@
 #include <common.h>
 #include <config.h>
 #include <fuse.h>
+#include <mapmem.h>
+#include <image.h>
 #include <asm/io.h>
 #include <asm/system.h>
 #include <asm/arch/clock.h>
@@ -302,18 +304,41 @@ static int do_hab_status(cmd_tbl_t *cmdtp, int flag, int argc,
 	return 0;
 }
 
+static ulong get_image_ivt_offset(ulong img_addr)
+{
+	const void *buf;
+
+	buf = map_sysmem(img_addr, 0);
+	switch (genimg_get_format(buf)) {
+#if defined(CONFIG_IMAGE_FORMAT_LEGACY)
+	case IMAGE_FORMAT_LEGACY:
+		return (image_get_image_size((image_header_t *)img_addr)
+			+ 0x1000 - 1)  & ~(0x1000 - 1);
+#endif
+#if IMAGE_ENABLE_FIT
+	case IMAGE_FORMAT_FIT:
+		return (fit_get_size(buf) + 0x1000 - 1)  & ~(0x1000 - 1);
+#endif
+	default:
+		return 0;
+	}
+}
+
 static int do_authenticate_image(cmd_tbl_t *cmdtp, int flag, int argc,
 				 char * const argv[])
 {
 	ulong	addr, length, ivt_offset;
 	int	rcode = 0;
 
-	if (argc < 4)
+	if (argc < 3)
 		return CMD_RET_USAGE;
 
 	addr = simple_strtoul(argv[1], NULL, 16);
 	length = simple_strtoul(argv[2], NULL, 16);
-	ivt_offset = simple_strtoul(argv[3], NULL, 16);
+	if (argc == 3)
+		ivt_offset = get_image_ivt_offset(addr);
+	else
+		ivt_offset = simple_strtoul(argv[3], NULL, 16);
 
 	rcode = imx_hab_authenticate_image(addr, length, ivt_offset);
 	if (rcode == 0)
-- 
2.39.5