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

[Ruoyu Wang <ruoyuw560@gmail.com>]

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 <ruoyuw560@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;