Re: [PATCH bpf-next v2 1/2] bpf: Fix task work scheduling and callback race

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

On 9/21/26 11:14 AM, Yun Lu wrote:
> From: Yun Lu <luyun@kylinos.cn>
> 
> bpf_task_work_irq() publishes a task_work callback before changing the
> context from SCHEDULING to SCHEDULED.  The callback can therefore run on
> another CPU while the irq_work handler is still using the same scheduling
> round.
> 
> Nothing currently orders the handler's use of ctx->task against
> bpf_task_work_ctx_reset().  The callback can reset ctx->task and publish
> STANDBY before the handler resumes, allowing map-value deletion to make the
> handler pass NULL to task_work_cancel().  Publishing STANDBY this early
> also allows a new round to reuse the context and irq_work while the old
> handler is still running.
> 
> One concrete interleaving is:
> 
>   1. CPU 0 changes PENDING to SCHEDULING in bpf_task_work_irq() and
>      successfully publishes ctx->work with task_work_add(), but has not yet
>      attempted the SCHEDULING-to-SCHEDULED transition.
> 
>   2. The target task on CPU 1 runs the callback.  It changes SCHEDULING to
>      RUNNING, executes the BPF subprogram, then bpf_task_work_ctx_reset()
>      releases ctx->task and sets it to NULL.  The callback changes RUNNING
>      to STANDBY and drops its ctx reference.
> 
>   3. CPU 2 deletes the map value.  bpf_task_work_cancel_and_free() changes
>      STANDBY to FREED.  Since the old state is not SCHEDULED, it does not
>      queue cancellation and drops the map's ctx reference.
> 
>   4. CPU 0 resumes.  Its SCHEDULING-to-SCHEDULED cmpxchg observes FREED, so
>      bpf_task_work_cancel() calls task_work_cancel(NULL, &ctx->work).
> 
> If deletion changes the state to FREED before the callback runs, the
> callback instead takes its FREED exit and may drop the last ctx reference.
> The destroy path then resets ctx->task, producing the same NULL dereference
> when CPU 0 resumes.
> 
> A controlled reproducer that delays CPU 0 between task_work_add() and the
> state cmpxchg triggered the following fault on a v7.3-rc3 based x86-64 KVM
> guest:
> 
>   BUG: kernel NULL pointer dereference, address: 0000000000000890
>   #PF: supervisor read access in kernel mode
>   RIP: task_work_cancel+0xd/0xa0
>   RDI: 0000000000000000
>   CR2: 0000000000000890
>   Call Trace:
>    <IRQ>
>    bpf_task_work_irq+0x95/0x100
>    irq_work_run_list+0x4f/0x90
>    irq_work_run+0x18/0x50
>    __sysvec_irq_work+0x18/0xb0
>    sysvec_irq_work+0x66/0x80
>    </IRQ>
> 
> 0x890 is the task_struct::task_works offset in that build.  It is read by
> task_work_pending() through the NULL task argument.  The rcu_read_lock() in
> bpf_task_work_irq() only keeps ctx memory alive; it does not retain
> ctx->task.  A NULL check would still race with releasing the task
> reference.
> 
> Fix this by making the callback claim RUNNING before synchronizing with the
> scheduling irq_work.  RUNNING prevents map-value deletion from
> reinitializing the irq_work for asynchronous cancellation, so
> irq_work_sync() waits for the scheduling invocation that published the
> callback.  Only after that invocation returns may the callback execute the
> BPF subprogram, reset task/prog and publish STANDBY.
> 
> Pin the context with one temporary reference held by the scheduling
> handler.  If deletion wins and the callback takes its FREED exit first,
> this reference prevents destruction from resetting ctx->task before the
> handler completes its cancellation attempt.
> 
> An add failure has no callback to perform the synchronization.  Initialize
> the scheduling irq_work when the context is created and reject a new round
> while the failed invocation remains BUSY.  This prevents overlapping
> invocations from sharing the single BUSY bit and making a later
> irq_work_sync() return before its scheduling handler has finished.
> 
> The callback runs in task context after task_work_run() has released
> task->pi_lock, and existing task_work callbacks such as ____fput() may
> sleep.  Keep the Tasks Trace RCU read-side section across irq_work_sync()
> so the map value remains live until the callback finishes.  No context
> fields or states are added, and the existing cancellation and destruction
> paths are retained.
> 
> Fixes: 38aa7003e369 ("bpf: task work scheduling kfuncs")
> Signed-off-by: Yun Lu <luyun@kylinos.cn>
> ---
>  kernel/bpf/helpers.c | 46 ++++++++++++++++++++++++++++++++++++++------
>  1 file changed, 40 insertions(+), 6 deletions(-)
> 
> diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
> index b3cc5c8fc875..e5d5683626ea 100644
> --- a/kernel/bpf/helpers.c
> +++ b/kernel/bpf/helpers.c
> @@ -4469,6 +4469,14 @@ static void bpf_task_work_callback(struct callback_head *cb)
>  		bpf_task_work_ctx_put(ctx);
>  		return;
>  	}
> +	if (WARN_ON_ONCE(state != BPF_TW_SCHEDULING &&
> +			 state != BPF_TW_SCHEDULED)) {
> +		bpf_task_work_ctx_put(ctx);
> +		return;
> +	}

If this is an impossible condition, we should remove this hunk. 
If it is possible, we should remove WARN_ON_ONCE.

> +
> +	/* Do not release this round's resources until its scheduler is done. */
> +	irq_work_sync(&ctx->irq_work);

This looks like a perf problem, we are blocking the task work callback waiting 
for the irq_work. Instead we should try to make the irq_work safe after the 
task_work_add(), maybe take a task reference.
>  
>  	key = (void *)map_key_from_value(ctx->map, ctx->map_val, &idx);
>  
> @@ -4495,6 +4503,12 @@ static void bpf_task_work_irq(struct irq_work *irq_work)
>  		bpf_task_work_ctx_put(ctx);
>  		return;
>  	}
> +	/*
> +	 * Pin the ctx until this handler is done. The callback may observe
> +	 * FREED and drop its ref first, and destroy must not reset ctx->task
> +	 * before the cancellation attempt below.
> +	 */
> +	refcount_inc(&ctx->refcnt);
>  
>  	err = task_work_add(ctx->task, &ctx->work, ctx->mode);
>  	if (err) {
> @@ -4504,20 +4518,26 @@ static void bpf_task_work_irq(struct irq_work *irq_work)
>  		 * gone to FREED already, which is fine as we already cleaned up after ourselves
>  		 */
>  		(void)cmpxchg(&ctx->state, BPF_TW_SCHEDULING, BPF_TW_STANDBY);
> +		/*
> +		 * No callback was published, so drop both refs owned by this
> +		 * failed round: the callback ref and the scheduler's temporary ref.
> +		 */
> +		bpf_task_work_ctx_put(ctx);
>  		bpf_task_work_ctx_put(ctx);
>  		return;
>  	}
>  
>  	/*
> -	 * It's technically possible for just scheduled task_work callback to
> -	 * complete running by now, going SCHEDULING -> RUNNING and then
> -	 * dropping its ctx refcount. Instead of capturing an extra ref just
> -	 * to protect below ctx->state access, we rely on rcu_read_lock
> -	 * above to prevent kfree_rcu from freeing ctx before we return.
> +	 * The callback may already be running on the target task's CPU, but
> +	 * it waits for this invocation to finish before resetting task/prog
> +	 * or publishing STANDBY, and the temporary reference above keeps the
> +	 * ctx alive no matter how the other references are dropped here.
>  	 */
>  	state = cmpxchg(&ctx->state, BPF_TW_SCHEDULING, BPF_TW_SCHEDULED);
>  	if (state == BPF_TW_FREED)
>  		bpf_task_work_cancel(ctx); /* clean up if we switched into FREED state */
> +
> +	bpf_task_work_ctx_put(ctx);
>  }
>  
>  static struct bpf_task_work_ctx *bpf_task_work_fetch_ctx(struct bpf_task_work *tw,
> @@ -4537,6 +4557,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) {
> @@ -4555,6 +4576,7 @@ static struct bpf_task_work_ctx *bpf_task_work_acquire_ctx(struct bpf_task_work
>  							   struct bpf_map *map)
>  {
>  	struct bpf_task_work_ctx *ctx;
> +	enum bpf_task_work_state state;
>  
>  	/*
>  	 * Sleepable BPF programs hold rcu_read_lock_trace but not
> @@ -4579,6 +4601,19 @@ 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);
>  	}
> +	/*
> +	 * An add failure publishes STANDBY before its irq_work handler
> +	 * returns. Do not let a new round requeue the same irq_work until that
> +	 * handler has cleared BUSY. Otherwise two invocations can overlap on
> +	 * different CPUs; either tail can clear the shared BUSY bit and let a
> +	 * later irq_work_sync() return while the other invocation still runs.
> +	 */
> +	if (unlikely(irq_work_is_busy(&ctx->irq_work))) {
> +		state = cmpxchg(&ctx->state, BPF_TW_PENDING, BPF_TW_STANDBY);
> +		WARN_ON_ONCE(state != BPF_TW_PENDING && state != BPF_TW_FREED);

Let's remove this WARN_ON_ONCE(), we do not have asserts for state machine
states in other places.
> +		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
> @@ -4628,7 +4663,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;