sched_ext: Convert scx_bpf_cid_override() to __arena array arguments

scx_bpf_cid_override() predates the cid-form arena transition and takes its
arrays as verifier-checked mem+size buffers, forcing scx_qmap to keep the
cpu_to_cid and shard_start arrays in writable bss while the rest of its
state lives in the arena. Unify on arena arguments before cid-form
schedulers start seeing real use.

BPF now translates between BPF and kernel arena addresses for __arena
arguments. Take the arrays as __arena arguments, with the counts passed in
entries. The counts now size the snapshot copies and are bounds-checked
before them.

scx_qmap moves the arrays into struct qmap_arena. As the arena is mmapped at
load, the loader populates them between load and attach instead of before
load.

The arena argument address translation is currently implemented only on
x86-64. Schedulers calling this kfunc load only there for now.

Signed-off-by: Tejun Heo <tj@kernel.org>

1 files changed, 28 insertions, 21 deletions

diff --git a/kernel/sched/ext/cid.c b/kernel/sched/ext/cid.c
index 11fe9be80f1e..39f88deb94bc 100644
--- a/kernel/sched/ext/cid.c
+++ b/kernel/sched/ext/cid.c
@@ -449,10 +449,10 @@ __bpf_kfunc_start_defs();
 
 /**
  * scx_bpf_cid_override - Install an explicit cpu->cid mapping with shard info
- * @cpu_to_cid_src: array of nr_cpu_ids s32 entries (cid for each cpu)
- * @cpu_to_cid_src__sz: must be nr_cpu_ids * sizeof(s32) bytes
- * @shard_start_src: array of first-cid-of-each-shard, strictly increasing from 0
- * @shard_start_src__sz: nr_shards * sizeof(s32) bytes
+ * @cpu_to_cid__arena: array of nr_cpu_ids s32 entries (cid for each cpu)
+ * @cpu_to_cid_cnt: number of entries, must be nr_cpu_ids
+ * @shard_start__arena: array of first-cid-of-each-shard, one entry per shard
+ * @shard_start_cnt: number of shards
  * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
  *
  * May only be called from ops.init_cids() of the root scheduler. Replace the
@@ -464,9 +464,9 @@ __bpf_kfunc_start_defs();
  * (core/LLC/node) is cleared and the shard layout is set from the input. On
  * invalid input, abort the scheduler.
  */
-__bpf_kfunc void scx_bpf_cid_override(const s32 *cpu_to_cid_src, u32 cpu_to_cid_src__sz,
-				       const s32 *shard_start_src, u32 shard_start_src__sz,
-				       const struct bpf_prog_aux *aux)
+__bpf_kfunc void scx_bpf_cid_override(const s32 *cpu_to_cid__arena, u32 cpu_to_cid_cnt,
+				      const s32 *shard_start__arena, u32 shard_start_cnt,
+				      const struct bpf_prog_aux *aux)
 {
 	cpumask_var_t seen __free(free_cpumask_var) = CPUMASK_VAR_NULL;
 	u32 *node_counts __free(kfree) = NULL;
@@ -475,19 +475,28 @@ __bpf_kfunc void scx_bpf_cid_override(const s32 *cpu_to_cid_src, u32 cpu_to_cid_
 	u32 npossible = num_possible_cpus();
 	struct scx_cid_tables *tbls;
 	struct scx_sched *sch;
-	u32 nr_shards;
+	u32 nr_shards = shard_start_cnt;
 	bool alloced;
 	s32 cpu, cid, si;
 
 	/*
 	 * GFP_KERNEL allocs must happen before the rcu read section. Snapshot
-	 * the BPF-supplied arrays so a concurrent map mutation can't change
+	 * the BPF-supplied arrays so a concurrent arena write can't change
 	 * them between validation and use.
+	 *
+	 * The BPF-supplied counts size the snapshots and thus the arena reads.
+	 * Gate the copies on the count bounds, reported below once @sch is
+	 * available. The bounded reads, at most 32KB, stay within the guard
+	 * region that arena fault recovery covers.
 	 */
 	alloced = zalloc_cpumask_var(&seen, GFP_KERNEL);
 	node_counts = kcalloc(nr_node_ids, sizeof(*node_counts), GFP_KERNEL);
-	cpu_to_cid = kmemdup(cpu_to_cid_src, cpu_to_cid_src__sz, GFP_KERNEL);
-	shard_start = kmemdup(shard_start_src, shard_start_src__sz, GFP_KERNEL);
+	if (cpu_to_cid_cnt == nr_cpu_ids)
+		cpu_to_cid = kmemdup(cpu_to_cid__arena, cpu_to_cid_cnt * sizeof(s32),
+				     GFP_KERNEL);
+	if (nr_shards && nr_shards <= npossible)
+		shard_start = kmemdup(shard_start__arena, nr_shards * sizeof(s32),
+				      GFP_KERNEL);
 
 	guard(rcu)();
 
@@ -499,25 +508,23 @@ __bpf_kfunc void scx_bpf_cid_override(const s32 *cpu_to_cid_src, u32 cpu_to_cid_
 	lockdep_assert_held(&scx_enable_mutex);
 	tbls = scx_cid_tables;
 
-	if (!alloced || !node_counts || !cpu_to_cid || !shard_start) {
-		scx_error(sch, "scx_bpf_cid_override: allocation failed");
+	if (cpu_to_cid_cnt != nr_cpu_ids) {
+		scx_error(sch, "scx_bpf_cid_override: cpu_to_cid expected %u entries, got %u",
+			  nr_cpu_ids, cpu_to_cid_cnt);
 		return;
 	}
 
-	if (cpu_to_cid_src__sz != nr_cpu_ids * sizeof(s32)) {
-		scx_error(sch, "scx_bpf_cid_override: cpu_to_cid expected %zu bytes, got %u",
-			  nr_cpu_ids * sizeof(s32), cpu_to_cid_src__sz);
+	if (!nr_shards || nr_shards > npossible) {
+		scx_error(sch, "scx_bpf_cid_override: invalid shard_start count %u",
+			  nr_shards);
 		return;
 	}
 
-	if (!shard_start_src__sz || shard_start_src__sz % sizeof(s32)) {
-		scx_error(sch, "scx_bpf_cid_override: invalid shard_start size %u",
-			  shard_start_src__sz);
+	if (!alloced || !node_counts || !cpu_to_cid || !shard_start) {
+		scx_error(sch, "scx_bpf_cid_override: allocation failed");
 		return;
 	}
 
-	nr_shards = shard_start_src__sz / sizeof(s32);
-
 	/* validate shard_start[]: starts at 0, strictly increasing, in range */
 	if (shard_start[0] != 0) {
 		scx_error(sch, "scx_bpf_cid_override: shard_start[0] must be 0, got %d",