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 <[email protected]> --- 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
