stm32mp: stm32prog: add support of STM32IMAGE version 2

Add support of new header for the STM32IMAGE version V2
in command stm32prog command for STM32MP13x family.

Signed-off-by: Patrick Delaunay <patrick.delaunay@foss.st.com>
Reviewed-by: Patrice Chotard <patrice.chotard@foss.st.com>
Signed-off-by: Patrice Chotard <patrice.chotard@foss.st.com>
diff --git a/arch/arm/mach-stm32mp/cmd_stm32prog/cmd_stm32prog.c b/arch/arm/mach-stm32mp/cmd_stm32prog/cmd_stm32prog.c
index 41452b5..3957e06 100644
--- a/arch/arm/mach-stm32mp/cmd_stm32prog/cmd_stm32prog.c
+++ b/arch/arm/mach-stm32mp/cmd_stm32prog/cmd_stm32prog.c
@@ -73,15 +73,15 @@
 
 	/* check STM32IMAGE presence */
 	if (size == 0) {
-		stm32prog_header_check((struct raw_header_s *)addr, &header);
+		stm32prog_header_check(addr, &header);
 		if (header.type == HEADER_STM32IMAGE) {
-			size = header.image_length + BL_HEADER_SIZE;
+			size = header.image_length + header.length;
 
 #if defined(CONFIG_LEGACY_IMAGE_FORMAT)
 			/* uImage detected in STM32IMAGE, execute the script */
 			if (IMAGE_FORMAT_LEGACY ==
-			    genimg_get_format((void *)(addr + BL_HEADER_SIZE)))
-				return image_source_script(addr + BL_HEADER_SIZE, "script@1");
+			    genimg_get_format((void *)(addr + header.length)))
+				return image_source_script(addr + header.length, "script@1");
 #endif
 		}
 	}
diff --git a/arch/arm/mach-stm32mp/cmd_stm32prog/stm32prog.c b/arch/arm/mach-stm32mp/cmd_stm32prog/stm32prog.c
index 5d53e61..3e1fdee 100644
--- a/arch/arm/mach-stm32mp/cmd_stm32prog/stm32prog.c
+++ b/arch/arm/mach-stm32mp/cmd_stm32prog/stm32prog.c
@@ -205,52 +205,98 @@
 	return (header->name == FIP_TOC_HEADER_NAME) && header->serial_number;
 }
 
-void stm32prog_header_check(struct raw_header_s *raw_header,
-			    struct image_header_s *header)
+static bool stm32prog_is_stm32_header_v1(struct stm32_header_v1 *header)
 {
 	unsigned int i;
 
-	if (!raw_header || !header) {
-		log_debug("%s:no header data\n", __func__);
-		return;
+	if (header->magic_number !=
+		(('S' << 0) | ('T' << 8) | ('M' << 16) | (0x32 << 24))) {
+		log_debug("%s:invalid magic number : 0x%x\n",
+			  __func__, header->magic_number);
+		return false;
 	}
-
-	header->type = HEADER_NONE;
-	header->image_checksum = 0x0;
-	header->image_length = 0x0;
+	if (header->header_version != 0x00010000) {
+		log_debug("%s:invalid header version : 0x%x\n",
+			  __func__, header->header_version);
+		return false;
+	}
 
-	if (stm32prog_is_fip_header((struct fip_toc_header *)raw_header)) {
-		header->type = HEADER_FIP;
-		return;
+	if (header->reserved1 || header->reserved2) {
+		log_debug("%s:invalid reserved field\n", __func__);
+		return false;
+	}
+	for (i = 0; i < sizeof(header->padding); i++) {
+		if (header->padding[i] != 0) {
+			log_debug("%s:invalid padding field\n", __func__);
+			return false;
+		}
 	}
 
-	if (raw_header->magic_number !=
+	return true;
+}
+
+static bool stm32prog_is_stm32_header_v2(struct stm32_header_v2 *header)
+{
+	unsigned int i;
+
+	if (header->magic_number !=
 		(('S' << 0) | ('T' << 8) | ('M' << 16) | (0x32 << 24))) {
 		log_debug("%s:invalid magic number : 0x%x\n",
-			  __func__, raw_header->magic_number);
-		return;
+			  __func__, header->magic_number);
+		return false;
 	}
-	/* only header v1.0 supported */
-	if (raw_header->header_version != 0x00010000) {
+	if (header->header_version != 0x00020000) {
 		log_debug("%s:invalid header version : 0x%x\n",
-			  __func__, raw_header->header_version);
+			  __func__, header->header_version);
+		return false;
+	}
+	if (header->reserved1 || header->reserved2)
+		return false;
+
+	for (i = 0; i < sizeof(header->padding); i++) {
+		if (header->padding[i] != 0) {
+			log_debug("%s:invalid padding field\n", __func__);
+			return false;
+		}
+	}
+
+	return true;
+}
+
+void stm32prog_header_check(uintptr_t raw_header, struct image_header_s *header)
+{
+	struct stm32_header_v1 *v1_header = (struct stm32_header_v1 *)raw_header;
+	struct stm32_header_v2 *v2_header = (struct stm32_header_v2 *)raw_header;
+
+	if (!raw_header || !header) {
+		log_debug("%s:no header data\n", __func__);
 		return;
 	}
-	if (raw_header->reserved1 != 0x0 || raw_header->reserved2) {
-		log_debug("%s:invalid reserved field\n", __func__);
+
+	if (stm32prog_is_fip_header((struct fip_toc_header *)raw_header)) {
+		header->type = HEADER_FIP;
+		header->length = 0;
 		return;
 	}
-	for (i = 0; i < (sizeof(raw_header->padding) / 4); i++) {
-		if (raw_header->padding[i] != 0) {
-			log_debug("%s:invalid padding field\n", __func__);
-			return;
-		}
+	if (stm32prog_is_stm32_header_v1(v1_header)) {
+		header->type = HEADER_STM32IMAGE;
+		header->image_checksum = le32_to_cpu(v1_header->image_checksum);
+		header->image_length = le32_to_cpu(v1_header->image_length);
+		header->length = sizeof(struct stm32_header_v1);
+		return;
+	}
+	if (stm32prog_is_stm32_header_v2(v2_header)) {
+		header->type = HEADER_STM32IMAGE_V2;
+		header->image_checksum = le32_to_cpu(v2_header->image_checksum);
+		header->image_length = le32_to_cpu(v2_header->image_length);
+		header->length = sizeof(struct stm32_header_v1) +
+				 v2_header->extension_headers_length;
+		return;
 	}
-	header->type = HEADER_STM32IMAGE;
-	header->image_checksum = le32_to_cpu(raw_header->image_checksum);
-	header->image_length = le32_to_cpu(raw_header->image_length);
 
-	return;
+	header->type = HEADER_NONE;
+	header->image_checksum = 0x0;
+	header->image_length = 0x0;
 }
 
 static u32 stm32prog_header_checksum(u32 addr, struct image_header_s *header)
@@ -480,11 +526,11 @@
 	data->part_nb = 0;
 
 	/* check if STM32image is detected */
-	stm32prog_header_check((struct raw_header_s *)addr, &header);
+	stm32prog_header_check(addr, &header);
 	if (header.type == HEADER_STM32IMAGE) {
 		u32 checksum;
 
-		addr = addr + BL_HEADER_SIZE;
+		addr = addr + header.length;
 		size = header.image_length;
 
 		checksum = stm32prog_header_checksum(addr, &header);
@@ -1560,7 +1606,7 @@
 	int ret, i;
 	void *fsbl;
 	struct image_header_s header;
-	struct raw_header_s raw_header;
+	struct stm32_header_v2 raw_header; /* V2 size > v1 size */
 	struct dfu_entity *dfu;
 	long size, offset;
 
@@ -1572,17 +1618,18 @@
 
 	/* read header */
 	dfu_transaction_cleanup(dfu);
-	size = BL_HEADER_SIZE;
+	size = sizeof(raw_header);
 	ret = dfu->read_medium(dfu, 0, (void *)&raw_header, &size);
 	if (ret)
 		return ret;
 
-	stm32prog_header_check(&raw_header, &header);
-	if (header.type != HEADER_STM32IMAGE)
+	stm32prog_header_check((ulong)&raw_header, &header);
+	if (header.type != HEADER_STM32IMAGE &&
+	    header.type != HEADER_STM32IMAGE_V2)
 		return -ENOENT;
 
 	/* read header + payload */
-	size = header.image_length + BL_HEADER_SIZE;
+	size = header.image_length + header.length;
 	size = round_up(size, part->dev->mtd->erasesize);
 	fsbl = calloc(1, size);
 	if (!fsbl)
diff --git a/arch/arm/mach-stm32mp/cmd_stm32prog/stm32prog.h b/arch/arm/mach-stm32mp/cmd_stm32prog/stm32prog.h
index 928b7b3..90cdc2b 100644
--- a/arch/arm/mach-stm32mp/cmd_stm32prog/stm32prog.h
+++ b/arch/arm/mach-stm32mp/cmd_stm32prog/stm32prog.h
@@ -42,6 +42,7 @@
 enum stm32prog_header_t {
 	HEADER_NONE,
 	HEADER_STM32IMAGE,
+	HEADER_STM32IMAGE_V2,
 	HEADER_FIP,
 };
 
@@ -49,11 +50,12 @@
 	enum stm32prog_header_t type;
 	u32	image_checksum;
 	u32	image_length;
+	u32	length;
 };
 
-struct raw_header_s {
+struct stm32_header_v1 {
 	u32 magic_number;
-	u32 image_signature[64 / 4];
+	u8 image_signature[64];
 	u32 image_checksum;
 	u32 header_version;
 	u32 image_length;
@@ -64,12 +66,30 @@
 	u32 version_number;
 	u32 option_flags;
 	u32 ecdsa_algorithm;
-	u32 ecdsa_public_key[64 / 4];
-	u32 padding[83 / 4];
-	u32 binary_type;
+	u8 ecdsa_public_key[64];
+	u8 padding[83];
+	u8 binary_type;
 };
 
-#define BL_HEADER_SIZE	sizeof(struct raw_header_s)
+struct stm32_header_v2 {
+	u32 magic_number;
+	u8 image_signature[64];
+	u32 image_checksum;
+	u32 header_version;
+	u32 image_length;
+	u32 image_entry_point;
+	u32 reserved1;
+	u32 load_address;
+	u32 reserved2;
+	u32 version_number;
+	u32 extension_flags;
+	u32 extension_headers_length;
+	u32 binary_type;
+	u8 padding[16];
+	u32 extension_header_type;
+	u32 extension_header_length;
+	u8 extension_padding[376];
+};
 
 /* partition type in flashlayout file */
 enum stm32prog_part_type {
@@ -171,8 +191,7 @@
 int stm32prog_pmic_start(struct stm32prog_data *data);
 
 /* generic part*/
-void stm32prog_header_check(struct raw_header_s *raw_header,
-			    struct image_header_s *header);
+void stm32prog_header_check(uintptr_t raw_header, struct image_header_s *header);
 int stm32prog_dfu_init(struct stm32prog_data *data);
 void stm32prog_next_phase(struct stm32prog_data *data);
 void stm32prog_do_reset(struct stm32prog_data *data);