[PATCH v4 2/4] iommufd/viommu: Associate a kvm pointer to iommufd_viommu

["Aneesh Kumar K.V (Arm)" <aneesh.kumar@kernel.org>]

From: Nicolin Chen <nicolinc@nvidia.com>

For TSM use case, vIOMMU will need to get access to KVM to make different
TSM calls. Associate the kvm pointer to the vIOMMU and take kvm's refcount.

Based on an original patch by Shameer Kolothum <shameerali.kolothum.thodi@huawei.com>

[nicolinc: hold kvm's users_count]
Signed-off-by: Nicolin Chen <nicolinc@nvidia.com>
Signed-off-by: Aneesh Kumar K.V (Arm) <aneesh.kumar@kernel.org>
---
 drivers/iommu/iommufd/viommu.c | 20 +++++++++++++++++---
 include/linux/iommu.h          |  1 +
 include/linux/iommufd.h        |  1 +
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/drivers/iommu/iommufd/viommu.c b/drivers/iommu/iommufd/viommu.c
index 4081deda9b33..1e023d1cf8d5 100644
--- a/drivers/iommu/iommufd/viommu.c
+++ b/drivers/iommu/iommufd/viommu.c
@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: GPL-2.0-only
 /* Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES
  */
+#include <linux/kvm_host.h>
 #include "iommufd_private.h"
 
 void iommufd_viommu_destroy(struct iommufd_object *obj)
@@ -11,6 +12,8 @@ void iommufd_viommu_destroy(struct iommufd_object *obj)
 	if (viommu->ops && viommu->ops->destroy)
 		viommu->ops->destroy(viommu);
 	refcount_dec(&viommu->hwpt->common.obj.users);
+	if (viommu->kvm)
+		kvm_put_kvm(viommu->kvm);
 	xa_destroy(&viommu->vdevs);
 }
 
@@ -68,14 +71,22 @@ int iommufd_viommu_alloc_ioctl(struct iommufd_ucmd *ucmd)
 		goto out_put_hwpt;
 	}
 
+	if (idev->kvm && !kvm_get_kvm_safe(idev->kvm)) {
+		rc = -ENOENT;
+		goto out_put_hwpt;
+	}
+
 	viommu = (struct iommufd_viommu *)_iommufd_object_alloc_ucmd(
 		ucmd, viommu_size, IOMMUFD_OBJ_VIOMMU);
 	if (IS_ERR(viommu)) {
 		rc = PTR_ERR(viommu);
-		goto out_put_hwpt;
+		goto out_put_kvm;
 	}
 
 	xa_init(&viommu->vdevs);
+	viommu->kvm = idev->kvm;
+	if (viommu->kvm)
+		kvm_get_kvm(viommu->kvm);
 	viommu->type = cmd->type;
 	viommu->ictx = ucmd->ictx;
 	viommu->hwpt = hwpt_paging;
@@ -92,17 +103,20 @@ int iommufd_viommu_alloc_ioctl(struct iommufd_ucmd *ucmd)
 	rc = ops->viommu_init(viommu, hwpt_paging->common.domain,
 			      user_data.len ? &user_data : NULL);
 	if (rc)
-		goto out_put_hwpt;
+		goto out_put_kvm;
 
 	/* It is a driver bug that viommu->ops isn't filled */
 	if (WARN_ON_ONCE(!viommu->ops)) {
 		rc = -EOPNOTSUPP;
-		goto out_put_hwpt;
+		goto out_put_kvm;
 	}
 
 	cmd->out_viommu_id = viommu->obj.id;
 	rc = iommufd_ucmd_respond(ucmd, sizeof(*cmd));
 
+out_put_kvm:
+	if (idev->kvm)
+		kvm_put_kvm(idev->kvm);
 out_put_hwpt:
 	iommufd_put_object(ucmd->ictx, &hwpt_paging->common.obj);
 out_put_idev:
diff --git a/include/linux/iommu.h b/include/linux/iommu.h
index 555597b54083..300281dae01d 100644
--- a/include/linux/iommu.h
+++ b/include/linux/iommu.h
@@ -46,6 +46,7 @@ struct iommu_dma_msi_cookie;
 struct iommu_fault_param;
 struct iommufd_ctx;
 struct iommufd_viommu;
+struct kvm;
 struct msi_desc;
 struct msi_msg;
 
diff --git a/include/linux/iommufd.h b/include/linux/iommufd.h
index 5cdcb8d2f305..8e0505a9bd0c 100644
--- a/include/linux/iommufd.h
+++ b/include/linux/iommufd.h
@@ -103,6 +103,7 @@ struct iommufd_viommu {
 	struct iommufd_ctx *ictx;
 	struct iommu_device *iommu_dev;
 	struct iommufd_hwpt_paging *hwpt;
+	struct kvm *kvm;
 
 	const struct iommufd_viommu_ops *ops;
 
-- 
2.43.0