RDMA/uverbs: Fix DMABUF lifetime bugs

index 2411ebee69e2..d88d30f1e42f 100644
@@ -68,9 +68,10 @@ static void uverbs_dmabuf_release(struct dma_buf *dmabuf)
 	struct ib_uverbs_dmabuf_file *priv = dmabuf->priv;
 
 	/*
-	 * This can only happen if the fput came from alloc_abort_fd_uobject()
+	 * This can only happen if the fput came from the error unwind of
+	 * UVERBS_METHOD_DMABUF_ALLOC. The abort cleans up the uobject.
 	 */
-	if (!priv->uobj.context)
+	if (!priv)
 		return;
 
 	uverbs_uobject_release(&priv->uobj);

@@ -141,6 +142,8 @@ static int UVERBS_HANDLER(UVERBS_METHOD_DMABUF_ALLOC)(
 	return 0;
 
 err_revoked:
+	/* Tell uverbs_dmabuf_release() that the uobject is being aborted */
+	uverbs_dmabuf->dmabuf->priv = NULL;
 	dma_buf_put(uverbs_dmabuf->dmabuf);
 err:
 	rdma_user_mmap_entry_put(mmap_entry);

index 41c84ffe8c09..2bdb6972fe1a 100644
@@ -244,16 +244,8 @@ void rdma_user_mmap_entry_remove(struct rdma_user_mmap_entry *entry)
 	entry->driver_removed = true;
 	xa_unlock(&entry->ucontext->mmap_xa);
 	list_for_each_entry_safe(uverbs_dmabuf, tmp, &entry->dmabufs, dmabufs_elm) {
-		dma_resv_lock(uverbs_dmabuf->dmabuf->resv, NULL);
 		list_del(&uverbs_dmabuf->dmabufs_elm);
-		uverbs_dmabuf->revoked = true;
-		dma_buf_invalidate_mappings(uverbs_dmabuf->dmabuf);
-		dma_resv_wait_timeout(uverbs_dmabuf->dmabuf->resv,
-				      DMA_RESV_USAGE_BOOKKEEP, false,
-				      MAX_SCHEDULE_TIMEOUT);
-		dma_resv_unlock(uverbs_dmabuf->dmabuf->resv);
-		kref_put(&uverbs_dmabuf->kref, ib_uverbs_dmabuf_done);
-		wait_for_completion(&uverbs_dmabuf->comp);
+		ib_uverbs_dmabuf_revoke(uverbs_dmabuf);
 	}
 	mutex_unlock(&entry->dmabufs_lock);
 

index c64dd6b94e10..728e09b41992 100644
@@ -42,6 +42,8 @@
 #include <linux/mutex.h>
 #include <linux/completion.h>
 #include <linux/cdev.h>
+#include <linux/dma-buf.h>
+#include <linux/dma-resv.h>
 
 #include <rdma/ib_verbs.h>
 #include <rdma/ib_umem.h>

@@ -135,6 +137,13 @@ struct ib_uverbs_completion_event_file {
 
 struct ib_uverbs_dmabuf_file {
 	struct ib_uobject uobj;
+	/*
+	 * Serializes users of dmabuf with the dma_buf's release, which sets it
+	 * to NULL. The uobject can outlive the dma_buf while it waits to be
+	 * destroyed. The map and unmap callbacks use attachment->dmabuf
+	 * instead of taking this lock.
+	 */
+	struct mutex dmabuf_release_lock;
 	struct dma_buf *dmabuf;
 	struct list_head dmabufs_elm;
 	struct rdma_user_mmap_entry *mmap_entry;

@@ -359,6 +368,34 @@ static inline void ib_uverbs_dmabuf_done(struct kref *kref)
 	complete(&priv->comp);
 }
 
+/*
+ * Revoke the dmabuf and wait for all of its DMA mappings to be gone. Must be
+ * called with the mmap entry's dmabufs_lock held, after the dmabuf has been
+ * removed from the entry's dmabufs list, so it runs once per dmabuf. Lock
+ * order is dmabufs_lock, dmabuf_release_lock, then the dma_resv lock.
+ */
+static inline void ib_uverbs_dmabuf_revoke(struct ib_uverbs_dmabuf_file *priv)
+{
+	mutex_lock(&priv->dmabuf_release_lock);
+	if (priv->dmabuf) {
+		dma_resv_lock(priv->dmabuf->resv, NULL);
+		priv->revoked = true;
+		dma_buf_invalidate_mappings(priv->dmabuf);
+		dma_resv_wait_timeout(priv->dmabuf->resv,
+				      DMA_RESV_USAGE_BOOKKEEP, false,
+				      MAX_SCHEDULE_TIMEOUT);
+		dma_resv_unlock(priv->dmabuf->resv);
+	} else {
+		/* A released dma_buf has no attachments left to invalidate */
+		priv->revoked = true;
+	}
+	mutex_unlock(&priv->dmabuf_release_lock);
+
+	kref_put(&priv->kref, ib_uverbs_dmabuf_done);
+	/* Wait till all DMA unmaps are completed */
+	wait_for_completion(&priv->comp);
+}
+
 int __uverbs_cleanup_ufile(struct ib_uverbs_file *ufile,
 			   enum rdma_remove_reason reason);
 

index d88d30f1e42f..8d8c997a4883 100644
@@ -28,7 +28,7 @@ uverbs_dmabuf_map(struct dma_buf_attachment *attachment,
 	struct ib_uverbs_dmabuf_file *priv = attachment->dmabuf->priv;
 	struct sg_table *ret;
 
-	dma_resv_assert_held(priv->dmabuf->resv);
+	dma_resv_assert_held(attachment->dmabuf->resv);
 
 	if (priv->revoked)
 		return ERR_PTR(-ENODEV);

@@ -49,7 +49,7 @@ static void uverbs_dmabuf_unmap(struct dma_buf_attachment *attachment,
 {
 	struct ib_uverbs_dmabuf_file *priv = attachment->dmabuf->priv;
 
-	dma_resv_assert_held(priv->dmabuf->resv);
+	dma_resv_assert_held(attachment->dmabuf->resv);
 	dma_buf_free_sgt(attachment, sgt, dir);
 	kref_put(&priv->kref, ib_uverbs_dmabuf_done);
 }

@@ -74,7 +74,25 @@ static void uverbs_dmabuf_release(struct dma_buf *dmabuf)
 	if (!priv)
 		return;
 
+	/*
+	 * uverbs_uobject_release() may drop the last reference to the uobject,
+	 * keep it around until the dmabuf pointer below has been cleared.
+	 */
+	uverbs_uobject_get(&priv->uobj);
 	uverbs_uobject_release(&priv->uobj);
+
+	/*
+	 * The dma_buf is freed once this function returns. If the uobject was
+	 * not destroyed above, because the ufile is concurrently being cleaned
+	 * up, it is still reachable through the ufile's uobjects list and the
+	 * mmap entry's dmabufs list. Make sure those paths are not in the
+	 * middle of using the dma_buf and cannot use it afterwards.
+	 */
+	mutex_lock(&priv->dmabuf_release_lock);
+	priv->dmabuf = NULL;
+	mutex_unlock(&priv->dmabuf_release_lock);
+
+	uverbs_uobject_put(&priv->uobj);
 }
 
 static const struct dma_buf_ops uverbs_dmabuf_ops = {

@@ -100,6 +118,8 @@ static int UVERBS_HANDLER(UVERBS_METHOD_DMABUF_ALLOC)(
 	off_t pg_off;
 	int ret;
 
+	mutex_init(&uverbs_dmabuf->dmabuf_release_lock);
+
 	ret = uverbs_get_const(&pg_off, attrs, UVERBS_ATTR_ALLOC_DMABUF_PGOFF);
 	if (ret)
 		return ret;

@@ -165,24 +185,11 @@ static void uverbs_dmabuf_fd_destroy_uobj(struct ib_uobject *uobj,
 {
 	struct ib_uverbs_dmabuf_file *uverbs_dmabuf =
 		container_of(uobj, struct ib_uverbs_dmabuf_file, uobj);
-	bool wait_for_comp = false;
 
 	mutex_lock(&uverbs_dmabuf->mmap_entry->dmabufs_lock);
-	dma_resv_lock(uverbs_dmabuf->dmabuf->resv, NULL);
 	if (!uverbs_dmabuf->revoked) {
-		uverbs_dmabuf->revoked = true;
 		list_del(&uverbs_dmabuf->dmabufs_elm);
-		dma_buf_invalidate_mappings(uverbs_dmabuf->dmabuf);
-		dma_resv_wait_timeout(uverbs_dmabuf->dmabuf->resv,
-				      DMA_RESV_USAGE_BOOKKEEP, false,
-				      MAX_SCHEDULE_TIMEOUT);
-		wait_for_comp = true;
-	}
-	dma_resv_unlock(uverbs_dmabuf->dmabuf->resv);
-	if (wait_for_comp) {
-		kref_put(&uverbs_dmabuf->kref, ib_uverbs_dmabuf_done);
-		/* Let's wait till all DMA unmap are completed. */
-		wait_for_completion(&uverbs_dmabuf->comp);
+		ib_uverbs_dmabuf_revoke(uverbs_dmabuf);
 	}
 	mutex_unlock(&uverbs_dmabuf->mmap_entry->dmabufs_lock);