Re: [PATCH bpf-next v4 2/3] bpf: Prevent task work context reuse while irq_work is busy

[Mykyta Yatsenko <mykyta.yatsenko5@gmail.com>]

On 9/29/26 8:27 AM, Yun Lu wrote:
> From: Yun Lu <luyun@kylinos.cn>
> 
> A bpf_task_work context can return to STANDBY before its scheduling
> irq_work handler returns.  Both callback completion and task_work_add()
> failure can publish STANDBY while the irq_work is still BUSY.  A new
> round can then reinitialize and queue the same irq_work on another CPU
> while the old handler is still running.
> 
> One concrete interleaving is:
> 
>   1. CPU 0 successfully queues ctx->work with task_work_add() and pauses
>      before its SCHEDULING-to-SCHEDULED cmpxchg.
> 
>   2. The target task on CPU 1 completes the callback, resets ctx->task
>      and publishes STANDBY.
> 
>   3. CPU 2 starts a new round, changes STANDBY to PENDING and queues the
>      same irq_work.  The new handler changes PENDING to SCHEDULING, but
>      task_work_add() returns -ESRCH because its target task has exited.
>      The handler has not yet reset the context or rolled back the state.
> 
>   4. CPU 0 resumes and changes the new round's SCHEDULING to SCHEDULED,
>      even though no task work was queued for that round.
> 
>   5. The new handler resets ctx->task, but its SCHEDULING-to-STANDBY
>      cmpxchg fails.  The context remains SCHEDULED with a NULL task, and
>      subsequent attempts to acquire this context fail with -EBUSY.
> 
>   6. Map-value deletion observes SCHEDULED and queues asynchronous
>      cancellation, which passes the NULL ctx->task to task_work_cancel().
> 
> Overlapping invocations also share one BUSY bit.  An invocation can clear
> that bit while the other invocation is still running.
> 
> Initialize the scheduling irq_work when the context is created and stop
> reinitializing it for each round.  After claiming the context with the
> STANDBY-to-PENDING cmpxchg, check BUSY and reject reuse with -EBUSY if a
> previous invocation is still in progress.  Checking after the cmpxchg
> orders the BUSY read after the STANDBY publication.
> 
> No irq_work has been queued for the attempted round at this point, so
> the rejection can roll PENDING back to STANDBY.  Use cmpxchg for that
> rollback to preserve FREED if a concurrent deletion has published it.
> Once BUSY clears, scheduling can be retried.
> 
> Fixes: 38aa7003e369 ("bpf: task work scheduling kfuncs")
> Signed-off-by: Yun Lu <luyun@kylinos.cn>
> ---

Acked-by: Mykyta Yatsenko <yatsenko@meta.com>

>  kernel/bpf/helpers.c | 14 +++++++++++++-
>  1 file changed, 13 insertions(+), 1 deletion(-)
> 
> diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
> index 028fc6796b59..350269b8646e 100644
> --- a/kernel/bpf/helpers.c
> +++ b/kernel/bpf/helpers.c
> @@ -4553,6 +4553,7 @@ static struct bpf_task_work_ctx *bpf_task_work_fetch_ctx(struct bpf_task_work *t
>  	memset(ctx, 0, sizeof(*ctx));
>  	refcount_set(&ctx->refcnt, 1); /* map's own ref */
>  	ctx->state = BPF_TW_STANDBY;
> +	init_irq_work(&ctx->irq_work, bpf_task_work_irq);
>  
>  	old_ctx = cmpxchg(&twk->ctx, NULL, ctx);
>  	if (old_ctx) {
> @@ -4595,6 +4596,18 @@ static struct bpf_task_work_ctx *bpf_task_work_acquire_ctx(struct bpf_task_work
>  		bpf_task_work_ctx_put(ctx);
>  		return ERR_PTR(-EBUSY);
>  	}
> +	/*
> +	 * STANDBY can be published before the previous scheduling irq_work
> +	 * returns. Claim PENDING before checking BUSY to order the check after
> +	 * that publication. No irq_work has been queued for this attempt yet,
> +	 * so reject reuse and roll back to STANDBY while BUSY is set. Preserve
> +	 * FREED if deletion wins the rollback race.
> +	 */
> +	if (unlikely(irq_work_is_busy(&ctx->irq_work))) {
> +		(void)cmpxchg(&ctx->state, BPF_TW_PENDING, BPF_TW_STANDBY);
> +		bpf_task_work_ctx_put(ctx);
> +		return ERR_PTR(-EBUSY);
> +	}
>  
>  	/*
>  	 * If no process or bpffs is holding a reference to the map, no new callbacks should be
> @@ -4644,7 +4657,6 @@ static int bpf_task_work_schedule(struct task_struct *task, struct bpf_task_work
>  	ctx->map = map;
>  	ctx->map_val = (void *)tw - map->record->task_work_off;
>  	init_task_work(&ctx->work, bpf_task_work_callback);
> -	init_irq_work(&ctx->irq_work, bpf_task_work_irq);
>  
>  	irq_work_queue(&ctx->irq_work);
>  	return 0;