test: test the save restore logic for brbcr_el1

This patch tests the save/restore logic by enabling
branch recording at NS-EL2. Additionally this
patch also tests the trap logic when FEAT_FGT is enabled
and a Realm tries to access any FEAT_BRBE related registers.

Signed-off-by: Sona Mathew <sonarebecca.mathew@arm.com>
Change-Id: I176ea6feaf01d42cfd6231dc65a9470da8d1e37c
diff --git a/include/lib/aarch64/arch.h b/include/lib/aarch64/arch.h
index 1d5e75a..09835b5 100644
--- a/include/lib/aarch64/arch.h
+++ b/include/lib/aarch64/arch.h
@@ -837,6 +837,35 @@
 #define DISABLE_INTERRUPTS	(DAIF_FIQ_BIT | DAIF_IRQ_BIT)
 
 /*
+ * BRBCR_EL2/EL1 definitions
+ */
+#define BRBCR_EL1_EXCEPTION_EN	(U(1) << 23)
+#define BRBCR_EL1_ERTN_EN	(U(1) << 22)
+#define BRBCR_EL1_MPRED_EN	(U(1) << 4)
+#define BRBCR_EL1_CC_EN		(U(1) << 3)
+#define BRBCR_EL1_CC_SHIFT	3
+#define BRBCR_EL1_CC_WIDTH	U(1)
+#define BRBCR_EL1_E1BRE_EN	(U(1) << 1)
+#define BRBCR_EL1_E0BRE_EN	(U(1) << 0)
+#define BRBCR_EL1_INIT		(BRBCR_EL1_EXCEPTION_EN | BRBCR_EL1_ERTN_EN | \
+				 BRBCR_EL1_MPRED_EN | BRBCR_EL1_CC_EN | BRBCR_EL1_E1BRE_EN | \
+				 BRBCR_EL1_E0BRE_EN)
+
+#define BRBCR_EL1_E0BRE_SHIFT	U(0)
+#define BRBCR_EL1_E0BRE_WIDTH	U(1)
+#define BRBCR_EL1_E1BRE_SHIFT	U(1)
+#define BRBCR_EL1_E1BRE_WIDTH	U(1)
+#define BRBCR_EL2_E0HBRE_SHIFT	U(0)
+
+#define BRBCR_EL2_E2BRE_ENABLE		(U(1) << 1)
+#define BRBCR_EL2_CC_ENABLE		(U(1) << 3)
+#define BRBCR_EL2_MPRED_ENABLE		(U(1) << 4)
+#define BRBCR_EL2_ERTN_ENABLE		(U(1) << 22)
+#define BRBCR_EL2_EXCEPTION_ENABLE	(U(1) << 23)
+#define BRBCR_EL2_INIT			(BRBCR_EL2_E2BRE_ENABLE | BRBCR_EL2_CC_ENABLE | \
+					BRBCR_EL2_MPRED_ENABLE | BRBCR_EL2_ERTN_ENABLE | \
+					BRBCR_EL2_EXCEPTION_ENABLE)
+/*
  * RMR_EL3 definitions
  */
 #define RMR_EL3_RR_BIT		(U(1) << 1)
@@ -1549,14 +1578,17 @@
  * FEAT_BRBE - Branch Record Buffer Extension System Registers
  ******************************************************************************/
 
-#define BRBCR_EL1	S2_1_C9_C0_0
-#define BRBCR_EL2	S2_4_C9_C0_0
-#define BRBFCR_EL1	S2_1_C9_C0_1
-#define BRBTS_EL1	S2_1_C9_C0_2
-#define BRBINFINJ_EL1	S2_1_C9_C1_0
-#define BRBSRCINJ_EL1	S2_1_C9_C1_1
-#define BRBTGTINJ_EL1	S2_1_C9_C1_2
-#define BRBIDR0_EL1	S2_1_C9_C2_0
+#define BRBCR_EL1		S2_1_C9_C0_0
+#define BRBCR_EL2		S2_4_C9_C0_0
+#define BRBFCR_EL1		S2_1_C9_C0_1
+#define BRBTS_EL1		S2_1_C9_C0_2
+#define BRBINFINJ_EL1		S2_1_C9_C1_0
+#define BRBSRCINJ_EL1		S2_1_C9_C1_1
+#define BRBTGTINJ_EL1		S2_1_C9_C1_2
+#define BRBIDR0_EL1		S2_1_C9_C2_0
+#define BRBINF15_EL1		S2_1_C8_C15_0
+#define BRBSRC11_EL1		S2_1_C8_C11_1
+#define BRBTGT0_EL1		S2_1_C8_C0_2
 
 /*******************************************************************************
  * FEAT_TCR2 - Extended Translation Control Registers
diff --git a/include/lib/aarch64/arch_helpers.h b/include/lib/aarch64/arch_helpers.h
index da1794d..8e55c8a 100644
--- a/include/lib/aarch64/arch_helpers.h
+++ b/include/lib/aarch64/arch_helpers.h
@@ -636,6 +636,9 @@
 DEFINE_RENAME_SYSREG_RW_FUNCS(brbsrcinj_el1, BRBSRCINJ_EL1)
 DEFINE_RENAME_SYSREG_RW_FUNCS(brbtgtinj_el1, BRBTGTINJ_EL1)
 DEFINE_RENAME_SYSREG_READ_FUNC(brbidr0_el1, BRBIDR0_EL1)
+DEFINE_RENAME_SYSREG_READ_FUNC(brbtgt0_el1, BRBTGT0_EL1)
+DEFINE_RENAME_SYSREG_READ_FUNC(brbsrc11_el1, BRBSRC11_EL1)
+DEFINE_RENAME_SYSREG_READ_FUNC(brbinf15_el1, BRBINF15_EL1)
 
 /* Armv8.4 Trace filter control System Registers */
 DEFINE_RENAME_SYSREG_RW_FUNCS(trfcr_el1, TRFCR_EL1)
diff --git a/include/runtime_services/host_realm_managment/host_shared_data.h b/include/runtime_services/host_realm_managment/host_shared_data.h
index 763306f..787a343 100644
--- a/include/runtime_services/host_realm_managment/host_shared_data.h
+++ b/include/runtime_services/host_realm_managment/host_shared_data.h
@@ -79,7 +79,8 @@
 	REALM_ATTESTATION,
 	REALM_ATTESTATION_FAULT,
 	REALM_ENTER_PLANE_N_CMD,
-	REALM_PLANE_N_REG_RW_CMD
+	REALM_PLANE_N_REG_RW_CMD,
+	REALM_WRITE_BRBCR_EL1
 };
 
 /*
diff --git a/realm/include/realm_tests.h b/realm/include/realm_tests.h
index 6c1e6c3..c2d4102 100644
--- a/realm/include/realm_tests.h
+++ b/realm/include/realm_tests.h
@@ -31,6 +31,7 @@
 bool test_realm_attestation(void);
 bool test_realm_attestation_fault(void);
 bool test_realm_mpam_undef_abort(void);
+bool test_realm_write_brbcr_el1_reg(void);
 
 #endif /* REALM_TESTS_H */
 
diff --git a/realm/realm.mk b/realm/realm.mk
index 6281779..d9c1bec 100644
--- a/realm/realm.mk
+++ b/realm/realm.mk
@@ -42,6 +42,7 @@
 	realm_shared_data.c						\
 	realm_simd.c							\
 	realm_mpam.c							\
+	realm_brbe_tests.c						\
 	)
 
 REALM_SOURCES += lib/${ARCH}/cache_helpers.S				\
diff --git a/realm/realm_brbe_tests.c b/realm/realm_brbe_tests.c
new file mode 100644
index 0000000..52eeaf3
--- /dev/null
+++ b/realm/realm_brbe_tests.c
@@ -0,0 +1,43 @@
+/*
+ * Copyright (c) 2024, Arm Limited. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+
+#include <arch_helpers.h>
+#include <debug.h>
+#include <realm_tests.h>
+#include <realm_helpers.h>
+#include <sync.h>
+
+ /* Try to enable Branch recording in R-EL1 via BRBCR_EL1 */
+
+bool test_realm_write_brbcr_el1_reg(void)
+{
+	unsigned long brbcr_el1;
+	bool ret_result = false;
+	realm_reset_undef_abort_count();
+
+	/* Install exception handler to catch undefined abort */
+	register_custom_sync_exception_handler(realm_sync_exception_handler);
+
+	brbcr_el1 = read_brbcr_el1();
+	brbcr_el1 |= BRBCR_EL1_INIT;
+	write_brbcr_el1(brbcr_el1);
+
+	read_brbfcr_el1();
+	read_brbts_el1();
+	read_brbinfinj_el1();
+	read_brbsrcinj_el1();
+	read_brbtgtinj_el1();
+	read_brbidr0_el1();
+	read_brbsrc11_el1();
+	read_brbtgt0_el1();
+	read_brbinf15_el1();
+
+	if (realm_get_undef_abort_count() == 11UL) {
+		ret_result = true;
+	}
+	unregister_custom_sync_exception_handler();
+	return ret_result;
+
+}
diff --git a/realm/realm_helpers.c b/realm/realm_helpers.c
index f11e2f3..1887164 100644
--- a/realm/realm_helpers.c
+++ b/realm/realm_helpers.c
@@ -54,7 +54,6 @@
 			     esr_el1, read_elr_el1());
 		realm_got_undef_abort++;
 	}
-
 	return true;
 }
 
diff --git a/realm/realm_payload_main.c b/realm/realm_payload_main.c
index d364aa0..4052ba3 100644
--- a/realm/realm_payload_main.c
+++ b/realm/realm_payload_main.c
@@ -494,6 +494,9 @@
 		case REALM_ATTESTATION_FAULT:
 			test_succeed = test_realm_attestation_fault();
 			break;
+		case REALM_WRITE_BRBCR_EL1:
+			test_succeed = test_realm_write_brbcr_el1_reg();
+			break;
 		default:
 			realm_printf("%s() invalid cmd %u\n", __func__, cmd);
 			break;
diff --git a/tftf/tests/runtime_services/realm_payload/host_realm_brbe_tests.c b/tftf/tests/runtime_services/realm_payload/host_realm_brbe_tests.c
new file mode 100644
index 0000000..9f6246b
--- /dev/null
+++ b/tftf/tests/runtime_services/realm_payload/host_realm_brbe_tests.c
@@ -0,0 +1,95 @@
+/*
+ * Copyright (c) 2025, Arm Limited. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+
+#include <assert.h>
+#include <debug.h>
+#include <host_realm_helper.h>
+#include <host_realm_mem_layout.h>
+#include <host_shared_data.h>
+#include <test_helpers.h>
+
+/*
+ * This test verifies that the values of BRBCR_EL2 and
+ * BRBCR_EL1 registers are being preserved by EL3 and
+ * R-EL2 as expected.
+ */
+
+test_result_t host_realm_test_brbe_save_restore(void)
+{
+	unsigned long rec_flag[] = {RMI_RUNNABLE};
+	unsigned long feature_flag =  0UL;
+	unsigned long feature_flag1 =  0UL;
+	struct realm realm;
+	long sl = RTT_MIN_LEVEL;
+	static unsigned long brbcr_el2_saved;
+	static unsigned long brbcr_el1_saved;
+	test_result_t test_result = TEST_RESULT_SUCCESS;
+
+	SKIP_TEST_IF_RME_NOT_SUPPORTED_OR_RMM_IS_TRP();
+	SKIP_TEST_IF_BRBE_NOT_SUPPORTED();
+
+	if (!host_create_activate_realm_payload(&realm, (u_register_t)REALM_IMAGE_BASE,
+						feature_flag, feature_flag1, sl, rec_flag, 1U, 0U)) {
+		return TEST_RESULT_FAIL;
+	}
+
+	/*
+	 * Program BRBCR_EL2 to enable branch recording, now
+	 * since MDCR_EL3.SBRBE is set to 0b01 in EL3, accesses
+	 * to BRBCR_EL2 should not trap to EL3.
+	 */
+
+	INFO ("Enable Branch recording in NS-EL2\n");
+
+	/* SET BRBCR_EL2 */
+	unsigned long val = read_brbcr_el2();
+	val = val | BRBCR_EL2_INIT;
+	write_brbcr_el2(val);
+
+	/* SET BRBCR_EL1 */
+	unsigned long val_brbcr_el1 = read_brbcr_el1();
+	val_brbcr_el1 = BRBCR_EL1_INIT & (~MASK(BRBCR_EL1_CC));
+	write_brbcr_el1(val_brbcr_el1);
+
+	brbcr_el2_saved = read_brbcr_el2();
+	brbcr_el1_saved	= read_brbcr_el1();
+
+	/*
+	 * Enter Realm and try to access BRBCR_EL1 and enable
+	 * branch recording for R-EL1.
+	 */
+
+	if (!host_enter_realm_execute(&realm, REALM_WRITE_BRBCR_EL1,
+					RMI_EXIT_HOST_CALL, 0)) {
+		test_result = TEST_RESULT_FAIL;
+		goto destroy_realm;
+	}
+
+	/*
+	 * Verify if EL3 has saved and restored the value in
+	 * BRBCR_EL2 register.
+	 */
+	INFO("NS Read and compare BRBCR_EL2\n");
+	if (read_brbcr_el2() != brbcr_el2_saved) {
+		test_result = TEST_RESULT_FAIL;
+		INFO("brbcr_el2 did not match, read = %lx\n and saved = %lx\n", read_brbcr_el2(), \
+									brbcr_el2_saved);
+		goto destroy_realm;
+	}
+
+	if (read_brbcr_el1() != brbcr_el1_saved) {
+		test_result = TEST_RESULT_FAIL;
+		INFO("brbcr_el1 did not match, read = %lx\n and saved = %lx\n", read_brbcr_el1(), \
+									brbcr_el1_saved);
+		goto destroy_realm;
+	}
+
+destroy_realm:
+	if (!host_destroy_realm(&realm)) {
+		return TEST_RESULT_FAIL;
+	}
+
+	return test_result;
+}
diff --git a/tftf/tests/tests-realm-payload.mk b/tftf/tests/tests-realm-payload.mk
index 1fe01ed..8d9b8b1 100644
--- a/tftf/tests/tests-realm-payload.mk
+++ b/tftf/tests/tests-realm-payload.mk
@@ -17,6 +17,7 @@
 		host_realm_payload_simd_tests.c				\
 		host_realm_lpa2_tests.c					\
 		host_realm_mpam_tests.c					\
+		host_realm_brbe_tests.c					\
 	)
 
 TESTS_SOURCES	+=							\
diff --git a/tftf/tests/tests-realm-payload.xml b/tftf/tests/tests-realm-payload.xml
index 277e6db..84ffab7 100644
--- a/tftf/tests/tests-realm-payload.xml
+++ b/tftf/tests/tests-realm-payload.xml
@@ -153,5 +153,8 @@
 	  function="host_realm_hide_feat_mpam" />
 	  <testcase name="Test that access to FEAT_MPAM from Realm causes an undef abort taken to the Realm"
 	  function="host_realm_mpam_undef_abort" />
+	  <!-- Test case related to FEAT_BRBE -->
+	  <testcase name="Test if BRBE realted registers are preserved "
+	  function="host_realm_test_brbe_save_restore" />
   </testsuite>
 </testsuites>