test(rmm): add testcase for multiple rec on multiple cpu

Test creates and enter 8 recs on different CPUs.
Exercises CPU_ON, CPU_OFF and PSCI_AFFINITY_INFO
from realm.

Signed-off-by: Shruti Gupta <shruti.gupta@arm.com>
Change-Id: Iecc58ea79bfde28f307d1df99680d707e57a1d80
diff --git a/realm/include/realm_tests.h b/realm/include/realm_tests.h
index 5caea8c..3016a4d 100644
--- a/realm/include/realm_tests.h
+++ b/realm/include/realm_tests.h
@@ -23,6 +23,7 @@
 bool test_realm_sve_cmp_regs(void);
 bool test_realm_sve_undef_abort(void);
 bool test_realm_multiple_rec_psci_denied_cmd(void);
+bool test_realm_multiple_rec_multiple_cpu_cmd(void);
 bool test_realm_sme_read_id_registers(void);
 bool test_realm_sme_undef_abort(void);
 
diff --git a/realm/realm_multiple_rec.c b/realm/realm_multiple_rec.c
index abd166b..c584cd4 100644
--- a/realm/realm_multiple_rec.c
+++ b/realm/realm_multiple_rec.c
@@ -22,6 +22,7 @@
 
 #define CXT_ID_MAGIC 0x100
 static uint64_t is_secondary_cpu_booted;
+static spinlock_t lock;
 
 static void rec1_handler(u_register_t cxt_id)
 {
@@ -31,7 +32,9 @@
 		realm_printf("Wrong cxt_id\n");
 		rsi_exit_to_host(HOST_CALL_EXIT_FAILED_CMD);
 	}
+	spin_lock(&lock);
 	is_secondary_cpu_booted++;
+	spin_unlock(&lock);
 	realm_cpu_off();
 }
 
@@ -67,3 +70,48 @@
 	}
 	return true;
 }
+
+bool test_realm_multiple_rec_multiple_cpu_cmd(void)
+{
+	unsigned int i = 1U, rec_count;
+	u_register_t ret;
+
+	realm_printf("Realm: running on CPU = 0x%lx\n", read_mpidr_el1() & MPID_MASK);
+	rec_count = realm_shared_data_get_my_host_val(HOST_ARG1_INDEX);
+
+	/* Check CPU_ON is supported */
+	ret = realm_psci_features(SMC_PSCI_CPU_ON);
+	if (ret != PSCI_E_SUCCESS) {
+		realm_printf("SMC_PSCI_CPU_ON not supported\n");
+		return false;
+	}
+
+	for (unsigned int j = 1U; j < rec_count; j++) {
+		ret = realm_cpu_on(j, (uintptr_t)rec1_handler, CXT_ID_MAGIC + j);
+		if (ret != PSCI_E_SUCCESS) {
+			realm_printf("SMC_PSCI_CPU_ON failed %d.\n", j);
+			return false;
+		}
+	}
+
+	/* Exit to host to allow host to run all CPUs */
+	rsi_exit_to_host(HOST_CALL_EXIT_SUCCESS_CMD);
+	/* wait for all CPUs to come up */
+	while (is_secondary_cpu_booted != rec_count - 1U) {
+		waitms(200);
+	}
+
+	/* wait for all CPUs to turn off */
+	while (i < rec_count) {
+		ret = realm_psci_affinity_info(i, MPIDR_AFFLVL0);
+		if (ret != PSCI_STATE_OFF) {
+			/* wait and query again */
+			realm_printf(" CPU %d is not off\n", i);
+			waitms(200);
+			continue;
+		}
+		i++;
+	}
+	realm_printf("All CPU are off\n");
+	return true;
+}
diff --git a/realm/realm_payload_main.c b/realm/realm_payload_main.c
index 5c488ee..ddaa3cb 100644
--- a/realm/realm_payload_main.c
+++ b/realm/realm_payload_main.c
@@ -156,6 +156,8 @@
 			break;
 		case REALM_MULTIPLE_REC_PSCI_DENIED_CMD:
 			test_succeed = test_realm_multiple_rec_psci_denied_cmd();
+		case REALM_MULTIPLE_REC_MULTIPLE_CPU_CMD:
+			test_succeed = test_realm_multiple_rec_multiple_cpu_cmd();
 			break;
 		case REALM_PAUTH_SET_CMD:
 			test_succeed = test_realm_pauth_set_cmd();
diff --git a/realm/realm_psci.c b/realm/realm_psci.c
index 2a5b951..a4a287b 100644
--- a/realm/realm_psci.c
+++ b/realm/realm_psci.c
@@ -74,15 +74,16 @@
 
 void realm_secondary_entrypoint(u_register_t cxt_id)
 {
-	u_register_t my_mpidr;
+	u_register_t my_mpidr, id;
 	secondary_ep_t ep;
 
 	my_mpidr = read_mpidr_el1() & MPID_MASK;
 	ep = entrypoint[my_mpidr];
+	id = context_id[my_mpidr];
 	if (ep != NULL) {
 		entrypoint[my_mpidr] = NULL;
 		context_id[my_mpidr] = 0;
-		(ep)(context_id[my_mpidr]);
+		(ep)(id);
 	} else {
 		/*
 		 * Host can execute Rec directly without CPU_ON