fix(ff-a): memory reclaim after a donate

This change fixes the case in which memory reclaim is performed after
a successful memory donate. Prior to this change memory reclaim would
fail with FFA_DENIED.
Basic tests are added to make sure a memory reclaim can happen after
successful calls to any of the memory send interfaces.

Change-Id: I87ce3e113246bec8564dde6754a4d65cb6794170
Signed-off-by: J-Alves <joao.alves@arm.com>
diff --git a/src/ffa_memory.c b/src/ffa_memory.c
index bf2fe5f..51a73ec 100644
--- a/src/ffa_memory.c
+++ b/src/ffa_memory.c
@@ -667,10 +667,13 @@
 			MM_MODE_INVALID | MM_MODE_UNOWNED | MM_MODE_SHARED;
 		uint32_t orig_to_state = orig_to_mode & state_mask;
 
-		if (orig_to_state != MM_MODE_INVALID &&
-		    orig_to_state != MM_MODE_SHARED) {
-			return ffa_error(FFA_DENIED);
-		}
+		/*
+		 * If the original ffa memory send call has been processed
+		 * successfully, it is expected the orig_to_mode would overlay
+		 * with `state_mask`, as a result of the function
+		 * `ffa_send_check_transition`.
+		 */
+		assert(orig_to_state != 0U);
 	} else {
 		/*
 		 * Ensure the retriever has the expected state. We don't care
diff --git a/test/vmapi/primary_with_secondaries/memory_sharing.c b/test/vmapi/primary_with_secondaries/memory_sharing.c
index 90e6d5a..997f435 100644
--- a/test/vmapi/primary_with_secondaries/memory_sharing.c
+++ b/test/vmapi/primary_with_secondaries/memory_sharing.c
@@ -221,12 +221,115 @@
 	}
 }
 
+/**
+ * Base test function for a memory reclaim after a successful memory send call.
+ */
+static void memory_send_reclaim(uint32_t msg_size,
+				struct ffa_value (*mem_send_function)(uint32_t,
+								      uint32_t))
+{
+	struct ffa_value ret;
+	ffa_memory_handle_t handle;
+
+	/*
+	 * It is assumed that the same pages as for other mem share tests are
+	 * used.
+	 */
+	uint8_t *ptr = pages;
+	ret = mem_send_function(msg_size, msg_size);
+	EXPECT_EQ(ret.func, FFA_SUCCESS_32);
+
+	handle = ffa_mem_success_handle(ret);
+
+	EXPECT_EQ(ffa_mem_reclaim(handle, 0).func, FFA_SUCCESS_32);
+
+	/* Write to pages to validate access has been reestablished. */
+	for (uint8_t i = 0; i < 5; i++) {
+		*ptr = i;
+	}
+}
+
 TEAR_DOWN(memory_sharing)
 {
 	EXPECT_FFA_ERROR(ffa_rx_release(), FFA_DENIED);
 }
 
 /**
+ * Test memory reclaim after a donate.
+ */
+TEST(memory_sharing, donate_reclaim)
+{
+	struct mailbox_buffers mb = set_up_mailbox();
+	struct ffa_memory_region_constituent constituents[] = {
+		{.address = (uint64_t)pages, .page_count = 2},
+		{.address = (uint64_t)pages + PAGE_SIZE * 3, .page_count = 1},
+	};
+	uint32_t msg_size;
+
+	EXPECT_EQ(ffa_memory_region_init(
+			  mb.send, HF_MAILBOX_SIZE, HF_PRIMARY_VM_ID,
+			  SERVICE_VM1, constituents, ARRAY_SIZE(constituents),
+			  0, 0, FFA_DATA_ACCESS_NOT_SPECIFIED,
+			  FFA_INSTRUCTION_ACCESS_NOT_SPECIFIED,
+			  FFA_MEMORY_NORMAL_MEM, FFA_MEMORY_CACHE_WRITE_BACK,
+			  FFA_MEMORY_INNER_SHAREABLE, NULL, &msg_size),
+		  0);
+
+	/* Call base function's test. */
+	memory_send_reclaim(msg_size, ffa_mem_donate);
+}
+
+/**
+ * Test memory reclaim after a lend.
+ */
+TEST(memory_sharing, lend_reclaim)
+{
+	struct mailbox_buffers mb = set_up_mailbox();
+	struct ffa_memory_region_constituent constituents[] = {
+		{.address = (uint64_t)pages, .page_count = 2},
+		{.address = (uint64_t)pages + PAGE_SIZE * 3, .page_count = 1},
+	};
+	uint32_t msg_size;
+
+	EXPECT_EQ(ffa_memory_region_init(
+			  mb.send, HF_MAILBOX_SIZE, HF_PRIMARY_VM_ID,
+			  SERVICE_VM1, constituents, ARRAY_SIZE(constituents),
+			  0, 0, FFA_DATA_ACCESS_RW,
+			  FFA_INSTRUCTION_ACCESS_NOT_SPECIFIED,
+			  FFA_MEMORY_NORMAL_MEM, FFA_MEMORY_CACHE_WRITE_BACK,
+			  FFA_MEMORY_INNER_SHAREABLE, NULL, &msg_size),
+		  0);
+
+	/* Call base function's test. */
+	memory_send_reclaim(msg_size, ffa_mem_lend);
+}
+
+/**
+ * Test memory reclaim after a share.
+ */
+TEST(memory_sharing, share_reclaim)
+{
+	struct mailbox_buffers mb = set_up_mailbox();
+	struct ffa_memory_region_constituent constituents[] = {
+		{.address = (uint64_t)pages, .page_count = 2},
+		{.address = (uint64_t)pages + PAGE_SIZE * 3, .page_count = 1},
+	};
+	uint32_t msg_size;
+
+	EXPECT_EQ(ffa_memory_region_init(
+			  mb.send, HF_MAILBOX_SIZE, HF_PRIMARY_VM_ID,
+			  SERVICE_VM1, constituents, ARRAY_SIZE(constituents),
+			  0, 0, FFA_DATA_ACCESS_RW,
+			  FFA_INSTRUCTION_ACCESS_NOT_SPECIFIED,
+			  FFA_MEMORY_NORMAL_MEM, FFA_MEMORY_CACHE_WRITE_BACK,
+			  FFA_MEMORY_INNER_SHAREABLE, NULL, &msg_size),
+		  0);
+
+	/* Call base function's test. */
+	memory_send_reclaim(msg_size, ffa_mem_share);
+}
+
+/**
  * Sharing memory concurrently gives both VMs access to the memory so it can be
  * used for communication.
  */