Set correct mailbox state when sending messages to TEE.

This fixes a bug introduced in
https://hafnium-review.googlesource.com/c/hafnium/+/8226.

Bug: 132429380
Change-Id: I86ae9147664fa9ccc6858dd2eb73d7e8de7b59ff
diff --git a/src/api.c b/src/api.c
index d61bf45..d6779a8 100644
--- a/src/api.c
+++ b/src/api.c
@@ -936,12 +936,11 @@
 		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);
 		/*
 		 * Don't return to the primary VM in this case, as the TEE is
@@ -949,6 +948,8 @@
 		 */
 	}
 
+	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,