[PATCH RFCv1 08/14] iommufd: Add IOMMU_VIOMMU_SET_DEV_ID ioctl

[Nicolin Chen <nicolinc@nvidia.com>]

Introduce a new ioctl to set a per-viommu device virtual id that should be
linked to the physical device id (or just a struct device pointer).

Since a viommu (user space IOMMU instance) can have multiple devices while
it's not ideal to confine a device to one single user space IOMMU instance
either, these two shouldn't just do a 1:1 mapping. Add two xarrays in their
structures to bind them bidirectionally.

Signed-off-by: Nicolin Chen <nicolinc@nvidia.com>
---
 drivers/iommu/iommufd/device.c          | 11 +++++
 drivers/iommu/iommufd/iommufd_private.h | 10 +++++
 drivers/iommu/iommufd/main.c            |  2 +
 drivers/iommu/iommufd/viommu.c          | 58 +++++++++++++++++++++++++
 include/linux/iommufd.h                 |  2 +
 include/uapi/linux/iommufd.h            | 20 +++++++++
 6 files changed, 103 insertions(+)

diff --git a/drivers/iommu/iommufd/device.c b/drivers/iommu/iommufd/device.c
index 873630c111c1..68086f3074b6 100644
--- a/drivers/iommu/iommufd/device.c
+++ b/drivers/iommu/iommufd/device.c
@@ -135,7 +135,16 @@ void iommufd_device_destroy(struct iommufd_object *obj)
 {
 	struct iommufd_device *idev =
 		container_of(obj, struct iommufd_device, obj);
+	struct iommufd_viommu *viommu;
+	unsigned long index;
 
+	xa_for_each(&idev->viommus, index, viommu) {
+		if (viommu->ops->unset_dev_id)
+			viommu->ops->unset_dev_id(viommu, idev->dev);
+		xa_erase(&viommu->idevs, idev->obj.id);
+		xa_erase(&idev->viommus, index);
+	}
+	xa_destroy(&idev->viommus);
 	iommu_device_release_dma_owner(idev->dev);
 	iommufd_put_group(idev->igroup);
 	if (!iommufd_selftest_is_mock_dev(idev->dev))
@@ -216,6 +225,8 @@ struct iommufd_device *iommufd_device_bind(struct iommufd_ctx *ictx,
 	/* igroup refcount moves into iommufd_device */
 	idev->igroup = igroup;
 
+	xa_init_flags(&idev->viommus, XA_FLAGS_ALLOC1);
+
 	/*
 	 * If the caller fails after this success it must call
 	 * iommufd_unbind_device() which is safe since we hold this refcount.
diff --git a/drivers/iommu/iommufd/iommufd_private.h b/drivers/iommu/iommufd/iommufd_private.h
index ae90b4493109..9ba8f4ecc221 100644
--- a/drivers/iommu/iommufd/iommufd_private.h
+++ b/drivers/iommu/iommufd/iommufd_private.h
@@ -392,6 +392,7 @@ struct iommufd_device {
 	struct list_head group_item;
 	/* always the physical device */
 	struct device *dev;
+	struct xarray viommus;
 	bool enforce_cache_coherency;
 };
 
@@ -424,8 +425,17 @@ void iopt_remove_access(struct io_pagetable *iopt,
 			u32 iopt_access_list_id);
 void iommufd_access_destroy_object(struct iommufd_object *obj);
 
+static inline struct iommufd_viommu *
+iommufd_get_viommu(struct iommufd_ucmd *ucmd, u32 id)
+{
+	return container_of(iommufd_get_object(ucmd->ictx, id,
+					       IOMMUFD_OBJ_VIOMMU),
+			    struct iommufd_viommu, obj);
+}
+
 int iommufd_viommu_alloc_ioctl(struct iommufd_ucmd *ucmd);
 void iommufd_viommu_destroy(struct iommufd_object *obj);
+int iommufd_viommu_set_device_id(struct iommufd_ucmd *ucmd);
 
 #ifdef CONFIG_IOMMUFD_TEST
 int iommufd_test(struct iommufd_ucmd *ucmd);
diff --git a/drivers/iommu/iommufd/main.c b/drivers/iommu/iommufd/main.c
index 9de7e3e63ce4..16efc3346a2a 100644
--- a/drivers/iommu/iommufd/main.c
+++ b/drivers/iommu/iommufd/main.c
@@ -381,6 +381,8 @@ static const struct iommufd_ioctl_op iommufd_ioctl_ops[] = {
 		 __reserved),
 	IOCTL_OP(IOMMU_VIOMMU_ALLOC, iommufd_viommu_alloc_ioctl,
 		 struct iommu_viommu_alloc, out_viommu_id),
+	IOCTL_OP(IOMMU_VIOMMU_SET_DEV_ID, iommufd_viommu_set_device_id,
+		 struct iommu_viommu_set_dev_id, id),
 #ifdef CONFIG_IOMMUFD_TEST
 	IOCTL_OP(IOMMU_TEST_CMD, iommufd_test, struct iommu_test_cmd, last),
 #endif
diff --git a/drivers/iommu/iommufd/viommu.c b/drivers/iommu/iommufd/viommu.c
index 079e0ff79942..71baca0c75de 100644
--- a/drivers/iommu/iommufd/viommu.c
+++ b/drivers/iommu/iommufd/viommu.c
@@ -69,6 +69,7 @@ int iommufd_viommu_alloc_ioctl(struct iommufd_ucmd *ucmd)
 		rc = PTR_ERR(viommu);
 		goto out_put_hwpt;
 	}
+	xa_init_flags(&viommu->idevs, XA_FLAGS_ALLOC1);
 
 	/* iommufd_object_finalize will store the viommu->obj.id */
 	rc = xa_alloc(&ucmd->ictx->objects, &viommu->obj.id, XA_ZERO_ENTRY,
@@ -102,3 +103,60 @@ int iommufd_viommu_alloc_ioctl(struct iommufd_ucmd *ucmd)
 	iommufd_put_object(ucmd->ictx, &idev->obj);
 	return rc;
 }
+
+int iommufd_viommu_set_device_id(struct iommufd_ucmd *ucmd)
+{
+	struct iommu_viommu_set_dev_id *cmd = ucmd->cmd;
+	unsigned int dev_id, viommu_id;
+	struct iommufd_viommu *viommu;
+	struct iommufd_device *idev;
+	int rc;
+
+	idev = iommufd_get_device(ucmd, cmd->dev_id);
+	if (IS_ERR(idev))
+		return PTR_ERR(idev);
+	dev_id = idev->obj.id;
+
+	viommu = iommufd_get_viommu(ucmd, cmd->viommu_id);
+	if (IS_ERR(viommu)) {
+		rc = PTR_ERR(viommu);
+		goto out_put_idev;
+	}
+
+	if (viommu->iommu_dev != idev->dev->iommu->iommu_dev) {
+		rc = -EINVAL;
+		goto out_put_viommu;
+	}
+
+	if (!viommu->ops->set_dev_id || !viommu->ops->unset_dev_id) {
+		rc = -EOPNOTSUPP;
+		goto out_put_viommu;
+	}
+
+	rc = xa_alloc(&idev->viommus, &viommu_id, viommu,
+		      XA_LIMIT(viommu->obj.id, viommu->obj.id),
+		      GFP_KERNEL_ACCOUNT);
+	if (rc)
+		goto out_put_viommu;
+
+	rc = xa_alloc(&viommu->idevs, &dev_id, idev,
+		      XA_LIMIT(dev_id, dev_id), GFP_KERNEL_ACCOUNT);
+	if (rc)
+		goto out_xa_erase_viommu;
+
+	rc = viommu->ops->set_dev_id(viommu, idev->dev, cmd->id);
+	if (rc)
+		goto out_xa_erase_idev;
+
+	goto out_put_viommu;
+
+out_xa_erase_idev:
+	xa_erase(&viommu->idevs, idev->obj.id);
+out_xa_erase_viommu:
+	xa_erase(&idev->viommus, viommu->obj.id);
+out_put_viommu:
+	iommufd_put_object(ucmd->ictx, &viommu->obj);
+out_put_idev:
+	iommufd_put_object(ucmd->ictx, &idev->obj);
+	return rc;
+}
diff --git a/include/linux/iommufd.h b/include/linux/iommufd.h
index ca6ac8a1ffd0..2be302b82f47 100644
--- a/include/linux/iommufd.h
+++ b/include/linux/iommufd.h
@@ -10,6 +10,7 @@
 #include <linux/errno.h>
 #include <linux/err.h>
 #include <linux/refcount.h>
+#include <linux/xarray.h>
 
 struct device;
 struct iommufd_device;
@@ -88,6 +89,7 @@ struct iommufd_viommu {
 	const struct iommufd_viommu_ops *ops;
 
 	unsigned int type;
+	struct xarray idevs;
 };
 
 /**
diff --git a/include/uapi/linux/iommufd.h b/include/uapi/linux/iommufd.h
index 2b0825d69846..eaa192de63d3 100644
--- a/include/uapi/linux/iommufd.h
+++ b/include/uapi/linux/iommufd.h
@@ -51,6 +51,7 @@ enum {
 	IOMMUFD_CMD_HWPT_GET_DIRTY_BITMAP,
 	IOMMUFD_CMD_HWPT_INVALIDATE,
 	IOMMUFD_CMD_VIOMMU_ALLOC,
+	IOMMUFD_CMD_VIOMMU_SET_DEV_ID,
 };
 
 /**
@@ -722,4 +723,23 @@ struct iommu_viommu_alloc {
 	__u32 out_viommu_id;
 };
 #define IOMMU_VIOMMU_ALLOC _IO(IOMMUFD_TYPE, IOMMUFD_CMD_VIOMMU_ALLOC)
+
+/**
+ * struct iommu_viommu_set_dev_id - ioctl(IOMMU_VIOMMU_SET_DEV_ID)
+ * @size: sizeof(struct iommu_viommu_set_dev_id)
+ * @viommu_id: viommu ID to associate with the device to store its virtual ID
+ * @dev_id: device ID to set a device virtual ID
+ * @__reserved: Must be 0
+ * @id: Device virtual ID
+ *
+ * Set a viommu-specific virtual ID of a device
+ */
+struct iommu_viommu_set_dev_id {
+	__u32 size;
+	__u32 viommu_id;
+	__u32 dev_id;
+	__u32 __reserved;
+	__aligned_u64 id;
+};
+#define IOMMU_VIOMMU_SET_DEV_ID _IO(IOMMUFD_TYPE, IOMMUFD_CMD_VIOMMU_SET_DEV_ID)
 #endif
-- 
2.43.0