feat(cm): context switch MDCR_EL3 register

Currently MDCR_EL3 register value is same for all the
worlds(Non-secure, Secure, Realm and Root).

With this approach, features enable/disable settings
remain same across all the worlds. This is not ideal as
there must be flexibility in controlling feature as per
the requirements for individual world.

The patch addresses this by providing MDCR_EL3 a per world
value. Features with identical values for all the worlds are
grouped under ``manage_extensions_common`` API.

Change-Id: Ibc068d985fe165d8cb6d0ffb84119bffd743b3d1
Signed-off-by: Jayanth Dodderi Chidanand <jayanthdodderi.chidanand@arm.com>
diff --git a/include/arch/aarch64/el3_common_macros.S b/include/arch/aarch64/el3_common_macros.S
index 1a3e9b6..b4c5c1b 100644
--- a/include/arch/aarch64/el3_common_macros.S
+++ b/include/arch/aarch64/el3_common_macros.S
@@ -88,21 +88,8 @@
 	 * MDCR_EL3.SDD: Set to one to disable AArch64 Secure self-hosted debug.
 	 *  Debug exceptions, other than Breakpoint Instruction exceptions, are
 	 *  disabled from all ELs in Secure state.
-	 *
-	 * MDCR_EL3.SPD32: Set to 0b10 to disable AArch32 Secure self-hosted
-	 *  privileged debug from S-EL1.
-	 *
-	 * MDCR_EL3.TDOSA: Set to zero so that EL2 and EL2 System register
-	 *  access to the powerdown debug registers do not trap to EL3.
-	 *
-	 * MDCR_EL3.TDA: Set to zero to allow EL0, EL1 and EL2 access to the
-	 *  debug registers, other than those registers that are controlled by
-	 *  MDCR_EL3.TDOSA.
 	 */
-	mov_imm	x0, ((MDCR_EL3_RESET_VAL | MDCR_SDD_BIT | \
-		      MDCR_SPD32(MDCR_SPD32_DISABLE)) & \
-		    ~(MDCR_TDOSA_BIT | MDCR_TDA_BIT))
-
+	mov_imm	x0, (MDCR_EL3_RESET_VAL | MDCR_SDD_BIT)
 	msr	mdcr_el3, x0
 
 	/* ---------------------------------------------------------------------
diff --git a/include/lib/el3_runtime/aarch64/context.h b/include/lib/el3_runtime/aarch64/context.h
index 6f97bed..acf111b 100644
--- a/include/lib/el3_runtime/aarch64/context.h
+++ b/include/lib/el3_runtime/aarch64/context.h
@@ -63,19 +63,20 @@
 #define CTX_ELR_EL3		U(0x20)
 #define CTX_PMCR_EL0		U(0x28)
 #define CTX_IS_IN_EL3		U(0x30)
+#define CTX_MDCR_EL3		U(0x38)
 /* Constants required in supporting nested exception in EL3 */
-#define CTX_SAVED_ELR_EL3	U(0x38)
+#define CTX_SAVED_ELR_EL3	U(0x40)
 /*
  * General purpose flag, to save various EL3 states
  * FFH mode : Used to identify if handling nested exception
  * KFH mode : Used as counter value
  */
-#define CTX_NESTED_EA_FLAG	U(0x40)
+#define CTX_NESTED_EA_FLAG	U(0x48)
 #if FFH_SUPPORT
- #define CTX_SAVED_ESR_EL3	U(0x48)
- #define CTX_SAVED_SPSR_EL3	U(0x50)
- #define CTX_SAVED_GPREG_LR	U(0x58)
- #define CTX_EL3STATE_END	U(0x60) /* Align to the next 16 byte boundary */
+ #define CTX_SAVED_ESR_EL3	U(0x50)
+ #define CTX_SAVED_SPSR_EL3	U(0x58)
+ #define CTX_SAVED_GPREG_LR	U(0x60)
+ #define CTX_EL3STATE_END	U(0x70) /* Align to the next 16 byte boundary */
 #else
  #define CTX_EL3STATE_END	U(0x50) /* Align to the next 16 byte boundary */
 #endif /* FFH_SUPPORT */
diff --git a/include/lib/extensions/brbe.h b/include/lib/extensions/brbe.h
index 194efba..425a037 100644
--- a/include/lib/extensions/brbe.h
+++ b/include/lib/extensions/brbe.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, Arm Limited. All rights reserved.
+ * Copyright (c) 2022-2024, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -7,10 +7,12 @@
 #ifndef BRBE_H
 #define BRBE_H
 
+#include <context.h>
+
 #if ENABLE_BRBE_FOR_NS
-void brbe_init_el3(void);
+void brbe_enable(cpu_context_t *ctx);
 #else
-static inline void brbe_init_el3(void)
+static inline void brbe_enable(cpu_context_t *ctx)
 {
 }
 #endif /* ENABLE_BRBE_FOR_NS */
diff --git a/include/lib/extensions/spe.h b/include/lib/extensions/spe.h
index 7b39037..c6e44f9 100644
--- a/include/lib/extensions/spe.h
+++ b/include/lib/extensions/spe.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017-2023, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2017-2024, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -8,13 +8,14 @@
 #define SPE_H
 
 #include <stdbool.h>
+#include <context.h>
 
 #if ENABLE_SPE_FOR_NS
-void spe_init_el3(void);
+void spe_enable(cpu_context_t *ctx);
 void spe_init_el2_unused(void);
 void spe_disable(void);
 #else
-static inline void spe_init_el3(void)
+static inline void spe_enable(cpu_context_t *ctx)
 {
 }
 static inline void spe_init_el2_unused(void)
diff --git a/include/lib/extensions/trbe.h b/include/lib/extensions/trbe.h
index 0bed433..5db3316 100644
--- a/include/lib/extensions/trbe.h
+++ b/include/lib/extensions/trbe.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2023, Arm Limited. All rights reserved.
+ * Copyright (c) 2021-2024, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -7,11 +7,13 @@
 #ifndef TRBE_H
 #define TRBE_H
 
+#include <context.h>
+
 #if ENABLE_TRBE_FOR_NS
-void trbe_init_el3(void);
+void trbe_enable(cpu_context_t *ctx);
 void trbe_init_el2_unused(void);
 #else
-static inline void trbe_init_el3(void)
+static inline void trbe_enable(cpu_context_t *ctx)
 {
 }
 static inline void trbe_init_el2_unused(void)
diff --git a/include/lib/extensions/trf.h b/include/lib/extensions/trf.h
index 1ac7cda..f0a946b 100644
--- a/include/lib/extensions/trf.h
+++ b/include/lib/extensions/trf.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2023, Arm Limited. All rights reserved.
+ * Copyright (c) 2021-2024, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -7,16 +7,32 @@
 #ifndef TRF_H
 #define TRF_H
 
+#include <context.h>
+
 #if ENABLE_TRF_FOR_NS
-void trf_init_el3(void);
+
+#if __aarch64__
+void trf_enable(cpu_context_t *ctx);
 void trf_init_el2_unused(void);
-#else
-static inline void trf_init_el3(void)
+#else /* !__aarch64 */
+void trf_init_el3(void);
+#endif /* __aarch64__ */
+
+#else /* ENABLE_TRF_FOR_NS=0 */
+
+#if __aarch64__
+static inline void trf_enable(cpu_context_t *ctx)
 {
 }
 static inline void trf_init_el2_unused(void)
 {
 }
+#else /* !__aarch64 */
+static inline void trf_init_el3(void)
+{
+}
+#endif /* __aarch64__*/
+
 #endif /* ENABLE_TRF_FOR_NS */
 
 #endif /* TRF_H */
diff --git a/lib/el3_runtime/aarch64/context.S b/lib/el3_runtime/aarch64/context.S
index 76aebf9..1fce1bf 100644
--- a/lib/el3_runtime/aarch64/context.S
+++ b/lib/el3_runtime/aarch64/context.S
@@ -454,15 +454,17 @@
 	synchronize_errors
 #endif /* IMAGE_BL31 */
 
-	/* ----------------------------------------------------------
-	 * Restore SPSR_EL3, ELR_EL3 and SCR_EL3 prior to ERET
-	 * ----------------------------------------------------------
+	/* --------------------------------------------------------------
+	 * Restore MDCR_EL3, SPSR_EL3, ELR_EL3 and SCR_EL3 prior to ERET
+	 * --------------------------------------------------------------
 	 */
-	ldr	x18, [sp, #CTX_EL3STATE_OFFSET + CTX_SCR_EL3]
 	ldp	x16, x17, [sp, #CTX_EL3STATE_OFFSET + CTX_SPSR_EL3]
-	msr	scr_el3, x18
+	ldr	x18, [sp, #CTX_EL3STATE_OFFSET + CTX_SCR_EL3]
+	ldr	x19, [sp, #CTX_EL3STATE_OFFSET + CTX_MDCR_EL3]
 	msr	spsr_el3, x16
 	msr	elr_el3, x17
+	msr	scr_el3, x18
+	msr	mdcr_el3, x19
 
 	restore_ptw_el1_sys_regs
 
diff --git a/lib/el3_runtime/aarch64/context_mgmt.c b/lib/el3_runtime/aarch64/context_mgmt.c
index 1937c30..058f7fa 100644
--- a/lib/el3_runtime/aarch64/context_mgmt.c
+++ b/lib/el3_runtime/aarch64/context_mgmt.c
@@ -42,6 +42,7 @@
 per_world_context_t per_world_context[CPU_DATA_CONTEXT_NUM];
 static bool has_secure_perworld_init;
 
+static void manage_extensions_common(cpu_context_t *ctx);
 static void manage_extensions_nonsecure(cpu_context_t *ctx);
 static void manage_extensions_secure(cpu_context_t *ctx);
 static void manage_extensions_secure_per_world(void);
@@ -312,6 +313,7 @@
 static void setup_context_common(cpu_context_t *ctx, const entry_point_info_t *ep)
 {
 	u_register_t scr_el3;
+	u_register_t mdcr_el3;
 	el3_state_t *state;
 	gp_regs_t *gp_regs;
 
@@ -484,6 +486,37 @@
 	write_ctx_reg(state, CTX_ELR_EL3, ep->pc);
 	write_ctx_reg(state, CTX_SPSR_EL3, ep->spsr);
 
+	/* Start with a clean MDCR_EL3 copy as all relevant values are set */
+	mdcr_el3 = MDCR_EL3_RESET_VAL;
+
+	/* ---------------------------------------------------------------------
+	 * Initialise MDCR_EL3, setting all fields rather than relying on hw.
+	 * Some fields are architecturally UNKNOWN on reset.
+	 *
+	 * MDCR_EL3.SDD: Set to one to disable AArch64 Secure self-hosted debug.
+	 *  Debug exceptions, other than Breakpoint Instruction exceptions, are
+	 *  disabled from all ELs in Secure state.
+	 *
+	 * MDCR_EL3.SPD32: Set to 0b10 to disable AArch32 Secure self-hosted
+	 *  privileged debug from S-EL1.
+	 *
+	 * MDCR_EL3.TDOSA: Set to zero so that EL2 and EL2 System register
+	 *  access to the powerdown debug registers do not trap to EL3.
+	 *
+	 * MDCR_EL3.TDA: Set to zero to allow EL0, EL1 and EL2 access to the
+	 *  debug registers, other than those registers that are controlled by
+	 *  MDCR_EL3.TDOSA.
+	 */
+	mdcr_el3 |= ((MDCR_SDD_BIT | MDCR_SPD32(MDCR_SPD32_DISABLE))
+			& ~(MDCR_TDA_BIT | MDCR_TDOSA_BIT)) ;
+	write_ctx_reg(state, CTX_MDCR_EL3, mdcr_el3);
+
+	/*
+	 * Configure MDCR_EL3 register as applicable for each world
+	 * (NS/Secure/Realm) context.
+	 */
+	manage_extensions_common(ctx);
+
 	/*
 	 * Store the X0-X7 value from the entrypoint into the context
 	 * Use memcpy as we are in control of the layout of the structures
@@ -560,10 +593,6 @@
 #if IMAGE_BL31
 void cm_manage_extensions_el3(void)
 {
-	if (is_feat_spe_supported()) {
-		spe_init_el3();
-	}
-
 	if (is_feat_amu_supported()) {
 		amu_init_el3();
 	}
@@ -572,18 +601,6 @@
 		sme_init_el3();
 	}
 
-	if (is_feat_trbe_supported()) {
-		trbe_init_el3();
-	}
-
-	if (is_feat_brbe_supported()) {
-		brbe_init_el3();
-	}
-
-	if (is_feat_trf_supported()) {
-		trf_init_el3();
-	}
-
 	pmuv3_init_el3();
 }
 #endif /* IMAGE_BL31 */
@@ -704,6 +721,48 @@
 }
 
 /*******************************************************************************
+ * Enable architecture extensions on first entry to Non-secure world only
+ * and disable for secure world.
+ *
+ * NOTE: Arch features which have been provided with the capability of getting
+ * enabled only for non-secure world and being disabled for secure world are
+ * grouped here, as the MDCR_EL3 context value remains same across the worlds.
+ ******************************************************************************/
+static void manage_extensions_common(cpu_context_t *ctx)
+{
+#if IMAGE_BL31
+	if (is_feat_spe_supported()) {
+		/*
+		 * Enable FEAT_SPE for Non-Secure and prohibit for Secure state.
+		 */
+		spe_enable(ctx);
+	}
+
+	if (is_feat_trbe_supported()) {
+		/*
+		 * Enable FEAT_SPE for Non-Secure and prohibit for Secure and
+		 * Realm state.
+		 */
+		trbe_enable(ctx);
+	}
+
+	if (is_feat_trf_supported()) {
+		/*
+		 * Enable FEAT_SPE for Non-Secure and prohibit for Secure state.
+		 */
+		trf_enable(ctx);
+	}
+
+	if (is_feat_brbe_supported()) {
+		/*
+		 * Enable FEAT_SPE for Non-Secure and prohibit for Secure state.
+		 */
+		brbe_enable(ctx);
+	}
+#endif /* IMAGE_BL31 */
+}
+
+/*******************************************************************************
  * Enable architecture extensions on first entry to Non-secure world.
  ******************************************************************************/
 static void manage_extensions_nonsecure(cpu_context_t *ctx)
diff --git a/lib/extensions/brbe/brbe.c b/lib/extensions/brbe/brbe.c
index 37bd834..dde0266 100644
--- a/lib/extensions/brbe/brbe.c
+++ b/lib/extensions/brbe/brbe.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, Arm Limited. All rights reserved.
+ * Copyright (c) 2022-2024, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -9,9 +9,10 @@
 #include <arch_helpers.h>
 #include <lib/extensions/brbe.h>
 
-void brbe_init_el3(void)
+void brbe_enable(cpu_context_t *ctx)
 {
-	uint64_t val;
+	el3_state_t *state = get_el3state_ctx(ctx);
+	u_register_t mdcr_el3_val = read_ctx_reg(state, CTX_MDCR_EL3);
 
 	/*
 	 * MDCR_EL3.SBRBE = 0b01
@@ -19,8 +20,7 @@
 	 * Allows BRBE usage in non-secure world and prohibited in
 	 * secure world.
 	 */
-	val = read_mdcr_el3();
-	val &= ~(MDCR_SBRBE_MASK << MDCR_SBRBE_SHIFT);
-	val |= (0x1UL << MDCR_SBRBE_SHIFT);
-	write_mdcr_el3(val);
+	mdcr_el3_val &= ~(MDCR_SBRBE_MASK << MDCR_SBRBE_SHIFT);
+	mdcr_el3_val |= (0x1UL << MDCR_SBRBE_SHIFT);
+	write_ctx_reg(state, CTX_MDCR_EL3, mdcr_el3_val);
 }
diff --git a/lib/extensions/spe/spe.c b/lib/extensions/spe/spe.c
index d1fb182..c6076fe 100644
--- a/lib/extensions/spe/spe.c
+++ b/lib/extensions/spe/spe.c
@@ -29,9 +29,10 @@
 	__asm__ volatile("hint #17");
 }
 
-void spe_init_el3(void)
+void spe_enable(cpu_context_t *ctx)
 {
-	uint64_t v;
+	el3_state_t *state = get_el3state_ctx(ctx);
+	u_register_t mdcr_el3_val = read_ctx_reg(state, CTX_MDCR_EL3);
 
 	/*
 	 * MDCR_EL3.NSPB (ARM v8.2): SPE enabled in Non-secure state
@@ -46,10 +47,9 @@
 	 * Setting this bit to 1 doesn't have any effect on it when
 	 * FEAT_SPEv1p2 not implemented.
 	 */
-	v = read_mdcr_el3();
-	v |= MDCR_NSPB(MDCR_NSPB_EL1) | MDCR_EnPMSN_BIT;
-	v &= ~(MDCR_NSPBE_BIT);
-	write_mdcr_el3(v);
+	mdcr_el3_val |= MDCR_NSPB(MDCR_NSPB_EL1) | MDCR_EnPMSN_BIT;
+	mdcr_el3_val &= ~(MDCR_NSPBE_BIT);
+	write_ctx_reg(state, CTX_MDCR_EL3, mdcr_el3_val);
 }
 
 void spe_init_el2_unused(void)
diff --git a/lib/extensions/trbe/trbe.c b/lib/extensions/trbe/trbe.c
index d4fbdfb..9157734 100644
--- a/lib/extensions/trbe/trbe.c
+++ b/lib/extensions/trbe/trbe.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2023, Arm Limited. All rights reserved.
+ * Copyright (c) 2021-2024, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -19,9 +19,10 @@
 	__asm__ volatile("hint #18");
 }
 
-void trbe_init_el3(void)
+void trbe_enable(cpu_context_t *ctx)
 {
-	u_register_t val;
+	el3_state_t *state = get_el3state_ctx(ctx);
+	u_register_t mdcr_el3_val = read_ctx_reg(state, CTX_MDCR_EL3);
 
 	/*
 	 * MDCR_EL3.NSTBE = 0b0
@@ -33,10 +34,9 @@
 	 *  NS-EL2, tracing is prohibited in Secure and Realm state (if
 	 *  implemented).
 	 */
-	val = read_mdcr_el3();
-	val |= MDCR_NSTB(MDCR_NSTB_EL1);
-	val &= ~(MDCR_NSTBE_BIT);
-	write_mdcr_el3(val);
+	mdcr_el3_val |= MDCR_NSTB(MDCR_NSTB_EL1);
+	mdcr_el3_val &= ~(MDCR_NSTBE_BIT);
+	write_ctx_reg(state, CTX_MDCR_EL3, mdcr_el3_val);
 }
 
 void trbe_init_el2_unused(void)
diff --git a/lib/extensions/trf/aarch64/trf.c b/lib/extensions/trf/aarch64/trf.c
index 83fbf85..d36853a 100644
--- a/lib/extensions/trf/aarch64/trf.c
+++ b/lib/extensions/trf/aarch64/trf.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2023, Arm Limited. All rights reserved.
+ * Copyright (c) 2021-2024, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -9,9 +9,10 @@
 #include <arch_helpers.h>
 #include <lib/extensions/trf.h>
 
-void trf_init_el3(void)
+void trf_enable(cpu_context_t *ctx)
 {
-	u_register_t val;
+	el3_state_t *state = get_el3state_ctx(ctx);
+	u_register_t mdcr_el3_val = read_ctx_reg(state, CTX_MDCR_EL3);
 
 	/*
 	 * MDCR_EL3.STE = b0
@@ -22,9 +23,8 @@
 	 * Allow access of trace filter control registers from NS-EL2
 	 * and NS-EL1 when NS-EL2 is implemented but not used
 	 */
-	val = read_mdcr_el3();
-	val &= ~(MDCR_STE_BIT | MDCR_TTRF_BIT);
-	write_mdcr_el3(val);
+	mdcr_el3_val &= ~(MDCR_STE_BIT | MDCR_TTRF_BIT);
+	write_ctx_reg(state, CTX_MDCR_EL3, mdcr_el3_val);
 }
 
 void trf_init_el2_unused(void)