feat(rme): add context management changes for FEAT_RME

This patch adds a new context for realm world and realm world
awareness in context management.

Signed-off-by: Zelalem Aweke <zelalem.aweke@arm.com>
Signed-off-by: Subhasish Ghosh <subhasish.ghosh@arm.com>
Change-Id: Ic17469393603e789d7adc025880346bc3d6233d7
diff --git a/bl31/bl31_context_mgmt.c b/bl31/bl31_context_mgmt.c
index 9175ee3..34f69ad 100644
--- a/bl31/bl31_context_mgmt.c
+++ b/bl31/bl31_context_mgmt.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013-2020, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2013-2021, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -19,9 +19,9 @@
  ******************************************************************************/
 void *cm_get_context(uint32_t security_state)
 {
-	assert(security_state <= NON_SECURE);
+	assert(sec_state_is_valid(security_state));
 
-	return get_cpu_data(cpu_context[security_state]);
+	return get_cpu_data(cpu_context[get_cpu_context_index(security_state)]);
 }
 
 /*******************************************************************************
@@ -30,9 +30,10 @@
  ******************************************************************************/
 void cm_set_context(void *context, uint32_t security_state)
 {
-	assert(security_state <= NON_SECURE);
+	assert(sec_state_is_valid(security_state));
 
-	set_cpu_data(cpu_context[security_state], context);
+	set_cpu_data(cpu_context[get_cpu_context_index(security_state)],
+			context);
 }
 
 /*******************************************************************************
@@ -46,7 +47,8 @@
 {
 	assert(sec_state_is_valid(security_state));
 
-	return get_cpu_data_by_index(cpu_idx, cpu_context[security_state]);
+	return get_cpu_data_by_index(cpu_idx,
+			cpu_context[get_cpu_context_index(security_state)]);
 }
 
 /*******************************************************************************
@@ -58,5 +60,7 @@
 {
 	assert(sec_state_is_valid(security_state));
 
-	set_cpu_data_by_index(cpu_idx, cpu_context[security_state], context);
+	set_cpu_data_by_index(cpu_idx,
+			cpu_context[get_cpu_context_index(security_state)],
+			context);
 }
diff --git a/include/lib/el3_runtime/aarch64/context.h b/include/lib/el3_runtime/aarch64/context.h
index c3f4117..698e208 100644
--- a/include/lib/el3_runtime/aarch64/context.h
+++ b/include/lib/el3_runtime/aarch64/context.h
@@ -405,13 +405,12 @@
 					 = (uint64_t) (val))
 
 /*
- * Top-level context structure which is used by EL3 firmware to
- * preserve the state of a core at EL1 in one of the two security
- * states and save enough EL3 meta data to be able to return to that
- * EL and security state. The context management library will be used
- * to ensure that SP_EL3 always points to an instance of this
- * structure at exception entry and exit. Each instance will
- * correspond to either the secure or the non-secure state.
+ * Top-level context structure which is used by EL3 firmware to preserve
+ * the state of a core at the next lower EL in a given security state and
+ * save enough EL3 meta data to be able to return to that EL and security
+ * state. The context management library will be used to ensure that
+ * SP_EL3 always points to an instance of this structure at exception
+ * entry and exit.
  */
 typedef struct cpu_context {
 	gp_regs_t gpregs_ctx;
diff --git a/include/lib/el3_runtime/cpu_data.h b/include/lib/el3_runtime/cpu_data.h
index 3d57a5c..2c7b619 100644
--- a/include/lib/el3_runtime/cpu_data.h
+++ b/include/lib/el3_runtime/cpu_data.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2014-2021, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2014-2021, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -19,16 +19,25 @@
 /* 8-bytes aligned size of psci_cpu_data structure */
 #define PSCI_CPU_DATA_SIZE_ALIGNED	((PSCI_CPU_DATA_SIZE + 7) & ~7)
 
+#if ENABLE_RME
+/* Size of cpu_context array */
+#define CPU_DATA_CONTEXT_NUM		3
 /* Offset of cpu_ops_ptr, size 8 bytes */
+#define CPU_DATA_CPU_OPS_PTR		0x18
+#else /* ENABLE_RME */
+#define CPU_DATA_CONTEXT_NUM		2
 #define CPU_DATA_CPU_OPS_PTR		0x10
+#endif /* ENABLE_RME */
 
 #if ENABLE_PAUTH
 /* 8-bytes aligned offset of apiakey[2], size 16 bytes */
-#define	CPU_DATA_APIAKEY_OFFSET		(0x18 + PSCI_CPU_DATA_SIZE_ALIGNED)
-#define CPU_DATA_CRASH_BUF_OFFSET	(CPU_DATA_APIAKEY_OFFSET + 0x10)
-#else
-#define CPU_DATA_CRASH_BUF_OFFSET	(0x18 + PSCI_CPU_DATA_SIZE_ALIGNED)
-#endif	/* ENABLE_PAUTH */
+#define	CPU_DATA_APIAKEY_OFFSET		(0x8 + PSCI_CPU_DATA_SIZE_ALIGNED \
+					     + CPU_DATA_CPU_OPS_PTR)
+#define CPU_DATA_CRASH_BUF_OFFSET	(0x10 + CPU_DATA_APIAKEY_OFFSET)
+#else /* ENABLE_PAUTH */
+#define CPU_DATA_CRASH_BUF_OFFSET	(0x8 + PSCI_CPU_DATA_SIZE_ALIGNED \
+					     + CPU_DATA_CPU_OPS_PTR)
+#endif /* ENABLE_PAUTH */
 
 /* need enough space in crash buffer to save 8 registers */
 #define CPU_DATA_CRASH_BUF_SIZE		64
@@ -65,11 +74,14 @@
 
 #ifndef __ASSEMBLER__
 
+#include <assert.h>
+#include <stdint.h>
+
 #include <arch_helpers.h>
 #include <lib/cassert.h>
 #include <lib/psci/psci.h>
+
 #include <platform_def.h>
-#include <stdint.h>
 
 /* Offsets for the cpu_data structure */
 #define CPU_DATA_PSCI_LOCK_OFFSET	__builtin_offsetof\
@@ -80,27 +92,34 @@
 		(cpu_data_t, platform_cpu_data)
 #endif
 
+typedef enum context_pas {
+	CPU_CONTEXT_SECURE = 0,
+	CPU_CONTEXT_NS,
+#if ENABLE_RME
+	CPU_CONTEXT_REALM,
+#endif
+	CPU_CONTEXT_NUM
+} context_pas_t;
+
 /*******************************************************************************
  * Function & variable prototypes
  ******************************************************************************/
 
 /*******************************************************************************
  * Cache of frequently used per-cpu data:
- *   Pointers to non-secure and secure security state contexts
+ *   Pointers to non-secure, realm, and secure security state contexts
  *   Address of the crash stack
  * It is aligned to the cache line boundary to allow efficient concurrent
  * manipulation of these pointers on different cpus
  *
- * TODO: Add other commonly used variables to this (tf_issues#90)
- *
  * The data structure and the _cpu_data accessors should not be used directly
  * by components that have per-cpu members. The member access macros should be
  * used for this.
  ******************************************************************************/
 typedef struct cpu_data {
 #ifdef __aarch64__
-	void *cpu_context[2];
-#endif
+	void *cpu_context[CPU_DATA_CONTEXT_NUM];
+#endif /* __aarch64__ */
 	uintptr_t cpu_ops_ptr;
 	struct psci_cpu_data psci_svc_cpu_data;
 #if ENABLE_PAUTH
@@ -122,6 +141,11 @@
 
 extern cpu_data_t percpu_data[PLATFORM_CORE_COUNT];
 
+#ifdef __aarch64__
+CASSERT(CPU_DATA_CONTEXT_NUM == CPU_CONTEXT_NUM,
+		assert_cpu_data_context_num_mismatch);
+#endif
+
 #if ENABLE_PAUTH
 CASSERT(CPU_DATA_APIAKEY_OFFSET == __builtin_offsetof
 	(cpu_data_t, apiakey),
@@ -160,6 +184,31 @@
 struct cpu_data *_cpu_data(void);
 #endif
 
+/*
+ * Returns the index of the cpu_context array for the given security state.
+ * All accesses to cpu_context should be through this helper to make sure
+ * an access is not out-of-bounds. The function assumes security_state is
+ * valid.
+ */
+static inline context_pas_t get_cpu_context_index(uint32_t security_state)
+{
+	if (security_state == SECURE) {
+		return CPU_CONTEXT_SECURE;
+	} else {
+#if ENABLE_RME
+		if (security_state == NON_SECURE) {
+			return CPU_CONTEXT_NS;
+		} else {
+			assert(security_state == REALM);
+			return CPU_CONTEXT_REALM;
+		}
+#else
+		assert(security_state == NON_SECURE);
+		return CPU_CONTEXT_NS;
+#endif
+	}
+}
+
 /**************************************************************************
  * APIs for initialising and accessing per-cpu data
  *************************************************************************/
diff --git a/lib/el3_runtime/aarch64/context_mgmt.c b/lib/el3_runtime/aarch64/context_mgmt.c
index 08022d4..0ec7e7e 100644
--- a/lib/el3_runtime/aarch64/context_mgmt.c
+++ b/lib/el3_runtime/aarch64/context_mgmt.c
@@ -93,24 +93,41 @@
 	scr_el3 = read_scr();
 	scr_el3 &= ~(SCR_NS_BIT | SCR_RW_BIT | SCR_FIQ_BIT | SCR_IRQ_BIT |
 			SCR_ST_BIT | SCR_HCE_BIT);
+
+#if ENABLE_RME
+	/* When RME support is enabled, clear the NSE bit as well. */
+	scr_el3 &= ~SCR_NSE_BIT;
+#endif /* ENABLE_RME */
+
 	/*
 	 * SCR_NS: Set the security state of the next EL.
 	 */
-	if (security_state != SECURE)
+	if (security_state == NON_SECURE) {
 		scr_el3 |= SCR_NS_BIT;
+	}
+
+#if ENABLE_RME
+	/* Check for realm state if RME support enabled. */
+	if (security_state == REALM) {
+		scr_el3 |= SCR_NS_BIT | SCR_NSE_BIT | SCR_EnSCXT_BIT;
+	}
+#endif /* ENABLE_RME */
+
 	/*
 	 * SCR_EL3.RW: Set the execution state, AArch32 or AArch64, for next
 	 *  Exception level as specified by SPSR.
 	 */
-	if (GET_RW(ep->spsr) == MODE_RW_64)
+	if (GET_RW(ep->spsr) == MODE_RW_64) {
 		scr_el3 |= SCR_RW_BIT;
+	}
 	/*
 	 * SCR_EL3.ST: Traps Secure EL1 accesses to the Counter-timer Physical
 	 *  Secure timer registers to EL3, from AArch64 state only, if specified
 	 *  by the entrypoint attributes.
 	 */
-	if (EP_GET_ST(ep->h.attr) != 0U)
+	if (EP_GET_ST(ep->h.attr) != 0U) {
 		scr_el3 |= SCR_ST_BIT;
+	}
 
 	/*
 	 * If FEAT_HCX is enabled, enable access to HCRX_EL2 by setting
@@ -152,8 +169,9 @@
 	 * If the Secure world wants to use pointer authentication,
 	 * CTX_INCLUDE_PAUTH_REGS must be set to 1.
 	 */
-	if (security_state == NON_SECURE)
+	if (security_state == NON_SECURE) {
 		scr_el3 |= SCR_API_BIT | SCR_APK_BIT;
+	}
 #endif /* !CTX_INCLUDE_PAUTH_REGS */
 
 #if !CTX_INCLUDE_MTE_REGS || ENABLE_ASSERTIONS
@@ -188,8 +206,14 @@
 	/*
 	 * SCR_EL3.IRQ, SCR_EL3.FIQ: Enable the physical FIQ and IRQ routing as
 	 *  indicated by the interrupt routing model for BL31.
+	 *
+	 * TODO: The interrupt routing model code is not updated for REALM
+	 * state. Use the default values of IRQ = FIQ = 0 for REALM security
+	 * state for now.
 	 */
-	scr_el3 |= get_scr_el3_from_routing_model(security_state);
+	if (security_state != REALM) {
+		scr_el3 |= get_scr_el3_from_routing_model(security_state);
+	}
 #endif
 
 	/* Save the initialized value of CPTR_EL3 register */
@@ -256,9 +280,9 @@
 	 *  required by PSCI specification)
 	 */
 	sctlr_elx = (EP_GET_EE(ep->h.attr) != 0U) ? SCTLR_EE_BIT : 0U;
-	if (GET_RW(ep->spsr) == MODE_RW_64)
+	if (GET_RW(ep->spsr) == MODE_RW_64) {
 		sctlr_elx |= SCTLR_EL1_RES1;
-	else {
+	} else {
 		/*
 		 * If the target execution state is AArch32 then the following
 		 * fields need to be set.
@@ -413,7 +437,8 @@
 }
 
 /*******************************************************************************
- * Prepare the CPU system registers for first entry into secure or normal world
+ * Prepare the CPU system registers for first entry into realm, secure, or
+ * normal world.
  *
  * If execution is requested to EL2 or hyp mode, SCTLR_EL2 is initialized
  * If execution is requested to non-secure EL1 or svc mode, and the CPU supports
@@ -497,7 +522,7 @@
 			 * architecturally UNKNOWN on reset and are set to zero
 			 * except for field(s) listed below.
 			 *
-			 * CNTHCTL_EL2.EL1PCEN: Set to one to disable traps to
+			 * CNTHCTL_EL2.EL1PTEN: Set to one to disable traps to
 			 *  Hyp mode of Non-secure EL0 and EL1 accesses to the
 			 *  physical timer registers.
 			 *
@@ -645,10 +670,10 @@
 	u_register_t scr_el3 = read_scr();
 
 	/*
-	 * Always save the non-secure EL2 context, only save the
+	 * Always save the non-secure and realm EL2 context, only save the
 	 * S-EL2 context if S-EL2 is enabled.
 	 */
-	if ((security_state == NON_SECURE) ||
+	if ((security_state != SECURE) ||
 	    ((security_state == SECURE) && ((scr_el3 & SCR_EEL2_BIT) != 0U))) {
 		cpu_context_t *ctx;
 
@@ -667,10 +692,10 @@
 	u_register_t scr_el3 = read_scr();
 
 	/*
-	 * Always restore the non-secure EL2 context, only restore the
+	 * Always restore the non-secure and realm EL2 context, only restore the
 	 * S-EL2 context if S-EL2 is enabled.
 	 */
-	if ((security_state == NON_SECURE) ||
+	if ((security_state != SECURE) ||
 	    ((security_state == SECURE) && ((scr_el3 & SCR_EEL2_BIT) != 0U))) {
 		cpu_context_t *ctx;