[PATCH 0/2] sched_ext: two more context-safety fixes found in the NMI kfunc audit

[PATCH 0/2] sched_ext: two more context-safety fixes found in the NMI kfunc audit

[Wanwu Li]

Continuing the audit that the sashiko bot kicked off on the NMI-reject
series, I went through every kfunc exposed to BPF_PROG_TYPE_TRACING
(the any / idle / cid context-filter sets). Two more context-safety
issues came out of it; this pair addresses both.

 - Patch 1 fixes the scx_locked_rq() class (three "any"-category kfuncs
   read scx_locked_rq() and take an unsafe fast path on a non-NULL
   return from NMI). scx_locked_rq() gains an in_nmi() test at the
   source that returns NULL from NMI, so the three callers fall onto
   their unlocked paths. This supersedes guarding them individually.

 - Patch 2 fixes an IRQ re-entrancy race on the per-CPU
   per_cpu_unvisited nodemask scratch in the idle path. The search
   scratch is protected only by preempt_disable(), which doesn't mask
   IRQs; switch to irqsave instead. The NMI case is left as-is (a
   misuse with no legitimate use and no crash risk).

Wanwu Li (2):
  sched_ext: Make scx_locked_rq() return NULL from NMI
  sched_ext: Protect the idle-search scratch nodemask with irqsave

 kernel/sched/ext/idle.c     | 14 ++++++++++++--
 kernel/sched/ext/internal.h |  9 +++++++++
 2 files changed, 21 insertions(+), 2 deletions(-)

-- 
2.25.1

[PATCH 1/2] sched_ext: Make scx_locked_rq() return NULL from NMI

[Wanwu Li]

scx_locked_rq() reads the per-CPU scx_locked_rq_state, which tracks the
rq locked by the context running on the CPU. Tracing progs can call
kfuncs from NMI, and an NMI interrupts - rather than replaces - the
context that set scx_locked_rq_state, so a non-NULL read from NMI
falsely tells the caller that it holds the interrupted context's rq
lock.

Three "any"-category kfuncs read scx_locked_rq() on their success path
and take an unsafe fast path on a non-NULL return:

 - scx_bpf_task_set_slice() writes p->scx.slice directly, racing
   update_curr_scx()'s non-atomic read-modify-write of the same field.
 - scx_bpf_dsq_nr_queued() resolves %SCX_DSQ_LOCAL to
   (scx_locked_rq() ?: this_rq()) and can report the interrupted
   context's local DSQ length instead of the caller's.
 - scx_bpf_locked_rq() hands the interrupted context's rq to the BPF
   program, which may then operate on it as if it owned the rq lock.

Make scx_locked_rq() return NULL from NMI so that all three take their
unlocked paths: scx_bpf_task_set_slice() stashes the request into the
atomic p->scx.slice_oob for application under the rq lock,
scx_bpf_dsq_nr_queued() falls back to this_rq(), and
scx_bpf_locked_rq() reports an error and aborts the scheduler through
the NMI-safe exit path.

The kfuncs that take scheduler locks reject NMI calls through
scx_kf_allowed_ctx() before reaching scx_locked_rq(), and the internal
callers only run from struct_ops callbacks, which never run in NMI, so
no other caller is affected.

Suggested-by: Tejun Heo <tj@kernel.org>
Link: https://lore.kernel.org/r/d84b31727f04e1ed0d40042ba1c09e61@kernel.org
Signed-off-by: Wanwu Li <liwanwu@kylinos.cn>
---
 kernel/sched/ext/internal.h | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
index a67277b0fee6..809e0ee0fd5f 100644
--- a/kernel/sched/ext/internal.h
+++ b/kernel/sched/ext/internal.h
@@ -2218,6 +2218,15 @@ static inline void scx_schedule_reenq_local(struct rq *rq, u64 reenq_flags)
  */
 static inline struct rq *scx_locked_rq(void)
 {
+	/*
+	 * Tracing progs can call kfuncs from NMI. scx_locked_rq_state tracks
+	 * the rq locked by the interrupted context, so a non-NULL read from
+	 * NMI would falsely claim its lock. Return NULL from NMI so that
+	 * callers take their unlocked paths.
+	 */
+	if (unlikely(in_nmi()))
+		return NULL;
+
 	return __this_cpu_read(scx_locked_rq_state);
 }
 
-- 
2.25.1

[PATCH 2/2] sched_ext: Protect the idle-search scratch nodemask with irqsave

[Wanwu Li]

pick_idle_cpu_from_online_nodes() uses the per-CPU per_cpu_unvisited
nodemask as scratch while walking the online nodes, protected only by
preempt_disable(). preempt_disable() does not mask IRQs and the idle
kfuncs are callable from IRQ-enabled contexts, so a nested invocation
on the same CPU can overwrite the mask with nodes_copy() while the
interrupted invocation is still iterating it, leading to a wrong node
traversal and a wrong idle CPU pick.

Switch to irqsave so a nested invocation can't run on the same CPU.
A stack-allocated nodemask would also close the race, but that would
enlarge the diff to fix a race that is already rare (per-node idle,
CONFIG_NUMA and a cross-node search all at once); irqsave is the
minimal fix for the context that actually triggers it. The NMI case is
deliberately not addressed: there is no legitimate reason to
call pick_idle from NMI and doing so poses no crash risk, so such a
caller is on its own.

Suggested-by: Tejun Heo <tj@kernel.org>
Link: https://lore.kernel.org/r/d84b31727f04e1ed0d40042ba1c09e61@kernel.org
Signed-off-by: Wanwu Li <liwanwu@kylinos.cn>
---
 kernel/sched/ext/idle.c | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/kernel/sched/ext/idle.c b/kernel/sched/ext/idle.c
index d2973fb3af6d..93e2e0b2d1f8 100644
--- a/kernel/sched/ext/idle.c
+++ b/kernel/sched/ext/idle.c
@@ -153,7 +153,18 @@ static s32 pick_idle_cpu_from_online_nodes(const struct cpumask *cpus_allowed, i
 	nodemask_t *unvisited;
 	s32 cpu = -EBUSY;
 
-	preempt_disable();
+	/*
+	 * @per_cpu_unvisited is per-CPU scratch and the idle kfuncs can be
+	 * called from IRQ-enabled contexts, so mask IRQs to keep a nested
+	 * invocation from clobbering the mask an outer invocation is still
+	 * iterating.
+	 *
+	 * NMI nesting is not handled: there is no legitimate reason to call
+	 * pick_idle from NMI and doing so poses no crash risk, so such a
+	 * caller is on its own.
+	 */
+	guard(irqsave)();
+
 	unvisited = this_cpu_ptr(&per_cpu_unvisited);
 
 	/*
@@ -183,7 +194,6 @@ static s32 pick_idle_cpu_from_online_nodes(const struct cpumask *cpus_allowed, i
 		if (cpu >= 0)
 			break;
 	}
-	preempt_enable();
 
 	return cpu;
 }
-- 
2.25.1

Re: [PATCH 0/2] sched_ext: two more context-safety fixes found in the NMI kfunc audit

[Tejun Heo]

> Wanwu Li (2):
>   sched_ext: Make scx_locked_rq() return NULL from NMI
>   sched_ext: Protect the idle-search scratch nodemask with irqsave

Applied 1-2 to sched_ext/for-7.4.

Thanks.

--
tejun