fix(el3-runtime): set unset pstate bits to default

During a transition to a higher EL some of the PSTATE bits are not set
by hardware, this means that their state may be leaked from lower ELs.
This patch sets those bits to a default value upon entry to EL3.

This patch was tested using a debugger to check the PSTATE values
are correctly set. As well as adding a test in the next patch to
ensure the PSTATE in lower ELs is still maintained after this change.

Change-Id: Ie546acbca7b9aa3c86bd68185edded91b2a64ae5
Signed-off-by: Daniel Boulby <daniel.boulby@arm.com>
diff --git a/Makefile b/Makefile
index a238ee4..1baacf0 100644
--- a/Makefile
+++ b/Makefile
@@ -277,6 +277,10 @@
 ENABLE_FEAT_ECV		=	1
 endif
 
+ifeq "8.4" "$(word 1, $(sort 8.4 $(ARM_ARCH_MAJOR).$(ARM_ARCH_MINOR)))"
+ENABLE_FEAT_DIT		= 	1
+endif
+
 ifneq ($(findstring armclang,$(notdir $(CC))),)
 TF_CFLAGS_aarch32	=	-target arm-arm-none-eabi $(march32-directive)
 TF_CFLAGS_aarch64	=	-target aarch64-arm-none-eabi $(march64-directive)
@@ -1040,6 +1044,7 @@
         USE_SP804_TIMER \
         ENABLE_FEAT_RNG \
         ENABLE_FEAT_SB \
+        ENABLE_FEAT_DIT \
         PSA_FWU_SUPPORT \
         ENABLE_TRBE_FOR_NS \
         ENABLE_SYS_REG_TRACE_FOR_NS \
@@ -1154,6 +1159,7 @@
         USE_SP804_TIMER \
         ENABLE_FEAT_RNG \
         ENABLE_FEAT_SB \
+        ENABLE_FEAT_DIT \
         NR_OF_FW_BANKS \
         NR_OF_IMAGES_IN_FW_BANK \
         PSA_FWU_SUPPORT \
diff --git a/bl31/aarch64/ea_delegate.S b/bl31/aarch64/ea_delegate.S
index 03820bd..fa6ede8 100644
--- a/bl31/aarch64/ea_delegate.S
+++ b/bl31/aarch64/ea_delegate.S
@@ -92,6 +92,7 @@
 	 * Save general purpose and ARMv8.3-PAuth registers (if enabled).
 	 * If Secure Cycle Counter is not disabled in MDCR_EL3 when
 	 * ARMv8.5-PMU is implemented, save PMCR_EL0 and disable Cycle Counter.
+	 * Also set the PSTATE to a known state.
 	 */
 	bl	prepare_el3_entry
 
@@ -139,6 +140,7 @@
 	 * Save general purpose and ARMv8.3-PAuth registers (if enabled).
 	 * If Secure Cycle Counter is not disabled in MDCR_EL3 when
 	 * ARMv8.5-PMU is implemented, save PMCR_EL0 and disable Cycle Counter.
+	 * Also set the PSTATE to a known state.
 	 */
 	bl	prepare_el3_entry
 
diff --git a/bl31/aarch64/runtime_exceptions.S b/bl31/aarch64/runtime_exceptions.S
index dbdbe43..bf5bd8d 100644
--- a/bl31/aarch64/runtime_exceptions.S
+++ b/bl31/aarch64/runtime_exceptions.S
@@ -71,6 +71,7 @@
 	 * Save general purpose and ARMv8.3-PAuth registers (if enabled).
 	 * If Secure Cycle Counter is not disabled in MDCR_EL3 when
 	 * ARMv8.5-PMU is implemented, save PMCR_EL0 and disable Cycle Counter.
+	 * Also set the PSTATE to a known state.
 	 */
 	bl	prepare_el3_entry
 
@@ -209,6 +210,7 @@
 	 * Save general purpose and ARMv8.3-PAuth registers (if enabled).
 	 * If Secure Cycle Counter is not disabled in MDCR_EL3 when
 	 * ARMv8.5-PMU is implemented, save PMCR_EL0 and disable Cycle Counter.
+	 * Also set the PSTATE to a known state.
 	 */
 	bl	prepare_el3_entry
 
@@ -462,6 +464,7 @@
 	 * Save general purpose and ARMv8.3-PAuth registers (if enabled).
 	 * If Secure Cycle Counter is not disabled in MDCR_EL3 when
 	 * ARMv8.5-PMU is implemented, save PMCR_EL0 and disable Cycle Counter.
+	 * Also set the PSTATE to a known state.
 	 */
 	bl	prepare_el3_entry
 
diff --git a/include/arch/aarch64/el3_common_macros.S b/include/arch/aarch64/el3_common_macros.S
index f29def7..d47244e 100644
--- a/include/arch/aarch64/el3_common_macros.S
+++ b/include/arch/aarch64/el3_common_macros.S
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2015-2021, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2015-2022, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -9,6 +9,7 @@
 
 #include <arch.h>
 #include <asm_macros.S>
+#include <assert_macros.S>
 #include <context.h>
 #include <lib/xlat_tables/xlat_tables_defs.h>
 
@@ -237,15 +238,20 @@
 
 	/*
 	 * If Data Independent Timing (DIT) functionality is implemented,
-	 * always enable DIT in EL3
+	 * always enable DIT in EL3.
+	 * First assert that the FEAT_DIT build flag matches the feature id
+	 * register value for DIT.
 	 */
+#if ENABLE_FEAT_DIT
+#if ENABLE_ASSERTIONS
 	mrs	x0, id_aa64pfr0_el1
 	ubfx	x0, x0, #ID_AA64PFR0_DIT_SHIFT, #ID_AA64PFR0_DIT_LENGTH
 	cmp	x0, #ID_AA64PFR0_DIT_SUPPORTED
-	bne	1f
+	ASM_ASSERT(eq)
+#endif /* ENABLE_ASSERTIONS */
 	mov	x0, #DIT_BIT
 	msr	DIT, x0
-1:
+#endif
 	.endm
 
 /* -----------------------------------------------------------------------------
diff --git a/lib/el3_runtime/aarch64/context.S b/lib/el3_runtime/aarch64/context.S
index bea72c8..c1c0612 100644
--- a/lib/el3_runtime/aarch64/context.S
+++ b/lib/el3_runtime/aarch64/context.S
@@ -682,6 +682,22 @@
 endfunc fpregs_context_restore
 #endif /* CTX_INCLUDE_FPREGS */
 
+	/*
+	 * Set the PSTATE bits not set when the exception was taken as
+	 * described in the AArch64.TakeException() pseudocode function
+	 * in ARM DDI 0487F.c page J1-7635 to a default value.
+	 */
+	.macro set_unset_pstate_bits
+        /*
+         * If Data Independent Timing (DIT) functionality is implemented,
+         * always enable DIT in EL3
+         */
+#if ENABLE_FEAT_DIT
+        mov     x8, #DIT_BIT
+        msr     DIT, x8
+#endif /* ENABLE_FEAT_DIT */
+	.endm /* set_unset_pstate_bits */
+
 /* ------------------------------------------------------------------
  * The following macro is used to save and restore all the general
  * purpose and ARMv8.3-PAuth (if enabled) registers.
@@ -767,17 +783,26 @@
 	stp	x26, x27, [x19, #CTX_PACDBKEY_LO]
 	stp	x28, x29, [x19, #CTX_PACGAKEY_LO]
 #endif /* CTX_INCLUDE_PAUTH_REGS */
-
 	.endm /* save_gp_pmcr_pauth_regs */
 
 /* -----------------------------------------------------------------
- * This function saves the context preparing entry to el3.
+ * This function saves the context and sets the PSTATE to a known
+ * state, preparing entry to el3.
  * Save all the general purpose and ARMv8.3-PAuth (if enabled)
  * registers.
+ * Then set any of the PSTATE bits that are not set by hardware
+ * according to the Aarch64.TakeException pseudocode in the Arm
+ * Architecture Reference Manual to a default value for EL3.
+ * clobbers: x17
  * -----------------------------------------------------------------
  */
 func prepare_el3_entry
 	save_gp_pmcr_pauth_regs
+	/*
+	 * Set the PSTATE bits not described in the Aarch64.TakeException
+	 * pseudocode to their default values.
+	 */
+	set_unset_pstate_bits
 	ret
 endfunc prepare_el3_entry
 
diff --git a/make_helpers/defaults.mk b/make_helpers/defaults.mk
index 551b689..b39dcf4 100644
--- a/make_helpers/defaults.mk
+++ b/make_helpers/defaults.mk
@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2016-2021, Arm Limited. All rights reserved.
+# Copyright (c) 2016-2022, Arm Limited. All rights reserved.
 #
 # SPDX-License-Identifier: BSD-3-Clause
 #
@@ -145,6 +145,9 @@
 # Flag to enable access to the CNTPOFF_EL2 register
 ENABLE_FEAT_ECV			:= 0
 
+# Flag to enable use of the DIT feature.
+ENABLE_FEAT_DIT			:= 0
+
 # By default BL31 encryption disabled
 ENCRYPT_BL31			:= 0