Add helper function to check whether VM ID is in current world.

Assert that vCPU is in the right world before switching to it.

Change-Id: I5f994c63ec97058a52adf7dba04a88e6e69b1b25
Signed-off-by: Andrew Walbran <qwandor@google.com>
diff --git a/inc/hf/vm.h b/inc/hf/vm.h
index cf03b0c..01aaaca 100644
--- a/inc/hf/vm.h
+++ b/inc/hf/vm.h
@@ -148,6 +148,7 @@
 struct vcpu *vm_get_vcpu(struct vm *vm, ffa_vcpu_index_t vcpu_index);
 struct wait_entry *vm_get_wait_entry(struct vm *vm, ffa_vm_id_t for_vm);
 ffa_vm_id_t vm_id_for_wait_entry(struct vm *vm, struct wait_entry *entry);
+bool vm_id_is_current_world(ffa_vm_id_t vm_id);
 
 bool vm_identity_map(struct vm_locked vm_locked, paddr_t begin, paddr_t end,
 		     uint32_t mode, struct mpool *ppool, ipaddr_t *ipa);
diff --git a/inc/vmapi/hf/types.h b/inc/vmapi/hf/types.h
index 3243256..017c4b0 100644
--- a/inc/vmapi/hf/types.h
+++ b/inc/vmapi/hf/types.h
@@ -31,11 +31,19 @@
 
 #define HF_INVALID_VM_ID 0x7fff
 
+/**
+ * The bit of the VM ID which indicates whether the VM ID is allocated by the
+ * normal world or the secure world.
+ */
+#define HF_VM_ID_WORLD_MASK 0x8000
+
+/** The special VM ID reserved for the hypervisor in the normal world. */
 #define HF_HYPERVISOR_VM_ID 0
 
 /**
- * An offset to use when assigning VM IDs.
- * The offset is needed because VM ID 0 is reserved.
+ * An offset to use when assigning VM IDs within the current world.
+ * The offset from HF_VM_ID_BASE is needed because VM ID `HF_VM_ID_BASE + 0` is
+ * reserved for the hypervisor/SPM.
  */
 #define HF_VM_ID_OFFSET (HF_VM_ID_BASE + 1)
 
@@ -49,8 +57,8 @@
 #define HF_PRIMARY_VM_ID (HF_VM_ID_OFFSET + HF_PRIMARY_VM_INDEX)
 
 /**
- * The special VM ID reserved for the OS running in the trusted execution
- * environment, e.g. secure EL1 on AArch64.
+ * The special VM ID reserved for the OS or SPMC running in the trusted
+ * execution environment, e.g. secure EL1 or EL2 on AArch64.
  */
 #define HF_TEE_VM_ID 0x8000
 
diff --git a/src/arch/aarch64/hypervisor/handler.c b/src/arch/aarch64/hypervisor/handler.c
index e330176..c6909c6 100644
--- a/src/arch/aarch64/hypervisor/handler.c
+++ b/src/arch/aarch64/hypervisor/handler.c
@@ -448,6 +448,7 @@
 	struct ffa_value other_world_args =
 		arch_regs_get_args(&other_world_vcpu->regs);
 
+	CHECK(!vm_id_is_current_world(other_world_vcpu->vm->id));
 	CHECK(*next == NULL);
 
 	while (*next == NULL) {
@@ -472,6 +473,14 @@
 	}
 
 	/*
+	 * ffa_handler set *next to something, which means it wants to switch
+	 * back to an SP in EL1. It must be something in this world though, as
+	 * if it wanted to return back to the other world (where the last FF-A
+	 * call came from) it wouldn't have set *next at all.
+	 */
+	CHECK(vm_id_is_current_world((*next)->vm->id));
+
+	/*
 	 * Store the return value on the other world vCPU, ready for next time
 	 * we switch to it (in case they aren't overwritten at that point by
 	 * whatever API function decides to make the switch).
@@ -498,6 +507,8 @@
 			vcpu->interrupts.enabled_and_pending_count > 0);
 		sl_unlock(&vcpu->lock);
 	} else {
+		CHECK(vm_id_is_current_world(next->vm->id));
+
 		/*
 		 * About to switch vCPUs, set the bit for the vCPU to which we
 		 * are switching in the saved copy of the register.
diff --git a/src/vm.c b/src/vm.c
index 2ad3516..fbbdc9f 100644
--- a/src/vm.c
+++ b/src/vm.c
@@ -201,6 +201,17 @@
 }
 
 /**
+ * Return whether the given VM ID represents an entity in the current world:
+ * i.e. the hypervisor or a normal world VM when running in the normal world, or
+ * the SPM or an SP when running in the secure world.
+ */
+bool vm_id_is_current_world(ffa_vm_id_t vm_id)
+{
+	return (vm_id & HF_VM_ID_WORLD_MASK) !=
+	       (HF_OTHER_WORLD_ID & HF_VM_ID_WORLD_MASK);
+}
+
+/**
  * Map a range of addresses to the VM in both the MMU and the IOMMU.
  *
  * mm_vm_defrag should always be called after a series of page table updates,