test(el3-runtime): check DIT is retained on exception

Add a test to check that the PSTATE bits not set in the
Aarch64.TakeException but set to a default when taking an exception to
EL3 are maintained after an exception and that changes in TSP do not
effect the PSTATE in TFTF and vice versa.

Signed-off-by: Daniel Boulby <daniel.boulby@arm.com>
Change-Id: Id4d625c7e9cbb565ac236f844274319cc02c2335
diff --git a/include/common/test_helpers.h b/include/common/test_helpers.h
index 0881f00..74b71a7 100644
--- a/include/common/test_helpers.h
+++ b/include/common/test_helpers.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018-2021, Arm Limited. All rights reserved.
+ * Copyright (c) 2018-2022, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -90,6 +90,15 @@
 		}								\
 	} while (0)
 
+#define SKIP_TEST_IF_DIT_NOT_SUPPORTED()					\
+	do {									\
+		if (!is_armv8_4_dit_present()) {				\
+			tftf_testcase_printf(					\
+				"DIT not supported\n");				\
+			return TEST_RESULT_SKIPPED;				\
+		}								\
+	} while (0)
+
 #define SKIP_TEST_IF_PAUTH_NOT_SUPPORTED()					\
 	do {									\
 		if (!is_armv8_3_pauth_present()) {				\
diff --git a/include/lib/aarch32/arch_features.h b/include/lib/aarch32/arch_features.h
index 1addc44..3c6a338 100644
--- a/include/lib/aarch32/arch_features.h
+++ b/include/lib/aarch32/arch_features.h
@@ -42,6 +42,12 @@
 		ID_DFR0_TRACEFILT_SUPPORTED;
 }
 
+static inline bool is_armv8_4_dit_present(void)
+{
+	return ((read_id_pfr0() >> ID_PFR0_DIT_SHIFT) &
+		ID_PFR0_DIT_MASK) != 0;
+}
+
 static inline bool get_armv8_0_sys_reg_trace_support(void)
 {
 	return ((read_id_dfr0() >> ID_DFR0_COPTRC_SHIFT) &
diff --git a/include/lib/aarch32/arch_helpers.h b/include/lib/aarch32/arch_helpers.h
index 6e1097d..aca3952 100644
--- a/include/lib/aarch32/arch_helpers.h
+++ b/include/lib/aarch32/arch_helpers.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2021, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2016-2022, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -422,6 +422,9 @@
 #define read_daif()		read_cpsr()
 #define write_daif(flags)	write_cpsr(flags)
 
+#define read_dit()		read_cpsr()
+#define write_dit(flags)	write_cpsr(flags)
+
 #define read_cnthp_cval_el2()	read64_cnthp_cval_el2()
 #define write_cnthp_cval_el2(v)	write64_cnthp_cval_el2(v)
 
diff --git a/include/runtime_services/secure_el1_payloads/tsp.h b/include/runtime_services/secure_el1_payloads/tsp.h
index 87ee6f7..19db911 100644
--- a/include/runtime_services/secure_el1_payloads/tsp.h
+++ b/include/runtime_services/secure_el1_payloads/tsp.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018, Arm Limited. All rights reserved.
+ * Copyright (c) 2018-2022, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -32,6 +32,7 @@
 #define TSP_MUL		0x2002
 #define TSP_DIV		0x2003
 #define TSP_HANDLE_SEL1_INTR_AND_RETURN	0x2004
+#define TSP_CHECK_DIT	0x2005
 
 /*
  * Identify a TSP service from function ID filtering the last 16 bits from the
diff --git a/tftf/tests/runtime_services/trusted_os/tsp/test_pstate_after_exception.c b/tftf/tests/runtime_services/trusted_os/tsp/test_pstate_after_exception.c
new file mode 100644
index 0000000..a17a54e
--- /dev/null
+++ b/tftf/tests/runtime_services/trusted_os/tsp/test_pstate_after_exception.c
@@ -0,0 +1,75 @@
+/*
+ * Copyright (c) 2022, Arm Limited. All rights reserved.
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+
+#include <test_helpers.h>
+
+/*
+ * Test that the PSTATE bits not set in Aarch64.TakeException but
+ * set to a default when taking an exception to EL3 are maintained
+ * after an exception and that changes in TSP do not effect the PSTATE
+ * in TFTF and vice versa.
+ */
+test_result_t tsp_check_pstate_maintained_on_exception(void)
+{
+	smc_args tsp_svc_params;
+	smc_ret_values ret;
+	u_register_t dit;
+	u_register_t dit_bit;
+
+	SKIP_TEST_IF_TSP_NOT_PRESENT();
+	SKIP_TEST_IF_DIT_NOT_SUPPORTED();
+
+#ifdef __aarch64__
+	dit_bit = DIT_BIT;
+#else
+	dit_bit = CPSR_DIT_BIT;
+#endif
+
+	write_dit(dit_bit);
+
+	/* Standard SMC */
+	tsp_svc_params.fid = TSP_STD_FID(TSP_CHECK_DIT);
+	tsp_svc_params.arg1 = 0;
+	tsp_svc_params.arg2 = 0;
+	ret = tftf_smc(&tsp_svc_params);
+	if (ret.ret1 == 0) {
+		if (ret.ret2 == 0xffff) {
+			tftf_testcase_printf("DIT bit not supported by TSP");
+			return TEST_RESULT_SKIPPED;
+		}
+		tftf_testcase_printf("DIT bit in the TSP is not 0.\n");
+		return TEST_RESULT_FAIL;
+	}
+
+	dit = read_dit();
+	if (dit != dit_bit) {
+		tftf_testcase_printf("DIT bit in TFTF was not maintained.\n"
+				     "Expected: 0x%x, Actual: 0x%x",
+				     (uint32_t) dit_bit, (uint32_t) dit);
+		return TEST_RESULT_FAIL;
+	}
+
+	tsp_svc_params.fid = TSP_STD_FID(TSP_CHECK_DIT);
+	tsp_svc_params.arg1 = dit_bit;
+	tsp_svc_params.arg2 = 0;
+	ret = tftf_smc(&tsp_svc_params);
+	if (ret.ret1 == 0) {
+		tftf_testcase_printf("DIT bit in the TSP was not maintained\n"
+				     "Expected: 0x%x, Actual: 0x%x",
+				     (uint32_t) dit_bit, (uint32_t) ret.ret2);
+		return TEST_RESULT_FAIL;
+	}
+
+	dit = read_dit();
+	if (dit != dit_bit) {
+		tftf_testcase_printf("DIT bit in TFTF was not maintained.\n"
+				     "Expected: 0x%x, Actual: 0x%x",
+				     (uint32_t) dit_bit, (uint32_t) dit);
+		return TEST_RESULT_FAIL;
+	}
+
+	return TEST_RESULT_SUCCESS;
+}
diff --git a/tftf/tests/tests-tsp.mk b/tftf/tests/tests-tsp.mk
index b1d8b15..35ef02a 100644
--- a/tftf/tests/tests-tsp.mk
+++ b/tftf/tests/tests-tsp.mk
@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2018, Arm Limited. All rights reserved.
+# Copyright (c) 2018-2022, Arm Limited. All rights reserved.
 #
 # SPDX-License-Identifier: BSD-3-Clause
 #
@@ -11,4 +11,5 @@
 		test_smc_tsp_std_fn_call.c					\
 		test_tsp_fast_smc.c						\
 		test_normal_int_switch.c					\
+		test_pstate_after_exception.c					\
 	)
diff --git a/tftf/tests/tests-tsp.xml b/tftf/tests/tests-tsp.xml
index 7e1018e..55dfbea 100644
--- a/tftf/tests/tests-tsp.xml
+++ b/tftf/tests/tests-tsp.xml
@@ -1,13 +1,12 @@
 <?xml version="1.0" encoding="utf-8"?>
 
 <!--
-  Copyright (c) 2018, Arm Limited. All rights reserved.
+  Copyright (c) 2018-2022, Arm Limited. All rights reserved.
 
   SPDX-License-Identifier: BSD-3-Clause
 -->
 
 <testsuites>
-
   <testsuite name="IRQ support in TSP" description="Test the normal IRQ preemption support in TSP.">
     <testcase name="TSP preempt by IRQ and resume" function="tsp_int_and_resume" />
     <testcase name="Fast SMC while TSP preempted" function="test_fast_smc_when_tsp_preempted" />
@@ -32,4 +31,8 @@
      <testcase name="Stress test TSP functionality" function="test_tsp_fast_smc_operations" />
   </testsuite>
 
+  <testsuite name="TSP PSTATE test" description="Test PSTATE bits are maintained during exception">
+     <testcase name="Test PSTATE bits maintained on exception" function="tsp_check_pstate_maintained_on_exception" />
+  </testsuite>
+
 </testsuites>