[PATCH 3/3] sched_ext: Guard the cid kfuncs against unallocated cid tables

[Tejun Heo <tj@kernel.org>]

The cid tables are allocated by the first enable's scx_cid_init(), which
runs after scx_alloc_and_add_sched() has published ops->priv. A TRACING or
SYSCALL program associated with the enabling struct_ops map gets a non-NULL
sched from scx_prog_sched() as soon as ops->priv is set, so a cid kfunc
called in that window dereferences the still-NULL table pointer. Only the
first enable since boot is exposed as the tables are never freed.

scx_bpf_this_cid() and scx_bpf_task_cid() already handle the window by
testing the table pointer. Do the same in scx_cid_to_cpu(),
scx_cpu_to_cid() and scx_bpf_cid_topo(), returning -EINVAL / all-(-1) topo
as before any scheduler is enabled. __scx_cid_to_cpu() and
__scx_cpu_to_cid() stay unchecked for callers with the tables guaranteed
allocated - ops invocations on a live scheduler and the enable path itself.

Fixes: e9b55af47edf ("sched_ext: Add topological CPU IDs (cids)")
Signed-off-by: Tejun Heo <tj@kernel.org>
---
 kernel/sched/ext/cid.c |  5 +++--
 kernel/sched/ext/cid.h | 23 +++++++++++++++--------
 2 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/kernel/sched/ext/cid.c b/kernel/sched/ext/cid.c
index 9dfd242be34f..699cb8db228d 100644
--- a/kernel/sched/ext/cid.c
+++ b/kernel/sched/ext/cid.c
@@ -875,17 +875,18 @@ bool scx_cmask_empty(const struct scx_cmask *m)
 __bpf_kfunc void scx_bpf_cid_topo(s32 cid, struct scx_cid_topo *out__uninit,
 				  const struct bpf_prog_aux *aux)
 {
+	struct scx_cid_topo *topo = READ_ONCE(scx_cid_topo);
 	struct scx_sched *sch;
 
 	guard(rcu)();
 
 	sch = scx_prog_sched(aux);
-	if (unlikely(!sch) || !cid_valid(sch, cid)) {
+	if (unlikely(!sch) || !cid_valid(sch, cid) || unlikely(!topo)) {
 		*out__uninit = SCX_CID_TOPO_NEG;
 		return;
 	}
 
-	*out__uninit = READ_ONCE(scx_cid_topo)[cid];
+	*out__uninit = topo[cid];
 }
 
 __bpf_kfunc_end_defs();
diff --git a/kernel/sched/ext/cid.h b/kernel/sched/ext/cid.h
index b36a1a28eac8..1498efb81549 100644
--- a/kernel/sched/ext/cid.h
+++ b/kernel/sched/ext/cid.h
@@ -90,9 +90,10 @@ static inline bool cid_valid(struct scx_sched *sch, s32 cid)
  * __scx_cid_to_cpu - Unchecked cid->cpu table lookup
  * @cid: cid to look up. Must be in [0, num_possible_cpus()).
  *
- * Intended for callsites that have already validated @cid and that hold a
- * non-NULL @sch from scx_prog_sched() - a live sched implies the table has
- * been allocated, so no NULL check is needed here.
+ * Intended for callsites that have already validated @cid and where the
+ * tables are guaranteed allocated - ops invocations on a live scheduler or
+ * the enable path itself. Prog-facing kfuncs, which can run while the first
+ * enable is still allocating the tables, use the checked wrappers instead.
  */
 static inline s32 __scx_cid_to_cpu(s32 cid)
 {
@@ -119,13 +120,17 @@ static inline s32 __scx_cpu_to_cid(s32 cpu)
  * Return the cpu for @cid or a negative errno on failure. Invalid cid triggers
  * scx_error() on @sch. The cid arrays are allocated on first scheduler enable
  * and never freed, so the returned cpu is stable for the lifetime of the loaded
- * scheduler.
+ * scheduler. Return -EINVAL without triggering scx_error() if the tables are
+ * not allocated yet, which a prog-facing kfunc can observe while racing the
+ * first enable.
  */
 static inline s32 scx_cid_to_cpu(struct scx_sched *sch, s32 cid)
 {
-	if (!cid_valid(sch, cid))
+	s16 *tbl = READ_ONCE(scx_cid_to_cpu_tbl);
+
+	if (!cid_valid(sch, cid) || unlikely(!tbl))
 		return -EINVAL;
-	return __scx_cid_to_cpu(cid);
+	return tbl[cid];
 }
 
 /**
@@ -138,9 +143,11 @@ static inline s32 scx_cid_to_cpu(struct scx_sched *sch, s32 cid)
  */
 static inline s32 scx_cpu_to_cid(struct scx_sched *sch, s32 cpu)
 {
-	if (!scx_cpu_valid(sch, cpu, NULL))
+	s16 *tbl = READ_ONCE(scx_cpu_to_cid_tbl);
+
+	if (!scx_cpu_valid(sch, cpu, NULL) || unlikely(!tbl))
 		return -EINVAL;
-	return __scx_cpu_to_cid(cpu);
+	return tbl[cpu];
 }
 
 /**
-- 
2.55.0