diff --git a/arch/arm64/kvm/hyp/nvhe/ffa.c b/arch/arm64/kvm/hyp/nvhe/ffa.c
index feaba59af32e..b595d5d80d3c 100644
--- a/arch/arm64/kvm/hyp/nvhe/ffa.c
+++ b/arch/arm64/kvm/hyp/nvhe/ffa.c
@@ -113,33 +113,27 @@ static struct kvm_ffa_buffers *ffa_get_buffers(struct pkvm_hyp_vcpu *hyp_vcpu)
 
 DECLARE_STATIC_KEY_FALSE(kvm_ffa_unmap_on_lend);
 
-static int ffa_host_store_handle(u64 ffa_handle, bool is_lend)
+static struct ffa_handle *ffa_host_alloc_handle(void)
 {
 	u32 i;
-	struct ffa_handle *free_handle = NULL;
-
-	if (!static_branch_unlikely(&kvm_ffa_unmap_on_lend))
-		return 0;
+	struct ffa_handle *handle;
 
 	if (spm_free_handle) {
 		WARN_ON(spm_free_handle < spm_handles ||
 			spm_free_handle >= (spm_handles + num_spm_handles));
-		free_handle = spm_free_handle;
+		handle = spm_free_handle;
 		spm_free_handle = NULL;
-	} else {
-		for (i = 0; i < num_spm_handles; i++)
-			if (spm_handles[i].handle == FFA_INVALID_SPM_HANDLE)
-				break;
-
-		if (i == num_spm_handles)
-			return -ENOSPC;
-
-		free_handle = &spm_handles[i];
+		return handle;
 	}
 
-	free_handle->handle = ffa_handle;
-	free_handle->is_lend = is_lend;
-	return 0;
+	for (i = 0; i < num_spm_handles; i++)
+		if (spm_handles[i].handle == FFA_INVALID_SPM_HANDLE)
+			break;
+
+	if (i == num_spm_handles)
+		return NULL;
+
+	return &spm_handles[i];
 }
 
 static struct ffa_handle *ffa_host_get_handle(u64 ffa_handle)
@@ -909,6 +903,7 @@ static int __do_ffa_mem_xfer(const u64 func_id,
 	struct ffa_mem_transfer *transfer = NULL;
 	u64 ffa_handle;
 	bool is_lend = func_id == FFA_FN64_MEM_LEND;
+	struct ffa_handle *handle = NULL;
 
 	if (addr_mbz || npages_mbz || fraglen > len ||
 	    fraglen > KVM_FFA_MBOX_NR_PAGES * PAGE_SIZE) {
@@ -1009,6 +1004,14 @@ static int __do_ffa_mem_xfer(const u64 func_id,
 	if (ret)
 		goto out_unlock;
 
+	if (!hyp_vcpu && static_branch_unlikely(&kvm_ffa_unmap_on_lend)) {
+		handle = ffa_host_alloc_handle();
+		if (!handle) {
+			ret = -ENOSPC;
+			goto out_unlock;
+		}
+	}
+
 	ffa_mem_xfer(res, func_id, len, fraglen);
 	if (fraglen != len) {
 		if (res->a0 != FFA_MEM_FRAG_RX)
@@ -1027,10 +1030,9 @@ static int __do_ffa_mem_xfer(const u64 func_id,
 	if (hyp_vcpu && transfer) {
 		transfer->ffa_handle = ffa_handle;
 		list_add(&transfer->node, &ffa_buf->xfer_list);
-	} else if (!hyp_vcpu) {
-		ret = ffa_host_store_handle(ffa_handle, is_lend);
-		if (ret)
-			goto err_unshare;
+	} else if (handle) {
+		handle->handle = ffa_handle;
+		handle->is_lend = is_lend;
 	}
 	hyp_spin_unlock(&kvm_ffa_hyp_lock);
 	return 0;