[PATCH 2/2] vfio/type1: Update group->domain after aux attach and detach

[Lu Baolu <baolu.lu@xxxxxxxxxxxxxxx>]

Update group->domain whenever an aux-domain is attached to or detached
from a mediated device. Without this change, iommu_get_domain_for_dev()
will be broken for mdev devices.

Fixes: 7bd50f0cd2fd5 ("vfio/type1: Add domain at(de)taching group helpers")
Signed-off-by: Lu Baolu <baolu.lu@xxxxxxxxxxxxxxx>
---
 drivers/vfio/vfio_iommu_type1.c | 37 ++++++++++++++++++++++++++++-----
 1 file changed, 32 insertions(+), 5 deletions(-)

diff --git a/drivers/vfio/vfio_iommu_type1.c b/drivers/vfio/vfio_iommu_type1.c
index 5e556ac9102a..e0d8802ce0c9 100644
--- a/drivers/vfio/vfio_iommu_type1.c
+++ b/drivers/vfio/vfio_iommu_type1.c
@@ -1634,10 +1634,28 @@ static int vfio_mdev_attach_domain(struct device *dev, void *data)
 
 	iommu_device = vfio_mdev_get_iommu_device(dev);
 	if (iommu_device) {
-		if (iommu_dev_feature_enabled(iommu_device, IOMMU_DEV_FEAT_AUX))
-			return iommu_aux_attach_device(domain, iommu_device);
-		else
+		if (iommu_dev_feature_enabled(iommu_device,
+					      IOMMU_DEV_FEAT_AUX)) {
+			struct iommu_group *group = iommu_group_get(dev);
+			int ret;
+
+			if (!group)
+				return -EINVAL;
+
+			if (iommu_group_get_domain(group)) {
+				iommu_group_put(group);
+				return -EBUSY;
+			}
+
+			ret = iommu_aux_attach_device(domain, iommu_device);
+			if (!ret)
+				iommu_group_set_domain(group, domain);
+
+			iommu_group_put(group);
+			return ret;
+		} else {
 			return iommu_attach_device(domain, iommu_device);
+		}
 	}
 
 	return -EINVAL;
@@ -1650,10 +1668,19 @@ static int vfio_mdev_detach_domain(struct device *dev, void *data)
 
 	iommu_device = vfio_mdev_get_iommu_device(dev);
 	if (iommu_device) {
-		if (iommu_dev_feature_enabled(iommu_device, IOMMU_DEV_FEAT_AUX))
+		if (iommu_dev_feature_enabled(iommu_device,
+					      IOMMU_DEV_FEAT_AUX)) {
+			struct iommu_group *group;
+
 			iommu_aux_detach_device(domain, iommu_device);
-		else
+			group = iommu_group_get(dev);
+			if (group) {
+				iommu_group_set_domain(group, NULL);
+				iommu_group_put(group);
+			}
+		} else {
 			iommu_detach_device(domain, iommu_device);
+		}
 	}
 
 	return 0;
-- 
2.17.1