fix(mpam): refine MPAM initialization and enablement process

Restricts MPAM to only NS world and enables trap to EL3 for access of
MPAM registers from lower ELs of Secure and Realm world.

This patch removes MPAM enablement from global context and adds it to
EL3 State context which enables/disables MPAM during world switches.
Renamed ENABLE_MPAM_FOR_LOWER_ELS to ENABLE_FEAT_MPAM and
removed mpam_init_el3() as RESET behaviour is trapping.

Signed-off-by: Arvind Ram Prakash <arvind.ramprakash@arm.com>
Change-Id: I131f9dba5df236a71959b2d425ee11af7f3c38c4
diff --git a/lib/el3_runtime/aarch64/context.S b/lib/el3_runtime/aarch64/context.S
index 290a93a..758355a 100644
--- a/lib/el3_runtime/aarch64/context.S
+++ b/lib/el3_runtime/aarch64/context.S
@@ -350,6 +350,43 @@
 #endif /* ENABLE_FEAT_DIT */
 	.endm /* set_unset_pstate_bits */
 
+/*-------------------------------------------------------------------------
+ * This macro checks the ENABLE_FEAT_MPAM state, performs ID register
+ * check to see if the platform supports MPAM extension and restores MPAM3
+ * register value if it is FEAT_STATE_ENABLED/FEAT_STATE_CHECKED.
+ *
+ * This is particularly more complicated because we can't check
+ * if the platform supports MPAM  by looking for status of a particular bit
+ * in the MDCR_EL3 or CPTR_EL3 register like other extensions.
+ * ------------------------------------------------------------------------
+ */
+
+	.macro	restore_mpam3_el3
+#if ENABLE_FEAT_MPAM
+#if ENABLE_FEAT_MPAM == 2
+
+	mrs x8, id_aa64pfr0_el1
+	lsr x8, x8, #(ID_AA64PFR0_MPAM_SHIFT)
+	and x8, x8, #(ID_AA64PFR0_MPAM_MASK)
+	mrs x7, id_aa64pfr1_el1
+	lsr x7, x7, #(ID_AA64PFR1_MPAM_FRAC_SHIFT)
+	and x7, x7, #(ID_AA64PFR1_MPAM_FRAC_MASK)
+	orr x7, x7, x8
+	cbz x7, no_mpam
+#endif
+	/* -----------------------------------------------------------
+	 * Restore MPAM3_EL3 register as per context state
+	 * Currently we only enable MPAM for NS world and trap to EL3
+	 * for MPAM access in lower ELs of Secure and Realm world
+	 * -----------------------------------------------------------
+	 */
+	ldr	x17, [sp, #CTX_EL3STATE_OFFSET + CTX_MPAM3_EL3]
+	msr	S3_6_C10_C5_0, x17 /* mpam3_el3 */
+
+no_mpam:
+#endif
+	.endm /* restore_mpam3_el3 */
+
 /* ------------------------------------------------------------------
  * The following macro is used to save and restore all the general
  * purpose and ARMv8.3-PAuth (if enabled) registers.
@@ -573,6 +610,9 @@
 	isb
 	msr	S3_6_C1_C2_0, x20 /* zcr_el3 */
 sve_not_enabled:
+
+	restore_mpam3_el3
+
 #endif /* IMAGE_BL31 */
 
 #if IMAGE_BL31 && DYNAMIC_WORKAROUND_CVE_2018_3639
diff --git a/lib/el3_runtime/aarch64/context_mgmt.c b/lib/el3_runtime/aarch64/context_mgmt.c
index 3a7e50f..98cee16 100644
--- a/lib/el3_runtime/aarch64/context_mgmt.c
+++ b/lib/el3_runtime/aarch64/context_mgmt.c
@@ -482,6 +482,11 @@
 	}
 #endif /* (IMAGE_BL31 && defined(SPD_spmd) && SPMD_SPM_AT_SEL2) */
 
+	if (is_feat_mpam_supported()) {
+		write_ctx_reg(get_el3state_ctx(ctx), CTX_MPAM3_EL3, \
+				MPAM3_EL3_RESET_VAL);
+	}
+
 	/*
 	 * Populate EL3 state so that we've the right context
 	 * before doing ERET
@@ -578,10 +583,6 @@
 		sme_init_el3();
 	}
 
-	if (is_feat_mpam_supported()) {
-		mpam_init_el3();
-	}
-
 	if (is_feat_trbe_supported()) {
 		trbe_init_el3();
 	}
@@ -621,6 +622,9 @@
 		sys_reg_trace_enable(ctx);
 	}
 
+	if (is_feat_mpam_supported()) {
+		mpam_enable(ctx);
+	}
 	pmuv3_enable(ctx);
 #endif /* IMAGE_BL31 */
 }
diff --git a/lib/extensions/mpam/mpam.c b/lib/extensions/mpam/mpam.c
index 6462c97..875ad9c 100644
--- a/lib/extensions/mpam/mpam.c
+++ b/lib/extensions/mpam/mpam.c
@@ -11,14 +11,19 @@
 #include <arch_helpers.h>
 #include <lib/extensions/mpam.h>
 
-void mpam_init_el3(void)
+void mpam_enable(cpu_context_t *context)
 {
+	u_register_t mpam3_el3;
+
+	mpam3_el3 = read_ctx_reg(get_el3state_ctx(context), CTX_MPAM3_EL3);
+
 	/*
 	 * Enable MPAM, and disable trapping to EL3 when lower ELs access their
-	 * own MPAM registers.
+	 * own MPAM registers
 	 */
-	write_mpam3_el3(MPAM3_EL3_MPAMEN_BIT);
-
+	mpam3_el3 = (mpam3_el3 | MPAM3_EL3_MPAMEN_BIT) &
+				~(MPAM3_EL3_TRAPLOWER_BIT);
+	write_ctx_reg(get_el3state_ctx(context), CTX_MPAM3_EL3, mpam3_el3);
 }
 
 /*