[Intel-xe] [PATCH] drm/xe: Fix vm refcount races

[Thomas Hellström]

Fix a race in xe_vm_lookup() where the vm could disappear after
the lookup mutex unlock but before the get. The xe_vm_get() call
must be inside the lookup mutex.

Also fix a vm close race where multiple callers could potentially
succeed in calling xe_vm_close_and_put().

Reported-by: Oded Gabbay <ogabbay at kernel.org>
Link: https://lists.freedesktop.org/archives/intel-xe/2023-May/004704.html
Signed-off-by: Thomas Hellström <thomas.hellstrom at linux.intel.com>
---
 drivers/gpu/drm/xe/xe_vm.c | 26 ++++++++++++--------------
 1 file changed, 12 insertions(+), 14 deletions(-)

diff --git a/drivers/gpu/drm/xe/xe_vm.c b/drivers/gpu/drm/xe/xe_vm.c
index a0306526b269..ac25afac89fd 100644
--- a/drivers/gpu/drm/xe/xe_vm.c
+++ b/drivers/gpu/drm/xe/xe_vm.c
@@ -1463,10 +1463,9 @@ struct xe_vm *xe_vm_lookup(struct xe_file *xef, u32 id)
 
 	mutex_lock(&xef->vm.lock);
 	vm = xa_load(&xef->vm.xa, id);
-	mutex_unlock(&xef->vm.lock);
-
 	if (vm)
 		xe_vm_get(vm);
+	mutex_unlock(&xef->vm.lock);
 
 	return vm;
 }
@@ -1940,26 +1939,25 @@ int xe_vm_destroy_ioctl(struct drm_device *dev, void *data,
 	struct xe_file *xef = to_xe_file(file);
 	struct drm_xe_vm_destroy *args = data;
 	struct xe_vm *vm;
+	int err = 0;
 
 	if (XE_IOCTL_ERR(xe, args->pad))
 		return -EINVAL;
 
-	vm = xe_vm_lookup(xef, args->vm_id);
-	if (XE_IOCTL_ERR(xe, !vm))
-		return -ENOENT;
-	xe_vm_put(vm);
-
-	/* FIXME: Extend this check to non-compute mode VMs */
-	if (XE_IOCTL_ERR(xe, vm->preempt.num_engines))
-		return -EBUSY;
-
 	mutex_lock(&xef->vm.lock);
-	xa_erase(&xef->vm.xa, args->vm_id);
+	vm = xa_load(&xef->vm.xa, args->vm_id);
+	if (XE_IOCTL_ERR(xe, !vm))
+		err = -ENOENT;
+	else if (XE_IOCTL_ERR(xe, vm->preempt.num_engines))
+		err = -EBUSY;
+	else
+		xa_erase(&xef->vm.xa, args->vm_id);
 	mutex_unlock(&xef->vm.lock);
 
-	xe_vm_close_and_put(vm);
+	if (!err)
+		xe_vm_close_and_put(vm);
 
-	return 0;
+	return err;
 }
 
 static const u32 region_to_mem_type[] = {
-- 
2.39.2