FFA: Set and verify allocator of memory handle

Change-Id: I19dff61311228399cdf8f26a855e91981473a8a7
Signed-off-by: J-Alves <joao.alves@arm.com>
diff --git a/src/api.c b/src/api.c
index f0766a8..41bc1a3 100644
--- a/src/api.c
+++ b/src/api.c
@@ -9,6 +9,7 @@
 #include "hf/api.h"
 
 #include "hf/arch/cpu.h"
+#include "hf/arch/ffa_memory_handle.h"
 #include "hf/arch/mm.h"
 #include "hf/arch/other_world.h"
 #include "hf/arch/timer.h"
@@ -2103,8 +2104,7 @@
 	struct vm *to = current->vm;
 	struct ffa_value ret;
 
-	if ((handle & FFA_MEMORY_HANDLE_ALLOCATOR_MASK) ==
-	    FFA_MEMORY_HANDLE_ALLOCATOR_HYPERVISOR) {
+	if (ffa_memory_handle_allocated_by_current_world(handle)) {
 		struct vm_locked to_locked = vm_lock(to);
 
 		ret = ffa_memory_reclaim(to_locked, handle, flags,
diff --git a/src/arch/aarch64/hypervisor/BUILD.gn b/src/arch/aarch64/hypervisor/BUILD.gn
index 0f9c15e..8f57598 100644
--- a/src/arch/aarch64/hypervisor/BUILD.gn
+++ b/src/arch/aarch64/hypervisor/BUILD.gn
@@ -17,6 +17,7 @@
 source_set("other_world") {
   public_configs = [ "//src/arch/aarch64:config" ]
   sources = [
+    "ffa_memory_handle.c",
     "other_world.c",
   ]
   deps = [
diff --git a/src/arch/aarch64/hypervisor/ffa_memory_handle.c b/src/arch/aarch64/hypervisor/ffa_memory_handle.c
new file mode 100644
index 0000000..647186d
--- /dev/null
+++ b/src/arch/aarch64/hypervisor/ffa_memory_handle.c
@@ -0,0 +1,29 @@
+/*
+ * Copyright 2021 The Hafnium Authors.
+ *
+ * Use of this source code is governed by a BSD-style
+ * license that can be found in the LICENSE file or at
+ * https://opensource.org/licenses/BSD-3-Clause.
+ */
+
+#include "hf/arch/ffa_memory_handle.h"
+
+ffa_memory_handle_t ffa_memory_handle_make(uint64_t index)
+{
+#if SECURE_WORLD == 1
+	return (index & ~FFA_MEMORY_HANDLE_ALLOCATOR_MASK) |
+	       FFA_MEMORY_HANDLE_ALLOCATOR_SPMC;
+#else
+	return index | FFA_MEMORY_HANDLE_ALLOCATOR_HYPERVISOR;
+#endif
+}
+
+bool ffa_memory_handle_allocated_by_current_world(ffa_memory_handle_t handle)
+{
+	return (handle & FFA_MEMORY_HANDLE_ALLOCATOR_MASK) ==
+#if SECURE_WORLD == 1
+	       FFA_MEMORY_HANDLE_ALLOCATOR_SPMC;
+#else
+	       FFA_MEMORY_HANDLE_ALLOCATOR_HYPERVISOR;
+#endif
+}
diff --git a/src/arch/fake/hypervisor/BUILD.gn b/src/arch/fake/hypervisor/BUILD.gn
index 97a82f6..8f28802 100644
--- a/src/arch/fake/hypervisor/BUILD.gn
+++ b/src/arch/fake/hypervisor/BUILD.gn
@@ -16,6 +16,7 @@
 
 source_set("other_world") {
   sources = [
+    "ffa_memory_handle.c",
     "other_world.c",
   ]
 }
diff --git a/src/arch/fake/hypervisor/ffa_memory_handle.c b/src/arch/fake/hypervisor/ffa_memory_handle.c
new file mode 100644
index 0000000..9091295
--- /dev/null
+++ b/src/arch/fake/hypervisor/ffa_memory_handle.c
@@ -0,0 +1,20 @@
+/*
+ * Copyright 2021 The Hafnium Authors.
+ *
+ * Use of this source code is governed by a BSD-style
+ * license that can be found in the LICENSE file or at
+ * https://opensource.org/licenses/BSD-3-Clause.
+ */
+
+#include "hf/arch/ffa_memory_handle.h"
+
+ffa_memory_handle_t ffa_memory_handle_make(uint64_t index)
+{
+	return index;
+}
+
+bool ffa_memory_handle_allocated_by_current_world(ffa_memory_handle_t handle)
+{
+	(void)handle;
+	return true;
+}
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) {