[PATCH bpf-next v3 0/2] bpf: Fix task work scheduling race

[PATCH bpf-next v3 0/2] bpf: Fix task work scheduling race

[Yun Lu]

From: Yun Lu <luyun@kylinos.cn>

This series fixes a race in the bpf_task_work scheduling kfuncs and
adds a regression test for it.

bpf_task_work_irq() publishes a task_work callback before it finishes
using the scheduling round.  A target task running on another CPU can
execute the callback, reset ctx->task and publish STANDBY before the
irq_work handler performs its final state transition.  Concurrent map
value deletion can then make the handler call task_work_cancel() with a
NULL task.  The early STANDBY transition can also let a new round reuse
the context while the old irq_work handler is still running.

Patch 1 captures an independent task reference in the scheduling
handler before task_work_add() and uses it for the post-add
cancellation, so the handler never reads ctx->task after the callback
may have reset it.  A BUSY check in acquire keeps a new round from
starting while a previous handler is still running, which prevents a
delayed handler from consuming the new round's SCHEDULING state.  The
scheduling irq_work is initialized once at context creation.  No
fields or states are added to the context, the callback stays
non-blocking, and the cancellation and destruction paths are unchanged.

Patch 2 adds a 1500-round cross-CPU regression test covering context
reuse, callback completion followed by deletion, and deletion racing
with scheduling.

Changes since v2:

- Kernel: drop irq_work_sync() from the callback and both
  WARN_ON_ONCE() checks, as requested in review.  Following the
  suggested direction, the scheduling handler now takes its own task
  reference around task_work_add() and the post-add cancellation; the
  temporary ctx refcount from v2 is gone as well.  The callback ends
  up unchanged from upstream.
- The BUSY check in acquire is kept: without it a handler delayed
  between task_work_add() and its state cmpxchg can flip the next
  round into a bogus SCHEDULED, which later routes map-value deletion
  into task_work_cancel() with an already-reset ctx->task.
- Selftest is unchanged apart from the rebase.

v2: https://lore.kernel.org/all/20260921101453.69273-1-luyun_611@163.com/

Changes since v1 :

- Kernel: rework the fix per review feedback.  Replace the round_refs
  ownership counter and the extra irq_work objects with a completion
  edge: the callback claims RUNNING and waits for the scheduling
  irq_work via irq_work_sync(), plus one temporary reference from the
  existing ctx refcount and a BUSY check covering the add-failure
  path.  No context fields or states are added.
- Selftest: address review feedback (single packed result store,
  bounded thread-start wait, tolerate the expected deletion-race
  outcomes).

v1: https://lore.kernel.org/bpf/20260918080014.54012-1-luyun_611@163.com/

Testing:

- Built with KASAN, lockdep, PROVE_RCU, DEBUG_ATOMIC_SLEEP and
  KMEMLEAK; test_progs -t task_work and 900 stress rounds in QEMU are
  clean, zero leaks reported.
- Replayed the deterministic reproducer (debug-delay only, not part of
  the series): without the fix it panics reliably in
  task_work_cancel(); with the fix the callback staying non-blocking
  and the guest stays healthy.

Yun Lu (2):
  bpf: Fix task work scheduling and callback race
  selftests/bpf: Add task work scheduling race test

 kernel/bpf/helpers.c                          |  39 ++-
 .../selftests/bpf/prog_tests/test_task_work.c | 283 ++++++++++++++++++
 .../selftests/bpf/progs/task_work_race.c      | 126 ++++++++
 3 files changed, 437 insertions(+), 11 deletions(-)
 create mode 100644 tools/testing/selftests/bpf/progs/task_work_race.c


base-commit: 5dd1818b15d98d4a20806cd00b1b40320b06004f
-- 
2.43.0

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

[Yun Lu]

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 taking a task reference in the scheduling irq_work before
publishing the callback.  Use the captured task for cancellation at the
handler tail, so the task remains valid even if the callback or destroy
path has reset ctx->task.  Release the reference on every add success and
failure path.

The asynchronous cancellation path is queued only after deletion wins
SCHEDULED-to-FREED.  In that case the callback observes FREED and cannot
reset ctx->task, while the map's transferred ctx reference prevents the
destroy path from doing so.  It can therefore safely pass ctx->task to the
same cancellation helper.  The two cancellation paths are mutually
exclusive: the scheduling handler handles deletion from SCHEDULING, and
the asynchronous path handles deletion from SCHEDULED.

The callback or add-failure path can publish STANDBY before the scheduling
irq_work returns.  Initialize irq_work when the context is created and
reject a new round while it is still BUSY.  This prevents the old handler
from consuming the new round's SCHEDULING state and publishing a false
SCHEDULED state, and prevents invocations on different CPUs from
overlapping on the same context.

This keeps the task-work callback non-blocking and adds no context fields
or states.

Fixes: 38aa7003e369 ("bpf: task work scheduling kfuncs")
Reported-by: Jackie Liu <liuyun01@kylinos.cn>
Suggested-by: Mykyta Yatsenko <mykyta.yatsenko5@gmail.com>
Signed-off-by: Yun Lu <luyun@kylinos.cn>
---
 kernel/bpf/helpers.c | 39 ++++++++++++++++++++++++++++-----------
 1 file changed, 28 insertions(+), 11 deletions(-)

diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
index b3cc5c8fc875..d8bea7538de5 100644
--- a/kernel/bpf/helpers.c
+++ b/kernel/bpf/helpers.c
@@ -4437,7 +4437,8 @@ static void bpf_task_work_ctx_put(struct bpf_task_work_ctx *ctx)
 	}
 }
 
-static void bpf_task_work_cancel(struct bpf_task_work_ctx *ctx)
+static void bpf_task_work_cancel(struct bpf_task_work_ctx *ctx,
+				 struct task_struct *task)
 {
 	/*
 	 * Scheduled task_work callback holds ctx ref, so if we successfully
@@ -4445,7 +4446,7 @@ static void bpf_task_work_cancel(struct bpf_task_work_ctx *ctx)
 	 * cancel, callback will inevitably run or has already completed
 	 * running, and it would have taken care of its ctx ref itself.
 	 */
-	if (task_work_cancel(ctx->task, &ctx->work))
+	if (task_work_cancel(task, &ctx->work))
 		bpf_task_work_ctx_put(ctx);
 }
 
@@ -4486,6 +4487,7 @@ static void bpf_task_work_callback(struct callback_head *cb)
 static void bpf_task_work_irq(struct irq_work *irq_work)
 {
 	struct bpf_task_work_ctx *ctx = container_of(irq_work, struct bpf_task_work_ctx, irq_work);
+	struct task_struct *task;
 	enum bpf_task_work_state state;
 	int err;
 
@@ -4496,7 +4498,8 @@ static void bpf_task_work_irq(struct irq_work *irq_work)
 		return;
 	}
 
-	err = task_work_add(ctx->task, &ctx->work, ctx->mode);
+	task = get_task_struct(ctx->task);
+	err = task_work_add(task, &ctx->work, ctx->mode);
 	if (err) {
 		bpf_task_work_ctx_reset(ctx);
 		/*
@@ -4505,19 +4508,20 @@ static void bpf_task_work_irq(struct irq_work *irq_work)
 		 */
 		(void)cmpxchg(&ctx->state, BPF_TW_SCHEDULING, BPF_TW_STANDBY);
 		bpf_task_work_ctx_put(ctx);
+		put_task_struct(task);
 		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 have completed and reset ctx->task already. Keep the
+	 * captured task alive until this invocation is done with cancellation.
+	 * The RCU read-side section keeps ctx memory alive until then.
 	 */
 	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_cancel(ctx, task); /* clean up if we switched into FREED state */
+
+	put_task_struct(task);
 }
 
 static struct bpf_task_work_ctx *bpf_task_work_fetch_ctx(struct bpf_task_work *tw,
@@ -4537,6 +4541,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) {
@@ -4579,6 +4584,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);
 	}
+	/*
+	 * The callback or add-failure path can publish STANDBY before the
+	 * scheduling irq_work returns. Do not let a new round reuse the context
+	 * until BUSY clears. Otherwise the old handler can consume the new
+	 * round's SCHEDULING state and publish a false SCHEDULED state, and two
+	 * invocations can overlap on different CPUs.
+	 */
+	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
@@ -4628,7 +4645,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;
@@ -4777,7 +4793,8 @@ static void bpf_task_work_cancel_scheduled(struct irq_work *irq_work)
 {
 	struct bpf_task_work_ctx *ctx = container_of(irq_work, struct bpf_task_work_ctx, irq_work);
 
-	bpf_task_work_cancel(ctx); /* this might put task_work callback's ref */
+	/* SCHEDULED -> FREED keeps ctx->task stable until this runs. */
+	bpf_task_work_cancel(ctx, ctx->task); /* this might put task_work callback's ref */
 	bpf_task_work_ctx_put(ctx); /* and here we put map's own ref that was transferred to us */
 }
 

-- 
2.43.0

[PATCH bpf-next v3 2/2] selftests/bpf: Add task work scheduling race test

[Yun Lu]

From: Yun Lu <luyun@kylinos.cn>

Exercise task-work scheduling, callback completion and map-value deletion
on separate execution contexts. Rotate through context reuse, deletion
after callback completion and deletion immediately before scheduling.
These variants stress the SCHEDULING/SCHEDULED cancellation window and
reuse of the same context across rounds.

Use READY and a single packed RESULT containing the generation and error,
so scheduling errors cannot be missed without relying on ordering between
separate map slots. Give every callback a unique generation so a late
callback cannot satisfy a later assertion. Retry transient -EBUSY results
while a previous callback wrapper is finishing. In the immediate-delete
variant, accept -EBUSY and a follow-up -ENOENT when deletion wins the race.

Select two CPUs from the process affinity mask, bound the thread-start
wait and use an atomic stop handshake. Skip when fewer than two CPUs are
available, while reporting thread creation or affinity setup errors as
test failures.

Signed-off-by: Yun Lu <luyun@kylinos.cn>
---
 .../selftests/bpf/prog_tests/test_task_work.c | 283 ++++++++++++++++++
 .../selftests/bpf/progs/task_work_race.c      | 126 ++++++++
 2 files changed, 409 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/progs/task_work_race.c

diff --git a/tools/testing/selftests/bpf/prog_tests/test_task_work.c b/tools/testing/selftests/bpf/prog_tests/test_task_work.c
index 774b31a5f6ca..ff56976f4226 100644
--- a/tools/testing/selftests/bpf/prog_tests/test_task_work.c
+++ b/tools/testing/selftests/bpf/prog_tests/test_task_work.c
@@ -5,10 +5,14 @@
 #include <stdio.h>
 #include "task_work.skel.h"
 #include "task_work_fail.skel.h"
+#include "task_work_race.skel.h"
 #include <linux/bpf.h>
 #include <linux/perf_event.h>
 #include <sys/syscall.h>
 #include <time.h>
+#include <pthread.h>
+#include <sched.h>
+#include <unistd.h>
 
 static int perf_event_open(__u32 type, __u64 config, int pid)
 {
@@ -155,3 +159,282 @@ void test_task_work(void)
 
 	RUN_TESTS(task_work_fail);
 }
+
+#define TASK_WORK_RACE_ROUNDS	1500
+
+/* Must match progs/task_work_race.c. */
+enum task_work_race_status {
+	RACE_ARM_SEQ,
+	RACE_READY_SEQ,
+	/* High 32 bits hold the sequence, low 32 bits the signed error. */
+	RACE_RESULT,
+};
+
+struct task_work_race_value {
+	__u32 seq;
+	char data[60];
+	struct bpf_task_work tw;
+};
+
+struct task_work_race_ctx {
+	int stop;
+	int setup_err;
+	int trigger_tid;
+	int target_tid;
+	int trigger_cpu;
+	int target_cpu;
+};
+
+static int task_work_race_pin_cpu(int cpu)
+{
+	cpu_set_t set;
+
+	CPU_ZERO(&set);
+	CPU_SET(cpu, &set);
+	return pthread_setaffinity_np(pthread_self(), sizeof(set), &set);
+}
+
+static void *task_work_race_trigger(void *arg)
+{
+	struct task_work_race_ctx *ctx = arg;
+	int err;
+
+	err = task_work_race_pin_cpu(ctx->trigger_cpu);
+	if (err)
+		__atomic_store_n(&ctx->setup_err, err, __ATOMIC_RELEASE);
+	__atomic_store_n(&ctx->trigger_tid, syscall(__NR_gettid),
+			 __ATOMIC_RELEASE);
+	while (!__atomic_load_n(&ctx->stop, __ATOMIC_ACQUIRE))
+		getppid();
+	return NULL;
+}
+
+static void *task_work_race_target(void *arg)
+{
+	struct task_work_race_ctx *ctx = arg;
+	int err;
+
+	err = task_work_race_pin_cpu(ctx->target_cpu);
+	if (err)
+		__atomic_store_n(&ctx->setup_err, err, __ATOMIC_RELEASE);
+	__atomic_store_n(&ctx->target_tid, syscall(__NR_gettid),
+			 __ATOMIC_RELEASE);
+	while (!__atomic_load_n(&ctx->stop, __ATOMIC_ACQUIRE))
+		getppid();
+	return NULL;
+}
+
+static int task_work_race_status(struct task_work_race *skel, int idx,
+				 __s64 *value)
+{
+	return bpf_map_lookup_elem(bpf_map__fd(skel->maps.status), &idx,
+				   value);
+}
+
+static int task_work_race_set_status(struct task_work_race *skel, int idx,
+				     __s64 value)
+{
+	return bpf_map_update_elem(bpf_map__fd(skel->maps.status), &idx,
+				   &value, BPF_ANY);
+}
+
+static int task_work_race_wait_ready(struct task_work_race *skel, __u32 seq)
+{
+	int i;
+
+	for (i = 0; i < 2000000 / 100; i++) {
+		__s64 value = 0;
+
+		if (!task_work_race_status(skel, RACE_READY_SEQ, &value) &&
+		    value == seq)
+			return 0;
+		if (!task_work_race_status(skel, RACE_RESULT, &value) &&
+		    (__u32)((__u64)value >> 32) == seq)
+			/*
+			 * The BPF program finished between the two reads:
+			 * RESULT makes waiting for READY unnecessary. Any error
+			 * that occurred before READY is checked by the caller.
+			 */
+			return 0;
+		usleep(100);
+	}
+	return -ETIMEDOUT;
+}
+
+static int task_work_race_prepare_elem(struct task_work_race *skel)
+{
+	struct task_work_race_value value = {};
+	int key = 0;
+
+	if (!bpf_map_lookup_elem(bpf_map__fd(skel->maps.hmap), &key, &value))
+		return 0;
+	if (errno != ENOENT)
+		return -errno;
+	if (bpf_map_update_elem(bpf_map__fd(skel->maps.hmap), &key, &value,
+				BPF_NOEXIST))
+		return -errno;
+	return 0;
+}
+
+static int task_work_race_arm(struct task_work_race *skel, __u32 seq)
+{
+	if (task_work_race_set_status(skel, RACE_ARM_SEQ, seq))
+		return -errno;
+	return task_work_race_wait_ready(skel, seq);
+}
+
+static int task_work_race_check_done(struct task_work_race *skel, __u32 seq)
+{
+	int i;
+
+	for (i = 0; i < 2000000 / 100; i++) {
+		__s64 result = 0;
+
+		if (!task_work_race_status(skel, RACE_RESULT, &result) &&
+		    (__u32)((__u64)result >> 32) == seq)
+			return (__s32)result;
+		usleep(100);
+	}
+	return -ETIMEDOUT;
+}
+
+static int task_work_race_wait_callback(struct task_work_race *skel,
+					__u32 seq)
+{
+	int i, key = seq;
+
+	for (i = 0; i < 2000000 / 100; i++) {
+		__u64 completed = 0;
+
+		if (!bpf_map_lookup_elem(bpf_map__fd(skel->maps.completed),
+					 &key, &completed) && completed)
+			return 0;
+		usleep(100);
+	}
+	return -ETIMEDOUT;
+}
+
+void serial_test_task_work_race(void)
+{
+	struct task_work_race_ctx ctx = {};
+	struct task_work_race *skel;
+	pthread_t trigger, target;
+	cpu_set_t allowed;
+	int cpu, cpu_count = 0;
+	bool target_started = false, trigger_started = false;
+	int err, i, key = 0;
+
+	if (sched_getaffinity(0, sizeof(allowed), &allowed)) {
+		ASSERT_OK(-errno, "sched_getaffinity");
+		return;
+	}
+	for (cpu = 0; cpu < CPU_SETSIZE && cpu_count < 2; cpu++) {
+		if (!CPU_ISSET(cpu, &allowed))
+			continue;
+		if (!cpu_count)
+			ctx.trigger_cpu = cpu;
+		else
+			ctx.target_cpu = cpu;
+		cpu_count++;
+	}
+	if (cpu_count < 2) {
+		printf("%s:SKIP:need two CPUs in the process affinity mask\n",
+		       __func__);
+		test__skip();
+		return;
+	}
+
+	skel = task_work_race__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "open_and_load"))
+		return;
+	if (!ASSERT_OK(task_work_race__attach(skel), "attach"))
+		goto cleanup;
+
+	err = pthread_create(&target, NULL, task_work_race_target, &ctx);
+	if (!ASSERT_OK(err, "pthread_create target"))
+		goto cleanup;
+	target_started = true;
+	err = pthread_create(&trigger, NULL, task_work_race_trigger, &ctx);
+	if (!ASSERT_OK(err, "pthread_create trigger"))
+		goto stop;
+	trigger_started = true;
+
+	for (i = 0; i < 5000000 / 100; i++) {
+		if (__atomic_load_n(&ctx.trigger_tid, __ATOMIC_ACQUIRE) &&
+		    __atomic_load_n(&ctx.target_tid, __ATOMIC_ACQUIRE))
+			break;
+		usleep(100);
+	}
+	if (!ASSERT_TRUE(__atomic_load_n(&ctx.trigger_tid, __ATOMIC_ACQUIRE) &&
+			 __atomic_load_n(&ctx.target_tid, __ATOMIC_ACQUIRE),
+			 "thread start"))
+		goto stop;
+	if (!ASSERT_OK(__atomic_load_n(&ctx.setup_err, __ATOMIC_ACQUIRE),
+		       "thread affinity"))
+		goto stop;
+
+	skel->bss->trigger_tid = __atomic_load_n(&ctx.trigger_tid, __ATOMIC_ACQUIRE);
+	skel->bss->target_tid = __atomic_load_n(&ctx.target_tid, __ATOMIC_ACQUIRE);
+
+	for (i = 0; i < TASK_WORK_RACE_ROUNDS; i++) {
+		__u32 seq = i + 1;
+		int variant = i % 3;
+
+		if (!ASSERT_OK(task_work_race_prepare_elem(skel), "prepare elem") ||
+		    !ASSERT_OK(task_work_race_arm(skel, seq), "round ready"))
+			goto stop;
+
+		if (variant == 2 &&
+		    !ASSERT_OK(bpf_map_delete_elem(bpf_map__fd(skel->maps.hmap),
+						   &key), "early delete"))
+			goto stop;
+
+		err = task_work_race_check_done(skel, seq);
+		/*
+		 * Variant 2 deletes the element right after READY, so losing
+		 * the race against the scheduling kfunc is a valid outcome:
+		 * the schedule itself fails with -EBUSY (ctx already FREED),
+		 * and the retry on the next tracepoint invocation fails with
+		 * -ENOENT because the element is gone.  Either way the round
+		 * exercised the deletion paths; a later variant 0 round
+		 * recreates the element.
+		 */
+		if (variant == 2 && (err == -EBUSY || err == -ENOENT))
+			err = 0;
+		if (!ASSERT_OK(err, "schedule")) {
+			fprintf(stderr, "round %d schedule failed: %d\n", i, err);
+			goto stop;
+		}
+
+		if (variant != 2 &&
+		    !ASSERT_OK(task_work_race_wait_callback(skel, seq),
+			       "callback"))
+			goto stop;
+
+		if (variant == 1 &&
+		    !ASSERT_OK(bpf_map_delete_elem(bpf_map__fd(skel->maps.hmap),
+						   &key), "late delete"))
+			goto stop;
+	}
+
+	/* A distinct completion generation prevents an old callback satisfying this. */
+	if (!ASSERT_OK(task_work_race_prepare_elem(skel), "final prepare") ||
+	    !ASSERT_OK(task_work_race_arm(skel, TASK_WORK_RACE_ROUNDS + 1),
+		       "final ready") ||
+	    !ASSERT_OK(task_work_race_check_done(skel,
+					TASK_WORK_RACE_ROUNDS + 1),
+		       "final schedule") ||
+	    !ASSERT_OK(task_work_race_wait_callback(skel,
+					     TASK_WORK_RACE_ROUNDS + 1),
+		       "final callback"))
+		goto stop;
+
+stop:
+	__atomic_store_n(&ctx.stop, 1, __ATOMIC_RELEASE);
+	if (trigger_started)
+		pthread_join(trigger, NULL);
+	if (target_started)
+		pthread_join(target, NULL);
+cleanup:
+	task_work_race__destroy(skel);
+}
diff --git a/tools/testing/selftests/bpf/progs/task_work_race.c b/tools/testing/selftests/bpf/progs/task_work_race.c
new file mode 100644
index 000000000000..c94eca510cb3
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/task_work_race.c
@@ -0,0 +1,126 @@
+// SPDX-License-Identifier: GPL-2.0
+/* Copyright (c) 2026 KylinSoft Corporation. */
+#include <vmlinux.h>
+#include <stdbool.h>
+#include <errno.h>
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+#include "bpf_misc.h"
+
+char _license[] SEC("license") = "GPL";
+
+#define TASK_WORK_RACE_MAX_SEQ	2048
+
+enum task_work_race_status {
+	RACE_ARM_SEQ,
+	RACE_READY_SEQ,
+	/* High 32 bits hold the sequence, low 32 bits the signed error. */
+	RACE_RESULT,
+	RACE_STATUS_MAX,
+};
+
+int trigger_tid;
+int target_tid;
+
+struct {
+	__uint(type, BPF_MAP_TYPE_ARRAY);
+	__uint(max_entries, RACE_STATUS_MAX);
+	__type(key, int);
+	__type(value, __s64);
+} status SEC(".maps");
+
+struct {
+	__uint(type, BPF_MAP_TYPE_ARRAY);
+	__uint(max_entries, TASK_WORK_RACE_MAX_SEQ);
+	__type(key, int);
+	__type(value, __u64);
+} completed SEC(".maps");
+
+struct task_work_race_value {
+	__u32 seq;
+	char data[60];
+	struct bpf_task_work tw;
+};
+
+struct {
+	__uint(type, BPF_MAP_TYPE_HASH);
+	__uint(map_flags, BPF_F_NO_PREALLOC);
+	__uint(max_entries, 4);
+	__type(key, int);
+	__type(value, struct task_work_race_value);
+} hmap SEC(".maps");
+
+static __always_inline void set_status(int key, __s64 value)
+{
+	__s64 *slot;
+
+	slot = bpf_map_lookup_elem(&status, &key);
+	if (slot)
+		*slot = value;
+}
+
+static int process_work(struct bpf_map *map, void *key, void *value)
+{
+	struct task_work_race_value *work = value;
+	__u64 *done;
+	int seq = work->seq;
+
+	if (seq <= 0 || seq >= TASK_WORK_RACE_MAX_SEQ)
+		return 0;
+	done = bpf_map_lookup_elem(&completed, &seq);
+	if (done)
+		*done = 1;
+	return 0;
+}
+
+SEC("tracepoint/syscalls/sys_enter_getppid")
+int race_sched_work(void *ctx)
+{
+	struct task_work_race_value *work;
+	struct task_struct *task;
+	__s64 *arm, *result;
+	__u32 tid = (__u32)bpf_get_current_pid_tgid();
+	int key = 0, err;
+	__u32 seq;
+
+	if (tid != trigger_tid)
+		return 0;
+
+	key = RACE_ARM_SEQ;
+	arm = bpf_map_lookup_elem(&status, &key);
+	key = RACE_RESULT;
+	result = bpf_map_lookup_elem(&status, &key);
+	if (!arm || !result || *arm <= 0 ||
+	    (__u32)((__u64)*result >> 32) == *arm)
+		return 0;
+	seq = *arm;
+
+	task = bpf_task_from_pid(target_tid);
+	if (!task) {
+		err = -ESRCH;
+		goto out_done;
+	}
+
+	key = 0;
+	work = bpf_map_lookup_elem(&hmap, &key);
+	if (!work) {
+		err = -ENOENT;
+		goto out_task;
+	}
+
+	work->seq = seq;
+	set_status(RACE_READY_SEQ, seq);
+	err = bpf_task_work_schedule_signal(task, &work->tw, &hmap,
+					    process_work);
+	if (err == -EBUSY) {
+		bpf_task_release(task);
+		return 0;
+	}
+
+out_task:
+	bpf_task_release(task);
+out_done:
+	/* Publish the sequence and its error in one store. */
+	*result = ((__u64)seq << 32) | (__u32)err;
+	return 0;
+}
-- 
2.43.0

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

[Mykyta Yatsenko]

On 9/22/26 11:00 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 taking a task reference in the scheduling irq_work before
> publishing the callback.  Use the captured task for cancellation at the
> handler tail, so the task remains valid even if the callback or destroy
> path has reset ctx->task.  Release the reference on every add success and
> failure path.
> 
> The asynchronous cancellation path is queued only after deletion wins
> SCHEDULED-to-FREED.  In that case the callback observes FREED and cannot
> reset ctx->task, while the map's transferred ctx reference prevents the
> destroy path from doing so.  It can therefore safely pass ctx->task to the
> same cancellation helper.  The two cancellation paths are mutually
> exclusive: the scheduling handler handles deletion from SCHEDULING, and
> the asynchronous path handles deletion from SCHEDULED.
> 
> The callback or add-failure path can publish STANDBY before the scheduling
> irq_work returns.  Initialize irq_work when the context is created and
> reject a new round while it is still BUSY.  This prevents the old handler
> from consuming the new round's SCHEDULING state and publishing a false
> SCHEDULED state, and prevents invocations on different CPUs from
> overlapping on the same context.
> 
> This keeps the task-work callback non-blocking and adds no context fields
> or states.
> 
> Fixes: 38aa7003e369 ("bpf: task work scheduling kfuncs")
> Reported-by: Jackie Liu <liuyun01@kylinos.cn>
> Suggested-by: Mykyta Yatsenko <mykyta.yatsenko5@gmail.com>
> Signed-off-by: Yun Lu <luyun@kylinos.cn>
> ---

Looks good, the only thing that bothers a bit is
PENDING -> STANDBY transition when irq_work_is_busy(),
but overall looks like the cleanest solution.
Acked-by: Mykyta Yatsenko <yatsenko@meta.com>

>  kernel/bpf/helpers.c | 39 ++++++++++++++++++++++++++++-----------
>  1 file changed, 28 insertions(+), 11 deletions(-)
> 
> diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
> index b3cc5c8fc875..d8bea7538de5 100644
> --- a/kernel/bpf/helpers.c
> +++ b/kernel/bpf/helpers.c
> @@ -4437,7 +4437,8 @@ static void bpf_task_work_ctx_put(struct bpf_task_work_ctx *ctx)
>  	}
>  }
>  
> -static void bpf_task_work_cancel(struct bpf_task_work_ctx *ctx)
> +static void bpf_task_work_cancel(struct bpf_task_work_ctx *ctx,
> +				 struct task_struct *task)
>  {
>  	/*
>  	 * Scheduled task_work callback holds ctx ref, so if we successfully
> @@ -4445,7 +4446,7 @@ static void bpf_task_work_cancel(struct bpf_task_work_ctx *ctx)
>  	 * cancel, callback will inevitably run or has already completed
>  	 * running, and it would have taken care of its ctx ref itself.
>  	 */
> -	if (task_work_cancel(ctx->task, &ctx->work))
> +	if (task_work_cancel(task, &ctx->work))
>  		bpf_task_work_ctx_put(ctx);
>  }
>  
> @@ -4486,6 +4487,7 @@ static void bpf_task_work_callback(struct callback_head *cb)
>  static void bpf_task_work_irq(struct irq_work *irq_work)
>  {
>  	struct bpf_task_work_ctx *ctx = container_of(irq_work, struct bpf_task_work_ctx, irq_work);
> +	struct task_struct *task;
>  	enum bpf_task_work_state state;
>  	int err;
>  
> @@ -4496,7 +4498,8 @@ static void bpf_task_work_irq(struct irq_work *irq_work)
>  		return;
>  	}
>  
> -	err = task_work_add(ctx->task, &ctx->work, ctx->mode);
> +	task = get_task_struct(ctx->task);
> +	err = task_work_add(task, &ctx->work, ctx->mode);
>  	if (err) {
>  		bpf_task_work_ctx_reset(ctx);
>  		/*
> @@ -4505,19 +4508,20 @@ static void bpf_task_work_irq(struct irq_work *irq_work)
>  		 */
>  		(void)cmpxchg(&ctx->state, BPF_TW_SCHEDULING, BPF_TW_STANDBY);
>  		bpf_task_work_ctx_put(ctx);
> +		put_task_struct(task);
>  		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 have completed and reset ctx->task already. Keep the
> +	 * captured task alive until this invocation is done with cancellation.
> +	 * The RCU read-side section keeps ctx memory alive until then.
>  	 */
>  	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_cancel(ctx, task); /* clean up if we switched into FREED state */
> +
> +	put_task_struct(task);
>  }
>  
>  static struct bpf_task_work_ctx *bpf_task_work_fetch_ctx(struct bpf_task_work *tw,
> @@ -4537,6 +4541,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) {
> @@ -4579,6 +4584,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);
>  	}
> +	/*
> +	 * The callback or add-failure path can publish STANDBY before the
> +	 * scheduling irq_work returns. Do not let a new round reuse the context
> +	 * until BUSY clears. Otherwise the old handler can consume the new
> +	 * round's SCHEDULING state and publish a false SCHEDULED state, and two
> +	 * invocations can overlap on different CPUs.
> +	 */
> +	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
> @@ -4628,7 +4645,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;
> @@ -4777,7 +4793,8 @@ static void bpf_task_work_cancel_scheduled(struct irq_work *irq_work)
>  {
>  	struct bpf_task_work_ctx *ctx = container_of(irq_work, struct bpf_task_work_ctx, irq_work);
>  
> -	bpf_task_work_cancel(ctx); /* this might put task_work callback's ref */
> +	/* SCHEDULED -> FREED keeps ctx->task stable until this runs. */
> +	bpf_task_work_cancel(ctx, ctx->task); /* this might put task_work callback's ref */
>  	bpf_task_work_ctx_put(ctx); /* and here we put map's own ref that was transferred to us */
>  }
>  
> 

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

[Alexei Starovoitov]

On Tue, Sep 22, 2026 at 06:00 PM Yun Lu <luyun_611@163.com> wrote:
> -	err = task_work_add(ctx->task, &ctx->work, ctx->mode);
> +	task = get_task_struct(ctx->task);
> +	err = task_work_add(task, &ctx->work, ctx->mode);

Why?
There should be no need for get_task_struct().
rcu_read_lock() keeps task_struct around.
bpf_task_release() is put_task_struct_rcu_user(), so the task is
freed after rcu gp and guard(rcu)() covers the whole handler.
  task = ctx->task;
is enough.
The commit log still says that rcu_read_lock() "does not retain
ctx->task". why?

[...]

> +	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);
> +	}

This is a fix for a different bug.
Pls split it into a separate patch and describe the sequence in
its commit log the way it's done for the first one.
"publishing a false SCHEDULED state" doesn't say what breaks.

pw-bot: cr

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

[luyun]

在 2026/9/24 07:04, Alexei Starovoitov 写道:
> On Tue, Sep 22, 2026 at 06:00 PM Yun Lu <luyun_611@163.com> wrote:
>> -	err = task_work_add(ctx->task, &ctx->work, ctx->mode);
>> +	task = get_task_struct(ctx->task);
>> +	err = task_work_add(task, &ctx->work, ctx->mode);
> Why?
> There should be no need for get_task_struct().
> rcu_read_lock() keeps task_struct around.
> bpf_task_release() is put_task_struct_rcu_user(), so the task is
> freed after rcu gp and guard(rcu)() covers the whole handler.
>    task = ctx->task;
> is enough.
> The commit log still says that rcu_read_lock() "does not retain
> ctx->task". why?

Hi, Alexei

Thanks for your reviewing.

You are right.
I added get_task_struct() to keep the task valid for the handler's post-add
cancellation, but overlooked that guard(rcu)() already covers the handler
and bpf_task_release() defers freeing the task past the RCU grace period.
v4  will drop the get/put_task_struct and captures the pointer, also
correct the commit log.


>
> [...]
>
>> +	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);
>> +	}
> This is a fix for a different bug.
> Pls split it into a separate patch and describe the sequence in
> its commit log the way it's done for the first one.
> "publishing a false SCHEDULED state" doesn't say what breaks.

OK, v4 will split the irq_work_is_busy() check into a separate patch,
and describe the interleaving and the concrete failure in commit log.

---

Thanks.

Yun Lu


>
> pw-bot: cr