sched_ext: Add scx_cmask_ref for validated arena cmask access

kfuncs taking struct scx_cmask * from BPF arena memory have two problems. The pointer can be any value the BPF prog hands in, and the header (@base, @nr_cids, @alloc_words) can be mutated by the prog concurrently with kernel access. Add scx_cmask_ref, a validated handle. _init() normalizes the input pointer into the arena's kern_vm range via scx_arena_to_kaddr() and snapshots the header, rejecting a range outside the machine or a nr_cids whose words exceed the declared @alloc_words. Downstream sizing uses the snapshot, not the live header. _shard() reads slices while _or() and _copy() write back, all bounded by the snapshot. No callers yet. struct scx_cmask's bits[] carried __counted_by(alloc_words), so UBSAN_BOUNDS and FORTIFY_SOURCE bound accesses to the array. That bound is read from @alloc_words at the access. For an arena cmask @alloc_words is BPF-writable. A prog that sets it larger than the real allocation makes the check pass on a genuine overrun, so the annotation catches nothing, and it only runs under those debug configs. Drop it - _init() validates @alloc_words explicitly, and kernel-owned cmasks set it themselves. v2: Validate @alloc_words in _init(), drop __counted_by. (Andrea, sashiko AI)

Tejun Heo committed 4 months ago

Commit cc7c2541 parent 8dba3bb3 files · +178 −1Browse files

Showing the changes to kernel (3 of 3 files). Show all

Unified · Split

3 files

kernel/sched/ext/cid.c +133 −0
kernel/sched/ext/cid.h +7 −0
kernel/sched/ext/types.h +38 −1

kernel/sched/ext/cid.c+133−0

@@ -641,6 +641,12 @@ enum cmask_op2 { · 640 unmodified lines
641	641	/* predicates - short-circuit when the per-word result is true */
642	642	CMASK_OP2_SUBSET,
643	643	CMASK_OP2_INTERSECTS,
	644	+ /*
	645	+ * @a is a BPF-arena cmask. Words on @a use READ_ONCE/WRITE_ONCE since
	646	+ * BPF may read/write concurrently. See scx_cmask_ref_or() / _copy().
	647	+ */
	648	+ CMASK_OP2_REF_OR,
	649	+ CMASK_OP2_REF_COPY,
644	650	};
645	651
646	652	static __always_inline bool cmask_op2_is_pred(const enum cmask_op2 op)
@@ -669,6 +675,12 @@ static __always_inline bool cmask_word_op2(u64 av, const u64 bp, u64 mask, · 22 unmodified lines
669	675	return (READ_ONCE(bp) & ~READ_ONCE(av)) & mask;
670	676	case CMASK_OP2_INTERSECTS:
671	677	return (READ_ONCE(av) & READ_ONCE(bp)) & mask;
	678	+ case CMASK_OP2_REF_OR:
	679	+ WRITE_ONCE(av, READ_ONCE(av) \| (READ_ONCE(*bp) & mask));
	680	+ return false;
	681	+ case CMASK_OP2_REF_COPY:
	682	+ WRITE_ONCE(av, (READ_ONCE(av) & ~mask) \| (READ_ONCE(*bp) & mask));
	683	+ return false;
672	684	}
673	685	unreachable();
674	686	}
@@ -899,6 +911,127 @@ static const struct btf_kfunc_id_set scx_kfunc_set_cid = { · 224 unmodified lines
899	911	.set = &scx_kfunc_ids_cid,
900	912	};
901	913
	914	+/**
	915	+ * scx_cmask_ref_init - Bind a scx_cmask_ref to a BPF-arena cmask
	916	+ * @sch: scheduler whose arena hosts @src
	917	+ * @src: BPF-supplied cmask pointer
	918	+ * @ref: output ref
	919	+ *
	920	+ * Snapshot @src's @base, @nr_cids and @alloc_words. The snapshot is necessary
	921	+ * because BPF may mutate the live header asynchronously.
	922	+ *
	923	+ * Return 0 on success, -EINVAL if the range is out of bounds or @alloc_words
	924	+ * doesn't cover it.
	925	+ */
	926	+int scx_cmask_ref_init(struct scx_sched sch, const struct scx_cmask src,
	927	+ struct scx_cmask_ref *ref)
	928	+{
	929	+ struct scx_cmask *kern_src = scx_arena_to_kaddr(sch, src);
	930	+ u32 base, nr_cids, alloc_words, npossible = num_possible_cpus();
	931	+
	932	+ base = READ_ONCE(kern_src->base);
	933	+ nr_cids = READ_ONCE(kern_src->nr_cids);
	934	+ alloc_words = READ_ONCE(kern_src->alloc_words);
	935	+
	936	+ if (unlikely(base >= npossible \|\| nr_cids > npossible - base \|\|
	937	+ SCX_CMASK_NR_WORDS(nr_cids) > alloc_words))
	938	+ return -EINVAL;
	939	+
	940	+ ref->sch = sch;
	941	+ ref->src = kern_src;
	942	+ ref->base = base;
	943	+ ref->nr_cids = nr_cids;
	944	+
	945	+ ref->shard_first = scx_cid_to_shard[base];
	946	+ if (likely(nr_cids))
	947	+ ref->shard_end = scx_cid_to_shard[base + nr_cids - 1] + 1;
	948	+ else
	949	+ ref->shard_end = ref->shard_first;
	950	+
	951	+ return 0;
	952	+}
	953	+
	954	+/**
	955	+ * scx_cmask_ref_shard - Read one shard from @ref into @out
	956	+ * @ref: validated ref
	957	+ * @shard_idx: target shard, in [@ref->shard_first, @ref->shard_end)
	958	+ * @out: output cmask whose @out->alloc_words must hold the shard
	959	+ *
	960	+ * Set @out to the intersection of @ref's range with @shard_idx's cid range,
	961	+ * with bits[] read from @ref->src via READ_ONCE. Empty intersection sets
	962	+ * @out->nr_cids to 0. scx_error()s on @ref's sched if @out can't hold the
	963	+ * shard.
	964	+ */
	965	+void scx_cmask_ref_shard(const struct scx_cmask_ref *ref, s32 shard_idx,
	966	+ struct scx_cmask *out)
	967	+{
	968	+ const struct scx_cid_shard *shard = &scx_cid_shard_ranges[shard_idx];
	969	+ u32 shard_base = shard->base_cid;
	970	+ u32 shard_end = shard_base + shard->nr_cids;
	971	+ u32 isect_base, isect_end, nr_words, src_off, wi;
	972	+ u64 head_mask, tail_mask;
	973	+
	974	+ isect_base = max(ref->base, shard_base);
	975	+ isect_end = min(ref->base + ref->nr_cids, shard_end);
	976	+
	977	+ if (isect_base >= isect_end) {
	978	+ out->base = shard_base;
	979	+ out->nr_cids = 0;
	980	+ return;
	981	+ }
	982	+
	983	+ nr_words = ((isect_end - 1) / 64) - (isect_base / 64) + 1;
	984	+ if (nr_words > out->alloc_words) {
	985	+ scx_error(ref->sch, "scx_cmask_ref_shard: out alloc_words=%u < %u for shard %d",
	986	+ out->alloc_words, nr_words, shard_idx);
	987	+ out->base = shard_base;
	988	+ out->nr_cids = 0;
	989	+ return;
	990	+ }
	991	+
	992	+ out->base = isect_base;
	993	+ out->nr_cids = isect_end - isect_base;
	994	+ src_off = (isect_base / 64) - (ref->base / 64);
	995	+
	996	+ for (wi = 0; wi < nr_words; wi++)
	997	+ out->bits[wi] = READ_ONCE(ref->src->bits[src_off + wi]);
	998	+
	999	+ head_mask = GENMASK_U64(63, isect_base & 63);
	1000	+ out->bits[0] &= head_mask;
	1001	+ tail_mask = GENMASK_U64((isect_end - 1) & 63, 0);
	1002	+ out->bits[nr_words - 1] &= tail_mask;
	1003	+}
	1004	+
	1005	+/**
	1006	+ * scx_cmask_ref_or - OR @src into the arena cmask referenced by @ref
	1007	+ * @ref: validated ref
	1008	+ * @src: stable kernel cmask
	1009	+ *
	1010	+ * Bits inside the intersection of @ref's snapshotted range with @src's range
	1011	+ * are OR'd into @ref->src and bits outside are left unchanged. Stores on
	1012	+ * @ref->src use WRITE_ONCE since BPF may read/write concurrently.
	1013	+ */
	1014	+void scx_cmask_ref_or(const struct scx_cmask_ref ref, const struct scx_cmask src)
	1015	+{
	1016	+ cmask_walk_op2(ref->src->bits, ref->base, ref->nr_cids,
	1017	+ src->bits, src->base, src->nr_cids, CMASK_OP2_REF_OR);
	1018	+}
	1019	+
	1020	+/**
	1021	+ * scx_cmask_ref_copy - Copy @src into the arena cmask referenced by @ref
	1022	+ * @ref: validated ref
	1023	+ * @src: stable kernel cmask
	1024	+ *
	1025	+ * Bits inside the intersection of @ref's snapshotted range with @src's range
	1026	+ * take @src's values and bits outside are left unchanged. Stores on @ref->src
	1027	+ * use WRITE_ONCE since BPF may read/write concurrently.
	1028	+ */
	1029	+void scx_cmask_ref_copy(const struct scx_cmask_ref ref, const struct scx_cmask src)
	1030	+{
	1031	+ cmask_walk_op2(ref->src->bits, ref->base, ref->nr_cids,
	1032	+ src->bits, src->base, src->nr_cids, CMASK_OP2_REF_COPY);
	1033	+}
	1034	+
902	1035	int scx_cid_kfunc_init(void)
903	1036	{
904	1037	return register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS, &scx_kfunc_set_init_cids) ?:

kernel/sched/ext/cid.h+7−0

@@ -293,4 +293,11 @@ static inline s32 scx_cpu_ret(struct scx_sched *sch, s32 cpu_or_cid) · 292 unmodified lines
293	293	return scx_cid_to_cpu(sch, cpu_or_cid);
294	294	}
295	295
	296	+int scx_cmask_ref_init(struct scx_sched sch, const struct scx_cmask src,
	297	+ struct scx_cmask_ref *ref);
	298	+void scx_cmask_ref_shard(const struct scx_cmask_ref *ref, s32 shard_idx,
	299	+ struct scx_cmask *out);
	300	+void scx_cmask_ref_or(const struct scx_cmask_ref ref, const struct scx_cmask src);
	301	+void scx_cmask_ref_copy(const struct scx_cmask_ref ref, const struct scx_cmask src);
	302	+
296	303	#endif /* _KERNEL_SCHED_EXT_CID_H */

kernel/sched/ext/types.h+38−1

@@ -116,7 +116,7 @@ struct scx_cmask { · 115 unmodified lines
116	116	u32 base;
117	117	u32 nr_cids;
118	118	u32 alloc_words;
119		− u64 bits[] __counted_by(alloc_words);
	119	+ u64 bits[];
120	120	};
121	121
122	122	/*
@@ -172,4 +172,41 @@ struct scx_cmask { · 49 unmodified lines
172	172	#define SCX_CMASK_DEFINE_SHARD(NAME, BASE, NR_CIDS) \
173	173	__SCX_CMASK_DEFINE(NAME, BASE, NR_CIDS, SCX_CID_SHARD_MAX_CPUS)
174	174
	175	+/*
	176	+ * scx_cmask_ref: validated reference to a BPF-arena cmask.
	177	+ *
	178	+ * scx_cmask_ref_init() normalizes the pointer into the arena and snapshots
	179	+ * @base/@nr_cids. The snapshot is what downstream code uses for sizing - the
	180	+ * live header can be mutated concurrently by BPF.
	181	+ *
	182	+ * scx_cmask_ref_shard() reads one shard into a cmask. scx_cmask_ref_or() and
	183	+ * scx_cmask_ref_copy() write back into the referenced arena cmask, bounded by
	184	+ * the snapshot.
	185	+ *
	186	+ * Typical input use:
	187	+ *
	188	+ * struct scx_cmask_ref ref;
	189	+ * SCX_CMASK_DEFINE(shard, 0, SCX_CID_SHARD_MAX_CPUS);
	190	+ * s32 idx, ret;
	191	+ *
	192	+ * ret = scx_cmask_ref_init(sch, src, &ref);
	193	+ * if (ret < 0)
	194	+ * return ret;
	195	+ *
	196	+ * for (idx = ref.shard_first; idx < ref.shard_end; idx++) {
	197	+ * scx_cmask_ref_shard(&ref, idx, shard);
	198	+ * if (!shard->nr_cids)
	199	+ * continue;
	200	+ * ... use idx and shard ...
	201	+ * }
	202	+ */
	203	+struct scx_cmask_ref {
	204	+ struct scx_sched *sch;
	205	+ struct scx_cmask *src;
	206	+ u32 base;
	207	+ u32 nr_cids;
	208	+ s32 shard_first;
	209	+ s32 shard_end;
	210	+};
	211	+
175	212	#endif /* _KERNEL_SCHED_EXT_TYPES_H */