[PATCH V3 3/3] accel/amdxdna: Fix iommu domain lifetime race during device removal

[Lizhi Hou]

When force_iova mode is enabled, amdxdna_remove() frees xdna->domain. If
amdxdna_gem_obj_free() is called after device removal, it may attempt to
access xdna->domain, resulting in a use-after-free.

Fix the race by adding freeing xdna->domain as a managed release action,
so its lifetime is managed by DRM and remains valid until all managed
resources are released.

Fixes: ece3e8980907 ("accel/amdxdna: Allow forcing IOVA-based DMA via module parameter")
Signed-off-by: Lizhi Hou <lizhi.hou@xxxxxxx>
---
 drivers/accel/amdxdna/amdxdna_iommu.c | 43 ++++++++++++++++++---------
 1 file changed, 29 insertions(+), 14 deletions(-)

diff --git a/drivers/accel/amdxdna/amdxdna_iommu.c b/drivers/accel/amdxdna/amdxdna_iommu.c
index eff00131d0f8..4f245b969eef 100644
--- a/drivers/accel/amdxdna/amdxdna_iommu.c
+++ b/drivers/accel/amdxdna/amdxdna_iommu.c
@@ -4,6 +4,7 @@
  */
 
 #include <drm/amdxdna_accel.h>
+#include <drm/drm_managed.h>
 #include <linux/iommu.h>
 #include <linux/iova.h>
 
@@ -153,10 +154,30 @@ void amdxdna_iommu_free(struct amdxdna_dev *xdna, size_t size,
 	free_pages((unsigned long)cpu_addr, get_order(size));
 }
 
+static void amdxdna_cleanup_force_iova(struct drm_device *dev, void *res)
+{
+	struct amdxdna_dev *xdna = to_xdna_dev(dev);
+
+	if (xdna->domain) {
+		iommu_detach_group(xdna->domain, xdna->group);
+		put_iova_domain(&xdna->iovad);
+		iova_cache_put();
+		iommu_domain_free(xdna->domain);
+	}
+
+	iommu_group_put(xdna->group);
+}
+
+void amdxdna_iommu_fini(struct amdxdna_dev *xdna)
+{
+	if (xdna->group && !xdna->domain)
+		iommu_group_put(xdna->group);
+}
+
 int amdxdna_iommu_init(struct amdxdna_dev *xdna)
 {
 	unsigned long order;
-	int ret;
+	int ret = 0;
 
 	xdna->group = iommu_group_get(xdna->ddev.dev);
 	if (!xdna->group || !force_iova)
@@ -182,8 +203,14 @@ int amdxdna_iommu_init(struct amdxdna_dev *xdna)
 	if (ret)
 		goto put_iova;
 
+	ret = drmm_add_action(&xdna->ddev, amdxdna_cleanup_force_iova, NULL);
+	if (ret)
+		goto detach_group;
+
 	return 0;
 
+detach_group:
+	iommu_detach_group(xdna->domain, xdna->group);
 put_iova:
 	put_iova_domain(&xdna->iovad);
 	iova_cache_put();
@@ -191,20 +218,8 @@ int amdxdna_iommu_init(struct amdxdna_dev *xdna)
 	iommu_domain_free(xdna->domain);
 put_group:
 	iommu_group_put(xdna->group);
+	xdna->group = NULL;
 	xdna->domain = NULL;
 
 	return ret;
 }
-
-void amdxdna_iommu_fini(struct amdxdna_dev *xdna)
-{
-	if (xdna->domain) {
-		iommu_detach_group(xdna->domain, xdna->group);
-		put_iova_domain(&xdna->iovad);
-		iova_cache_put();
-		iommu_domain_free(xdna->domain);
-	}
-
-	if (xdna->group)
-		iommu_group_put(xdna->group);
-}
-- 
2.34.1