diff --git a/kernel/ethosu_mailbox.c b/kernel/ethosu_mailbox.c
index 4c64f17..3e7284b 100644
--- a/kernel/ethosu_mailbox.c
+++ b/kernel/ethosu_mailbox.c
@@ -79,12 +79,16 @@
 }
 
 struct ethosu_mailbox_msg *ethosu_mailbox_find(struct ethosu_mailbox *mbox,
-					       int msg_id)
+					       int msg_id,
+					       uint32_t msg_type)
 {
 	struct ethosu_mailbox_msg *ptr = (struct ethosu_mailbox_msg *)idr_find(
 		&mbox->msg_idr, msg_id);
 
 	if (ptr == NULL)
+		return ERR_PTR(-ENOENT);
+
+	if (ptr->type != msg_type)
 		return ERR_PTR(-EINVAL);
 
 	return ptr;
@@ -147,6 +151,8 @@
 		}
 	};
 
+	msg->type = rpmsg.header.type;
+
 	return rpmsg_send(mbox->ept, &rpmsg, sizeof(rpmsg.header));
 }
 
@@ -172,6 +178,8 @@
 	struct ethosu_core_msg_inference_req *inf_req = &rpmsg.inf_req;
 	uint32_t i;
 
+	msg->type = rpmsg.header.type;
+
 	/* Verify that the uapi and core has the same number of pmus */
 	if (pmu_event_config_count != ETHOSU_CORE_PMU_MAX) {
 		dev_err(mbox->dev, "PMU count misconfigured.");
@@ -218,6 +226,8 @@
 	};
 	struct ethosu_core_msg_network_info_req *info_req = &rpmsg.net_info_req;
 
+	msg->type = rpmsg.header.type;
+
 	if (network != NULL) {
 		info_req->network.type = ETHOSU_CORE_NETWORK_BUFFER;
 		ethosu_core_set_size(network, &info_req->network.buffer);
@@ -246,6 +256,8 @@
 		}
 	};
 
+	msg->type = rpmsg.header.type;
+
 	return rpmsg_send(mbox->ept, &rpmsg,
 			  sizeof(rpmsg.header) + sizeof(rpmsg.cancel_req));
 }