Merge tag 'u-boot-stm32-20220510' of https://source.denx.de/u-boot/custodians/u-boot-stm

Add new STM32 MCU boards and Documentation
STM32 programmer improvements
video: support several LTDC HW versions and fix data enable polarity
board: fix stboard error message, consider USB cable connected when boot device is USB
configs: stm32mp1: set console variable for extlinux.conf
configs: stm32mp1: add support for baudrate higher than 115200 for ST-Link
ARM: stm32mp: Fix Silicon version handling and ft_system_setup()
phy: stm32-usbphyc: Add DT phy tuning support
arm: dts: stm32mp15: alignment with v5.18
ram: Conditionally enable ASR
mach-stm32mp: psci: retain MCUDIVR, PLL3CR, PLL4CR, MSSCKSELR across suspend
configs: Use TFTP_TSIZE on DHSOM and STMicroelectronics boards
ARM: stm32: Use default CONFIG_TFTP_BLOCKSIZE on DHSOM
pinctrl: stm32: rework GPIO holes management

1 files changed, 324 insertions, 94 deletions

diff --git a/arch/arm/mach-stm32mp/cmd_stm32prog/stm32prog.c b/arch/arm/mach-stm32mp/cmd_stm32prog/stm32prog.c
index 61cba157fdc..b7111123ba9 100644
--- a/arch/arm/mach-stm32mp/cmd_stm32prog/stm32prog.c
+++ b/arch/arm/mach-stm32mp/cmd_stm32prog/stm32prog.c
@@ -6,12 +6,15 @@
 #include <command.h>
 #include <console.h>
 #include <dfu.h>
+#include <image.h>
 #include <malloc.h>
 #include <misc.h>
 #include <mmc.h>
 #include <part.h>
+#include <tee.h>
 #include <asm/arch/stm32mp1_smc.h>
 #include <asm/global_data.h>
+#include <dm/device_compat.h>
 #include <dm/uclass.h>
 #include <jffs2/load_kernel.h>
 #include <linux/list.h>
@@ -46,7 +49,7 @@
 	EFI_GUID(0xFD58F1C7, 0xBE0D, 0x4338, \
 		 0x88, 0xE9, 0xAD, 0x8F, 0x05, 0x0A, 0xEB, 0x18)
 
-/* RAW parttion (binary / bootloader) used Linux - reserved UUID */
+/* RAW partition (binary / bootloader) used Linux - reserved UUID */
 #define LINUX_RESERVED_UUID "8DA63339-0007-60C0-C436-083AC8230908"
 
 /*
@@ -60,6 +63,28 @@ static const efi_guid_t uuid_mmc[3] = {
 	ROOTFS_MMC2_UUID
 };
 
+/* FIP type partition UUID used by TF-A*/
+#define FIP_TYPE_UUID "19D5DF83-11B0-457B-BE2C-7559C13142A5"
+
+/* unique partition guid (uuid) for FIP partitions A/B */
+#define FIP_A_UUID \
+	EFI_GUID(0x4FD84C93, 0x54EF, 0x463F, \
+		 0xA7, 0xEF, 0xAE, 0x25, 0xFF, 0x88, 0x70, 0x87)
+
+#define FIP_B_UUID \
+	EFI_GUID(0x09C54952, 0xD5BF, 0x45AF, \
+		 0xAC, 0xEE, 0x33, 0x53, 0x03, 0x76, 0x6F, 0xB3)
+
+static const char * const fip_part_name[] = {
+	"fip-a",
+	"fip-b"
+};
+
+static const efi_guid_t fip_part_uuid[] = {
+	FIP_A_UUID,
+	FIP_B_UUID
+};
+
 /* order of column in flash layout file */
 enum stm32prog_col_t {
 	COL_OPTION,
@@ -79,8 +104,110 @@ struct fip_toc_header {
 	u64	flags;
 };
 
+#define TA_NVMEM_UUID { 0x1a8342cc, 0x81a5, 0x4512, \
+		{ 0x99, 0xfe, 0x9e, 0x2b, 0x3e, 0x37, 0xd6, 0x26 } }
+
+/*
+ * Read NVMEM memory for STM32CubeProgrammer
+ *
+ * [in]		value[0].a:		Type (0 for OTP access)
+ * [out]	memref[1].buffer	Output buffer to return all read values
+ * [out]	memref[1].size		Size of buffer to be read
+ *
+ * Return codes:
+ * TEE_SUCCESS - Invoke command success
+ * TEE_ERROR_BAD_PARAMETERS - Incorrect input param
+ */
+#define TA_NVMEM_READ		0x0
+
+/*
+ * Write NVMEM memory for STM32CubeProgrammer
+ *
+ * [in]	     value[0].a		Type (0 for OTP access)
+ * [in]      memref[1].buffer	Input buffer with the values to write
+ * [in]      memref[1].size	Size of buffer to be written
+ *
+ * Return codes:
+ * TEE_SUCCESS - Invoke command success
+ * TEE_ERROR_BAD_PARAMETERS - Incorrect input param
+ */
+#define TA_NVMEM_WRITE		0x1
+
+/* value of TA_NVMEM type = value[in] a */
+#define NVMEM_OTP		0
+
 DECLARE_GLOBAL_DATA_PTR;
 
+/* OPTEE TA NVMEM open helper */
+static int optee_ta_open(struct stm32prog_data *data)
+{
+	const struct tee_optee_ta_uuid uuid = TA_NVMEM_UUID;
+	struct tee_open_session_arg arg;
+	struct udevice *tee = NULL;
+	int rc;
+
+	if (data->tee)
+		return 0;
+
+	tee = tee_find_device(NULL, NULL, NULL, NULL);
+	if (!tee)
+		return -ENODEV;
+
+	memset(&arg, 0, sizeof(arg));
+	tee_optee_ta_uuid_to_octets(arg.uuid, &uuid);
+	rc = tee_open_session(tee, &arg, 0, NULL);
+	if (rc < 0)
+		return -ENODEV;
+
+	data->tee = tee;
+	data->tee_session = arg.session;
+
+	return 0;
+}
+
+/* OPTEE TA NVMEM invoke helper */
+static int optee_ta_invoke(struct stm32prog_data *data, int cmd, int type,
+			   void *buff, ulong size)
+{
+	struct tee_invoke_arg arg;
+	struct tee_param param[2];
+	struct tee_shm *buff_shm;
+	int rc;
+
+	rc = tee_shm_register(data->tee, buff, size, 0, &buff_shm);
+	if (rc)
+		return rc;
+
+	memset(&arg, 0, sizeof(arg));
+	arg.func = cmd;
+	arg.session = data->tee_session;
+
+	memset(param, 0, sizeof(param));
+	param[0].attr = TEE_PARAM_ATTR_TYPE_VALUE_INPUT;
+	param[0].u.value.a = type;
+
+	if (cmd == TA_NVMEM_WRITE)
+		param[1].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INPUT;
+	else
+		param[1].attr = TEE_PARAM_ATTR_TYPE_MEMREF_OUTPUT;
+
+	param[1].u.memref.shm = buff_shm;
+	param[1].u.memref.size = size;
+
+	rc = tee_invoke_func(data->tee, &arg, 2, param);
+	if (rc < 0 || arg.ret != 0) {
+		dev_err(data->tee,
+			"TA_NVMEM invoke failed TEE err: %x, err:%x\n",
+			arg.ret, rc);
+		if (!rc)
+			rc = -EIO;
+	}
+
+	tee_shm_free(buff_shm);
+
+	return rc;
+}
+
 /* partition handling routines : CONFIG_CMD_MTDPARTS */
 int mtdparts_init(void);
 int find_dev_and_part(const char *id, struct mtd_device **dev,
@@ -101,52 +228,98 @@ static bool stm32prog_is_fip_header(struct fip_toc_header *header)
 	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;
+	}
+	if (header->header_version != 0x00010000) {
+		log_debug("%s:invalid header version : 0x%x\n",
+			  __func__, header->header_version);
+		return false;
 	}
 
-	header->type = HEADER_NONE;
-	header->image_checksum = 0x0;
-	header->image_length = 0x0;
-
-	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)
@@ -255,6 +428,8 @@ static int parse_type(struct stm32prog_data *data,
 				part->bin_nb =
 					dectoul(&p[7], NULL);
 		}
+	} else if (!strcmp(p, "FIP")) {
+		part->part_type = PART_FIP;
 	} else if (!strcmp(p, "System")) {
 		part->part_type = PART_SYSTEM;
 	} else if (!strcmp(p, "FileSystem")) {
@@ -376,11 +551,11 @@ static int parse_flash_layout(struct stm32prog_data *data,
 	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);
@@ -906,9 +1081,10 @@ static int create_gpt_partitions(struct stm32prog_data *data)
 	char uuid[UUID_STR_LEN + 1];
 	unsigned char *uuid_bin;
 	unsigned int mmc_id;
-	int i;
+	int i, j;
 	bool rootfs_found;
 	struct stm32prog_part_t *part;
+	const char *type_str;
 
 	buf = malloc(buflen);
 	if (!buf)
@@ -950,33 +1126,46 @@ static int create_gpt_partitions(struct stm32prog_data *data)
 					   part->addr,
 					   part->size);
 
-			if (part->part_type == PART_BINARY)
-				offset += snprintf(buf + offset,
-						   buflen - offset,
-						   ",type="
-						   LINUX_RESERVED_UUID);
-			else
-				offset += snprintf(buf + offset,
-						   buflen - offset,
-						   ",type=linux");
+			switch (part->part_type) {
+			case PART_BINARY:
+				type_str = LINUX_RESERVED_UUID;
+				break;
+			case PART_FIP:
+				type_str = FIP_TYPE_UUID;
+				break;
+			default:
+				type_str = "linux";
+				break;
+			}
+			offset += snprintf(buf + offset,
+					   buflen - offset,
+					   ",type=%s", type_str);
 
 			if (part->part_type == PART_SYSTEM)
 				offset += snprintf(buf + offset,
 						   buflen - offset,
 						   ",bootable");
 
+			/* partition UUID */
+			uuid_bin = NULL;
 			if (!rootfs_found && !strcmp(part->name, "rootfs")) {
 				mmc_id = part->dev_id;
 				rootfs_found = true;
-				if (mmc_id < ARRAY_SIZE(uuid_mmc)) {
-					uuid_bin =
-					  (unsigned char *)uuid_mmc[mmc_id].b;
-					uuid_bin_to_str(uuid_bin, uuid,
-							UUID_STR_FORMAT_GUID);
-					offset += snprintf(buf + offset,
-							   buflen - offset,
-							   ",uuid=%s", uuid);
-				}
+				if (mmc_id < ARRAY_SIZE(uuid_mmc))
+					uuid_bin = (unsigned char *)uuid_mmc[mmc_id].b;
+			}
+			if (part->part_type == PART_FIP) {
+				for (j = 0; j < ARRAY_SIZE(fip_part_name); j++)
+					if (!strcmp(part->name, fip_part_name[j])) {
+						uuid_bin = (unsigned char *)fip_part_uuid[j].b;
+						break;
+					}
+			}
+			if (uuid_bin) {
+				uuid_bin_to_str(uuid_bin, uuid, UUID_STR_FORMAT_GUID);
+				offset += snprintf(buf + offset,
+						   buflen - offset,
+						   ",uuid=%s", uuid);
 			}
 
 			offset += snprintf(buf + offset, buflen - offset, ";");
@@ -1154,7 +1343,9 @@ static int dfu_init_entities(struct stm32prog_data *data)
 	struct dfu_entity *dfu;
 	int alt_nb;
 
-	alt_nb = 2; /* number of virtual = CMD, OTP*/
+	alt_nb = 1; /* number of virtual = CMD*/
+	if (IS_ENABLED(CONFIG_CMD_STM32PROG_OTP))
+		alt_nb++; /* OTP*/
 	if (CONFIG_IS_ENABLED(DM_PMIC))
 		alt_nb++; /* PMIC NVMEM*/
 
@@ -1205,8 +1396,12 @@ static int dfu_init_entities(struct stm32prog_data *data)
 	if (!ret)
 		ret = stm32prog_alt_add_virt(dfu, "virtual", PHASE_CMD, CMD_SIZE);
 
-	if (!ret)
-		ret = stm32prog_alt_add_virt(dfu, "OTP", PHASE_OTP, OTP_SIZE);
+	if (!ret && IS_ENABLED(CONFIG_CMD_STM32PROG_OTP)) {
+		ret = optee_ta_open(data);
+		log_debug("optee_ta result %d\n", ret);
+		ret = stm32prog_alt_add_virt(dfu, "OTP", PHASE_OTP,
+					     data->tee ? OTP_SIZE_TA : OTP_SIZE_SMC);
+	}
 
 	if (!ret && CONFIG_IS_ENABLED(DM_PMIC))
 		ret = stm32prog_alt_add_virt(dfu, "PMIC", PHASE_PMIC, PMIC_SIZE);
@@ -1224,19 +1419,26 @@ static int dfu_init_entities(struct stm32prog_data *data)
 int stm32prog_otp_write(struct stm32prog_data *data, u32 offset, u8 *buffer,
 			long *size)
 {
+	u32 otp_size = data->tee ? OTP_SIZE_TA : OTP_SIZE_SMC;
 	log_debug("%s: %x %lx\n", __func__, offset, *size);
 
+	if (!IS_ENABLED(CONFIG_CMD_STM32PROG_OTP)) {
+		stm32prog_err("OTP update not supported");
+
+		return -EOPNOTSUPP;
+	}
+
 	if (!data->otp_part) {
-		data->otp_part = memalign(CONFIG_SYS_CACHELINE_SIZE, OTP_SIZE);
+		data->otp_part = memalign(CONFIG_SYS_CACHELINE_SIZE, otp_size);
 		if (!data->otp_part)
 			return -ENOMEM;
 	}
 
 	if (!offset)
-		memset(data->otp_part, 0, OTP_SIZE);
+		memset(data->otp_part, 0, otp_size);
 
-	if (offset + *size > OTP_SIZE)
-		*size = OTP_SIZE - offset;
+	if (offset + *size > otp_size)
+		*size = otp_size - offset;
 
 	memcpy((void *)((u32)data->otp_part + offset), buffer, *size);
 
@@ -1246,12 +1448,13 @@ int stm32prog_otp_write(struct stm32prog_data *data, u32 offset, u8 *buffer,
 int stm32prog_otp_read(struct stm32prog_data *data, u32 offset, u8 *buffer,
 		       long *size)
 {
+	u32 otp_size = data->tee ? OTP_SIZE_TA : OTP_SIZE_SMC;
 	int result = 0;
 
-	if (!IS_ENABLED(CONFIG_ARM_SMCCC)) {
+	if (!IS_ENABLED(CONFIG_CMD_STM32PROG_OTP)) {
 		stm32prog_err("OTP update not supported");
 
-		return -1;
+		return -EOPNOTSUPP;
 	}
 
 	log_debug("%s: %x %lx\n", __func__, offset, *size);
@@ -1259,7 +1462,7 @@ int stm32prog_otp_read(struct stm32prog_data *data, u32 offset, u8 *buffer,
 	if (!offset) {
 		if (!data->otp_part)
 			data->otp_part =
-				memalign(CONFIG_SYS_CACHELINE_SIZE, OTP_SIZE);
+				memalign(CONFIG_SYS_CACHELINE_SIZE, otp_size);
 
 		if (!data->otp_part) {
 			result = -ENOMEM;
@@ -1267,11 +1470,16 @@ int stm32prog_otp_read(struct stm32prog_data *data, u32 offset, u8 *buffer,
 		}
 
 		/* init struct with 0 */
-		memset(data->otp_part, 0, OTP_SIZE);
+		memset(data->otp_part, 0, otp_size);
 
 		/* call the service */
-		result = stm32_smc_exec(STM32_SMC_BSEC, STM32_SMC_READ_ALL,
-					(u32)data->otp_part, 0);
+		result = -EOPNOTSUPP;
+		if (data->tee && CONFIG_IS_ENABLED(OPTEE))
+			result = optee_ta_invoke(data, TA_NVMEM_READ, NVMEM_OTP,
+						 data->otp_part, OTP_SIZE_TA);
+		else if (IS_ENABLED(CONFIG_ARM_SMCCC))
+			result = stm32_smc_exec(STM32_SMC_BSEC, STM32_SMC_READ_ALL,
+						(u32)data->otp_part, 0);
 		if (result)
 			goto end_otp_read;
 	}
@@ -1281,8 +1489,8 @@ int stm32prog_otp_read(struct stm32prog_data *data, u32 offset, u8 *buffer,
 		goto end_otp_read;
 	}
 
-	if (offset + *size > OTP_SIZE)
-		*size = OTP_SIZE - offset;
+	if (offset + *size > otp_size)
+		*size = otp_size - offset;
 	memcpy(buffer, (void *)((u32)data->otp_part + offset), *size);
 
 end_otp_read:
@@ -1296,10 +1504,10 @@ int stm32prog_otp_start(struct stm32prog_data *data)
 	int result = 0;
 	struct arm_smccc_res res;
 
-	if (!IS_ENABLED(CONFIG_ARM_SMCCC)) {
+	if (!IS_ENABLED(CONFIG_CMD_STM32PROG_OTP)) {
 		stm32prog_err("OTP update not supported");
 
-		return -1;
+		return -EOPNOTSUPP;
 	}
 
 	if (!data->otp_part) {
@@ -1307,28 +1515,34 @@ int stm32prog_otp_start(struct stm32prog_data *data)
 		return -1;
 	}
 
-	arm_smccc_smc(STM32_SMC_BSEC, STM32_SMC_WRITE_ALL,
-		      (u32)data->otp_part, 0, 0, 0, 0, 0, &res);
-
-	if (!res.a0) {
-		switch (res.a1) {
-		case 0:
-			result = 0;
-			break;
-		case 1:
-			stm32prog_err("Provisioning");
-			result = 0;
-			break;
-		default:
-			log_err("%s: OTP incorrect value (err = %ld)\n",
-				__func__, res.a1);
+	result = -EOPNOTSUPP;
+	if (data->tee && CONFIG_IS_ENABLED(OPTEE)) {
+		result = optee_ta_invoke(data, TA_NVMEM_WRITE, NVMEM_OTP,
+					 data->otp_part, OTP_SIZE_TA);
+	} else if (IS_ENABLED(CONFIG_ARM_SMCCC)) {
+		arm_smccc_smc(STM32_SMC_BSEC, STM32_SMC_WRITE_ALL,
+			      (u32)data->otp_part, 0, 0, 0, 0, 0, &res);
+
+		if (!res.a0) {
+			switch (res.a1) {
+			case 0:
+				result = 0;
+				break;
+			case 1:
+				stm32prog_err("Provisioning");
+				result = 0;
+				break;
+			default:
+				log_err("%s: OTP incorrect value (err = %ld)\n",
+					__func__, res.a1);
+				result = -EINVAL;
+				break;
+			}
+		} else {
+			log_err("%s: Failed to exec svc=%x op=%x in secure mode (err = %ld)\n",
+				__func__, STM32_SMC_BSEC, STM32_SMC_WRITE_ALL, res.a0);
 			result = -EINVAL;
-			break;
 		}
-	} else {
-		log_err("%s: Failed to exec svc=%x op=%x in secure mode (err = %ld)\n",
-			__func__, STM32_SMC_BSEC, STM32_SMC_WRITE_ALL, res.a0);
-		result = -EINVAL;
 	}
 
 	free(data->otp_part);
@@ -1431,7 +1645,7 @@ static int stm32prog_copy_fsbl(struct stm32prog_part_t *part)
 	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;
 
@@ -1443,17 +1657,18 @@ static int stm32prog_copy_fsbl(struct stm32prog_part_t *part)
 
 	/* 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)
@@ -1483,7 +1698,16 @@ error:
 static void stm32prog_end_phase(struct stm32prog_data *data, u64 offset)
 {
 	if (data->phase == PHASE_FLASHLAYOUT) {
-		if (parse_flash_layout(data, STM32_DDR_BASE, 0))
+#if defined(CONFIG_LEGACY_IMAGE_FORMAT)
+		if (genimg_get_format((void *)STM32_DDR_BASE) == IMAGE_FORMAT_LEGACY) {
+			data->script = STM32_DDR_BASE;
+			data->phase = PHASE_END;
+			log_notice("U-Boot script received\n");
+			return;
+		}
+#endif
+		log_notice("\nFlashLayout received, size = %lld\n", offset);
+		if (parse_flash_layout(data, STM32_DDR_BASE, offset))
 			stm32prog_err("Layout: invalid FlashLayout");
 		return;
 	}
@@ -1739,6 +1963,12 @@ void stm32prog_clean(struct stm32prog_data *data)
 	free(data->part_array);
 	free(data->otp_part);
 	free(data->buffer);
+
+	if (CONFIG_IS_ENABLED(OPTEE) && data->tee) {
+		tee_close_session(data->tee, data->tee_session);
+		data->tee = NULL;
+		data->tee_session = 0x0;
+	}
 }
 
 /* DFU callback: used after serial and direct DFU USB access */