FFA: Set and verify allocator of memory handle

Change-Id: I19dff61311228399cdf8f26a855e91981473a8a7
Signed-off-by: J-Alves <joao.alves@arm.com>
diff --git a/src/ffa_memory.c b/src/ffa_memory.c
index 00ec68c..8b28191 100644
--- a/src/ffa_memory.c
+++ b/src/ffa_memory.c
@@ -8,6 +8,7 @@
 
 #include "hf/ffa_memory.h"
 
+#include "hf/arch/ffa_memory_handle.h"
 #include "hf/arch/other_world.h"
 
 #include "hf/api.h"
@@ -121,6 +122,14 @@
 	tee_retrieve_buffer[HF_MAILBOX_SIZE * MAX_FRAGMENTS];
 
 /**
+ * Extracts the index from a memory handle allocated by Hafnium's current world.
+ */
+uint64_t ffa_memory_handle_get_index(ffa_memory_handle_t handle)
+{
+	return handle & ~FFA_MEMORY_HANDLE_ALLOCATOR_MASK;
+}
+
+/**
  * Initialises the next available `struct ffa_memory_share_state` and sets
  * `share_state_ret` to a pointer to it. If `handle` is
  * `FFA_MEMORY_HANDLE_INVALID` then allocates an appropriate handle, otherwise
@@ -150,8 +159,7 @@
 
 			if (handle == FFA_MEMORY_HANDLE_INVALID) {
 				memory_region->handle =
-					i |
-					FFA_MEMORY_HANDLE_ALLOCATOR_HYPERVISOR;
+					ffa_memory_handle_make(i);
 			} else {
 				memory_region->handle = handle;
 			}
@@ -205,7 +213,7 @@
 			    struct ffa_memory_share_state **share_state_ret)
 {
 	struct ffa_memory_share_state *share_state;
-	uint32_t index;
+	uint64_t index;
 
 	CHECK(share_states.share_states != NULL);
 	CHECK(share_state_ret != NULL);
@@ -214,9 +222,8 @@
 	 * First look for a share_state allocated by us, in which case the
 	 * handle is based on the index.
 	 */
-	if ((handle & FFA_MEMORY_HANDLE_ALLOCATOR_MASK) ==
-	    FFA_MEMORY_HANDLE_ALLOCATOR_HYPERVISOR) {
-		index = handle & ~FFA_MEMORY_HANDLE_ALLOCATOR_MASK;
+	if (ffa_memory_handle_allocated_by_current_world(handle)) {
+		index = ffa_memory_handle_get_index(handle);
 		if (index < MAX_MEM_SHARES) {
 			share_state = &share_states.share_states[index];
 			if (share_state->share_func != 0) {