test(ff-a): test FFA_MSG_WAIT with retain RX buffer flag

Add helper functions for VMs/SPs to call FFA_MSG_WAIT with flags to
optionally retain RX buffer ownership.

Add test showing verifying that SP keeps ownership of its RX buffer when
calling FFA_MSG_WAIT with retain RX buffer flag set.

Signed-off-by: Kathleen Capella <kathleen.capella@arm.com>
Change-Id: I518a082462fbaf5d3f16811d1d61a6c36ba2aa2f
diff --git a/inc/vmapi/hf/call.h b/inc/vmapi/hf/call.h
index 5e68e2e..cbabee3 100644
--- a/inc/vmapi/hf/call.h
+++ b/inc/vmapi/hf/call.h
@@ -297,6 +297,10 @@
  *
  * The mailbox must be cleared before a new message can be received.
  *
+ * By default, FFA_MSG_WAIT will release the mailbox back to the SPMC. The
+ * FFA_MSG_WAIT_FLAG_RETAIN_RX flag can be used with `ffa_msg_wait_with_flags`
+ * function to override this default and allow the VM to retain the RX buffer.
+ *
  * If no message is immediately available and there are no enabled and pending
  * interrupts (irrespective of whether interrupts are enabled globally), then
  * this will block until a message is available or an enabled interrupt becomes
@@ -304,13 +308,18 @@
  * that a message becoming available is also treated like a wake-up event.
  *
  * Returns:
- *  - FFA_MSG_SEND if a message is successfully received.
  *  - FFA_ERROR FFA_NOT_SUPPORTED if called from the primary VM.
  *  - FFA_ERROR FFA_INTERRUPTED if an interrupt happened during the call.
  */
+static inline struct ffa_value ffa_msg_wait_with_flags(uint32_t flags)
+{
+	return ffa_call_ext(
+		(struct ffa_value){.func = FFA_MSG_WAIT_32, .arg2 = flags});
+}
+
 static inline struct ffa_value ffa_msg_wait(void)
 {
-	return ffa_call_ext((struct ffa_value){.func = FFA_MSG_WAIT_32});
+	return ffa_msg_wait_with_flags(0);
 }
 
 /**
diff --git a/test/vmapi/primary_with_secondaries/rx_ownership.c b/test/vmapi/primary_with_secondaries/rx_ownership.c
index f9330f7..7b5607e 100644
--- a/test/vmapi/primary_with_secondaries/rx_ownership.c
+++ b/test/vmapi/primary_with_secondaries/rx_ownership.c
@@ -53,3 +53,42 @@
 	ret = ffa_run(service1_info->vm_id, 0);
 	EXPECT_EQ(ret.func, FFA_YIELD_32);
 }
+
+/**
+ * Test to verify that when an SP calls FFA_MSG_WAIT with the
+ * FFA_MSG_WAIT_FLAG_RETAIN_RX flag it does not relinquish ownership of its RX
+ * buffer back to the SPMC. Test sequence is as follows:
+ * 	1. PVM runs SP via FFA_RUN.
+ * 	2. SP calls FFA_PARTITION_INFO_GET to gain ownership of its RX buffer
+ * from the SPMC.
+ * 	3. SP calls FFA_MSG_WAIT with FFA_MSG_WAIT_FLAG_RETAIN_RX flag
+ * returning control back to PVM but keeping ownership of the RX buffer.
+ * 	4. PVM runs SP via FFA_RUN.
+ * 	5. SP calls FFA_PARTITION_INFO_GET and checks that if fails because the
+ * 	SPMC does not have ownership of the buffer and cannot fill it.
+ */
+TEST_PRECONDITION(rx_ownership, ffa_msg_wait_retain_buffer, service1_is_not_vm)
+{
+	struct mailbox_buffers mb = set_up_mailbox();
+	struct ffa_value ret;
+
+	struct ffa_partition_info *service1_info = service1(mb.recv);
+
+	SERVICE_SELECT(service1_info->vm_id, "test_ffa_msg_wait_retain_buffer",
+		       mb.send);
+
+	/*
+	 * Run service to call FFA_PARTITION_INFO_GET, then FFA_MSG_WAIT with
+	 * the flag to retain the RX buffer set.
+	 */
+	ret = ffa_run(service1_info->vm_id, 0);
+	EXPECT_EQ(ret.func, FFA_MSG_WAIT_32);
+
+	/*
+	 * Run service to attempt to call FFA_PARTITION_INFO_GET again and check
+	 * that it fails since FFA_MSG_WAIT did not release RX buffer to the
+	 * SPMC.
+	 */
+	ret = ffa_run(service1_info->vm_id, 0);
+	EXPECT_EQ(ret.func, FFA_YIELD_32);
+}
diff --git a/test/vmapi/primary_with_secondaries/services/rx_ownership.c b/test/vmapi/primary_with_secondaries/services/rx_ownership.c
index 837acdf..8e55118 100644
--- a/test/vmapi/primary_with_secondaries/services/rx_ownership.c
+++ b/test/vmapi/primary_with_secondaries/services/rx_ownership.c
@@ -52,3 +52,42 @@
 
 	ffa_yield();
 }
+
+TEST_SERVICE(test_ffa_msg_wait_retain_buffer)
+{
+	struct ffa_value ret;
+	struct ffa_uuid uuid;
+
+	/* A Null UUID requests information for all partitions. */
+	ffa_uuid_init(0, 0, 0, 0, &uuid);
+	ret = ffa_partition_info_get(&uuid, 0);
+	EXPECT_EQ(ret.func, FFA_SUCCESS_32);
+
+	dlog_verbose("FFA_PARTITION_INFO_GET put the RX buffer FULL.");
+
+	/*
+	 * Subsequent call to FFA_PARTITION_INFO_GET should fail because buffer
+	 * is busy.
+	 */
+	ret = ffa_partition_info_get(&uuid, 0);
+	EXPECT_FFA_ERROR(ret, FFA_BUSY);
+
+	dlog_verbose("FFA_PARTITION_INFO_GET attested the RX buffer is FULL.");
+
+	/* FFA_MSG_WAIT should retain buffer. */
+	ret = ffa_msg_wait_with_flags(FFA_MSG_WAIT_FLAG_RETAIN_RX);
+	EXPECT_EQ(ret.func, FFA_RUN_32);
+	dlog_verbose("FFA_MSG_WAIT returned");
+
+	/*
+	 * Additional call to ffa_partition_info_get should fail since buffer
+	 * was retained.
+	 */
+	ret = ffa_partition_info_get(&uuid, 0);
+	EXPECT_FFA_ERROR(ret, FFA_BUSY);
+	dlog_verbose(
+		"FFA_PARTITION_INFO_GET attested the RX buffer is FULL after "
+		"retention.");
+
+	ffa_yield();
+}