[PATCH] accel/habanalabs: publish mmap buffers after allocation

[Ruoyu Wang]

hl_mmap_mem_buf_alloc() inserts a new buffer into the memory manager
IDR before initializing the descriptor and before running the behavior
allocation callback.

Handles from this IDR are used by hl_mmap_mem_buf_get() and the mmap
path. A racing lookup can therefore take the buffer kref or inspect
behavior-private state before the buffer is ready.

Initialize the descriptor, run the behavior allocation callback, and
only then publish the object in the IDR. If IDR allocation fails,
release the behavior-private resources before freeing the descriptor.

Signed-off-by: Ruoyu Wang <ruoyuw560@xxxxxxxxx>
---
 drivers/accel/habanalabs/common/memory_mgr.c | 30 ++++++++++----------
 1 file changed, 15 insertions(+), 15 deletions(-)

diff --git a/drivers/accel/habanalabs/common/memory_mgr.c b/drivers/accel/habanalabs/common/memory_mgr.c
index 9fdd34acf389f..b5b92dbbe5ddd 100644
--- a/drivers/accel/habanalabs/common/memory_mgr.c
+++ b/drivers/accel/habanalabs/common/memory_mgr.c
@@ -156,6 +156,17 @@ hl_mmap_mem_buf_alloc(struct hl_mem_mgr *mmg,
 	if (!buf)
 		return NULL;
 
+	buf->mmg = mmg;
+	buf->behavior = behavior;
+	kref_init(&buf->refcount);
+
+	rc = buf->behavior->alloc(buf, gfp, args);
+	if (rc) {
+		dev_err(mmg->dev, "%s: Failure in buffer alloc callback %d\n",
+			behavior->topic, rc);
+		goto free_buf;
+	}
+
 	spin_lock(&mmg->lock);
 	rc = idr_alloc(&mmg->handles, buf, 1, 0, GFP_ATOMIC);
 	spin_unlock(&mmg->lock);
@@ -163,27 +174,16 @@ hl_mmap_mem_buf_alloc(struct hl_mem_mgr *mmg,
 		dev_err(mmg->dev,
 			"%s: Failed to allocate IDR for a new buffer, rc=%d\n",
 			behavior->topic, rc);
-		goto free_buf;
+		goto destroy_buf;
 	}
 
-	buf->mmg = mmg;
-	buf->behavior = behavior;
 	buf->handle = (((u64)rc | buf->behavior->mem_id) << PAGE_SHIFT);
-	kref_init(&buf->refcount);
-
-	rc = buf->behavior->alloc(buf, gfp, args);
-	if (rc) {
-		dev_err(mmg->dev, "%s: Failure in buffer alloc callback %d\n",
-			behavior->topic, rc);
-		goto remove_idr;
-	}
 
 	return buf;
 
-remove_idr:
-	spin_lock(&mmg->lock);
-	idr_remove(&mmg->handles, lower_32_bits(buf->handle >> PAGE_SHIFT));
-	spin_unlock(&mmg->lock);
+destroy_buf:
+	hl_mmap_mem_buf_destroy(buf);
+	return NULL;
 free_buf:
 	kfree(buf);
 	return NULL;
-- 
2.51.0