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

[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 <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;
 

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

[sashiko-bot]

Thank you for your contribution! Sashiko AI review found 2 potential issue(s) to consider:

Pre-existing issues:
- [Critical] cs_ioctl_unreserve_signals directly frees the reference-counted encapsulated signal handle using kfree(), leading to a reliable Use-After-Free.
- [High] cs_ioctl_reserve_signals dereferences the signal handle after making it globally visible in the IDR, opening a Use-After-Free race condition.
--

commit e876f03fdbac8aca21d399cce254b40bfee1df02
Author: Ruoyu Wang <ruoyuw560@gmail.com>

accel/habanalabs: publish signal handle after SOB setup

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.

> diff --git a/drivers/accel/habanalabs/common/command_submission.c b/drivers/accel/habanalabs/common/command_submission.c
> index d51b4cf332b2a..0ed0950fee652 100644
> --- a/drivers/accel/habanalabs/common/command_submission.c
> +++ b/drivers/accel/habanalabs/common/command_submission.c
[ ... ]
> @@ -2102,11 +2104,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;

[Severity: High]
This is a pre-existing issue, but does dereferencing the handle here 
expose us to a use-after-free?

Once idr_replace() publishes the handle to mgr->handles and the lock is 
dropped, a concurrent thread could guess the IDR allocation ID and call 
cs_ioctl_unreserve_signals(). This would free the handle memory before 
this thread dereferences it to read sob_addr.

Since hw_sob is available as a local variable, could we assign this using 
the local hw_sob pointer instead of accessing the handle post-publication?

>  	*handle_id = hdl_id;
>  

While reviewing the reservation changes, I also noticed a related 
pre-existing issue in the unreserve path:

drivers/accel/habanalabs/common/command_submission.c:cs_ioctl_unreserve_signals() {
    ...
    kfree(encaps_sig_hdl);
    ...
}

[Severity: Critical]
This isn't a bug introduced by this patch, but does 
cs_ioctl_unreserve_signals() skip dropping the reference count when freeing 
the encapsulated signal handle?

The handle is initialized with kref_init(), and concurrent wait command 
submissions can acquire a reference to it via kref_get_unless_zero(). 
Because wait submissions do not modify next_sob_val or sob_addr, the 
validation checks in cs_ioctl_unreserve_signals() can succeed while a wait 
submission is active.

Since the unreserve path calls kfree() directly instead of using 
kref_put(), could this lead to a use-after-free when the wait submission 
eventually completes and tries to drop its reference?

-- 
Sashiko AI review · https://sashiko.dev/#/patchset/20260620155313.79464-1-ruoyuw560@gmail.com?part=1