[PATCH v2] RDMA/rxe: Fix race in do_task() when draining

[PATCH v2] RDMA/rxe: Fix race in do_task() when draining

[Gui-Dong Han]

When do_task() exhausts its iteration budget (!ret), it sets the state
to TASK_STATE_IDLE to reschedule, without a secondary check on the
current task->state. This can overwrite the TASK_STATE_DRAINING state
set by a concurrent call to rxe_cleanup_task() or rxe_disable_task().

While state changes are protected by a spinlock, both rxe_cleanup_task()
and rxe_disable_task() release the lock while waiting for the task to
finish draining in the while(!is_done(task)) loop. The race occurs if
do_task() hits its iteration limit and acquires the lock in this window.
The cleanup logic may then proceed while the task incorrectly
reschedules itself, leading to a potential use-after-free.

This bug was introduced during the migration from tasklets to workqueues,
where the special handling for the draining case was lost.

Fix this by restoring the original pre-migration behavior. If the state is
TASK_STATE_DRAINING when iterations are exhausted, set cont to 1 to
force a new loop iteration. This allows the task to finish its work, so
that a subsequent iteration can reach the switch statement and correctly
transition the state to TASK_STATE_DRAINED, stopping the task as intended.

Fixes: 9b4b7c1f9f54 ("RDMA/rxe: Add workqueue support for rxe tasks")
Cc: stable@vger.kernel.org
Reviewed-by: Zhu Yanjun <yanjun.zhu@linux.dev>
Signed-off-by: Gui-Dong Han <hanguidong02@gmail.com>
---
v2:
* Rewrite commit message for clarity. Thanks to Zhu Yanjun for the review.
---
 drivers/infiniband/sw/rxe/rxe_task.c | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/drivers/infiniband/sw/rxe/rxe_task.c b/drivers/infiniband/sw/rxe/rxe_task.c
index 6f8f353e9583..f522820b950c 100644
--- a/drivers/infiniband/sw/rxe/rxe_task.c
+++ b/drivers/infiniband/sw/rxe/rxe_task.c
@@ -132,8 +132,12 @@ static void do_task(struct rxe_task *task)
 		 * yield the cpu and reschedule the task
 		 */
 		if (!ret) {
-			task->state = TASK_STATE_IDLE;
-			resched = 1;
+			if (task->state != TASK_STATE_DRAINING) {
+				task->state = TASK_STATE_IDLE;
+				resched = 1;
+			} else {
+				cont = 1;
+			}
 			goto exit;
 		}
 
-- 
2.25.1

Re: [PATCH v2] RDMA/rxe: Fix race in do_task() when draining

[Yanjun.Zhu]

On 9/18/25 7:52 PM, Gui-Dong Han wrote:
> When do_task() exhausts its iteration budget (!ret), it sets the state
> to TASK_STATE_IDLE to reschedule, without a secondary check on the
> current task->state. This can overwrite the TASK_STATE_DRAINING state
> set by a concurrent call to rxe_cleanup_task() or rxe_disable_task().
>
> While state changes are protected by a spinlock, both rxe_cleanup_task()
> and rxe_disable_task() release the lock while waiting for the task to
> finish draining in the while(!is_done(task)) loop. The race occurs if
> do_task() hits its iteration limit and acquires the lock in this window.
> The cleanup logic may then proceed while the task incorrectly
> reschedules itself, leading to a potential use-after-free.
>
> This bug was introduced during the migration from tasklets to workqueues,
> where the special handling for the draining case was lost.
>
> Fix this by restoring the original pre-migration behavior. If the state is
> TASK_STATE_DRAINING when iterations are exhausted, set cont to 1 to
> force a new loop iteration. This allows the task to finish its work, so
> that a subsequent iteration can reach the switch statement and correctly
> transition the state to TASK_STATE_DRAINED, stopping the task as intended.
>
> Fixes: 9b4b7c1f9f54 ("RDMA/rxe: Add workqueue support for rxe tasks")
> Cc: stable@vger.kernel.org
> Reviewed-by: Zhu Yanjun <yanjun.zhu@linux.dev>

Thanks a lot. I am fine with this.

Yanjun.Zhu

> Signed-off-by: Gui-Dong Han <hanguidong02@gmail.com>
> ---
> v2:
> * Rewrite commit message for clarity. Thanks to Zhu Yanjun for the review.
> ---
>   drivers/infiniband/sw/rxe/rxe_task.c | 8 ++++++--
>   1 file changed, 6 insertions(+), 2 deletions(-)
>
> diff --git a/drivers/infiniband/sw/rxe/rxe_task.c b/drivers/infiniband/sw/rxe/rxe_task.c
> index 6f8f353e9583..f522820b950c 100644
> --- a/drivers/infiniband/sw/rxe/rxe_task.c
> +++ b/drivers/infiniband/sw/rxe/rxe_task.c
> @@ -132,8 +132,12 @@ static void do_task(struct rxe_task *task)
>   		 * yield the cpu and reschedule the task
>   		 */
>   		if (!ret) {
> -			task->state = TASK_STATE_IDLE;
> -			resched = 1;
> +			if (task->state != TASK_STATE_DRAINING) {
> +				task->state = TASK_STATE_IDLE;
> +				resched = 1;
> +			} else {
> +				cont = 1;
> +			}
>   			goto exit;
>   		}
>   

Re: [PATCH v2] RDMA/rxe: Fix race in do_task() when draining

[Leon Romanovsky]

On Fri, 19 Sep 2025 02:52:12 +0000, Gui-Dong Han wrote:
> When do_task() exhausts its iteration budget (!ret), it sets the state
> to TASK_STATE_IDLE to reschedule, without a secondary check on the
> current task->state. This can overwrite the TASK_STATE_DRAINING state
> set by a concurrent call to rxe_cleanup_task() or rxe_disable_task().
> 
> While state changes are protected by a spinlock, both rxe_cleanup_task()
> and rxe_disable_task() release the lock while waiting for the task to
> finish draining in the while(!is_done(task)) loop. The race occurs if
> do_task() hits its iteration limit and acquires the lock in this window.
> The cleanup logic may then proceed while the task incorrectly
> reschedules itself, leading to a potential use-after-free.
> 
> [...]

Applied, thanks!

[1/1] RDMA/rxe: Fix race in do_task() when draining
      https://git.kernel.org/rdma/rdma/c/8ca7eada62fcfa

Best regards,
-- 
Leon Romanovsky <leon@kernel.org>

Re: [PATCH v2] RDMA/rxe: Fix race in do_task() when draining

[Leon Romanovsky]

On Fri, 19 Sep 2025 02:52:12 +0000, Gui-Dong Han wrote:
> When do_task() exhausts its iteration budget (!ret), it sets the state
> to TASK_STATE_IDLE to reschedule, without a secondary check on the
> current task->state. This can overwrite the TASK_STATE_DRAINING state
> set by a concurrent call to rxe_cleanup_task() or rxe_disable_task().
> 
> While state changes are protected by a spinlock, both rxe_cleanup_task()
> and rxe_disable_task() release the lock while waiting for the task to
> finish draining in the while(!is_done(task)) loop. The race occurs if
> do_task() hits its iteration limit and acquires the lock in this window.
> The cleanup logic may then proceed while the task incorrectly
> reschedules itself, leading to a potential use-after-free.
> 
> [...]

Applied, thanks!

[1/1] RDMA/rxe: Fix race in do_task() when draining
      https://git.kernel.org/rdma/rdma/c/8ca7eada62fcfa

Best regards,
-- 
Leon Romanovsky <leon@kernel.org>