Merge "feat(cm): context switch MDCR_EL3 register" into integration
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 dafc3be..981fddc 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);
@@ -308,6 +309,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;
 
@@ -480,6 +482,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
@@ -556,10 +589,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();
 	}
@@ -568,18 +597,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 */
@@ -700,6 +717,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)