[PATCH] accel/habanalabs: publish signal handle after SOB setup

[Ruoyu Wang]

cs_ioctl_reserve_signals() makes the encapsulated signal handle visible
in the context IDR before the SOB pointer and pre-reserve SOB value are
set. Concurrent unreserve and wait paths dereference those fields after
IDR lookup.

Reserve the IDR slot with a NULL entry, initialize the handle including
the SOB fields, and replace the slot with the handle only after the
visible state is ready.

Fixes: dadf17abb724 ("habanalabs: add support for encapsulated signals 
reservation")
Signed-off-by: Ruoyu Wang <ruoyuw...@gmail.com>
---
 .../habanalabs/common/command_submission.c    | 25 +++++++++++++------
 1 file changed, 17 insertions(+), 8 deletions(-)

diff --git a/drivers/accel/habanalabs/common/command_submission.c 
b/drivers/accel/habanalabs/common/command_submission.c
index ba4257bda77b..791fc01e24c5 100644
--- a/drivers/accel/habanalabs/common/command_submission.c
+++ b/drivers/accel/habanalabs/common/command_submission.c
@@ -2009,6 +2009,7 @@ static int cs_ioctl_reserve_signals(struct hl_fpriv 
*hpriv,
        struct hl_cs_encaps_sig_handle *handle;
        struct hl_encaps_signals_mgr *mgr;
        struct hl_hw_sob *hw_sob;
+       void *old;
        int hdl_id;
        int rc = 0;
 
@@ -2045,13 +2046,19 @@ static int cs_ioctl_reserve_signals(struct hl_fpriv 
*hpriv,
        }
 
        handle->count = count;
+       handle->q_idx = q_idx;
+       handle->hdev = hdev;
+       handle->cs_seq = ULLONG_MAX;
+       kref_init(&handle->refcount);
 
        hl_ctx_get(hpriv->ctx);
        handle->ctx = hpriv->ctx;
        mgr = &hpriv->ctx->sig_mgr;
 
        spin_lock(&mgr->lock);
-       hdl_id = idr_alloc(&mgr->handles, handle, 1, 0, GFP_ATOMIC);
+       hdl_id = idr_alloc(&mgr->handles, NULL, 1, 0, GFP_ATOMIC);
+       if (hdl_id >= 0)
+               handle->id = hdl_id;
        spin_unlock(&mgr->lock);
 
        if (hdl_id < 0) {
@@ -2060,11 +2067,6 @@ static int cs_ioctl_reserve_signals(struct hl_fpriv 
*hpriv,
                goto put_ctx;
        }
 
-       handle->id = hdl_id;
-       handle->q_idx = q_idx;
-       handle->hdev = hdev;
-       kref_init(&handle->refcount);
-
        hdev->asic_funcs->hw_queues_lock(hdev);
 
        hw_sob = &prop->hw_sob[prop->curr_sob_offset];
@@ -2093,11 +2095,18 @@ static int cs_ioctl_reserve_signals(struct hl_fpriv 
*hpriv,
         */
        handle->pre_sob_val = prop->next_sob_val - handle->count;
 
-       handle->cs_seq = ULLONG_MAX;
-
        *signals_count = prop->next_sob_val;
        hdev->asic_funcs->hw_queues_unlock(hdev);
 
+       spin_lock(&mgr->lock);
+       old = idr_replace(&mgr->handles, handle, hdl_id);
+       spin_unlock(&mgr->lock);
+
+       if (WARN_ON(IS_ERR(old))) {
+               rc = PTR_ERR(old);
+               goto remove_idr;
+       }
+
        *sob_addr = handle->hw_sob->sob_addr;
        *handle_id = hdl_id;