[PATCH v2 2/3] iommufd/viommu: Publish a vDEVICE only after vdevice_init() succeeds

[Nicolin Chen]

iommufd_vdevice_alloc_ioctl() adds the vDEVICE to the viommu->vdevs xarray
with xa_cmpxchg() before the driver's vdevice_init() op runs. That op is
where a driver validates the device and may reject it, but the xarray entry
is already live by then: a concurrent IOMMU_HWPT_INVALIDATE can look it up
with iommufd_viommu_find_dev() and run the driver invalidation path against
a device that vdevice_init() would have refused.

Reserve the index with xa_insert(): it stores a zero entry that reads back
as NULL, and returns -EBUSY on a duplicate virt_id. Run vdevice_init() and
store the vDEVICE pointer only once it succeeds. A failed vdevice_init()
releases the reservation, so lookups observe the vDEVICE only after it is
fully initialized and accepted.

Fixes: ed42eee797ff3 ("iommufd/viommu: Add driver-defined vDEVICE support")
Cc: stable at vger.kernel.org
Reviewed-by: Kevin Tian <kevin.tian at intel.com>
Assisted-by: Claude:claude-opus-4-8
Signed-off-by: Nicolin Chen <nicolinc at nvidia.com>
---
 drivers/iommu/iommufd/viommu.c | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/drivers/iommu/iommufd/viommu.c b/drivers/iommu/iommufd/viommu.c
index 0c12c7e352a14..fc13cf4737ea3 100644
--- a/drivers/iommu/iommufd/viommu.c
+++ b/drivers/iommu/iommufd/viommu.c
@@ -143,7 +143,7 @@ void iommufd_vdevice_destroy(struct iommufd_object *obj)
 int iommufd_vdevice_alloc_ioctl(struct iommufd_ucmd *ucmd)
 {
 	struct iommu_vdevice_alloc *cmd = ucmd->cmd;
-	struct iommufd_vdevice *vdev, *curr;
+	struct iommufd_vdevice *vdev;
 	size_t vdev_size = sizeof(*vdev);
 	struct iommufd_viommu *viommu;
 	struct iommufd_device *idev;
@@ -218,18 +218,28 @@ int iommufd_vdevice_alloc_ioctl(struct iommufd_ucmd *ucmd)
 	 */
 	idev->vdev = vdev;
 
-	curr = xa_cmpxchg(&viommu->vdevs, virt_id, NULL, vdev, GFP_KERNEL);
-	if (curr) {
-		rc = xa_err(curr) ?: -EEXIST;
+	/*
+	 * Reserve the slot with a zero entry (reads back as NULL) until the
+	 * vdevice_init() op accepts the vDEVICE. Only the xa_* helpers hide a
+	 * reserved entry, so never use a raw xas_* iterator on this xarray.
+	 */
+	rc = xa_insert(&viommu->vdevs, virt_id, NULL, GFP_KERNEL);
+	if (rc) {
+		if (rc == -EBUSY)
+			rc = -EEXIST;
 		goto out_abort;
 	}
 
 	if (viommu->ops && viommu->ops->vdevice_init) {
 		rc = viommu->ops->vdevice_init(vdev);
-		if (rc)
+		if (rc) {
+			xa_release(&viommu->vdevs, virt_id);
 			goto out_abort;
+		}
 	}
 
+	xa_store(&viommu->vdevs, virt_id, vdev, GFP_KERNEL);
+
 	cmd->out_vdevice_id = vdev->obj.id;
 	rc = iommufd_ucmd_respond(ucmd, sizeof(*cmd));
 	if (rc)
-- 
2.43.0