[PATCH RFC 09/11] iommu: Add iommu_get_domain_for_dev_pasid()

[Lu Baolu]

This adds a helper to retrieve the iommu_domain which has been attached to
a {device, pasid}. One usage scenario of this helper exists in the I/O page
fault handling framework, where {device, pasid} was reported by hardware,
and software needs to retrieve the attached domain for further routing.

Signed-off-by: Lu Baolu <baolu.lu@xxxxxxxxxxxxxxx>
---
 include/linux/iommu.h |  8 ++++++++
 drivers/iommu/iommu.c | 31 +++++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/include/linux/iommu.h b/include/linux/iommu.h
index 47c9aa5aa9c8..803e7b07605e 100644
--- a/include/linux/iommu.h
+++ b/include/linux/iommu.h
@@ -676,6 +676,8 @@ int iommu_attach_device_pasid(struct iommu_domain *domain,
 			      struct device *dev, ioasid_t pasid);
 void iommu_detach_device_pasid(struct iommu_domain *domain,
 			       struct device *dev, ioasid_t pasid);
+struct iommu_domain *
+iommu_get_domain_for_dev_pasid(struct device *dev, ioasid_t pasid);
 
 #else /* CONFIG_IOMMU_API */
 
@@ -1041,6 +1043,12 @@ static inline void iommu_detach_device_pasid(struct iommu_domain *domain,
 					     struct device *dev, ioasid_t pasid)
 {
 }
+
+static inline struct iommu_domain *
+iommu_get_domain_for_dev_pasid(struct device *dev, ioasid_t pasid)
+{
+	return NULL;
+}
 #endif /* CONFIG_IOMMU_API */
 
 #ifdef CONFIG_IOMMU_SVA
diff --git a/drivers/iommu/iommu.c b/drivers/iommu/iommu.c
index 4f90b71c6f6e..508fdcabda5c 100644
--- a/drivers/iommu/iommu.c
+++ b/drivers/iommu/iommu.c
@@ -40,6 +40,7 @@ struct iommu_group {
 	struct kobject kobj;
 	struct kobject *devices_kobj;
 	struct list_head devices;
+	struct xarray pasid_array;
 	struct mutex mutex;
 	void *iommu_data;
 	void (*iommu_data_release)(void *iommu_data);
@@ -632,6 +633,7 @@ struct iommu_group *iommu_group_alloc(void)
 	mutex_init(&group->mutex);
 	INIT_LIST_HEAD(&group->devices);
 	INIT_LIST_HEAD(&group->entry);
+	xa_init(&group->pasid_array);
 
 	ret = ida_simple_get(&iommu_group_ida, 0, 0, GFP_KERNEL);
 	if (ret < 0) {
@@ -3086,6 +3088,7 @@ int iommu_attach_device_pasid(struct iommu_domain *domain,
 {
 	struct iommu_group *group;
 	int ret = -EINVAL;
+	void *curr;
 
 	if (!domain->ops->attach_dev_pasid)
 		return -EINVAL;
@@ -3098,7 +3101,16 @@ int iommu_attach_device_pasid(struct iommu_domain *domain,
 	if (iommu_group_device_count(group) != 1)
 		goto out_unlock;
 
+	xa_lock(&group->pasid_array);
+	curr = __xa_cmpxchg(&group->pasid_array, pasid, NULL,
+			    domain, GFP_KERNEL);
+	xa_unlock(&group->pasid_array);
+	if (curr)
+		goto out_unlock;
+
 	ret = domain->ops->attach_dev_pasid(domain, dev, pasid);
+	if (ret)
+		xa_erase(&group->pasid_array, pasid);
 
 out_unlock:
 	mutex_unlock(&group->mutex);
@@ -3118,6 +3130,25 @@ void iommu_detach_device_pasid(struct iommu_domain *domain,
 
 	mutex_lock(&group->mutex);
 	domain->ops->detach_dev_pasid(domain, dev, pasid);
+	xa_erase(&group->pasid_array, pasid);
+	mutex_unlock(&group->mutex);
+	iommu_group_put(group);
+}
+
+struct iommu_domain *
+iommu_get_domain_for_dev_pasid(struct device *dev, ioasid_t pasid)
+{
+	struct iommu_domain *domain;
+	struct iommu_group *group;
+
+	group = iommu_group_get(dev);
+	if (!group)
+		return NULL;
+
+	mutex_lock(&group->mutex);
+	domain = xa_load(&group->pasid_array, pasid);
 	mutex_unlock(&group->mutex);
 	iommu_group_put(group);
+
+	return domain;
 }
-- 
2.25.1