[RFC PATCH 1/1] iommu: Fix one concurrency issue happens on non-iommu binding device

[Zhenhua Huang]

Thread (1)
During SMMU probe:
iommu_device_register
        bus_iommu_probe
                probe_iommu_group for all devices in bus
                        _probe_iommu_device
(non-iommu binding device)	first allocate dev->iommu then free

Thread (2)
client device probing:
dma_configure
        of_iommu_configure
                get dev->iommu->fwspec and call fwspec->ops

There may be a time window that dev->iommu allocated even for non iommu
binding device in thread (1). It would be possible thread (2) meet Use-
after-free errors at that window.

Fix it by closing the time window, set dev->iommu only once
ops->probe_device successfully.
Previous discussion details:
https://lore.kernel.org/linux-arm-kernel/20231017163337.GE282036@ziepe.ca/T/#mee0d7bdc375541934a571ae69f43b9660f8e7312

Suggested-by: Jason Gunthorpe <jgg at ziepe.ca>
Signed-off-by: Zhenhua Huang <quic_zhenhuah at quicinc.com>
---
Hi Jason, Robin,
Shall we first address non-iommu binding device crash with below patch? It's
seen a lot especially when smmu probing and other driver probing at same time.
Although it doesn't comprehensively fixes race with of_xlate VS probe, it
indeed solves issue for non-iommu binding device.

 drivers/iommu/iommu.c | 42 +++++++++++++++++++++++++++++-------------
 include/linux/iommu.h |  5 +----
 2 files changed, 30 insertions(+), 17 deletions(-)

diff --git a/drivers/iommu/iommu.c b/drivers/iommu/iommu.c
index f17a111..88ce802 100644
--- a/drivers/iommu/iommu.c
+++ b/drivers/iommu/iommu.c
@@ -386,6 +386,17 @@ static u32 dev_iommu_get_max_pasids(struct device *dev)
 	return min_t(u32, max_pasids, dev->iommu->iommu_dev->max_pasids);
 }
 
+void dev_iommu_priv_set(struct device *dev, void *priv)
+{
+	struct dev_iommu *dev_iommu;
+
+	dev_iommu = dev_iommu_get(dev);
+	if (WARN_ON(!dev_iommu))
+		return;
+
+	dev->iommu->priv = priv;
+}
+
 /*
  * Init the dev->iommu and dev->iommu_group in the struct device and get the
  * driver probed
@@ -393,23 +404,26 @@ static u32 dev_iommu_get_max_pasids(struct device *dev)
 static int iommu_init_device(struct device *dev, const struct iommu_ops *ops)
 {
 	struct iommu_device *iommu_dev;
+	struct dev_iommu *dev_iommu;
 	struct iommu_group *group;
 	int ret;
 
-	if (!dev_iommu_get(dev))
-		return -ENOMEM;
-
-	if (!try_module_get(ops->owner)) {
-		ret = -EINVAL;
-		goto err_free;
-	}
+	if (!try_module_get(ops->owner))
+		return -EINVAL;
 
 	iommu_dev = ops->probe_device(dev);
 	if (IS_ERR(iommu_dev)) {
 		ret = PTR_ERR(iommu_dev);
 		goto err_module_put;
 	}
-	dev->iommu->iommu_dev = iommu_dev;
+
+	dev_iommu = dev_iommu_get(dev);
+	if (WARN_ON(!dev_iommu)) {
+		ret = -ENOMEM;
+		goto err_release;
+	}
+
+	dev_iommu->iommu_dev = iommu_dev;
 
 	ret = iommu_device_link(iommu_dev, dev);
 	if (ret)
@@ -424,9 +438,9 @@ static int iommu_init_device(struct device *dev, const struct iommu_ops *ops)
 	}
 	dev->iommu_group = group;
 
-	dev->iommu->max_pasids = dev_iommu_get_max_pasids(dev);
+	dev_iommu->max_pasids = dev_iommu_get_max_pasids(dev);
 	if (ops->is_attach_deferred)
-		dev->iommu->attach_deferred = ops->is_attach_deferred(dev);
+		dev_iommu->attach_deferred = ops->is_attach_deferred(dev);
 	return 0;
 
 err_unlink:
@@ -436,9 +450,11 @@ static int iommu_init_device(struct device *dev, const struct iommu_ops *ops)
 		ops->release_device(dev);
 err_module_put:
 	module_put(ops->owner);
-err_free:
-	dev->iommu->iommu_dev = NULL;
-	dev_iommu_free(dev);
+	/*
+	 * If probe_device allocated a dev->iommu and things failed later
+	 * we just leave it. We don't yet have robust locking, there
+	 * could be concurrent users.
+	 */
 	return ret;
 }
 
diff --git a/include/linux/iommu.h b/include/linux/iommu.h
index 73e58df..62e9c86 100644
--- a/include/linux/iommu.h
+++ b/include/linux/iommu.h
@@ -840,10 +840,7 @@ static inline void *dev_iommu_priv_get(struct device *dev)
 		return NULL;
 }
 
-static inline void dev_iommu_priv_set(struct device *dev, void *priv)
-{
-	dev->iommu->priv = priv;
-}
+void dev_iommu_priv_set(struct device *dev, void *priv);
 
 int iommu_probe_device(struct device *dev);
 
-- 
2.7.4