diff --git a/arch/arm64/kvm/hyp/nvhe/iommu/iommu.c b/arch/arm64/kvm/hyp/nvhe/iommu/iommu.c
index 198952c02bba..d40e111f8a41 100644
--- a/arch/arm64/kvm/hyp/nvhe/iommu/iommu.c
+++ b/arch/arm64/kvm/hyp/nvhe/iommu/iommu.c
@@ -219,14 +219,29 @@ handle_to_domain(pkvm_handle_t domain_id)
 static int domain_get(struct kvm_hyp_iommu_domain *domain)
 {
 	int old = atomic_fetch_inc_acquire(&domain->refs);
+	struct pkvm_hyp_vcpu *hyp_vcpu = __get_vcpu();
+	int ret = 0;
 
 	BUG_ON(!old || (old + 1 < 0));
-	return 0;
+
+	/* check done after refcount is elevated to avoid race with alloc_domain */
+	if (!hyp_vcpu && domain->vm)
+		ret = -EPERM;
+	if (hyp_vcpu && (domain->vm != pkvm_hyp_vcpu_to_hyp_vm(hyp_vcpu)))
+		ret = -EPERM;
+
+	if (ret)
+		atomic_dec_return_release(&domain->refs);
+	return ret;
 }
 
 static void domain_put(struct kvm_hyp_iommu_domain *domain)
 {
+	struct pkvm_hyp_vcpu *hyp_vcpu = __get_vcpu();
+
 	BUG_ON(!atomic_dec_return_release(&domain->refs));
+	WARN_ON(!hyp_vcpu && domain->vm);
+	WARN_ON(hyp_vcpu && (domain->vm != pkvm_hyp_vcpu_to_hyp_vm(hyp_vcpu)));
 }
 
 static int kvm_iommu_init_atomic_pool(struct kvm_hyp_memcache *atomic_mc)
@@ -285,6 +300,8 @@ int kvm_iommu_alloc_domain(pkvm_handle_t domain_id, int type)
 {
 	int ret = -EINVAL;
 	struct kvm_hyp_iommu_domain *domain;
+	struct pkvm_hyp_vcpu *hyp_vcpu = __get_vcpu();
+	struct pkvm_hyp_vm *vm;
 
 	domain = handle_to_domain(domain_id);
 	if (!domain)
@@ -299,6 +316,10 @@ int kvm_iommu_alloc_domain(pkvm_handle_t domain_id, int type)
 	if (ret)
 		goto out_unlock;
 
+	if (hyp_vcpu) {
+		vm = pkvm_hyp_vcpu_to_hyp_vm(hyp_vcpu);
+		domain->vm = vm;
+	}
 	atomic_set_release(&domain->refs, 1);
 out_unlock:
 	hyp_spin_unlock(&kvm_iommu_domain_lock);
@@ -309,13 +330,18 @@ int kvm_iommu_free_domain(pkvm_handle_t domain_id)
 {
 	int ret = 0;
 	struct kvm_hyp_iommu_domain *domain;
+	struct pkvm_hyp_vcpu *hyp_vcpu = __get_vcpu();
+	struct pkvm_hyp_vm *vm = NULL;
 
 	domain = handle_to_domain(domain_id);
 	if (!domain)
 		return -EINVAL;
 
 	hyp_spin_lock(&kvm_iommu_domain_lock);
-	if (WARN_ON(atomic_cmpxchg_acquire(&domain->refs, 1, 0) != 1)) {
+	if (hyp_vcpu)
+		vm = pkvm_hyp_vcpu_to_hyp_vm(hyp_vcpu);
+
+	if (WARN_ON(atomic_cmpxchg_acquire(&domain->refs, 1, 0) != 1) || domain->vm != vm) {
 		ret = -EINVAL;
 		goto out_unlock;
 	}
diff --git a/include/kvm/iommu.h b/include/kvm/iommu.h
index 734dca08ad5f..e571e0e1767b 100644
--- a/include/kvm/iommu.h
+++ b/include/kvm/iommu.h
@@ -21,6 +21,7 @@ struct kvm_hyp_iommu_domain {
 	atomic_t		refs;
 	pkvm_handle_t		domain_id;
 	void			*priv;
+	void			*vm;
 	ANDROID_KABI_RESERVE(1);
 	ANDROID_KABI_RESERVE(2);
 };