Set correct mailbox state after sending messages to TEE.

https://hafnium-review.googlesource.com/c/hafnium/+/8840 was still
wrong.

Bug: 132429380
Change-Id: Ife266a3bc4e962e50d07d3dbd1eb81cad867c040
diff --git a/src/api.c b/src/api.c
index d6779a8..39da4d8 100644
--- a/src/api.c
+++ b/src/api.c
@@ -936,20 +936,25 @@
 		return ret;
 	}
 
+	to.vm->mailbox.state = MAILBOX_STATE_RECEIVED;
+
 	/* Messages for the TEE are sent on via the dispatcher. */
 	if (to.vm->id == HF_TEE_VM_ID) {
 		struct spci_value call = spci_msg_recv_return(to.vm);
 
-		to.vm->mailbox.state = MAILBOX_STATE_READ;
-		return arch_tee_call(call);
+		ret = arch_tee_call(call);
+		/*
+		 * After the call to the TEE completes it must have finished
+		 * reading its RX buffer, so it is ready for another message.
+		 */
+		to.vm->mailbox.state = MAILBOX_STATE_EMPTY;
 		/*
 		 * Don't return to the primary VM in this case, as the TEE is
 		 * not (yet) scheduled via SPCI.
 		 */
+		return ret;
 	}
 
-	to.vm->mailbox.state = MAILBOX_STATE_RECEIVED;
-
 	/* Return to the primary VM directly or with a switch. */
 	if (from_id != HF_PRIMARY_VM_ID) {
 		*next = api_switch_to_primary(current, primary_ret,