Begin to abstract the ABI.

Different architectures or standards will require different ABIs but
this should be hidden from the Hafnium API. Further changes are required
to support returning multiple registers but that will be hidden from the
API.

Change-Id: I35bc674c35bd4bb4c8c30e02f1075024a3bc44db
diff --git a/build/BUILD.gn b/build/BUILD.gn
index 11b1d78..ef0cf9f 100644
--- a/build/BUILD.gn
+++ b/build/BUILD.gn
@@ -34,6 +34,7 @@
       current_toolchain == arch_toolchain) {
     include_dirs = [
       "//inc",
+      "//inc/vmapi",
       "//src/arch/${arch}/inc",
     ]
 
diff --git a/driver/linux b/driver/linux
index 0973a2e..dc8cab5 160000
--- a/driver/linux
+++ b/driver/linux
@@ -1 +1 @@
-Subproject commit 0973a2ec678062affef66303b1e057b7c0be555e
+Subproject commit dc8cab578ec3c7ea61aefd74f2d059acdc4d94df
diff --git a/inc/hf/api.h b/inc/hf/api.h
index de27c35..514ac2f 100644
--- a/inc/hf/api.h
+++ b/inc/hf/api.h
@@ -19,13 +19,17 @@
 #include "hf/cpu.h"
 #include "hf/vm.h"
 
+#include "vmapi/hf/call.h"
+
 int64_t api_vm_get_count(void);
 int64_t api_vcpu_get_count(uint32_t vm_id);
-int64_t api_vcpu_run(uint32_t vm_id, uint32_t vcpu_idx, struct vcpu **next);
+struct hf_vcpu_run_return api_vcpu_run(uint32_t vm_id, uint32_t vcpu_idx,
+				       struct vcpu **next);
 int64_t api_vm_configure(ipaddr_t send, ipaddr_t recv);
 
 int64_t api_mailbox_send(uint32_t vm_id, size_t size, struct vcpu **next);
-int64_t api_mailbox_receive(bool block, struct vcpu **next);
+struct hf_mailbox_receive_return api_mailbox_receive(bool block,
+						     struct vcpu **next);
 int64_t api_mailbox_clear(void);
 
 struct vcpu *api_wait_for_interrupt(void);
diff --git a/inc/vmapi/hf/abi.h b/inc/vmapi/hf/abi.h
new file mode 100644
index 0000000..9866059
--- /dev/null
+++ b/inc/vmapi/hf/abi.h
@@ -0,0 +1,120 @@
+/*
+ * Copyright 2018 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "hf/types.h"
+
+enum hf_vcpu_run_code {
+	HF_VCPU_RUN_YIELD,
+	HF_VCPU_RUN_WAIT_FOR_INTERRUPT,
+	HF_VCPU_RUN_WAKE_UP,
+	HF_VCPU_RUN_MESSAGE,
+	HF_VCPU_RUN_SLEEP,
+};
+
+struct hf_vcpu_run_return {
+	enum hf_vcpu_run_code code;
+	union {
+		struct {
+			uint32_t vm_id;
+			uint16_t vcpu;
+		} wake_up;
+		struct {
+			uint32_t size;
+		} message;
+		struct {
+			uint64_t ns;
+		} sleep;
+	};
+};
+
+struct hf_mailbox_receive_return {
+	uint32_t vm_id;
+	uint32_t size;
+};
+
+/**
+ * Encode an hf_vcpu_run_return struct in the 64-bit packing ABI.
+ */
+static inline uint64_t hf_vcpu_run_return_encode(struct hf_vcpu_run_return res)
+{
+	uint64_t ret = res.code & 0xff;
+	switch (res.code) {
+	case HF_VCPU_RUN_WAKE_UP:
+		ret |= (uint64_t)res.wake_up.vm_id << 32;
+		ret |= (uint64_t)res.wake_up.vcpu << 16;
+		break;
+	case HF_VCPU_RUN_MESSAGE:
+		ret |= (uint64_t)res.message.size << 32;
+		break;
+	case HF_VCPU_RUN_SLEEP:
+		ret |= res.sleep.ns << 8;
+		break;
+	default:
+		break;
+	}
+	return ret;
+}
+
+/**
+ * Decode an hf_vcpu_run_return struct from the 64-bit packing ABI.
+ */
+static inline struct hf_vcpu_run_return hf_vcpu_run_return_decode(uint64_t res)
+{
+	struct hf_vcpu_run_return ret;
+
+	ret.code = (enum hf_vcpu_run_code)(res & 0xff);
+
+	/* Some codes include more data. */
+	switch (ret.code) {
+	case HF_VCPU_RUN_WAKE_UP:
+		ret.wake_up.vm_id = res >> 32;
+		ret.wake_up.vcpu = (res >> 16) & 0xffff;
+		break;
+	case HF_VCPU_RUN_MESSAGE:
+		ret.message.size = res >> 32;
+		break;
+	case HF_VCPU_RUN_SLEEP:
+		ret.sleep.ns = res >> 8;
+		break;
+	default:
+		break;
+	}
+
+	return ret;
+}
+
+/**
+ * Encode an hf_mailbox_receive_return struct in the 64-bit packing ABI.
+ */
+static inline uint64_t hf_mailbox_receive_return_encode(
+	struct hf_mailbox_receive_return res)
+{
+	return res.vm_id | ((uint64_t)res.size << 32);
+}
+
+/**
+ * Decode an hf_mailbox_receive_return struct from the 64-bit packing ABI.
+ */
+static inline struct hf_mailbox_receive_return hf_mailbox_receive_return_decode(
+	uint64_t res)
+{
+	return (struct hf_mailbox_receive_return){
+		.vm_id = (uint32_t)(res & 0xffffffff),
+		.size = (uint32_t)(res >> 32),
+	};
+}
diff --git a/inc/vmapi/hf/call.h b/inc/vmapi/hf/call.h
index 4e70395..771997a 100644
--- a/inc/vmapi/hf/call.h
+++ b/inc/vmapi/hf/call.h
@@ -16,21 +16,8 @@
 
 #pragma once
 
-#if defined(__linux__) && defined(__KERNEL__)
-
-#include <linux/types.h>
-
-typedef phys_addr_t hf_ipaddr_t;
-
-#else
-
-#include <stdbool.h>
-#include <stddef.h>
-#include <stdint.h>
-
-typedef uintptr_t hf_ipaddr_t;
-
-#endif
+#include "hf/abi.h"
+#include "hf/types.h"
 
 /* Keep macro alignment */
 /* clang-format off */
@@ -44,34 +31,9 @@
 #define HF_MAILBOX_RECEIVE  0xff05
 #define HF_MAILBOX_CLEAR    0xff06
 
-/* The ID of the primary VM which is responsile for scheduling. */
-#define HF_PRIMARY_VM_ID 0
-
-/* Return codes for hf_vcpu_run(). */
-#define HF_VCPU_RUN_YIELD              0x00
-#define HF_VCPU_RUN_WAIT_FOR_INTERRUPT 0x01
-#define HF_VCPU_RUN_WAKE_UP            0x02
-#define HF_VCPU_RUN_MESSAGE            0x03
-
-/* Construct and destruct the hf_vcpu_run() response. */
-#define HF_VCPU_RUN_RESPONSE(code, vm_id, data)               \
-	((code & 0xff) | ((uint64_t)(vm_id & 0xffff) << 16) | \
-	 ((uint64_t)data << 32))
-#define HF_VCPU_RUN_CODE(ret) (ret & 0xff)
-#define HF_VCPU_RUN_VM_ID(ret) ((ret >> 16) & 0xffff)
-#define HF_VCPU_RUN_DATA(ret) (ret >> 32)
-
-/* Construct and destruct the hf_mailbox_receive() response. */
-#define HF_MAILBOX_RECEIVE_RESPONSE(vm_id, size) \
-	((vm_id & 0xffff) | ((uint64_t)size << 32))
-#define HF_MAILBOX_RECEIVE_VM_ID(ret) (ret & 0xffff)
-#define HF_MAILBOX_RECEIVE_SIZE(ret) (ret >> 32)
-
+/* The amount of data that can be sent to a mailbox. */
 #define HF_MAILBOX_SIZE 4096
 
-#define HF_INVALID_VM_ID 0xffff
-#define HF_INVALID_VCPU  0xffff
-
 /* clang-format on */
 
 /**
@@ -83,9 +45,11 @@
 /**
  * Runs the given vcpu of the given vm.
  */
-static inline int64_t hf_vcpu_run(uint32_t vm_id, uint32_t vcpu_idx)
+static inline struct hf_vcpu_run_return hf_vcpu_run(uint32_t vm_id,
+						    uint32_t vcpu_idx)
 {
-	return hf_call(HF_VCPU_RUN, vm_id, vcpu_idx, 0);
+	return hf_vcpu_run_return_decode(
+		hf_call(HF_VCPU_RUN, vm_id, vcpu_idx, 0));
 }
 
 /**
@@ -125,11 +89,14 @@
  * Called by secondary VMs to receive a message. The call can optionally block
  * until a message is received.
  *
+ * If no message was received, the VM ID will be HF_INVALID_VM_ID.
+ *
  * The mailbox must be cleared before a new message can be received.
  */
-static inline int64_t hf_mailbox_receive(bool block)
+static inline struct hf_mailbox_receive_return hf_mailbox_receive(bool block)
 {
-	return hf_call(HF_MAILBOX_RECEIVE, block, 0, 0);
+	return hf_mailbox_receive_return_decode(
+		hf_call(HF_MAILBOX_RECEIVE, block, 0, 0));
 }
 
 /**
diff --git a/inc/vmapi/hf/types.h b/inc/vmapi/hf/types.h
index 1e4c7f8..55de5f1 100644
--- a/inc/vmapi/hf/types.h
+++ b/inc/vmapi/hf/types.h
@@ -16,14 +16,13 @@
 
 #pragma once
 
+/* Define the standard types for the platform. */
 #if defined(__linux__) && defined(__KERNEL__)
 
 #include <linux/types.h>
 
 typedef phys_addr_t hf_ipaddr_t;
 
-#define PRIu16 "hu"
-
 #else
 
 #include <stdbool.h>
@@ -34,10 +33,9 @@
 
 #endif
 
-typedef uint16_t hf_vm_id_t;
-
-#define HF_VM_ID_MAX UINT16_MAX
-#define HF_PRI_VM_ID PRIu16
-
 /* The ID of the primary VM which is responsile for scheduling. */
 #define HF_PRIMARY_VM_ID 0
+
+/* Invalid values for fields to indicate absence or errors. */
+#define HF_INVALID_VM_ID 0xffffffff
+#define HF_INVALID_VCPU 0xffff
diff --git a/src/BUILD.gn b/src/BUILD.gn
index eb0f76f..cad0a06 100644
--- a/src/BUILD.gn
+++ b/src/BUILD.gn
@@ -106,11 +106,15 @@
 executable("unit_tests") {
   testonly = true
   sources = [
+    "abi_test.cc",
     "fdt_handler_test.cc",
     "fdt_test.cc",
     "mm_test.cc",
   ]
-  cflags_cc = [ "-Wno-c99-extensions" ]
+  cflags_cc = [
+    "-Wno-c99-extensions",
+    "-Wno-nested-anon-types",
+  ]
   ldflags = [
     "-Xlinker",
     "-defsym=text_begin=0",
diff --git a/src/abi_test.cc b/src/abi_test.cc
new file mode 100644
index 0000000..d6bff21
--- /dev/null
+++ b/src/abi_test.cc
@@ -0,0 +1,194 @@
+/*
+ * Copyright 2018 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+extern "C" {
+#include "vmapi/hf/abi.h"
+}
+
+#include <gmock/gmock.h>
+
+namespace
+{
+using ::testing::Eq;
+
+/**
+ * Simulate an uninitialized hf_vcpu_run_return so it can be detected if any
+ * uninitialized fields make their way into the encoded form which would
+ * indicate a data leak.
+ */
+struct hf_vcpu_run_return dirty_vcpu_run_return()
+{
+	struct hf_vcpu_run_return res;
+	memset(&res, 0xc5, sizeof(res));
+	return res;
+}
+
+/**
+ * Simulate an uninitialized hf_mailbox_receive_return so it can be detected if
+ * any uninitialized fields make their way into the encoded form which would
+ * indicate a data leak.
+ */
+struct hf_mailbox_receive_return dirty_mailbox_receive_return()
+{
+	struct hf_mailbox_receive_return res;
+	memset(&res, 0xc5, sizeof(res));
+	return res;
+}
+
+/**
+ * Encode a yield response without leaking.
+ */
+TEST(abi, hf_vcpu_run_return_encode_yield)
+{
+	struct hf_vcpu_run_return res = dirty_vcpu_run_return();
+	res.code = HF_VCPU_RUN_YIELD;
+	EXPECT_THAT(hf_vcpu_run_return_encode(res), Eq(0));
+}
+
+/**
+ * Decode a yield response ignoring the irrelevant bits.
+ */
+TEST(abi, hf_vcpu_run_return_decode_yield)
+{
+	struct hf_vcpu_run_return res =
+		hf_vcpu_run_return_decode(0x1a1a1a1a2b2b2b00);
+	EXPECT_THAT(res.code, Eq(HF_VCPU_RUN_YIELD));
+}
+
+/**
+ * Encode wait-for-interrupt response without leaking.
+ */
+TEST(abi, hf_vcpu_run_return_encode_wait_for_interrupt)
+{
+	struct hf_vcpu_run_return res = dirty_vcpu_run_return();
+	res.code = HF_VCPU_RUN_WAIT_FOR_INTERRUPT;
+	EXPECT_THAT(hf_vcpu_run_return_encode(res), Eq(1));
+}
+
+/**
+ * Decode a wait-for-interrupt response ignoring the irrelevant bits.
+ */
+TEST(abi, hf_vcpu_run_return_decode_wait_for_interrupt)
+{
+	struct hf_vcpu_run_return res =
+		hf_vcpu_run_return_decode(0x1234abcdbadb0101);
+	EXPECT_THAT(res.code, Eq(HF_VCPU_RUN_WAIT_FOR_INTERRUPT));
+}
+
+/**
+ * Encode wake up response without leaking.
+ */
+TEST(abi, hf_vcpu_run_return_encode_wake_up)
+{
+	struct hf_vcpu_run_return res = dirty_vcpu_run_return();
+	res.code = HF_VCPU_RUN_WAKE_UP;
+	res.wake_up.vm_id = 0x12345678;
+	res.wake_up.vcpu = 0xabcd;
+	EXPECT_THAT(hf_vcpu_run_return_encode(res), Eq(0x12345678abcd0002));
+}
+
+/**
+ * Decode a wake up response ignoring the irrelevant bits.
+ */
+TEST(abi, hf_vcpu_run_return_decode_wake_up)
+{
+	struct hf_vcpu_run_return res =
+		hf_vcpu_run_return_decode(0xbeefd00df00daf02);
+	EXPECT_THAT(res.code, Eq(HF_VCPU_RUN_WAKE_UP));
+	EXPECT_THAT(res.wake_up.vm_id, Eq(0xbeefd00d));
+	EXPECT_THAT(res.wake_up.vcpu, Eq(0xf00d));
+}
+
+/**
+ * Encode message response without leaking.
+ */
+TEST(abi, hf_vcpu_run_return_encode_message)
+{
+	struct hf_vcpu_run_return res = dirty_vcpu_run_return();
+	res.code = HF_VCPU_RUN_MESSAGE;
+	res.message.size = 0xdeadbeef;
+	EXPECT_THAT(hf_vcpu_run_return_encode(res), Eq(0xdeadbeef00000003));
+}
+
+/**
+ * Decode a wake up response ignoring the irrelevant bits.
+ */
+TEST(abi, hf_vcpu_run_return_decode_message)
+{
+	struct hf_vcpu_run_return res =
+		hf_vcpu_run_return_decode(0x1123581314916203);
+	EXPECT_THAT(res.code, Eq(HF_VCPU_RUN_MESSAGE));
+	EXPECT_THAT(res.message.size, Eq(0x11235813));
+}
+
+/**
+ * Encode sleep response without leaking.
+ */
+TEST(abi, hf_vcpu_run_return_encode_sleep)
+{
+	struct hf_vcpu_run_return res = dirty_vcpu_run_return();
+	res.code = HF_VCPU_RUN_SLEEP;
+	res.sleep.ns = 0xcafed00dfeeded;
+	EXPECT_THAT(hf_vcpu_run_return_encode(res), Eq(0xcafed00dfeeded04));
+}
+
+/**
+ * Encoding a sleep response with too large a sleep duration will drop the top
+ * octet.
+ */
+TEST(abi, hf_vcpu_run_return_encode_sleep_too_long)
+{
+	struct hf_vcpu_run_return res = dirty_vcpu_run_return();
+	res.code = HF_VCPU_RUN_SLEEP;
+	res.sleep.ns = 0xcc88888888888888;
+	EXPECT_THAT(hf_vcpu_run_return_encode(res), Eq(0x8888888888888804));
+}
+
+/**
+ * Decode a sleep response.
+ */
+TEST(abi, hf_vcpu_run_return_decode_sleep)
+{
+	struct hf_vcpu_run_return res =
+		hf_vcpu_run_return_decode(0x1a2b3c4d5e6f7704);
+	EXPECT_THAT(res.code, Eq(HF_VCPU_RUN_SLEEP));
+	EXPECT_THAT(res.sleep.ns, Eq(0x1a2b3c4d5e6f77));
+}
+
+/**
+ * Encode a mailbox receive response without leaking.
+ */
+TEST(abi, hf_mailbox_receive_return_encode)
+{
+	struct hf_mailbox_receive_return res = dirty_mailbox_receive_return();
+	res.vm_id = 0x12345678;
+	res.size = 0xaabbccdd;
+	EXPECT_THAT(hf_mailbox_receive_return_encode(res),
+		    Eq(0xaabbccdd12345678));
+}
+
+/**
+ * Decode a mailbox receive response.
+ */
+TEST(abi, hf_mailbox_receive_return_decode)
+{
+	struct hf_mailbox_receive_return res =
+		hf_mailbox_receive_return_decode(0X8badf00d00ddba11);
+	EXPECT_THAT(res.vm_id, Eq(0X00ddba11));
+	EXPECT_THAT(res.size, Eq(0x8badf00d));
+}
+
+} /* namespace */
diff --git a/src/api.c b/src/api.c
index b2e2c36..3b8d921 100644
--- a/src/api.c
+++ b/src/api.c
@@ -34,7 +34,7 @@
  * to cause HF_VCPU_RUN to return and the primary VM to regain control of the
  * cpu.
  */
-static struct vcpu *api_switch_to_primary(size_t primary_retval,
+static struct vcpu *api_switch_to_primary(struct hf_vcpu_run_return primary_ret,
 					  enum vcpu_state secondary_state)
 {
 	struct vcpu *vcpu = cpu()->current;
@@ -45,9 +45,10 @@
 	vm_set_current(primary);
 
 	/*
-	 * Set the return valuefor the primary VM's call to HF_VCPU_RUN.
+	 * Set the return value for the primary VM's call to HF_VCPU_RUN.
 	 */
-	arch_regs_set_retval(&next->regs, primary_retval);
+	arch_regs_set_retval(&next->regs,
+			     hf_vcpu_run_return_encode(primary_ret));
 
 	/* Mark the vcpu as waiting. */
 	sl_lock(&vcpu->lock);
@@ -63,9 +64,10 @@
  */
 struct vcpu *api_yield(void)
 {
-	return api_switch_to_primary(
-		HF_VCPU_RUN_RESPONSE(HF_VCPU_RUN_YIELD, 0, 0),
-		vcpu_state_ready);
+	struct hf_vcpu_run_return ret = {
+		.code = HF_VCPU_RUN_YIELD,
+	};
+	return api_switch_to_primary(ret, vcpu_state_ready);
 }
 
 /**
@@ -74,9 +76,10 @@
  */
 struct vcpu *api_wait_for_interrupt(void)
 {
-	return api_switch_to_primary(
-		HF_VCPU_RUN_RESPONSE(HF_VCPU_RUN_WAIT_FOR_INTERRUPT, 0, 0),
-		vcpu_state_blocked_interrupt);
+	struct hf_vcpu_run_return ret = {
+		.code = HF_VCPU_RUN_WAIT_FOR_INTERRUPT,
+	};
+	return api_switch_to_primary(ret, vcpu_state_blocked_interrupt);
 }
 
 /**
@@ -110,51 +113,51 @@
 /**
  * Runs the given vcpu of the given vm.
  */
-int64_t api_vcpu_run(uint32_t vm_id, uint32_t vcpu_idx, struct vcpu **next)
+struct hf_vcpu_run_return api_vcpu_run(uint32_t vm_id, uint32_t vcpu_idx,
+				       struct vcpu **next)
 {
 	struct vm *vm;
 	struct vcpu *vcpu;
-	int64_t ret;
+	struct hf_vcpu_run_return ret = {
+		.code = HF_VCPU_RUN_WAIT_FOR_INTERRUPT,
+	};
 
 	/* Only the primary VM can switch vcpus. */
 	if (cpu()->current->vm->id != HF_PRIMARY_VM_ID) {
-		goto fail;
+		goto out;
 	}
 
 	/* Only secondary VM vcpus can be run. */
 	if (vm_id == HF_PRIMARY_VM_ID) {
-		goto fail;
+		goto out;
 	}
 
 	/* The requested VM must exist. */
 	vm = vm_get(vm_id);
 	if (vm == NULL) {
-		goto fail;
+		goto out;
 	}
 
 	/* The requested vcpu must exist. */
 	if (vcpu_idx >= vm->vcpu_count) {
-		goto fail;
+		goto out;
 	}
 
 	vcpu = &vm->vcpus[vcpu_idx];
 
 	sl_lock(&vcpu->lock);
 	if (vcpu->state != vcpu_state_ready) {
-		ret = HF_VCPU_RUN_RESPONSE(HF_VCPU_RUN_WAIT_FOR_INTERRUPT, 0,
-					   0);
+		ret.code = HF_VCPU_RUN_WAIT_FOR_INTERRUPT;
 	} else {
 		vcpu->state = vcpu_state_running;
 		vm_set_current(vm);
 		*next = vcpu;
-		ret = HF_VCPU_RUN_RESPONSE(HF_VCPU_RUN_YIELD, 0, 0);
+		ret.code = HF_VCPU_RUN_YIELD;
 	}
 	sl_unlock(&vcpu->lock);
 
+out:
 	return ret;
-
-fail:
-	return HF_VCPU_RUN_RESPONSE(HF_VCPU_RUN_WAIT_FOR_INTERRUPT, 0, 0);
 }
 
 /**
@@ -250,7 +253,9 @@
 	const void *from_buf;
 	uint16_t vcpu;
 	int64_t ret;
-	int64_t primary_ret;
+	struct hf_vcpu_run_return primary_ret = {
+		.code = HF_VCPU_RUN_WAIT_FOR_INTERRUPT,
+	};
 
 	/* Limit the size of transfer. */
 	if (size > HF_MAILBOX_SIZE) {
@@ -297,8 +302,8 @@
 
 	/* Messages for the primary VM are delivered directly. */
 	if (to->id == HF_PRIMARY_VM_ID) {
-		primary_ret =
-			HF_VCPU_RUN_RESPONSE(HF_VCPU_RUN_MESSAGE, 0, size);
+		primary_ret.code = HF_VCPU_RUN_MESSAGE;
+		primary_ret.message.size = size;
 		ret = 0;
 		/*
 		 * clang-tidy isn't able to prove that
@@ -338,8 +343,11 @@
 
 		/* Return from HF_MAILBOX_RECEIVE. */
 		arch_regs_set_retval(&to_vcpu->regs,
-				     HF_MAILBOX_RECEIVE_RESPONSE(
-					     to->mailbox.recv_from_id, size));
+				     hf_mailbox_receive_return_encode((
+					     struct hf_mailbox_receive_return){
+					     .vm_id = to->mailbox.recv_from_id,
+					     .size = size,
+				     }));
 
 		sl_unlock(&to_vcpu->lock);
 
@@ -347,7 +355,9 @@
 	}
 
 	/* Return to the primary VM directly or with a switch. */
-	primary_ret = HF_VCPU_RUN_RESPONSE(HF_VCPU_RUN_WAKE_UP, to->id, vcpu);
+	primary_ret.code = HF_VCPU_RUN_WAKE_UP;
+	primary_ret.wake_up.vm_id = to->id;
+	primary_ret.wake_up.vcpu = vcpu;
 	ret = 0;
 
 out:
@@ -364,7 +374,7 @@
 
 	/* If the sender is the primary, return the vcpu to schedule. */
 	if (from->id == HF_PRIMARY_VM_ID) {
-		return vcpu;
+		return primary_ret.wake_up.vcpu;
 	}
 
 	/* Switch to primary for scheduling and return success to the sender. */
@@ -378,18 +388,21 @@
  *
  * No new messages can be received until the mailbox has been cleared.
  */
-int64_t api_mailbox_receive(bool block, struct vcpu **next)
+struct hf_mailbox_receive_return api_mailbox_receive(bool block,
+						     struct vcpu **next)
 {
 	struct vcpu *vcpu = cpu()->current;
 	struct vm *vm = vcpu->vm;
-	int64_t ret = 0;
+	struct hf_mailbox_receive_return ret = {
+		.vm_id = HF_INVALID_VM_ID,
+	};
 
 	/*
 	 * The primary VM will receive messages as a status code from running
 	 * vcpus and must not call this function.
 	 */
 	if (vm->id == HF_PRIMARY_VM_ID) {
-		return -1;
+		return ret;
 	}
 
 	sl_lock(&vm->lock);
@@ -397,14 +410,13 @@
 	/* Return pending messages without blocking. */
 	if (vm->mailbox.state == mailbox_state_received) {
 		vm->mailbox.state = mailbox_state_read;
-		ret = HF_MAILBOX_RECEIVE_RESPONSE(vm->mailbox.recv_from_id,
-						  vm->mailbox.recv_bytes);
+		ret.vm_id = vm->mailbox.recv_from_id;
+		ret.size = vm->mailbox.recv_bytes;
 		goto out;
 	}
 
 	/* No pending message so fail if not allowed to block. */
 	if (!block) {
-		ret = -1;
 		goto out;
 	}
 
diff --git a/src/arch/aarch64/handler.c b/src/arch/aarch64/handler.c
index 53f8c88..8fa5733 100644
--- a/src/arch/aarch64/handler.c
+++ b/src/arch/aarch64/handler.c
@@ -198,7 +198,8 @@
 		break;
 
 	case HF_VCPU_RUN:
-		ret.user_ret = api_vcpu_run(arg1, arg2, &ret.new);
+		ret.user_ret = hf_vcpu_run_return_encode(
+			api_vcpu_run(arg1, arg2, &ret.new));
 		break;
 
 	case HF_VM_CONFIGURE:
@@ -210,7 +211,8 @@
 		break;
 
 	case HF_MAILBOX_RECEIVE:
-		ret.user_ret = api_mailbox_receive(arg1, &ret.new);
+		ret.user_ret = hf_mailbox_receive_return_encode(
+			api_mailbox_receive(arg1, &ret.new));
 		break;
 
 	case HF_MAILBOX_CLEAR:
diff --git a/test/vm/primary_only.c b/test/vm/primary_only.c
index 166fb0c..dd7ab48 100644
--- a/test/vm/primary_only.c
+++ b/test/vm/primary_only.c
@@ -37,3 +37,15 @@
 {
 	EXPECT_EQ(hf_vcpu_get_count(0xffffffff), -1);
 }
+
+TEST(hf_vcpu_run, cannot_run_primary)
+{
+	struct hf_vcpu_run_return res = hf_vcpu_run(HF_PRIMARY_VM_ID, 0);
+	EXPECT_EQ(res.code, HF_VCPU_RUN_WAIT_FOR_INTERRUPT);
+}
+
+TEST(hf_vcpu_run, cannot_run_absent_secondary)
+{
+	struct hf_vcpu_run_return res = hf_vcpu_run(1, 0);
+	EXPECT_EQ(res.code, HF_VCPU_RUN_WAIT_FOR_INTERRUPT);
+}
diff --git a/test/vm/primary_with_secondaries.c b/test/vm/primary_with_secondaries.c
index 146b576..6a37090 100644
--- a/test/vm/primary_with_secondaries.c
+++ b/test/vm/primary_with_secondaries.c
@@ -43,7 +43,7 @@
 /* clang-format on */
 
 /**
- * Confirm there are 3 secondary VMs.
+ * Confirm there are 3 secondary VMs as well as this primary VM.
  */
 TEST(hf_vm_get_count, three_secondary_vms)
 {
@@ -68,6 +68,33 @@
 }
 
 /**
+ * The primary can't be run by the hypervisor.
+ */
+TEST(hf_vcpu_run, cannot_run_primary)
+{
+	struct hf_vcpu_run_return res = hf_vcpu_run(HF_PRIMARY_VM_ID, 0);
+	EXPECT_EQ(res.code, HF_VCPU_RUN_WAIT_FOR_INTERRUPT);
+}
+
+/**
+ * Can only run a VM that exists.
+ */
+TEST(hf_vcpu_run, cannot_run_absent_secondary)
+{
+	struct hf_vcpu_run_return res = hf_vcpu_run(1234, 0);
+	EXPECT_EQ(res.code, HF_VCPU_RUN_WAIT_FOR_INTERRUPT);
+}
+
+/**
+ * Can only run a vcpu that exists.
+ */
+TEST(hf_vcpu_run, cannot_run_absent_vcpu)
+{
+	struct hf_vcpu_run_return res = hf_vcpu_run(ECHO_VM_ID, 1234);
+	EXPECT_EQ(res.code, HF_VCPU_RUN_WAIT_FOR_INTERRUPT);
+}
+
+/**
  * The configured send/receive addresses can't be unaligned.
  */
 TEST(hf_vm_configure, fails_with_unaligned_pointer)
@@ -112,23 +139,44 @@
 }
 
 /**
+ * The primary receives messages from hf_vcpu_run().
+ */
+TEST(hf_mailbox_receive, cannot_receive_from_primary_blocking)
+{
+	struct hf_mailbox_receive_return res = hf_mailbox_receive(true);
+	EXPECT_EQ(res.vm_id, HF_INVALID_VM_ID);
+	EXPECT_EQ(res.size, 0);
+}
+
+/**
+ * The primary receives messages from hf_vcpu_run().
+ */
+TEST(hf_mailbox_receive, cannot_receive_from_primary_non_blocking)
+{
+	struct hf_mailbox_receive_return res = hf_mailbox_receive(false);
+	EXPECT_EQ(res.vm_id, HF_INVALID_VM_ID);
+	EXPECT_EQ(res.size, 0);
+}
+
+/**
  * Send and receive the same message from the echo VM.
  */
 TEST(mailbox, echo)
 {
 	const char message[] = "Echo this back to me!";
+	struct hf_vcpu_run_return run_res;
 
 	/* Configure mailbox pages. */
 	EXPECT_EQ(hf_vm_configure(send_page_addr, recv_page_addr), 0);
-	EXPECT_EQ(hf_vcpu_run(ECHO_VM_ID, 0),
-		  HF_VCPU_RUN_RESPONSE(HF_VCPU_RUN_WAIT_FOR_INTERRUPT, 0, 0));
+	run_res = hf_vcpu_run(ECHO_VM_ID, 0);
+	EXPECT_EQ(run_res.code, HF_VCPU_RUN_WAIT_FOR_INTERRUPT);
 
 	/* Set the message, echo it and check it didn't change. */
 	memcpy(send_page, message, sizeof(message));
 	EXPECT_EQ(hf_mailbox_send(ECHO_VM_ID, sizeof(message)), 0);
-	EXPECT_EQ(
-		hf_vcpu_run(ECHO_VM_ID, 0),
-		HF_VCPU_RUN_RESPONSE(HF_VCPU_RUN_MESSAGE, 0, sizeof(message)));
+	run_res = hf_vcpu_run(ECHO_VM_ID, 0);
+	EXPECT_EQ(run_res.code, HF_VCPU_RUN_MESSAGE);
+	EXPECT_EQ(run_res.message.size, sizeof(message));
 	EXPECT_EQ(memcmp(recv_page, message, sizeof(message)), 0);
 	EXPECT_EQ(hf_mailbox_clear(), 0);
 }
@@ -140,13 +188,14 @@
 TEST(mailbox, relay)
 {
 	const char message[] = "Send this round the relay!";
+	struct hf_vcpu_run_return run_res;
 
 	/* Configure mailbox pages. */
 	EXPECT_EQ(hf_vm_configure(send_page_addr, recv_page_addr), 0);
-	EXPECT_EQ(hf_vcpu_run(RELAY_A_VM_ID, 0),
-		  HF_VCPU_RUN_RESPONSE(HF_VCPU_RUN_WAIT_FOR_INTERRUPT, 0, 0));
-	EXPECT_EQ(hf_vcpu_run(RELAY_B_VM_ID, 0),
-		  HF_VCPU_RUN_RESPONSE(HF_VCPU_RUN_WAIT_FOR_INTERRUPT, 0, 0));
+	run_res = hf_vcpu_run(RELAY_A_VM_ID, 0);
+	EXPECT_EQ(run_res.code, HF_VCPU_RUN_WAIT_FOR_INTERRUPT);
+	run_res = hf_vcpu_run(RELAY_B_VM_ID, 0);
+	EXPECT_EQ(run_res.code, HF_VCPU_RUN_WAIT_FOR_INTERRUPT);
 
 	/*
 	 * Send the message to relay_a which is then sent to relay_b before
@@ -154,11 +203,13 @@
 	 */
 	memcpy(send_page, message, sizeof(message));
 	EXPECT_EQ(hf_mailbox_send(RELAY_A_VM_ID, sizeof(message)), 0);
-	EXPECT_EQ(hf_vcpu_run(RELAY_A_VM_ID, 0),
-		  HF_VCPU_RUN_RESPONSE(HF_VCPU_RUN_WAKE_UP, RELAY_B_VM_ID, 0));
-	EXPECT_EQ(
-		hf_vcpu_run(RELAY_B_VM_ID, 0),
-		HF_VCPU_RUN_RESPONSE(HF_VCPU_RUN_MESSAGE, 0, sizeof(message)));
+	run_res = hf_vcpu_run(RELAY_A_VM_ID, 0);
+	EXPECT_EQ(run_res.code, HF_VCPU_RUN_WAKE_UP);
+	EXPECT_EQ(run_res.wake_up.vm_id, RELAY_B_VM_ID);
+	EXPECT_EQ(run_res.wake_up.vcpu, 0);
+	run_res = hf_vcpu_run(RELAY_B_VM_ID, 0);
+	EXPECT_EQ(run_res.code, HF_VCPU_RUN_MESSAGE);
+	EXPECT_EQ(run_res.message.size, sizeof(message));
 	EXPECT_EQ(memcmp(recv_page, message, sizeof(message)), 0);
 	EXPECT_EQ(hf_mailbox_clear(), 0);
 }
diff --git a/test/vm/secondaries/echo.c b/test/vm/secondaries/echo.c
index 7982757..d2bf32b 100644
--- a/test/vm/secondaries/echo.c
+++ b/test/vm/secondaries/echo.c
@@ -37,11 +37,9 @@
 
 	/* Loop, echo messages back to the sender. */
 	for (;;) {
-		uint64_t ret = hf_mailbox_receive(true);
-		uint32_t vm_id = HF_MAILBOX_RECEIVE_VM_ID(ret);
-		uint32_t size = HF_MAILBOX_RECEIVE_SIZE(ret);
-		memcpy(send_page, recv_page, size);
+		struct hf_mailbox_receive_return res = hf_mailbox_receive(true);
+		memcpy(send_page, recv_page, res.size);
 		hf_mailbox_clear();
-		hf_mailbox_send(vm_id, size);
+		hf_mailbox_send(res.vm_id, res.size);
 	}
 }
diff --git a/test/vm/secondaries/relay_a.c b/test/vm/secondaries/relay_a.c
index 8fa7bb5..a6f6a3a 100644
--- a/test/vm/secondaries/relay_a.c
+++ b/test/vm/secondaries/relay_a.c
@@ -39,10 +39,9 @@
 
 	/* Loop, forward messages to the next VM. */
 	for (;;) {
-		uint64_t ret = hf_mailbox_receive(true);
-		uint32_t size = HF_MAILBOX_RECEIVE_SIZE(ret);
-		memcpy(send_page, recv_page, size);
+		struct hf_mailbox_receive_return res = hf_mailbox_receive(true);
+		memcpy(send_page, recv_page, res.size);
 		hf_mailbox_clear();
-		hf_mailbox_send(FORWARD_VM_ID, size);
+		hf_mailbox_send(FORWARD_VM_ID, res.size);
 	}
 }
diff --git a/test/vm/secondaries/relay_b.c b/test/vm/secondaries/relay_b.c
index ae74e6c..2f9643b 100644
--- a/test/vm/secondaries/relay_b.c
+++ b/test/vm/secondaries/relay_b.c
@@ -37,10 +37,9 @@
 
 	/* Loop, forward messages to the primary. */
 	for (;;) {
-		uint64_t ret = hf_mailbox_receive(true);
-		uint32_t size = HF_MAILBOX_RECEIVE_SIZE(ret);
-		memcpy(send_page, recv_page, size);
+		struct hf_mailbox_receive_return res = hf_mailbox_receive(true);
+		memcpy(send_page, recv_page, res.size);
 		hf_mailbox_clear();
-		hf_mailbox_send(HF_PRIMARY_VM_ID, size);
+		hf_mailbox_send(HF_PRIMARY_VM_ID, res.size);
 	}
 }