3 files changed, 53 insertions, 9 deletions
diff --git a/lib/cpus/aarch64/neoverse_zeus.S b/lib/cpus/aarch64/neoverse_zeus.S
index 3d850137..44882b45 100644
--- a/lib/cpus/aarch64/neoverse_zeus.S
+++ b/lib/cpus/aarch64/neoverse_zeus.S
@@ -46,6 +46,16 @@ func neoverse_zeus_errata_report
 endfunc neoverse_zeus_errata_report
 #endif
 
+func neoverse_zeus_reset_func
+	mov	x19, x30
+
+	/* Disable speculative loads */
+	msr	SSBS, xzr
+
+	isb
+	ret	x19
+endfunc neoverse_zeus_reset_func
+
 	/* ---------------------------------------------
 	 * This function provides Neoverse-Zeus specific
 	 * register information for crash reporting.
@@ -66,5 +76,5 @@ func neoverse_zeus_cpu_reg_dump
 endfunc neoverse_zeus_cpu_reg_dump
 
 declare_cpu_ops neoverse_zeus, NEOVERSE_ZEUS_MIDR, \
-	CPU_NO_RESET_FUNC, \
+	neoverse_zeus_reset_func, \
 	neoverse_zeus_core_pwr_dwn
diff --git a/lib/el3_runtime/aarch64/context.S b/lib/el3_runtime/aarch64/context.S
index 53dc02e6..37bb12c8 100644
--- a/lib/el3_runtime/aarch64/context.S
+++ b/lib/el3_runtime/aarch64/context.S
@@ -145,6 +145,17 @@ func el1_sysregs_context_save
 	str	x14, [x0, #CTX_CNTKCTL_EL1]
 #endif
 
+	/* Save MTE system registers if the build has instructed so */
+#if CTX_INCLUDE_MTE_REGS
+	mrs	x15, TFSRE0_EL1
+	mrs	x16, TFSR_EL1
+	stp	x15, x16, [x0, #CTX_TFSRE0_EL1]
+
+	mrs	x9, RGSR_EL1
+	mrs	x10, GCR_EL1
+	stp	x9, x10, [x0, #CTX_RGSR_EL1]
+#endif
+
 	ret
 endfunc el1_sysregs_context_save
 
@@ -229,6 +240,16 @@ func el1_sysregs_context_restore
 	ldr	x14, [x0, #CTX_CNTKCTL_EL1]
 	msr	cntkctl_el1, x14
 #endif
+	/* Restore MTE system registers if the build has instructed so */
+#if CTX_INCLUDE_MTE_REGS
+	ldp	x11, x12, [x0, #CTX_TFSRE0_EL1]
+	msr	TFSRE0_EL1, x11
+	msr	TFSR_EL1, x12
+
+	ldp	x13, x14, [x0, #CTX_RGSR_EL1]
+	msr	RGSR_EL1, x13
+	msr	GCR_EL1, x14
+#endif
 
 	/* No explict ISB required here as ERET covers it */
 	ret
diff --git a/lib/el3_runtime/aarch64/context_mgmt.c b/lib/el3_runtime/aarch64/context_mgmt.c
index bd5b3aa6..446d9da9 100644
--- a/lib/el3_runtime/aarch64/context_mgmt.c
+++ b/lib/el3_runtime/aarch64/context_mgmt.c
@@ -137,17 +137,30 @@ void cm_setup_context(cpu_context_t *ctx, const entry_point_info_t *ep)
 		scr_el3 |= SCR_API_BIT | SCR_APK_BIT;
 #endif /* !CTX_INCLUDE_PAUTH_REGS */
 
-	unsigned int mte = get_armv8_5_mte_support();
-
 	/*
-	 * Enable MTE support unilaterally for normal world if the CPU supports
-	 * it.
+	 * Enable MTE support. Support is enabled unilaterally for the normal
+	 * world, and only for the secure world when CTX_INCLUDE_MTE_REGS is
+	 * set.
 	 */
-	if (mte != MTE_UNIMPLEMENTED) {
-		if (security_state == NON_SECURE) {
-			scr_el3 |= SCR_ATA_BIT;
-		}
+	unsigned int mte = get_armv8_5_mte_support();
+#if CTX_INCLUDE_MTE_REGS
+	assert(mte == MTE_IMPLEMENTED_ELX);
+	scr_el3 |= SCR_ATA_BIT;
+#else
+	if (mte == MTE_IMPLEMENTED_EL0) {
+		/*
+		 * Can enable MTE across both worlds as no MTE registers are
+		 * used
+		 */
+		scr_el3 |= SCR_ATA_BIT;
+	} else if (mte == MTE_IMPLEMENTED_ELX && security_state == NON_SECURE) {
+		/*
+		 * Can only enable MTE in Non-Secure world without register
+		 * saving
+		 */
+		scr_el3 |= SCR_ATA_BIT;
 	}
+#endif
 
 #ifdef IMAGE_BL31
 	/*