feat(optee): add loading OP-TEE image via an SMC

This adds the ability to load the OP-TEE image via an SMC called from
non-secure userspace rather than loading it during boot. This should
only be utilized on platforms that can ensure security is maintained up
until the point the SMC is invoked as it breaks the normal barrier
between the secure and non-secure world.

Signed-off-by: Jeffrey Kardatzke <jkardatzke@google.com>
Change-Id: I21cfa9699617c493fa4190f01d1cbb714e7449cc
diff --git a/docs/components/spd/optee-dispatcher.rst b/docs/components/spd/optee-dispatcher.rst
index 63baccc..81476f1 100644
--- a/docs/components/spd/optee-dispatcher.rst
+++ b/docs/components/spd/optee-dispatcher.rst
@@ -6,9 +6,26 @@
 To build and execute OP-TEE follow the instructions at
 `OP-TEE build.git`_
 
+There are two different modes for loading the OP-TEE OS. The default mode will
+load it as the BL32 payload during boot, and is the recommended technique for
+platforms to use. There is also another technique that will load OP-TEE OS after
+boot via an SMC call by enabling the option for OPTEE_ALLOW_SMC_LOAD that was
+specifically added for ChromeOS. Loading OP-TEE via an SMC call may be insecure
+depending upon the platform configuration. If using that option, be sure to
+understand the risks involved with allowing the Trusted OS to be loaded this
+way. ChromeOS uses a boot flow where it verifies the signature of the firmware
+before executing it, and then only if the signature is valid will the 'secrets'
+used by the TEE become accessible. The firmware then verifies the signature of
+the kernel using depthcharge, and the kernel verifies the rootfs using
+dm-verity.  The SMC call to load OP-TEE is then invoked immediately after the
+kernel finishes loading and before any attack vectors can be opened up by
+mounting writable filesystems or opening network/device connections. this
+ensures the platform is 'closed' and running signed code through the point where
+OP-TEE is loaded.
+
 --------------
 
-*Copyright (c) 2014-2018, Arm Limited and Contributors. All rights reserved.*
+*Copyright (c) 2014-2023, Arm Limited and Contributors. All rights reserved.*
 
 .. _OP-TEE OS: https://github.com/OP-TEE/build
 .. _OP-TEE build.git: https://github.com/OP-TEE/build
diff --git a/docs/threat_model/threat_model.rst b/docs/threat_model/threat_model.rst
index 99bbb3a..0e967ba 100644
--- a/docs/threat_model/threat_model.rst
+++ b/docs/threat_model/threat_model.rst
@@ -918,9 +918,54 @@
 | Mitigations            | | Yes / Platform specific                           |
 +------------------------+-----------------------------------------------------+
 
++------------------------+-----------------------------------------------------+
+| ID                     | 14                                                  |
++========================+=====================================================+
+| Threat                 | | **Security vulnerabilities in the Non-secure OS   |
+|                        |   can lead to secure world compromise if the option |
+|                        |   OPTEE_ALLOW_SMC_LOAD is enabled.**                |
+|                        |                                                     |
+|                        | | This option trusts the non-secure world up until  |
+|                        |   the point it issues the SMC call to load the      |
+|                        |   Secure BL32 payload. If a compromise occurs       |
+|                        |   before the SMC call is invoked, then arbitrary    |
+|                        |   code execution in S-EL1 can occur or arbitrary    |
+|                        |   memory in EL3 can be overwritten.                 |
++------------------------+-----------------------------------------------------+
+| Diagram Elements       | DF5                                                 |
++------------------------+-----------------------------------------------------+
+| Affected TF-A          | BL31, BL32                                          |
+| Components             |                                                     |
++------------------------+-----------------------------------------------------+
+| Assets                 | Code Execution, Sensitive Data                      |
++------------------------+-----------------------------------------------------+
+| Threat Agent           | NSCode                                              |
++------------------------+-----------------------------------------------------+
+| Threat Type            | Tampering, Information Disclosure,                  |
+|                        | Elevation of privilege                              |
++------------------------+-----------------+-----------------+-----------------+
+| Application            | Server          | IoT             | Mobile          |
++------------------------+-----------------+-----------------+-----------------+
+| Impact                 | Critical (5)    | Critical (5)    | Critical (5)    |
++------------------------+-----------------+-----------------+-----------------+
+| Likelihood             | Low (2)         | Low (2)         | Low (2)         |
++------------------------+-----------------+-----------------+-----------------+
+| Total Risk Rating      | Medium (10)     | Medium (10)     | Medium (10)     |
++------------------------+-----------------+-----------------+-----------------+
+| Mitigations            | When enabling the option OPTEE_ALLOW_SMC_LOAD,      |
+|                        | the non-secure OS must be considered a closed       |
+|                        | platform up until the point the SMC can be invoked  |
+|                        | to load OP-TEE.                                     |
++------------------------+-----------------------------------------------------+
+| Mitigations            | | None in TF-A itself. This option is only used by  |
+| implemented?           |   ChromeOS currently which has other mechanisms to  |
+|                        |   to mitigate this threat which are described in    |
+|                        |   `OP-TEE Dispatcher`_.                             |
++------------------------+-----------------------------------------------------+
+
 --------------
 
-*Copyright (c) 2021-2022, Arm Limited. All rights reserved.*
+*Copyright (c) 2021-2023, Arm Limited. All rights reserved.*
 
 
 .. _STRIDE threat analysis technique: https://docs.microsoft.com/en-us/azure/security/develop/threat-modeling-tool-threats#stride-model
@@ -932,3 +977,4 @@
 .. _TF-A error handling policy: https://trustedfirmware-a.readthedocs.io/en/latest/process/coding-guidelines.html#error-handling-and-robustness
 .. _Secure Development Guidelines: https://trustedfirmware-a.readthedocs.io/en/latest/process/security-hardening.html#secure-development-guidelines
 .. _Trusted Firmware-A Tests: https://git.trustedfirmware.org/TF-A/tf-a-tests.git/about/
+.. _OP-TEE Dispatcher: https://github.com/ARM-software/arm-trusted-firmware/blob/master/docs/components/spd/optee-dispatcher.rst
diff --git a/include/lib/optee_utils.h b/include/lib/optee_utils.h
index 06378eb..8224d50 100644
--- a/include/lib/optee_utils.h
+++ b/include/lib/optee_utils.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017-2021, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2017-2023, ARM Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -17,4 +17,40 @@
 	image_info_t *pager_image_info,
 	image_info_t *paged_image_info);
 
+/*
+ * load_addr_hi and load_addr_lo: image load address.
+ * image_id: 0 - pager, 1 - paged
+ * size: image size in bytes.
+ */
+typedef struct optee_image {
+	uint32_t load_addr_hi;
+	uint32_t load_addr_lo;
+	uint32_t image_id;
+	uint32_t size;
+} optee_image_t;
+
+#define OPTEE_PAGER_IMAGE_ID		0
+#define OPTEE_PAGED_IMAGE_ID		1
+
+#define OPTEE_MAX_NUM_IMAGES		2u
+
+#define TEE_MAGIC_NUM_OPTEE		0x4554504f
+/*
+ * magic: header magic number.
+ * version: OPTEE header version:
+ *		1 - not supported
+ *		2 - supported
+ * arch: OPTEE os architecture type: 0 - AARCH32, 1 - AARCH64.
+ * flags: unused currently.
+ * nb_images: number of images.
+ */
+typedef struct optee_header {
+	uint32_t magic;
+	uint8_t version;
+	uint8_t arch;
+	uint16_t flags;
+	uint32_t nb_images;
+	optee_image_t optee_image_list[];
+} optee_header_t;
+
 #endif /* OPTEE_UTILS_H */
diff --git a/lib/optee/optee_utils.c b/lib/optee/optee_utils.c
index 6c87b0d..25272fc 100644
--- a/lib/optee/optee_utils.c
+++ b/lib/optee/optee_utils.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017-2022, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2017-2023, ARM Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -11,42 +11,6 @@
 
 #include <platform_def.h>
 
-/*
- * load_addr_hi and load_addr_lo: image load address.
- * image_id: 0 - pager, 1 - paged
- * size: image size in bytes.
- */
-typedef struct optee_image {
-	uint32_t load_addr_hi;
-	uint32_t load_addr_lo;
-	uint32_t image_id;
-	uint32_t size;
-} optee_image_t;
-
-#define OPTEE_PAGER_IMAGE_ID		0
-#define OPTEE_PAGED_IMAGE_ID		1
-
-#define OPTEE_MAX_NUM_IMAGES		2u
-
-#define TEE_MAGIC_NUM_OPTEE		0x4554504f
-/*
- * magic: header magic number.
- * version: OPTEE header version:
- *		1 - not supported
- *		2 - supported
- * arch: OPTEE os architecture type: 0 - AARCH32, 1 - AARCH64.
- * flags: unused currently.
- * nb_images: number of images.
- */
-typedef struct optee_header {
-	uint32_t magic;
-	uint8_t version;
-	uint8_t arch;
-	uint16_t flags;
-	uint32_t nb_images;
-	optee_image_t optee_image_list[];
-} optee_header_t;
-
 /*******************************************************************************
  * Check if it is a valid tee header
  * Return true if valid
diff --git a/services/spd/opteed/opteed.mk b/services/spd/opteed/opteed.mk
index 643b054..477b45d 100644
--- a/services/spd/opteed/opteed.mk
+++ b/services/spd/opteed/opteed.mk
@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2013-2019, ARM Limited and Contributors. All rights reserved.
+# Copyright (c) 2013-2023, ARM Limited and Contributors. All rights reserved.
 #
 # SPDX-License-Identifier: BSD-3-Clause
 #
@@ -16,3 +16,19 @@
 
 # required so that optee code can control access to the timer registers
 NS_TIMER_SWITCH		:=	1
+
+# WARNING: This enables loading of OP-TEE via an SMC, which can be potentially
+# insecure. This removes the boundary between the startup of the secure and
+# non-secure worlds until the point where this SMC is invoked. Only use this
+# setting if you can ensure that the non-secure OS can remain trusted up until
+# the point where this SMC is invoked.
+OPTEE_ALLOW_SMC_LOAD		:=	0
+ifeq ($(OPTEE_ALLOW_SMC_LOAD),1)
+ifeq ($(PLAT_XLAT_TABLES_DYNAMIC),0)
+$(error When OPTEE_ALLOW_SMC_LOAD=1, PLAT_XLAT_TABLES_DYNAMIC must also be 1)
+endif
+$(warning "OPTEE_ALLOW_SMC_LOAD is enabled which may result in an insecure \
+	platform")
+$(eval $(call add_define,PLAT_XLAT_TABLES_DYNAMIC))
+$(eval $(call add_define,OPTEE_ALLOW_SMC_LOAD))
+endif
diff --git a/services/spd/opteed/opteed_main.c b/services/spd/opteed/opteed_main.c
index 160a693..ff2aee0 100644
--- a/services/spd/opteed/opteed_main.c
+++ b/services/spd/opteed/opteed_main.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013-2017, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2013-2023, ARM Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -16,6 +16,7 @@
  ******************************************************************************/
 #include <assert.h>
 #include <errno.h>
+#include <inttypes.h>
 #include <stddef.h>
 
 #include <arch_helpers.h>
@@ -24,12 +25,13 @@
 #include <common/debug.h>
 #include <common/runtime_svc.h>
 #include <lib/el3_runtime/context_mgmt.h>
+#include <lib/optee_utils.h>
+#include <lib/xlat_tables/xlat_tables_v2.h>
 #include <plat/common/platform.h>
 #include <tools_share/uuid.h>
 
 #include "opteed_private.h"
 #include "teesmc_opteed.h"
-#include "teesmc_opteed_macros.h"
 
 /*******************************************************************************
  * Address of the entrypoint vector table in OPTEE. It is
@@ -43,7 +45,16 @@
 optee_context_t opteed_sp_context[OPTEED_CORE_COUNT];
 uint32_t opteed_rw;
 
+#if OPTEE_ALLOW_SMC_LOAD
+static bool opteed_allow_load;
+#else
 static int32_t opteed_init(void);
+#endif
+
+uint64_t dual32to64(uint32_t high, uint32_t low)
+{
+	return ((uint64_t)high << 32) | low;
+}
 
 /*******************************************************************************
  * This function is the handler registered for S-EL1 interrupts by the
@@ -93,6 +104,11 @@
  ******************************************************************************/
 static int32_t opteed_setup(void)
 {
+#if OPTEE_ALLOW_SMC_LOAD
+	opteed_allow_load = true;
+	INFO("Delaying OP-TEE setup until we receive an SMC call to load it\n");
+	return 0;
+#else
 	entry_point_info_t *optee_ep_info;
 	uint32_t linear_id;
 	uint64_t opteed_pageable_part;
@@ -142,6 +158,7 @@
 	bl31_register_bl32_init(&opteed_init);
 
 	return 0;
+#endif  /* OPTEE_ALLOW_SMC_LOAD */
 }
 
 /*******************************************************************************
@@ -153,18 +170,12 @@
  * non-secure state. This function performs a synchronous entry into
  * OPTEE. OPTEE passes control back to this routine through a SMC.
  ******************************************************************************/
-static int32_t opteed_init(void)
+static int32_t
+opteed_init_with_entry_point(entry_point_info_t *optee_entry_point)
 {
 	uint32_t linear_id = plat_my_core_pos();
 	optee_context_t *optee_ctx = &opteed_sp_context[linear_id];
-	entry_point_info_t *optee_entry_point;
 	uint64_t rc;
-
-	/*
-	 * Get information about the OPTEE (BL32) image. Its
-	 * absence is a critical failure.
-	 */
-	optee_entry_point = bl31_plat_get_next_image_ep_info(SECURE);
 	assert(optee_entry_point);
 
 	cm_init_my_context(optee_entry_point);
@@ -179,7 +190,116 @@
 	return rc;
 }
 
+#if !OPTEE_ALLOW_SMC_LOAD
+static int32_t opteed_init(void)
+{
+	entry_point_info_t *optee_entry_point;
+	/*
+	 * Get information about the OP-TEE (BL32) image. Its
+	 * absence is a critical failure.
+	 */
+	optee_entry_point = bl31_plat_get_next_image_ep_info(SECURE);
+	return opteed_init_with_entry_point(optee_entry_point);
+}
+#endif  /* !OPTEE_ALLOW_SMC_LOAD */
 
+#if OPTEE_ALLOW_SMC_LOAD
+/*******************************************************************************
+ * This function is responsible for handling the SMC that loads the OP-TEE
+ * binary image via a non-secure SMC call. It takes the size and physical
+ * address of the payload as parameters.
+ ******************************************************************************/
+static int32_t opteed_handle_smc_load(uint64_t data_size, uint32_t data_pa)
+{
+	uintptr_t data_va = data_pa;
+	uint64_t mapped_data_pa;
+	uintptr_t mapped_data_va;
+	uint64_t data_map_size;
+	int32_t rc;
+	optee_header_t *image_header;
+	uint8_t *image_ptr;
+	uint64_t target_pa;
+	uint64_t target_end_pa;
+	uint64_t image_pa;
+	uintptr_t image_va;
+	optee_image_t *curr_image;
+	uintptr_t target_va;
+	uint64_t target_size;
+	entry_point_info_t optee_ep_info;
+	uint32_t linear_id = plat_my_core_pos();
+
+	mapped_data_pa = page_align(data_pa, DOWN);
+	mapped_data_va = mapped_data_pa;
+	data_map_size = page_align(data_size + (mapped_data_pa - data_pa), UP);
+
+	rc = mmap_add_dynamic_region(mapped_data_pa, mapped_data_va,
+				     data_map_size, MT_MEMORY | MT_RO | MT_NS);
+	if (rc != 0) {
+		return rc;
+	}
+
+	image_header = (optee_header_t *)data_va;
+	if (image_header->magic != TEE_MAGIC_NUM_OPTEE ||
+	    image_header->version != 2 || image_header->nb_images != 1) {
+		mmap_remove_dynamic_region(mapped_data_va, data_map_size);
+		return -EINVAL;
+	}
+
+	image_ptr = (uint8_t *)data_va + sizeof(optee_header_t) +
+			sizeof(optee_image_t);
+	if (image_header->arch == 1) {
+		opteed_rw = OPTEE_AARCH64;
+	} else {
+		opteed_rw = OPTEE_AARCH32;
+	}
+
+	curr_image = &image_header->optee_image_list[0];
+	image_pa = dual32to64(curr_image->load_addr_hi,
+			      curr_image->load_addr_lo);
+	image_va = image_pa;
+	target_end_pa = image_pa + curr_image->size;
+
+	/* Now also map the memory we want to copy it to. */
+	target_pa = page_align(image_pa, DOWN);
+	target_va = target_pa;
+	target_size = page_align(target_end_pa, UP) - target_pa;
+
+	rc = mmap_add_dynamic_region(target_pa, target_va, target_size,
+				     MT_MEMORY | MT_RW | MT_SECURE);
+	if (rc != 0) {
+		mmap_remove_dynamic_region(mapped_data_va, data_map_size);
+		return rc;
+	}
+
+	INFO("Loaded OP-TEE via SMC: size %d addr 0x%" PRIx64 "\n",
+	     curr_image->size, image_va);
+
+	memcpy((void *)image_va, image_ptr, curr_image->size);
+	flush_dcache_range(target_pa, target_size);
+
+	mmap_remove_dynamic_region(mapped_data_va, data_map_size);
+	mmap_remove_dynamic_region(target_va, target_size);
+
+	/* Save the non-secure state */
+	cm_el1_sysregs_context_save(NON_SECURE);
+
+	opteed_init_optee_ep_state(&optee_ep_info,
+				   opteed_rw,
+				   image_pa,
+				   0,
+				   0,
+				   0,
+				   &opteed_sp_context[linear_id]);
+	rc = opteed_init_with_entry_point(&optee_ep_info);
+
+	/* Restore non-secure state */
+	cm_el1_sysregs_context_restore(NON_SECURE);
+	cm_set_next_eret_context(NON_SECURE);
+
+	return rc;
+}
+#endif  /* OPTEE_ALLOW_SMC_LOAD */
+
 /*******************************************************************************
  * This function is responsible for handling all SMCs in the Trusted OS/App
  * range from the non-secure state as defined in the SMC Calling Convention
@@ -207,6 +327,34 @@
 	 */
 
 	if (is_caller_non_secure(flags)) {
+#if OPTEE_ALLOW_SMC_LOAD
+		if (smc_fid == NSSMC_OPTEED_CALL_LOAD_IMAGE) {
+			/*
+			 * TODO: Consider wiping the code for SMC loading from
+			 * memory after it has been invoked similar to what is
+			 * done under RECLAIM_INIT, but extended to happen
+			 * later.
+			 */
+			if (!opteed_allow_load) {
+				SMC_RET1(handle, -EPERM);
+			}
+
+			opteed_allow_load = false;
+			uint64_t data_size = dual32to64(x1, x2);
+			uint64_t data_pa = dual32to64(x3, x4);
+			if (!data_size || !data_pa) {
+				/*
+				 * This is invoked when the OP-TEE image didn't
+				 * load correctly in the kernel but we want to
+				 * block off loading of it later for security
+				 * reasons.
+				 */
+				SMC_RET1(handle, -EINVAL);
+			}
+			SMC_RET1(handle, opteed_handle_smc_load(
+					data_size, data_pa));
+		}
+#endif  /* OPTEE_ALLOW_SMC_LOAD */
 		/*
 		 * This is a fresh request from the non-secure client.
 		 * The parameters are in x1 and x2. Figure out which
@@ -219,8 +367,18 @@
 
 		/*
 		 * We are done stashing the non-secure context. Ask the
-		 * OPTEE to do the work now.
+		 * OP-TEE to do the work now. If we are loading vi an SMC,
+		 * then we also need to init this CPU context if not done
+		 * already.
 		 */
+		if (optee_vector_table == NULL) {
+			SMC_RET1(handle, -EINVAL);
+		}
+
+		if (get_optee_pstate(optee_ctx->state) ==
+		    OPTEE_PSTATE_UNKNOWN) {
+			opteed_cpu_on_finish_handler(0);
+		}
 
 		/*
 		 * Verify if there is a valid context to use, copy the
diff --git a/services/spd/opteed/opteed_pm.c b/services/spd/opteed/opteed_pm.c
index 719eeb7..fa724a1 100644
--- a/services/spd/opteed/opteed_pm.c
+++ b/services/spd/opteed/opteed_pm.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013-2017, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2013-2023, ARM Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -32,6 +32,10 @@
 	uint32_t linear_id = plat_my_core_pos();
 	optee_context_t *optee_ctx = &opteed_sp_context[linear_id];
 
+	if (get_optee_pstate(optee_ctx->state) == OPTEE_PSTATE_UNKNOWN) {
+		return 0;
+	}
+
 	assert(optee_vector_table);
 	assert(get_optee_pstate(optee_ctx->state) == OPTEE_PSTATE_ON);
 
@@ -65,6 +69,10 @@
 	uint32_t linear_id = plat_my_core_pos();
 	optee_context_t *optee_ctx = &opteed_sp_context[linear_id];
 
+	if (get_optee_pstate(optee_ctx->state) == OPTEE_PSTATE_UNKNOWN) {
+		return;
+	}
+
 	assert(optee_vector_table);
 	assert(get_optee_pstate(optee_ctx->state) == OPTEE_PSTATE_ON);
 
@@ -92,7 +100,7 @@
  * after initialising minimal architectural state that guarantees safe
  * execution.
  ******************************************************************************/
-static void opteed_cpu_on_finish_handler(u_register_t unused)
+void opteed_cpu_on_finish_handler(u_register_t unused)
 {
 	int32_t rc = 0;
 	uint32_t linear_id = plat_my_core_pos();
@@ -100,7 +108,8 @@
 	entry_point_info_t optee_on_entrypoint;
 
 	assert(optee_vector_table);
-	assert(get_optee_pstate(optee_ctx->state) == OPTEE_PSTATE_OFF);
+	assert(get_optee_pstate(optee_ctx->state) == OPTEE_PSTATE_OFF ||
+	       get_optee_pstate(optee_ctx->state) == OPTEE_PSTATE_UNKNOWN);
 
 	opteed_init_optee_ep_state(&optee_on_entrypoint, opteed_rw,
 				(uint64_t)&optee_vector_table->cpu_on_entry,
@@ -134,6 +143,10 @@
 	uint32_t linear_id = plat_my_core_pos();
 	optee_context_t *optee_ctx = &opteed_sp_context[linear_id];
 
+	if (get_optee_pstate(optee_ctx->state) == OPTEE_PSTATE_UNKNOWN) {
+		return;
+	}
+
 	assert(optee_vector_table);
 	assert(get_optee_pstate(optee_ctx->state) == OPTEE_PSTATE_SUSPEND);
 
@@ -173,6 +186,14 @@
 	uint32_t linear_id = plat_my_core_pos();
 	optee_context_t *optee_ctx = &opteed_sp_context[linear_id];
 
+	/*
+	 * OP-TEE must have been initialized in order to reach this location so
+	 * it is safe to init the CPU context if not already done for this core.
+	 */
+	if (get_optee_pstate(optee_ctx->state) == OPTEE_PSTATE_UNKNOWN) {
+		opteed_cpu_on_finish_handler(0);
+	}
+
 	assert(optee_vector_table);
 	assert(get_optee_pstate(optee_ctx->state) == OPTEE_PSTATE_ON);
 
@@ -193,6 +214,14 @@
 	uint32_t linear_id = plat_my_core_pos();
 	optee_context_t *optee_ctx = &opteed_sp_context[linear_id];
 
+	/*
+	 * OP-TEE must have been initialized in order to reach this location so
+	 * it is safe to init the CPU context if not already done for this core.
+	 */
+	if (get_optee_pstate(optee_ctx->state) == OPTEE_PSTATE_UNKNOWN) {
+		opteed_cpu_on_finish_handler(0);
+	}
+
 	assert(optee_vector_table);
 	assert(get_optee_pstate(optee_ctx->state) == OPTEE_PSTATE_ON);
 
diff --git a/services/spd/opteed/opteed_private.h b/services/spd/opteed/opteed_private.h
index 242154f..ab6e4cd 100644
--- a/services/spd/opteed/opteed_private.h
+++ b/services/spd/opteed/opteed_private.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013-2017, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2013-2023, ARM Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -18,9 +18,10 @@
  * OPTEE PM state information e.g. OPTEE is suspended, uninitialised etc
  * and macros to access the state information in the per-cpu 'state' flags
  ******************************************************************************/
-#define OPTEE_PSTATE_OFF		0
-#define OPTEE_PSTATE_ON			1
-#define OPTEE_PSTATE_SUSPEND		2
+#define OPTEE_PSTATE_OFF		1
+#define OPTEE_PSTATE_ON			2
+#define OPTEE_PSTATE_SUSPEND		3
+#define OPTEE_PSTATE_UNKNOWN		0
 #define OPTEE_PSTATE_SHIFT		0
 #define OPTEE_PSTATE_MASK		0x3
 #define get_optee_pstate(state)	((state >> OPTEE_PSTATE_SHIFT) & \
@@ -153,6 +154,7 @@
 				uint64_t mem_limit,
 				uint64_t dt_addr,
 				optee_context_t *optee_ctx);
+void opteed_cpu_on_finish_handler(u_register_t unused);
 
 extern optee_context_t opteed_sp_context[OPTEED_CORE_COUNT];
 extern uint32_t opteed_rw;
diff --git a/services/spd/opteed/teesmc_opteed.h b/services/spd/opteed/teesmc_opteed.h
index c82b58a..eae3ed2 100644
--- a/services/spd/opteed/teesmc_opteed.h
+++ b/services/spd/opteed/teesmc_opteed.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2014-2019, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2014-2023, ARM Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -9,8 +9,10 @@
 #ifndef TEESMC_OPTEED_H
 #define TEESMC_OPTEED_H
 
+#include "teesmc_opteed_macros.h"
+
 /*
- * This file specifies SMC function IDs used when returning from TEE to the
+ * This section specifies SMC function IDs used when returning from TEE to the
  * secure monitor.
  *
  * All SMC Function IDs indicates SMC32 Calling Convention but will carry
@@ -120,4 +122,39 @@
 #define TEESMC_OPTEED_RETURN_SYSTEM_RESET_DONE \
 	TEESMC_OPTEED_RV(TEESMC_OPTEED_FUNCID_RETURN_SYSTEM_RESET_DONE)
 
+/*
+ * This section specifies SMC function IDs used when the secure monitor is
+ * invoked from the non-secure world.
+ */
+
+/*
+ * Load OP-TEE image from the payload specified in the registers.
+ *
+ * WARNING: Use this cautiously as it could lead to insecure loading of the
+ * Trusted OS. Further details are in opteed.mk.
+ *
+ * Call register usage:
+ * x0 SMC Function ID, OPTEE_SMC_CALL_LOAD_IMAGE
+ * x1 Upper 32bit of a 64bit size for the payload
+ * x2 Lower 32bit of a 64bit size for the payload
+ * x3 Upper 32bit of the physical address for the payload
+ * x4 Lower 32bit of the physical address for the payload
+ *
+ * The payload consists of a optee_header struct that contains optee_image
+ * structs in a flex array, immediately following that in memory is the data
+ * referenced by the optee_image structs.
+ * Example:
+ *
+ * struct optee_header (with n images specified)
+ * image 0 data
+ * image 1 data
+ * ...
+ * image n-1 data
+ *
+ * Returns 0 on success and an error code otherwise.
+ */
+#define NSSMC_OPTEED_FUNCID_LOAD_IMAGE 2
+#define NSSMC_OPTEED_CALL_LOAD_IMAGE \
+	NSSMC_OPTEED_CALL(NSSMC_OPTEED_FUNCID_LOAD_IMAGE)
+
 #endif /*TEESMC_OPTEED_H*/
diff --git a/services/spd/opteed/teesmc_opteed_macros.h b/services/spd/opteed/teesmc_opteed_macros.h
index 9d8a169..ad3ed75 100644
--- a/services/spd/opteed/teesmc_opteed_macros.h
+++ b/services/spd/opteed/teesmc_opteed_macros.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2014, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2014-2023, ARM Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -14,4 +14,10 @@
 		 (62 << FUNCID_OEN_SHIFT) | \
 		 ((func_num) & FUNCID_NUM_MASK))
 
+#define NSSMC_OPTEED_CALL(func_num) \
+		((SMC_TYPE_FAST << FUNCID_TYPE_SHIFT) | \
+		((SMC_32) << FUNCID_CC_SHIFT) | \
+		(50 << FUNCID_OEN_SHIFT) | \
+		((func_num) & FUNCID_NUM_MASK))
+
 #endif /* TEESMC_OPTEED_MACROS_H */