Merge changes from topic "bk/context_refactor" into integration

* changes:
  refactor(psci): extract cm_prepare_el3_exit_ns() to a common location
  refactor(cm): set MDCR_EL3/CPTR_EL3 bits in respective feat_init_el3() only
  fix(cm): set MDCR_EL3.{NSPBE, STE} explicitly
  refactor(cm): factor out EL2 register setting when EL2 is unused
diff --git a/include/arch/aarch32/arch.h b/include/arch/aarch32/arch.h
index dd2c0a6..da4ee7f 100644
--- a/include/arch/aarch32/arch.h
+++ b/include/arch/aarch32/arch.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2023, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2016-2023, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -197,11 +197,11 @@
 #define SDCR_SPD_LEGACY		U(0x0)
 #define SDCR_SPD_DISABLE	U(0x2)
 #define SDCR_SPD_ENABLE		U(0x3)
-#define SDCR_SCCD_BIT		(U(1) << 23)
-#define SDCR_TTRF_BIT		(U(1) << 19)
 #define SDCR_SPME_BIT		(U(1) << 17)
-#define SDCR_RESET_VAL		U(0x0)
+#define SDCR_TTRF_BIT		(U(1) << 19)
+#define SDCR_SCCD_BIT		(U(1) << 23)
 #define SDCR_MTPME_BIT		(U(1) << 28)
+#define SDCR_RESET_VAL		U(0x0)
 
 /* HSCTLR definitions */
 #define HSCTLR_RES1	((U(1) << 29) | (U(1) << 28) | (U(1) << 23) | \
diff --git a/include/arch/aarch64/arch.h b/include/arch/aarch64/arch.h
index 698fad4..b19e8af 100644
--- a/include/arch/aarch64/arch.h
+++ b/include/arch/aarch64/arch.h
@@ -596,7 +596,7 @@
 #define MDCR_SBRBE_MASK		ULL(0x3)
 #define MDCR_NSTB(x)		((x) << 24)
 #define MDCR_NSTB_EL1		ULL(0x3)
-#define MDCR_NSTBE		(ULL(1) << 26)
+#define MDCR_NSTBE_BIT		(ULL(1) << 26)
 #define MDCR_MTPME_BIT		(ULL(1) << 28)
 #define MDCR_TDCC_BIT		(ULL(1) << 27)
 #define MDCR_SCCD_BIT		(ULL(1) << 23)
@@ -612,6 +612,7 @@
 #define MDCR_SPD32_ENABLE	ULL(0x3)
 #define MDCR_NSPB(x)		((x) << 12)
 #define MDCR_NSPB_EL1		ULL(0x3)
+#define MDCR_NSPBE_BIT		(ULL(1) << 11)
 #define MDCR_TDOSA_BIT		(ULL(1) << 10)
 #define MDCR_TDA_BIT		(ULL(1) << 9)
 #define MDCR_TPM_BIT		(ULL(1) << 6)
diff --git a/include/arch/aarch64/el3_common_macros.S b/include/arch/aarch64/el3_common_macros.S
index 6360461..fa9310e 100644
--- a/include/arch/aarch64/el3_common_macros.S
+++ b/include/arch/aarch64/el3_common_macros.S
@@ -115,33 +115,11 @@
 	 * 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.TPM: Set to zero so that EL0, EL1, and EL2 System register
-	 *  accesses to all Performance Monitors registers do not trap to EL3.
-	 *
-	 * MDCR_EL3.NSTB, MDCR_EL3.NSTBE: Set to zero so that Trace Buffer
-	 *  owning security state is Secure state. If FEAT_TRBE is implemented,
-	 *  accesses to Trace Buffer control registers at EL2 and EL1 in any
-	 *  security state generates trap exceptions to EL3.
-	 *  If FEAT_TRBE is not implemented, these bits are RES0.
-	 *
-	 * MDCR_EL3.TTRF: Set to one so that access to trace filter control
-	 *  registers in non-monitor mode generate EL3 trap exception,
-	 *  unless the access generates a higher priority exception when trace
-	 *  filter control(FEAT_TRF) is implemented.
-	 *  When FEAT_TRF is not implemented, this bit is RES0.
-	 * ---------------------------------------------------------------------
 	 */
 	mov_imm	x0, ((MDCR_EL3_RESET_VAL | MDCR_SDD_BIT | \
 		      MDCR_SPD32(MDCR_SPD32_DISABLE)) & \
-		    ~(MDCR_TDOSA_BIT | MDCR_TDA_BIT | MDCR_TPM_BIT | \
-		      MDCR_NSTB(MDCR_NSTB_EL1) | MDCR_NSTBE | MDCR_TTRF_BIT))
+		    ~(MDCR_TDOSA_BIT | MDCR_TDA_BIT))
 
-	mrs	x1, id_aa64dfr0_el1
-	ubfx	x1, x1, #ID_AA64DFR0_TRACEFILT_SHIFT, #ID_AA64DFR0_TRACEFILT_LENGTH
-	cbz	x1, 1f
-	orr	x0, x0, #MDCR_TTRF_BIT
-1:
 	msr	mdcr_el3, x0
 
 	/* ---------------------------------------------------------------------
@@ -158,15 +136,6 @@
 	 * CPTR_EL3.TCPAC: Set to zero so that any accesses to CPACR_EL1,
 	 *  CPTR_EL2, CPACR, or HCPTR do not trap to EL3.
 	 *
-	 * CPTR_EL3.TTA: Set to one so that accesses to the trace system
-	 *  registers trap to EL3 from all exception levels and security
-	 *  states when system register trace is implemented.
-	 *  When system register trace is not implemented, this bit is RES0 and
-	 *  hence set to zero.
-	 *
-	 * CPTR_EL3.TTA: Set to zero so that System register accesses to the
-	 *  trace registers do not trap to EL3.
-	 *
 	 * CPTR_EL3.TFP: Set to zero so that accesses to the V- or Z- registers
 	 *  by Advanced SIMD, floating-point or SVE instructions (if implemented)
 	 *  do not trap to EL3.
@@ -181,12 +150,7 @@
 	 *  to EL3 by default.
 	 */
 
-	mov_imm x0, (CPTR_EL3_RESET_VAL & ~(TCPAC_BIT | TTA_BIT | TFP_BIT))
-	mrs	x1, id_aa64dfr0_el1
-	ubfx	x1, x1, #ID_AA64DFR0_TRACEVER_SHIFT, #ID_AA64DFR0_TRACEVER_LENGTH
-	cbz	x1, 1f
-	orr	x0, x0, #TTA_BIT
-1:
+	mov_imm x0, (CPTR_EL3_RESET_VAL & ~(TCPAC_BIT | TFP_BIT))
 	msr	cptr_el3, x0
 
 	/*
diff --git a/include/lib/extensions/sys_reg_trace.h b/include/lib/extensions/sys_reg_trace.h
index d9f7f1b..beda88a 100644
--- a/include/lib/extensions/sys_reg_trace.h
+++ b/include/lib/extensions/sys_reg_trace.h
@@ -13,6 +13,7 @@
 
 #if __aarch64__
 void sys_reg_trace_enable(cpu_context_t *context);
+void sys_reg_trace_disable(cpu_context_t *context);
 void sys_reg_trace_init_el2_unused(void);
 #else
 void sys_reg_trace_init_el3(void);
diff --git a/lib/el3_runtime/aarch64/context_mgmt.c b/lib/el3_runtime/aarch64/context_mgmt.c
index d9ff0b6..9d717bb 100644
--- a/lib/el3_runtime/aarch64/context_mgmt.c
+++ b/lib/el3_runtime/aarch64/context_mgmt.c
@@ -562,6 +562,20 @@
 #endif /* IMAGE_BL31 */
 }
 
+/* TODO: move to lib/extensions/pauth when it has been ported to FEAT_STATE */
+static __unused void enable_pauth_el2(void)
+{
+	u_register_t hcr_el2 = read_hcr_el2();
+	/*
+	 * For Armv8.3 pointer authentication feature, disable traps to EL2 when
+	 *  accessing key registers or using pointer authentication instructions
+	 *  from lower ELs.
+	 */
+	hcr_el2 |= (HCR_API_BIT | HCR_APK_BIT);
+
+	write_hcr_el2(hcr_el2);
+}
+
 /*******************************************************************************
  * Enable architecture extensions in-place at EL2 on first entry to Non-secure
  * world when EL2 is empty and unused.
@@ -602,6 +616,10 @@
 	if (is_feat_sme_supported()) {
 		sme_init_el2_unused();
 	}
+
+#if ENABLE_PAUTH
+	enable_pauth_el2();
+#endif /* ENABLE_PAUTH */
 #endif /* IMAGE_BL31 */
 }
 
@@ -644,6 +662,11 @@
 			sme_disable(ctx);
 		}
 	}
+
+	/* NS can access this but Secure shouldn't */
+	if (is_feat_sys_reg_trace_supported()) {
+		sys_reg_trace_disable(ctx);
+	}
 #endif /* IMAGE_BL31 */
 }
 
@@ -672,6 +695,109 @@
 	cm_setup_context(ctx, ep);
 }
 
+/* EL2 present but unused, need to disable safely. SCTLR_EL2 can be ignored */
+static __unused void init_nonsecure_el2_unused(cpu_context_t *ctx)
+{
+	u_register_t hcr_el2 = HCR_RESET_VAL;
+	u_register_t mdcr_el2;
+	u_register_t scr_el3;
+
+	scr_el3 = read_ctx_reg(get_el3state_ctx(ctx), CTX_SCR_EL3);
+
+	/* Set EL2 register width: Set HCR_EL2.RW to match SCR_EL3.RW */
+	if ((scr_el3 & SCR_RW_BIT) != 0U) {
+		hcr_el2 |= HCR_RW_BIT;
+	}
+
+	write_hcr_el2(hcr_el2);
+
+	/*
+	 * Initialise CPTR_EL2 setting all fields rather than relying on the hw.
+	 * All fields have architecturally UNKNOWN reset values.
+	 */
+	write_cptr_el2(CPTR_EL2_RESET_VAL);
+
+	/*
+	 * Initialise CNTHCTL_EL2. All fields are architecturally UNKNOWN on
+	 * reset and are set to zero except for field(s) listed below.
+	 *
+	 * CNTHCTL_EL2.EL1PTEN: Set to one to disable traps to Hyp mode of
+	 * Non-secure EL0 and EL1 accesses to the physical timer registers.
+	 *
+	 * CNTHCTL_EL2.EL1PCTEN: Set to one to disable traps to Hyp mode of
+	 * Non-secure EL0 and EL1 accesses to the physical counter registers.
+	 */
+	write_cnthctl_el2(CNTHCTL_RESET_VAL | EL1PCEN_BIT | EL1PCTEN_BIT);
+
+	/*
+	 * Initialise CNTVOFF_EL2 to zero as it resets to an architecturally
+	 * UNKNOWN value.
+	 */
+	write_cntvoff_el2(0);
+
+	/*
+	 * Set VPIDR_EL2 and VMPIDR_EL2 to match MIDR_EL1 and MPIDR_EL1
+	 * respectively.
+	 */
+	write_vpidr_el2(read_midr_el1());
+	write_vmpidr_el2(read_mpidr_el1());
+
+	/*
+	 * Initialise VTTBR_EL2. All fields are architecturally UNKNOWN on reset.
+	 *
+	 * VTTBR_EL2.VMID: Set to zero. Even though EL1&0 stage 2 address
+	 * translation is disabled, cache maintenance operations depend on the
+	 * VMID.
+	 *
+	 * VTTBR_EL2.BADDR: Set to zero as EL1&0 stage 2 address translation is
+	 * disabled.
+	 */
+	write_vttbr_el2(VTTBR_RESET_VAL &
+		     ~((VTTBR_VMID_MASK << VTTBR_VMID_SHIFT) |
+		       (VTTBR_BADDR_MASK << VTTBR_BADDR_SHIFT)));
+
+	/*
+	 * Initialise MDCR_EL2, setting all fields rather than relying on hw.
+	 * Some fields are architecturally UNKNOWN on reset.
+	 *
+	 * MDCR_EL2.TDRA: Set to zero so that Non-secure EL0 and EL1 System
+	 * register accesses to the Debug ROM registers are not trapped to EL2.
+	 *
+	 * MDCR_EL2.TDOSA: Set to zero so that Non-secure EL1 System register
+	 * accesses to the powerdown debug registers are not trapped to EL2.
+	 *
+	 * MDCR_EL2.TDA: Set to zero so that System register accesses to the
+	 * debug registers do not trap to EL2.
+	 *
+	 * MDCR_EL2.TDE: Set to zero so that debug exceptions are not routed to
+	 * EL2.
+	 */
+	mdcr_el2 = MDCR_EL2_RESET_VAL &
+		 ~(MDCR_EL2_TDRA_BIT | MDCR_EL2_TDOSA_BIT | MDCR_EL2_TDA_BIT |
+		   MDCR_EL2_TDE_BIT);
+
+	write_mdcr_el2(mdcr_el2);
+
+	/*
+	 * Initialise HSTR_EL2. All fields are architecturally UNKNOWN on reset.
+	 *
+	 * HSTR_EL2.T<n>: Set all these fields to zero so that Non-secure EL0 or
+	 * EL1 accesses to System registers do not trap to EL2.
+	 */
+	write_hstr_el2(HSTR_EL2_RESET_VAL & ~(HSTR_EL2_T_MASK));
+
+	/*
+	 * Initialise CNTHP_CTL_EL2. All fields are architecturally UNKNOWN on
+	 * reset.
+	 *
+	 * CNTHP_CTL_EL2:ENABLE: Set to zero to disable the EL2 physical timer
+	 * and prevent timer interrupts.
+	 */
+	write_cnthp_ctl_el2(CNTHP_CTL_RESET_VAL & ~(CNTHP_CTL_ENABLE_BIT));
+
+	manage_extensions_nonsecure_el2_unused();
+}
+
 /*******************************************************************************
  * Prepare the CPU system registers for first entry into realm, secure, or
  * normal world.
@@ -683,9 +809,8 @@
  ******************************************************************************/
 void cm_prepare_el3_exit(uint32_t security_state)
 {
-	u_register_t sctlr_elx, scr_el3, mdcr_el2;
+	u_register_t sctlr_elx, scr_el3;
 	cpu_context_t *ctx = cm_get_context(security_state);
-	uint64_t hcr_el2 = 0U;
 
 	assert(ctx != NULL);
 
@@ -722,121 +847,7 @@
 #endif
 			write_sctlr_el2(sctlr_elx);
 		} else if (el2_implemented != EL_IMPL_NONE) {
-			/*
-			 * EL2 present but unused, need to disable safely.
-			 * SCTLR_EL2 can be ignored in this case.
-			 *
-			 * Set EL2 register width appropriately: Set HCR_EL2
-			 * field to match SCR_EL3.RW.
-			 */
-			if ((scr_el3 & SCR_RW_BIT) != 0U)
-				hcr_el2 |= HCR_RW_BIT;
-
-			/*
-			 * For Armv8.3 pointer authentication feature, disable
-			 * traps to EL2 when accessing key registers or using
-			 * pointer authentication instructions from lower ELs.
-			 */
-			hcr_el2 |= (HCR_API_BIT | HCR_APK_BIT);
-
-			write_hcr_el2(hcr_el2);
-
-			/*
-			 * Initialise CPTR_EL2 setting all fields rather than
-			 * relying on the hw. All fields have architecturally
-			 * UNKNOWN reset values.
-			 */
-			write_cptr_el2(CPTR_EL2_RESET_VAL);
-
-			/*
-			 * Initialise CNTHCTL_EL2. All fields are
-			 * architecturally UNKNOWN on reset and are set to zero
-			 * except for field(s) listed below.
-			 *
-			 * CNTHCTL_EL2.EL1PTEN: Set to one to disable traps to
-			 *  Hyp mode of Non-secure EL0 and EL1 accesses to the
-			 *  physical timer registers.
-			 *
-			 * CNTHCTL_EL2.EL1PCTEN: Set to one to disable traps to
-			 *  Hyp mode of  Non-secure EL0 and EL1 accesses to the
-			 *  physical counter registers.
-			 */
-			write_cnthctl_el2(CNTHCTL_RESET_VAL |
-						EL1PCEN_BIT | EL1PCTEN_BIT);
-
-			/*
-			 * Initialise CNTVOFF_EL2 to zero as it resets to an
-			 * architecturally UNKNOWN value.
-			 */
-			write_cntvoff_el2(0);
-
-			/*
-			 * Set VPIDR_EL2 and VMPIDR_EL2 to match MIDR_EL1 and
-			 * MPIDR_EL1 respectively.
-			 */
-			write_vpidr_el2(read_midr_el1());
-			write_vmpidr_el2(read_mpidr_el1());
-
-			/*
-			 * Initialise VTTBR_EL2. All fields are architecturally
-			 * UNKNOWN on reset.
-			 *
-			 * VTTBR_EL2.VMID: Set to zero. Even though EL1&0 stage
-			 *  2 address translation is disabled, cache maintenance
-			 *  operations depend on the VMID.
-			 *
-			 * VTTBR_EL2.BADDR: Set to zero as EL1&0 stage 2 address
-			 *  translation is disabled.
-			 */
-			write_vttbr_el2(VTTBR_RESET_VAL &
-				~((VTTBR_VMID_MASK << VTTBR_VMID_SHIFT)
-				| (VTTBR_BADDR_MASK << VTTBR_BADDR_SHIFT)));
-
-			/*
-			 * Initialise MDCR_EL2, setting all fields rather than
-			 * relying on hw. Some fields are architecturally
-			 * UNKNOWN on reset.
-			 *
-			 * MDCR_EL2.TDRA: Set to zero so that Non-secure EL0 and
-			 *  EL1 System register accesses to the Debug ROM
-			 *  registers are not trapped to EL2.
-			 *
-			 * MDCR_EL2.TDOSA: Set to zero so that Non-secure EL1
-			 *  System register accesses to the powerdown debug
-			 *  registers are not trapped to EL2.
-			 *
-			 * MDCR_EL2.TDA: Set to zero so that System register
-			 *  accesses to the debug registers do not trap to EL2.
-			 *
-			 * MDCR_EL2.TDE: Set to zero so that debug exceptions
-			 *  are not routed to EL2.
-			 */
-			mdcr_el2 = ((MDCR_EL2_RESET_VAL) &
-				   ~(MDCR_EL2_TDRA_BIT | MDCR_EL2_TDOSA_BIT |
-				     MDCR_EL2_TDA_BIT | MDCR_EL2_TDE_BIT));
-
-			write_mdcr_el2(mdcr_el2);
-
-			/*
-			 * Initialise HSTR_EL2. All fields are architecturally
-			 * UNKNOWN on reset.
-			 *
-			 * HSTR_EL2.T<n>: Set all these fields to zero so that
-			 *  Non-secure EL0 or EL1 accesses to System registers
-			 *  do not trap to EL2.
-			 */
-			write_hstr_el2(HSTR_EL2_RESET_VAL & ~(HSTR_EL2_T_MASK));
-			/*
-			 * Initialise CNTHP_CTL_EL2. All fields are
-			 * architecturally UNKNOWN on reset.
-			 *
-			 * CNTHP_CTL_EL2:ENABLE: Set to zero to disable the EL2
-			 *  physical timer and prevent timer interrupts.
-			 */
-			write_cnthp_ctl_el2(CNTHP_CTL_RESET_VAL &
-						~(CNTHP_CTL_ENABLE_BIT));
-
-			manage_extensions_nonsecure_el2_unused();
+			init_nonsecure_el2_unused(ctx);
 		}
 	}
 
diff --git a/lib/extensions/pmuv3/aarch64/pmuv3.c b/lib/extensions/pmuv3/aarch64/pmuv3.c
index fda71aa..61fc47d 100644
--- a/lib/extensions/pmuv3/aarch64/pmuv3.c
+++ b/lib/extensions/pmuv3/aarch64/pmuv3.c
@@ -82,9 +82,12 @@
 	 *   0  |  1   |    enabled   | disabled
 	 *   1  |  1   |    enabled   | disabled only for counters 0 to
 	 *                              MDCR_EL2.HPMN - 1. Enabled for the rest
+	 *
+	 * MDCR_EL3.TPM: Set to zero so that EL0, EL1, and EL2 System register
+	 *  accesses to all Performance Monitors registers do not trap to EL3.
 	 */
 	mdcr_el3 = (mdcr_el3 | MDCR_SCCD_BIT | MDCR_MCCD_BIT) &
-		  ~(MDCR_MPMX_BIT | MDCR_SPME_BIT);
+		  ~(MDCR_MPMX_BIT | MDCR_SPME_BIT | MDCR_TPM_BIT);
 	mdcr_el3 = mtpmu_disable_el3(mdcr_el3);
 	write_mdcr_el3(mdcr_el3);
 
diff --git a/lib/extensions/spe/spe.c b/lib/extensions/spe/spe.c
index 236b102..2c25a9d 100644
--- a/lib/extensions/spe/spe.c
+++ b/lib/extensions/spe/spe.c
@@ -26,10 +26,13 @@
 	uint64_t v;
 
 	/*
-	 * MDCR_EL2.NSPB (ARM v8.2): SPE enabled in Non-secure state
+	 * MDCR_EL3.NSPB (ARM v8.2): SPE enabled in Non-secure state
 	 * and disabled in secure state. Accesses to SPE registers at
 	 * S-EL1 generate trap exceptions to EL3.
 	 *
+	 * MDCR_EL3.NSPBE: Profiling Buffer uses Non-secure Virtual Addresses.
+	 * When FEAT_RME is not implemented, this field is RES0.
+	 *
 	 * MDCR_EL3.EnPMSN (ARM v8.7): Do not trap access to PMSNEVFR_EL1
 	 * register at NS-EL1 or NS-EL2 to EL3 if FEAT_SPEv1p2 is implemented.
 	 * Setting this bit to 1 doesn't have any effect on it when
@@ -37,6 +40,7 @@
 	 */
 	v = read_mdcr_el3();
 	v |= MDCR_NSPB(MDCR_NSPB_EL1) | MDCR_EnPMSN_BIT;
+	v &= ~(MDCR_NSPBE_BIT);
 	write_mdcr_el3(v);
 }
 
diff --git a/lib/extensions/sys_reg_trace/aarch64/sys_reg_trace.c b/lib/extensions/sys_reg_trace/aarch64/sys_reg_trace.c
index 4b57f67..1349566 100644
--- a/lib/extensions/sys_reg_trace/aarch64/sys_reg_trace.c
+++ b/lib/extensions/sys_reg_trace/aarch64/sys_reg_trace.c
@@ -12,16 +12,26 @@
 
 void sys_reg_trace_enable(cpu_context_t *ctx)
 {
-	uint64_t val;
+	/*
+	 * CPTR_EL3.TTA: Set to zero so that System register accesses to the
+	 *  trace registers do not trap to EL3.
+	 */
+	uint64_t val = read_ctx_reg(get_el3state_ctx(ctx), CTX_CPTR_EL3);
+
+	val &= ~(TTA_BIT);
+	write_ctx_reg(get_el3state_ctx(ctx), CTX_CPTR_EL3, val);
+}
 
-	/* Retrieve CPTR_EL3 value from the given context 'ctx',
-	 * and update CPTR_EL3.TTA bit to 0.
-	 * This function is called while switching context to NS to
-	 * allow system trace register access to NS-EL2 and NS-EL1
-	 * when NS-EL2 is implemented but not used.
+void sys_reg_trace_disable(cpu_context_t *ctx)
+{
+	/*
+	 * CPTR_EL3.TTA: Set to one so that System register accesses to the
+	 *  trace registers trap to EL3, unless it is trapped by CPACR.TRCDIS,
+	 *  CPACR_EL1.TTA, or CPTR_EL2.TTA
 	 */
-	val = read_ctx_reg(get_el3state_ctx(ctx), CTX_CPTR_EL3);
-	val &= ~TTA_BIT;
+	uint64_t val = read_ctx_reg(get_el3state_ctx(ctx), CTX_CPTR_EL3);
+
+	val |= TTA_BIT;
 	write_ctx_reg(get_el3state_ctx(ctx), CTX_CPTR_EL3, val);
 }
 
diff --git a/lib/extensions/trbe/trbe.c b/lib/extensions/trbe/trbe.c
index 461ea73..d4fbdfb 100644
--- a/lib/extensions/trbe/trbe.c
+++ b/lib/extensions/trbe/trbe.c
@@ -24,13 +24,18 @@
 	u_register_t val;
 
 	/*
+	 * MDCR_EL3.NSTBE = 0b0
+	 *  Trace Buffer owning Security state is Non-secure state. If FEAT_RME
+	 *  is not implemented, this field is RES0.
+	 *
 	 * MDCR_EL3.NSTB = 0b11
-	 * Allow access of trace buffer control registers from NS-EL1
-	 * and NS-EL2, tracing is prohibited in Secure and Realm state
-	 * (if implemented).
+	 *  Allow access of trace buffer control registers from NS-EL1 and
+	 *  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);
 }
 
diff --git a/lib/extensions/trf/aarch32/trf.c b/lib/extensions/trf/aarch32/trf.c
index e13b4db..321c089 100644
--- a/lib/extensions/trf/aarch32/trf.c
+++ b/lib/extensions/trf/aarch32/trf.c
@@ -16,9 +16,9 @@
 
 	/*
 	 * Allow access of trace filter control registers from
-	 * non-monitor mode
+	 * non-monitor mode.
 	 */
 	val = read_sdcr();
-	val &= ~SDCR_TTRF_BIT;
+	val &= ~(SDCR_TTRF_BIT);
 	write_sdcr(val);
 }
diff --git a/lib/extensions/trf/aarch64/trf.c b/lib/extensions/trf/aarch64/trf.c
index f681b28..83fbf85 100644
--- a/lib/extensions/trf/aarch64/trf.c
+++ b/lib/extensions/trf/aarch64/trf.c
@@ -14,12 +14,16 @@
 	u_register_t val;
 
 	/*
+	 * MDCR_EL3.STE = b0
+	 * Trace prohibited in Secure state unless overridden by the
+	 * IMPLEMENTATION DEFINED authentication interface.
+	 *
 	 * MDCR_EL3.TTRF = b0
 	 * 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_TTRF_BIT;
+	val &= ~(MDCR_STE_BIT | MDCR_TTRF_BIT);
 	write_mdcr_el3(val);
 }
 
diff --git a/lib/psci/psci_common.c b/lib/psci/psci_common.c
index 8aa0cce..70bf77e 100644
--- a/lib/psci/psci_common.c
+++ b/lib/psci/psci_common.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013-2022, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2013-2023, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -1037,6 +1037,13 @@
 		psci_cpu_suspend_finish(cpu_idx, &state_info);
 
 	/*
+	 * Generic management: Now we just need to retrieve the
+	 * information that we had stashed away during the cpu_on
+	 * call to set this cpu on its way.
+	 */
+	cm_prepare_el3_exit_ns();
+
+	/*
 	 * Set the requested and target state of this CPU and all the higher
 	 * power domains which are ancestors of this CPU to run.
 	 */
diff --git a/lib/psci/psci_on.c b/lib/psci/psci_on.c
index 76eb50c..1f619ae 100644
--- a/lib/psci/psci_on.c
+++ b/lib/psci/psci_on.c
@@ -225,11 +225,4 @@
 	/* Populate the mpidr field within the cpu node array */
 	/* This needs to be done only once */
 	psci_cpu_pd_nodes[cpu_idx].mpidr = read_mpidr() & MPIDR_AFFINITY_MASK;
-
-	/*
-	 * Generic management: Now we just need to retrieve the
-	 * information that we had stashed away during the cpu_on
-	 * call to set this cpu on its way.
-	 */
-	cm_prepare_el3_exit_ns();
 }
diff --git a/lib/psci/psci_suspend.c b/lib/psci/psci_suspend.c
index d93e60d..cb12b83 100644
--- a/lib/psci/psci_suspend.c
+++ b/lib/psci/psci_suspend.c
@@ -367,11 +367,4 @@
 	psci_set_suspend_pwrlvl(PSCI_INVALID_PWR_LVL);
 
 	PUBLISH_EVENT(psci_suspend_pwrdown_finish);
-
-	/*
-	 * Generic management: Now we just need to retrieve the
-	 * information that we had stashed away during the suspend
-	 * call to set this cpu on its way.
-	 */
-	cm_prepare_el3_exit_ns();
 }
diff --git a/services/std_svc/rmmd/rmmd_main.c b/services/std_svc/rmmd/rmmd_main.c
index c80b524..a929ea2 100644
--- a/services/std_svc/rmmd/rmmd_main.c
+++ b/services/std_svc/rmmd/rmmd_main.c
@@ -128,6 +128,11 @@
 		sve_enable(ctx);
 	}
 
+	/* NS can access this but Realm shouldn't */
+	if (is_feat_sys_reg_trace_supported()) {
+		sys_reg_trace_disable(ctx);
+	}
+
 	pmuv3_enable(ctx);
 }