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/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)