refactor(context mgmt): add cm_prepare_el3_exit_ns function

As part of the RFC:
https://review.trustedfirmware.org/c/TF-A/trusted-firmware-a/+/13651,
this patch adds the 'cm_prepare_el3_exit_ns' function. The function is
a wrapper to 'cm_prepare_el3_exit' function for Non-secure state.

When EL2 sysregs context exists (CTX_INCLUDE_EL2_REGS is
enabled) EL1 and EL2 sysreg values are restored from the context
instead of directly updating the registers.

Signed-off-by: Zelalem Aweke <zelalem.aweke@arm.com>
Change-Id: I9b071030576bb05500d54090e2a03b3f125d1653
diff --git a/bl31/bl31_main.c b/bl31/bl31_main.c
index 797093e..2a3d838 100644
--- a/bl31/bl31_main.c
+++ b/bl31/bl31_main.c
@@ -259,7 +259,16 @@
 		(image_type == SECURE) ? "secure" : "normal");
 	print_entry_point_info(next_image_info);
 	cm_init_my_context(next_image_info);
-	cm_prepare_el3_exit(image_type);
+
+	/*
+	* If we are entering the Non-secure world, use
+	* 'cm_prepare_el3_exit_ns' to exit.
+	*/
+	if (image_type == NON_SECURE) {
+		cm_prepare_el3_exit_ns();
+	} else {
+		cm_prepare_el3_exit(image_type);
+	}
 }
 
 /*******************************************************************************
diff --git a/include/lib/el3_runtime/context_mgmt.h b/include/lib/el3_runtime/context_mgmt.h
index 2090687..1a76d8e 100644
--- a/include/lib/el3_runtime/context_mgmt.h
+++ b/include/lib/el3_runtime/context_mgmt.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013-2020, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2013-2022, ARM Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -34,6 +34,7 @@
 			      const struct entry_point_info *ep);
 void cm_setup_context(cpu_context_t *ctx, const struct entry_point_info *ep);
 void cm_prepare_el3_exit(uint32_t security_state);
+void cm_prepare_el3_exit_ns(void);
 
 #ifdef __aarch64__
 #if CTX_INCLUDE_EL2_REGS
diff --git a/lib/el3_runtime/aarch32/context_mgmt.c b/lib/el3_runtime/aarch32/context_mgmt.c
index 3ef378c..af8edf5 100644
--- a/lib/el3_runtime/aarch32/context_mgmt.c
+++ b/lib/el3_runtime/aarch32/context_mgmt.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2021, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2016-2022, ARM Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -332,3 +332,12 @@
 		enable_extensions_nonsecure(el2_unused);
 	}
 }
+
+/*******************************************************************************
+ * This function is used to exit to Non-secure world. It simply calls the
+ * cm_prepare_el3_exit function for AArch32.
+ ******************************************************************************/
+void cm_prepare_el3_exit_ns(void)
+{
+	cm_prepare_el3_exit(NON_SECURE);
+}
diff --git a/lib/el3_runtime/aarch64/context_mgmt.c b/lib/el3_runtime/aarch64/context_mgmt.c
index 72281f2..459ca2c 100644
--- a/lib/el3_runtime/aarch64/context_mgmt.c
+++ b/lib/el3_runtime/aarch64/context_mgmt.c
@@ -16,6 +16,7 @@
 #include <bl31/interrupt_mgmt.h>
 #include <common/bl_common.h>
 #include <context.h>
+#include <drivers/arm/gicv3.h>
 #include <lib/el3_runtime/context_mgmt.h>
 #include <lib/el3_runtime/pubsub_events.h>
 #include <lib/extensions/amu.h>
@@ -141,6 +142,31 @@
 	scr_el3 |= get_scr_el3_from_routing_model(NON_SECURE);
 #endif
 	write_ctx_reg(state, CTX_SCR_EL3, scr_el3);
+
+	/* Initialize EL2 context registers */
+#if CTX_INCLUDE_EL2_REGS
+
+	/*
+	 * Initialize SCTLR_EL2 context register using Endianness value
+	 * taken from the entrypoint attribute.
+	 */
+	u_register_t sctlr_el2 = (EP_GET_EE(ep->h.attr) != 0U) ? SCTLR_EE_BIT : 0UL;
+	sctlr_el2 |= SCTLR_EL2_RES1;
+	write_ctx_reg(get_el2_sysregs_ctx(ctx), CTX_SCTLR_EL2,
+			sctlr_el2);
+
+	/*
+	 * The GICv3 driver initializes the ICC_SRE_EL2 register during
+	 * platform setup. Use the same setting for the corresponding
+	 * context register to make sure the correct bits are set when
+	 * restoring NS context.
+	 */
+	u_register_t icc_sre_el2 = read_icc_sre_el2();
+	icc_sre_el2 |= (ICC_SRE_DIB_BIT | ICC_SRE_DFB_BIT);
+	icc_sre_el2 |= (ICC_SRE_EN_BIT | ICC_SRE_SRE_BIT);
+	write_ctx_reg(get_el2_sysregs_ctx(ctx), CTX_ICC_SRE_EL2,
+			icc_sre_el2);
+#endif /* CTX_INCLUDE_EL2_REGS */
 }
 
 /*******************************************************************************
@@ -792,6 +818,40 @@
 #endif /* CTX_INCLUDE_EL2_REGS */
 
 /*******************************************************************************
+ * This function is used to exit to Non-secure world. If CTX_INCLUDE_EL2_REGS
+ * is enabled, it restores EL1 and EL2 sysreg contexts instead of directly
+ * updating EL1 and EL2 registers. Otherwise, it calls the generic
+ * cm_prepare_el3_exit function.
+ ******************************************************************************/
+void cm_prepare_el3_exit_ns(void)
+{
+#if CTX_INCLUDE_EL2_REGS
+	cpu_context_t *ctx = cm_get_context(NON_SECURE);
+	assert(ctx != NULL);
+
+	/*
+	 * Currently some extensions are configured using
+	 * direct register updates. Therefore, do this here
+	 * instead of when setting up context.
+	 */
+	manage_extensions_nonsecure(0, ctx);
+
+	/*
+	 * Set the NS bit to be able to access the ICC_SRE_EL2
+	 * register when restoring context.
+	 */
+	write_scr_el3(read_scr_el3() | SCR_NS_BIT);
+
+	/* Restore EL2 and EL1 sysreg contexts */
+	cm_el2_sysregs_context_restore(NON_SECURE);
+	cm_el1_sysregs_context_restore(NON_SECURE);
+	cm_set_next_eret_context(NON_SECURE);
+#else
+	cm_prepare_el3_exit(NON_SECURE);
+#endif /* CTX_INCLUDE_EL2_REGS */
+}
+
+/*******************************************************************************
  * The next four functions are used by runtime services to save and restore
  * EL1 context on the 'cpu_context' structure for the specified security
  * state.
diff --git a/lib/psci/psci_on.c b/lib/psci/psci_on.c
index dd48e10..c70b377 100644
--- a/lib/psci/psci_on.c
+++ b/lib/psci/psci_on.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013-2020, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2013-2022, ARM Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -229,5 +229,5 @@
 	 * information that we had stashed away during the cpu_on
 	 * call to set this cpu on its way.
 	 */
-	cm_prepare_el3_exit(NON_SECURE);
+	cm_prepare_el3_exit_ns();
 }
diff --git a/lib/psci/psci_suspend.c b/lib/psci/psci_suspend.c
index da9f328..ffe3a91 100644
--- a/lib/psci/psci_suspend.c
+++ b/lib/psci/psci_suspend.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013-2020, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2013-2022, ARM Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -331,5 +331,5 @@
 	 * information that we had stashed away during the suspend
 	 * call to set this cpu on its way.
 	 */
-	cm_prepare_el3_exit(NON_SECURE);
+	cm_prepare_el3_exit_ns();
 }
diff --git a/plat/arm/common/aarch64/execution_state_switch.c b/plat/arm/common/aarch64/execution_state_switch.c
index bed929a..2353e6a 100644
--- a/plat/arm/common/aarch64/execution_state_switch.c
+++ b/plat/arm/common/aarch64/execution_state_switch.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017-2020, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2017-2022, ARM Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -162,7 +162,7 @@
 	 * calling EL.
 	 */
 	cm_init_my_context(&ep);
-	cm_prepare_el3_exit(NON_SECURE);
+	cm_prepare_el3_exit_ns();
 
 	/*
 	 * State switch success. The caller of SMC wouldn't see the SMC
diff --git a/services/std_svc/rmmd/rmmd_main.c b/services/std_svc/rmmd/rmmd_main.c
index cf5ff7b..746419e 100644
--- a/services/std_svc/rmmd/rmmd_main.c
+++ b/services/std_svc/rmmd/rmmd_main.c
@@ -60,10 +60,6 @@
 
 	cm_set_context(&(rmm_ctx->cpu_ctx), REALM);
 
-	/* Save the current el1/el2 context before loading realm context. */
-	cm_el1_sysregs_context_save(NON_SECURE);
-	cm_el2_sysregs_context_save(NON_SECURE);
-
 	/* Restore the realm context assigned above */
 	cm_el1_sysregs_context_restore(REALM);
 	cm_el2_sysregs_context_restore(REALM);
@@ -72,14 +68,15 @@
 	/* Enter RMM */
 	rc = rmmd_rmm_enter(&rmm_ctx->c_rt_ctx);
 
-	/* Save realm context */
+	/*
+	 * Save realm context. EL1 and EL2 Non-secure
+	 * contexts will be restored before exiting to
+	 * Non-secure world, therefore there is no need
+	 * to clear EL1 and EL2 context registers.
+	 */
 	cm_el1_sysregs_context_save(REALM);
 	cm_el2_sysregs_context_save(REALM);
 
-	/* Restore the el1/el2 context again. */
-	cm_el1_sysregs_context_restore(NON_SECURE);
-	cm_el2_sysregs_context_restore(NON_SECURE);
-
 	return rc;
 }