[RFC PATCH bpf-next 0/8] bpf, sched_ext: Make kfunc filters support struct_ops context to reduce runtime overhead

[RFC PATCH bpf-next 0/8] bpf, sched_ext: Make kfunc filters support struct_ops context to reduce runtime overhead

[Juntong Deng]

This patch series makes kfunc filters support the use of struct_ops
context information and adds corresponding filters for various kfunc
sets in SCX.

After improving kfunc filters, SCX no longer needs the mask-based
runtime kfuncs call restriction, so this patch removes the mask-based
runtime restriction and updates the corresponding test case.

I added *st_ops as part of the context information to avoid kfuncs being
incorrectly blocked when used in non-SCX scenarios where the member
offsets would have a different meaning (not sure if this is necessary).

Note that this is an RFC patch series and I am not sure if I
misunderstood something or broke something.

If I got something wrong please let me know.

Known issues:

scx_bpf_dsq_move appears in both scx_kfunc_ids_dispatch and
scx_kfunc_ids_unlocked.

This results in scx_bpf_dsq_move being incorrectly blocked by
the filter of scx_kfunc_ids_dispatch in the 'init' context.

The scx_qmap sample program fails because of this.

In the filter-based restrictions, each kfunc can only appear in
one kfunc set.

A possible solution is to add a scx_kfunc_set_unlocked_dispatch
kfunc set.

Signed-off-by: Juntong Deng <juntong.deng@outlook.com>

Juntong Deng (8):
  bpf: Add struct_ops context information to struct bpf_prog_aux
  sched_ext: Add filter for scx_kfunc_ids_select_cpu
  sched_ext: Add filter for scx_kfunc_ids_enqueue_dispatch
  sched_ext: Add filter for scx_kfunc_ids_dispatch
  sched_ext: Add filter for scx_kfunc_ids_cpu_release
  sched_ext: Add filter for scx_kfunc_ids_unlocked
  sched_ext: Removed mask-based runtime restrictions on calling kfuncs
    in different contexts
  selftests/sched_ext: Update enq_select_cpu_fails to adapt to
    struct_ops context filter

 include/linux/bpf.h                           |   2 +
 include/linux/sched/ext.h                     |  24 --
 kernel/bpf/verifier.c                         |   8 +-
 kernel/sched/ext.c                            | 354 +++++++++---------
 .../sched_ext/enq_select_cpu_fails.c          |  35 +-
 5 files changed, 184 insertions(+), 239 deletions(-)

-- 
2.39.5

[RFC PATCH bpf-next 1/8] bpf: Add struct_ops context information to struct bpf_prog_aux

[Juntong Deng]

This patch adds struct_ops context information to struct bpf_prog_aux.

This context information will be used in the kfunc filter.

Currently the added context information includes struct_ops member
offset and a pointer to struct bpf_struct_ops.

Signed-off-by: Juntong Deng <juntong.deng@outlook.com>
---
 include/linux/bpf.h   | 2 ++
 kernel/bpf/verifier.c | 8 ++++++--
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index f3f50e29d639..e06348a59dcf 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -1503,6 +1503,7 @@ struct bpf_prog_aux {
 	u32 real_func_cnt; /* includes hidden progs, only used for JIT and freeing progs */
 	u32 func_idx; /* 0 for non-func prog, the index in func array for func prog */
 	u32 attach_btf_id; /* in-kernel BTF type id to attach to */
+	u32 attach_st_ops_member_off;
 	u32 ctx_arg_info_size;
 	u32 max_rdonly_access;
 	u32 max_rdwr_access;
@@ -1547,6 +1548,7 @@ struct bpf_prog_aux {
 #endif
 	struct bpf_ksym ksym;
 	const struct bpf_prog_ops *ops;
+	const struct bpf_struct_ops *st_ops;
 	struct bpf_map **used_maps;
 	struct mutex used_maps_mutex; /* mutex for used_maps and used_map_cnt */
 	struct btf_mod_pair *used_btfs;
diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
index 9971c03adfd5..2dee3fd190a4 100644
--- a/kernel/bpf/verifier.c
+++ b/kernel/bpf/verifier.c
@@ -22384,7 +22384,7 @@ static int check_struct_ops_btf_id(struct bpf_verifier_env *env)
 	const struct bpf_struct_ops *st_ops;
 	const struct btf_member *member;
 	struct bpf_prog *prog = env->prog;
-	u32 btf_id, member_idx;
+	u32 btf_id, member_idx, member_off;
 	struct btf *btf;
 	const char *mname;
 	int err;
@@ -22435,7 +22435,8 @@ static int check_struct_ops_btf_id(struct bpf_verifier_env *env)
 		return -EINVAL;
 	}
 
-	err = bpf_struct_ops_supported(st_ops, __btf_member_bit_offset(t, member) / 8);
+	member_off = __btf_member_bit_offset(t, member) / 8;
+	err = bpf_struct_ops_supported(st_ops, member_off);
 	if (err) {
 		verbose(env, "attach to unsupported member %s of struct %s\n",
 			mname, st_ops->name);
@@ -22463,6 +22464,9 @@ static int check_struct_ops_btf_id(struct bpf_verifier_env *env)
 	prog->aux->ctx_arg_info_size =
 		st_ops_desc->arg_info[member_idx].cnt;
 
+	prog->aux->st_ops = st_ops;
+	prog->aux->attach_st_ops_member_off = member_off;
+
 	prog->aux->attach_func_proto = func_proto;
 	prog->aux->attach_func_name = mname;
 	env->ops = st_ops->verifier_ops;
-- 
2.39.5

[RFC PATCH bpf-next 2/8] sched_ext: Add filter for scx_kfunc_ids_select_cpu

[Juntong Deng]

This patch adds filter for scx_kfunc_ids_select_cpu.

The kfuncs in the scx_kfunc_ids_select_cpu set can be used in select_cpu
and other rq-locked operations.

Signed-off-by: Juntong Deng <juntong.deng@outlook.com>
---
 kernel/sched/ext.c | 42 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
index 8857c0709bdd..c92949aa23f6 100644
--- a/kernel/sched/ext.c
+++ b/kernel/sched/ext.c
@@ -6401,9 +6401,51 @@ BTF_KFUNCS_START(scx_kfunc_ids_select_cpu)
 BTF_ID_FLAGS(func, scx_bpf_select_cpu_dfl, KF_RCU)
 BTF_KFUNCS_END(scx_kfunc_ids_select_cpu)
 
+static int scx_kfunc_ids_other_rqlocked_filter(const struct bpf_prog *prog, u32 kfunc_id)
+{
+	u32 moff = prog->aux->attach_st_ops_member_off;
+
+	if (moff == offsetof(struct sched_ext_ops, runnable) ||
+	    moff == offsetof(struct sched_ext_ops, dequeue) ||
+	    moff == offsetof(struct sched_ext_ops, stopping) ||
+	    moff == offsetof(struct sched_ext_ops, quiescent) ||
+	    moff == offsetof(struct sched_ext_ops, yield) ||
+	    moff == offsetof(struct sched_ext_ops, cpu_acquire) ||
+	    moff == offsetof(struct sched_ext_ops, running) ||
+	    moff == offsetof(struct sched_ext_ops, core_sched_before) ||
+	    moff == offsetof(struct sched_ext_ops, set_cpumask) ||
+	    moff == offsetof(struct sched_ext_ops, update_idle) ||
+	    moff == offsetof(struct sched_ext_ops, tick) ||
+	    moff == offsetof(struct sched_ext_ops, enable) ||
+	    moff == offsetof(struct sched_ext_ops, set_weight) ||
+	    moff == offsetof(struct sched_ext_ops, disable) ||
+	    moff == offsetof(struct sched_ext_ops, exit_task) ||
+	    moff == offsetof(struct sched_ext_ops, dump_task) ||
+	    moff == offsetof(struct sched_ext_ops, dump_cpu))
+		return 0;
+
+	return -EACCES;
+}
+
+static int scx_kfunc_ids_select_cpu_filter(const struct bpf_prog *prog, u32 kfunc_id)
+{
+	u32 moff;
+
+	if (!btf_id_set8_contains(&scx_kfunc_ids_select_cpu, kfunc_id) ||
+	    prog->aux->st_ops != &bpf_sched_ext_ops)
+		return 0;
+
+	moff = prog->aux->attach_st_ops_member_off;
+	if (moff == offsetof(struct sched_ext_ops, select_cpu))
+		return 0;
+
+	return scx_kfunc_ids_other_rqlocked_filter(prog, kfunc_id);
+}
+
 static const struct btf_kfunc_id_set scx_kfunc_set_select_cpu = {
 	.owner			= THIS_MODULE,
 	.set			= &scx_kfunc_ids_select_cpu,
+	.filter			= scx_kfunc_ids_select_cpu_filter,
 };
 
 static bool scx_dsq_insert_preamble(struct task_struct *p, u64 enq_flags)
-- 
2.39.5

Re: [RFC PATCH bpf-next 2/8] sched_ext: Add filter for scx_kfunc_ids_select_cpu

[Andrea Righi]

Hi Juntong,

On Wed, Feb 05, 2025 at 07:30:14PM +0000, Juntong Deng wrote:
> This patch adds filter for scx_kfunc_ids_select_cpu.
> 
> The kfuncs in the scx_kfunc_ids_select_cpu set can be used in select_cpu
> and other rq-locked operations.

The only function in scx_kfunc_ids_select_cpu is scx_bpf_select_cpu_dfl(),
which should be called exclusively from ops.select_cpu() and not from any
rq-locked ops.

> 
> Signed-off-by: Juntong Deng <juntong.deng@outlook.com>
> ---
>  kernel/sched/ext.c | 42 ++++++++++++++++++++++++++++++++++++++++++
>  1 file changed, 42 insertions(+)
> 
> diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
> index 8857c0709bdd..c92949aa23f6 100644
> --- a/kernel/sched/ext.c
> +++ b/kernel/sched/ext.c
> @@ -6401,9 +6401,51 @@ BTF_KFUNCS_START(scx_kfunc_ids_select_cpu)
>  BTF_ID_FLAGS(func, scx_bpf_select_cpu_dfl, KF_RCU)
>  BTF_KFUNCS_END(scx_kfunc_ids_select_cpu)
>  
> +static int scx_kfunc_ids_other_rqlocked_filter(const struct bpf_prog *prog, u32 kfunc_id)
> +{
> +	u32 moff = prog->aux->attach_st_ops_member_off;
> +
> +	if (moff == offsetof(struct sched_ext_ops, runnable) ||
> +	    moff == offsetof(struct sched_ext_ops, dequeue) ||
> +	    moff == offsetof(struct sched_ext_ops, stopping) ||
> +	    moff == offsetof(struct sched_ext_ops, quiescent) ||
> +	    moff == offsetof(struct sched_ext_ops, yield) ||
> +	    moff == offsetof(struct sched_ext_ops, cpu_acquire) ||
> +	    moff == offsetof(struct sched_ext_ops, running) ||
> +	    moff == offsetof(struct sched_ext_ops, core_sched_before) ||
> +	    moff == offsetof(struct sched_ext_ops, set_cpumask) ||
> +	    moff == offsetof(struct sched_ext_ops, update_idle) ||
> +	    moff == offsetof(struct sched_ext_ops, tick) ||
> +	    moff == offsetof(struct sched_ext_ops, enable) ||
> +	    moff == offsetof(struct sched_ext_ops, set_weight) ||
> +	    moff == offsetof(struct sched_ext_ops, disable) ||
> +	    moff == offsetof(struct sched_ext_ops, exit_task) ||
> +	    moff == offsetof(struct sched_ext_ops, dump_task) ||
> +	    moff == offsetof(struct sched_ext_ops, dump_cpu))
> +		return 0;
> +
> +	return -EACCES;
> +}
> +
> +static int scx_kfunc_ids_select_cpu_filter(const struct bpf_prog *prog, u32 kfunc_id)
> +{
> +	u32 moff;
> +
> +	if (!btf_id_set8_contains(&scx_kfunc_ids_select_cpu, kfunc_id) ||
> +	    prog->aux->st_ops != &bpf_sched_ext_ops)
> +		return 0;
> +
> +	moff = prog->aux->attach_st_ops_member_off;
> +	if (moff == offsetof(struct sched_ext_ops, select_cpu))
> +		return 0;
> +
> +	return scx_kfunc_ids_other_rqlocked_filter(prog, kfunc_id);

So, I think we just need to return -EACCES here.

> +}
> +
>  static const struct btf_kfunc_id_set scx_kfunc_set_select_cpu = {
>  	.owner			= THIS_MODULE,
>  	.set			= &scx_kfunc_ids_select_cpu,
> +	.filter			= scx_kfunc_ids_select_cpu_filter,
>  };
>  
>  static bool scx_dsq_insert_preamble(struct task_struct *p, u64 enq_flags)
> -- 
> 2.39.5
> 

Thanks,
-Andrea

Re: [RFC PATCH bpf-next 2/8] sched_ext: Add filter for scx_kfunc_ids_select_cpu

[Andrea Righi]

On Wed, Feb 05, 2025 at 07:30:14PM +0000, Juntong Deng wrote:
...
> +static int scx_kfunc_ids_other_rqlocked_filter(const struct bpf_prog *prog, u32 kfunc_id)
> +{
> +	u32 moff = prog->aux->attach_st_ops_member_off;
> +
> +	if (moff == offsetof(struct sched_ext_ops, runnable) ||
> +	    moff == offsetof(struct sched_ext_ops, dequeue) ||
> +	    moff == offsetof(struct sched_ext_ops, stopping) ||
> +	    moff == offsetof(struct sched_ext_ops, quiescent) ||
> +	    moff == offsetof(struct sched_ext_ops, yield) ||
> +	    moff == offsetof(struct sched_ext_ops, cpu_acquire) ||
> +	    moff == offsetof(struct sched_ext_ops, running) ||
> +	    moff == offsetof(struct sched_ext_ops, core_sched_before) ||
> +	    moff == offsetof(struct sched_ext_ops, set_cpumask) ||
> +	    moff == offsetof(struct sched_ext_ops, update_idle) ||
> +	    moff == offsetof(struct sched_ext_ops, tick) ||
> +	    moff == offsetof(struct sched_ext_ops, enable) ||
> +	    moff == offsetof(struct sched_ext_ops, set_weight) ||
> +	    moff == offsetof(struct sched_ext_ops, disable) ||
> +	    moff == offsetof(struct sched_ext_ops, exit_task) ||
> +	    moff == offsetof(struct sched_ext_ops, dump_task) ||
> +	    moff == offsetof(struct sched_ext_ops, dump_cpu))
> +		return 0;
> +
> +	return -EACCES;

Actually, do we need this filter at all?

I think the other filters in your patch set should be sufficient to
establish the correct permissions for all kfuncs, as none of them need to
be called from any rq-locked operations. Or am I missing something?

-Andrea

Re: [RFC PATCH bpf-next 2/8] sched_ext: Add filter for scx_kfunc_ids_select_cpu

[Juntong Deng]

On 2025/2/6 23:39, Andrea Righi wrote:
> On Wed, Feb 05, 2025 at 07:30:14PM +0000, Juntong Deng wrote:
> ...
>> +static int scx_kfunc_ids_other_rqlocked_filter(const struct bpf_prog *prog, u32 kfunc_id)
>> +{
>> +	u32 moff = prog->aux->attach_st_ops_member_off;
>> +
>> +	if (moff == offsetof(struct sched_ext_ops, runnable) ||
>> +	    moff == offsetof(struct sched_ext_ops, dequeue) ||
>> +	    moff == offsetof(struct sched_ext_ops, stopping) ||
>> +	    moff == offsetof(struct sched_ext_ops, quiescent) ||
>> +	    moff == offsetof(struct sched_ext_ops, yield) ||
>> +	    moff == offsetof(struct sched_ext_ops, cpu_acquire) ||
>> +	    moff == offsetof(struct sched_ext_ops, running) ||
>> +	    moff == offsetof(struct sched_ext_ops, core_sched_before) ||
>> +	    moff == offsetof(struct sched_ext_ops, set_cpumask) ||
>> +	    moff == offsetof(struct sched_ext_ops, update_idle) ||
>> +	    moff == offsetof(struct sched_ext_ops, tick) ||
>> +	    moff == offsetof(struct sched_ext_ops, enable) ||
>> +	    moff == offsetof(struct sched_ext_ops, set_weight) ||
>> +	    moff == offsetof(struct sched_ext_ops, disable) ||
>> +	    moff == offsetof(struct sched_ext_ops, exit_task) ||
>> +	    moff == offsetof(struct sched_ext_ops, dump_task) ||
>> +	    moff == offsetof(struct sched_ext_ops, dump_cpu))
>> +		return 0;
>> +
>> +	return -EACCES;
> 
> Actually, do we need this filter at all?
> 
> I think the other filters in your patch set should be sufficient to
> establish the correct permissions for all kfuncs, as none of them need to
> be called from any rq-locked operations. Or am I missing something?
> 

Thanks for your reply.

I think I misunderstood SCX_KF_REST.

I incorrectly thought that all but SCX_KF_UNLOCKED belonged to
SCX_KF_REST (including SCX_KF_CPU_RELEASE, SCX_KF_DISPATCH, etc.).

I will remove scx_kfunc_ids_other_rqlocked_filter in the next version.

If you find any other mistakes, please let me know.

> -Andrea

[RFC PATCH bpf-next 3/8] sched_ext: Add filter for scx_kfunc_ids_enqueue_dispatch

[Juntong Deng]

This patch adds filter for scx_kfunc_ids_enqueue_dispatch.

The kfuncs in the scx_kfunc_ids_enqueue_dispatch set can be used in
enqueue, select_cpu, dispatch and other rq-locked operations.

Signed-off-by: Juntong Deng <juntong.deng@outlook.com>
---
 kernel/sched/ext.c | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
index c92949aa23f6..d782ee618d54 100644
--- a/kernel/sched/ext.c
+++ b/kernel/sched/ext.c
@@ -6605,9 +6605,27 @@ BTF_ID_FLAGS(func, scx_bpf_dispatch, KF_RCU)
 BTF_ID_FLAGS(func, scx_bpf_dispatch_vtime, KF_RCU)
 BTF_KFUNCS_END(scx_kfunc_ids_enqueue_dispatch)
 
+static int scx_kfunc_ids_enqueue_dispatch_filter(const struct bpf_prog *prog, u32 kfunc_id)
+{
+	u32 moff;
+
+	if (!btf_id_set8_contains(&scx_kfunc_ids_enqueue_dispatch, kfunc_id) ||
+	    prog->aux->st_ops != &bpf_sched_ext_ops)
+		return 0;
+
+	moff = prog->aux->attach_st_ops_member_off;
+	if (moff == offsetof(struct sched_ext_ops, enqueue) ||
+	    moff == offsetof(struct sched_ext_ops, select_cpu) ||
+	    moff == offsetof(struct sched_ext_ops, dispatch))
+		return 0;
+
+	return scx_kfunc_ids_other_rqlocked_filter(prog, kfunc_id);
+}
+
 static const struct btf_kfunc_id_set scx_kfunc_set_enqueue_dispatch = {
 	.owner			= THIS_MODULE,
 	.set			= &scx_kfunc_ids_enqueue_dispatch,
+	.filter			= scx_kfunc_ids_enqueue_dispatch_filter,
 };
 
 static bool scx_dsq_move(struct bpf_iter_scx_dsq_kern *kit,
-- 
2.39.5

[RFC PATCH bpf-next 4/8] sched_ext: Add filter for scx_kfunc_ids_dispatch

[Juntong Deng]

This patch adds filter for scx_kfunc_ids_dispatch.

The kfuncs in the scx_kfunc_ids_dispatch set can be used in dispatch
and other rq-locked operations.

Signed-off-by: Juntong Deng <juntong.deng@outlook.com>
---
 kernel/sched/ext.c | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
index d782ee618d54..caddcf41e5f1 100644
--- a/kernel/sched/ext.c
+++ b/kernel/sched/ext.c
@@ -6943,9 +6943,25 @@ BTF_ID_FLAGS(func, scx_bpf_dispatch_from_dsq, KF_RCU)
 BTF_ID_FLAGS(func, scx_bpf_dispatch_vtime_from_dsq, KF_RCU)
 BTF_KFUNCS_END(scx_kfunc_ids_dispatch)
 
+static int scx_kfunc_ids_dispatch_filter(const struct bpf_prog *prog, u32 kfunc_id)
+{
+	u32 moff;
+
+	if (!btf_id_set8_contains(&scx_kfunc_ids_dispatch, kfunc_id) ||
+	    prog->aux->st_ops != &bpf_sched_ext_ops)
+		return 0;
+
+	moff = prog->aux->attach_st_ops_member_off;
+	if (moff == offsetof(struct sched_ext_ops, dispatch))
+		return 0;
+
+	return scx_kfunc_ids_other_rqlocked_filter(prog, kfunc_id);
+}
+
 static const struct btf_kfunc_id_set scx_kfunc_set_dispatch = {
 	.owner			= THIS_MODULE,
 	.set			= &scx_kfunc_ids_dispatch,
+	.filter			= scx_kfunc_ids_dispatch_filter,
 };
 
 __bpf_kfunc_start_defs();
-- 
2.39.5

[RFC PATCH bpf-next 5/8] sched_ext: Add filter for scx_kfunc_ids_cpu_release

[Juntong Deng]

This patch adds filter for scx_kfunc_ids_cpu_release.

The kfuncs in the scx_kfunc_ids_cpu_release set can be used in
cpu_release and other rq-locked operations.

Signed-off-by: Juntong Deng <juntong.deng@outlook.com>
---
 kernel/sched/ext.c | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
index caddcf41e5f1..7f039a32f137 100644
--- a/kernel/sched/ext.c
+++ b/kernel/sched/ext.c
@@ -7026,9 +7026,25 @@ BTF_KFUNCS_START(scx_kfunc_ids_cpu_release)
 BTF_ID_FLAGS(func, scx_bpf_reenqueue_local)
 BTF_KFUNCS_END(scx_kfunc_ids_cpu_release)
 
+static int scx_kfunc_ids_cpu_release_filter(const struct bpf_prog *prog, u32 kfunc_id)
+{
+	u32 moff;
+
+	if (!btf_id_set8_contains(&scx_kfunc_ids_cpu_release, kfunc_id) ||
+	    prog->aux->st_ops != &bpf_sched_ext_ops)
+		return 0;
+
+	moff = prog->aux->attach_st_ops_member_off;
+	if (moff == offsetof(struct sched_ext_ops, cpu_release))
+		return 0;
+
+	return scx_kfunc_ids_other_rqlocked_filter(prog, kfunc_id);
+}
+
 static const struct btf_kfunc_id_set scx_kfunc_set_cpu_release = {
 	.owner			= THIS_MODULE,
 	.set			= &scx_kfunc_ids_cpu_release,
+	.filter			= scx_kfunc_ids_cpu_release_filter,
 };
 
 __bpf_kfunc_start_defs();
-- 
2.39.5

[RFC PATCH bpf-next 6/8] sched_ext: Add filter for scx_kfunc_ids_unlocked

[Juntong Deng]

This patch adds filter for scx_kfunc_ids_unlocked.

The kfuncs in the scx_kfunc_ids_unlocked set can be used in init, exit,
cpu_online, cpu_offline, init_task, dump, cgroup_init, cgroup_exit,
cgroup_prep_move, cgroup_cancel_move, cgroup_move, cgroup_set_weight
operations.

Signed-off-by: Juntong Deng <juntong.deng@outlook.com>
---
 kernel/sched/ext.c | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
index 7f039a32f137..955fb0f5fc5e 100644
--- a/kernel/sched/ext.c
+++ b/kernel/sched/ext.c
@@ -7079,9 +7079,39 @@ BTF_ID_FLAGS(func, scx_bpf_dispatch_from_dsq, KF_RCU)
 BTF_ID_FLAGS(func, scx_bpf_dispatch_vtime_from_dsq, KF_RCU)
 BTF_KFUNCS_END(scx_kfunc_ids_unlocked)
 
+static int scx_kfunc_ids_unlocked_filter(const struct bpf_prog *prog, u32 kfunc_id)
+{
+	u32 moff;
+
+	if (!btf_id_set8_contains(&scx_kfunc_ids_unlocked, kfunc_id) ||
+	    prog->aux->st_ops != &bpf_sched_ext_ops)
+		return 0;
+
+	moff = prog->aux->attach_st_ops_member_off;
+	if (moff == offsetof(struct sched_ext_ops, init) ||
+	    moff == offsetof(struct sched_ext_ops, exit) ||
+	    moff == offsetof(struct sched_ext_ops, cpu_online) ||
+	    moff == offsetof(struct sched_ext_ops, cpu_offline) ||
+	    moff == offsetof(struct sched_ext_ops, init_task) ||
+	    moff == offsetof(struct sched_ext_ops, dump))
+		return 0;
+
+#ifdef CONFIG_EXT_GROUP_SCHED
+	if (moff == offsetof(struct sched_ext_ops, cgroup_init) ||
+	    moff == offsetof(struct sched_ext_ops, cgroup_exit) ||
+	    moff == offsetof(struct sched_ext_ops, cgroup_prep_move) ||
+	    moff == offsetof(struct sched_ext_ops, cgroup_cancel_move) ||
+	    moff == offsetof(struct sched_ext_ops, cgroup_move) ||
+	    moff == offsetof(struct sched_ext_ops, cgroup_set_weight))
+		return 0;
+#endif
+	return -EACCES;
+}
+
 static const struct btf_kfunc_id_set scx_kfunc_set_unlocked = {
 	.owner			= THIS_MODULE,
 	.set			= &scx_kfunc_ids_unlocked,
+	.filter			= scx_kfunc_ids_unlocked_filter,
 };
 
 __bpf_kfunc_start_defs();
-- 
2.39.5

Re: [RFC PATCH bpf-next 6/8] sched_ext: Add filter for scx_kfunc_ids_unlocked

[Alexei Starovoitov]

On Wed, Feb 5, 2025 at 11:35 AM Juntong Deng <juntong.deng@outlook.com> wrote:
>
> This patch adds filter for scx_kfunc_ids_unlocked.
>
> The kfuncs in the scx_kfunc_ids_unlocked set can be used in init, exit,
> cpu_online, cpu_offline, init_task, dump, cgroup_init, cgroup_exit,
> cgroup_prep_move, cgroup_cancel_move, cgroup_move, cgroup_set_weight
> operations.
>
> Signed-off-by: Juntong Deng <juntong.deng@outlook.com>
> ---
>  kernel/sched/ext.c | 30 ++++++++++++++++++++++++++++++
>  1 file changed, 30 insertions(+)
>
> diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
> index 7f039a32f137..955fb0f5fc5e 100644
> --- a/kernel/sched/ext.c
> +++ b/kernel/sched/ext.c
> @@ -7079,9 +7079,39 @@ BTF_ID_FLAGS(func, scx_bpf_dispatch_from_dsq, KF_RCU)
>  BTF_ID_FLAGS(func, scx_bpf_dispatch_vtime_from_dsq, KF_RCU)
>  BTF_KFUNCS_END(scx_kfunc_ids_unlocked)
>
> +static int scx_kfunc_ids_unlocked_filter(const struct bpf_prog *prog, u32 kfunc_id)
> +{
> +       u32 moff;
> +
> +       if (!btf_id_set8_contains(&scx_kfunc_ids_unlocked, kfunc_id) ||
> +           prog->aux->st_ops != &bpf_sched_ext_ops)
> +               return 0;
> +
> +       moff = prog->aux->attach_st_ops_member_off;
> +       if (moff == offsetof(struct sched_ext_ops, init) ||
> +           moff == offsetof(struct sched_ext_ops, exit) ||
> +           moff == offsetof(struct sched_ext_ops, cpu_online) ||
> +           moff == offsetof(struct sched_ext_ops, cpu_offline) ||
> +           moff == offsetof(struct sched_ext_ops, init_task) ||
> +           moff == offsetof(struct sched_ext_ops, dump))
> +               return 0;
> +
> +#ifdef CONFIG_EXT_GROUP_SCHED
> +       if (moff == offsetof(struct sched_ext_ops, cgroup_init) ||
> +           moff == offsetof(struct sched_ext_ops, cgroup_exit) ||
> +           moff == offsetof(struct sched_ext_ops, cgroup_prep_move) ||
> +           moff == offsetof(struct sched_ext_ops, cgroup_cancel_move) ||
> +           moff == offsetof(struct sched_ext_ops, cgroup_move) ||
> +           moff == offsetof(struct sched_ext_ops, cgroup_set_weight))
> +               return 0;
> +#endif
> +       return -EACCES;
> +}
> +
>  static const struct btf_kfunc_id_set scx_kfunc_set_unlocked = {
>         .owner                  = THIS_MODULE,
>         .set                    = &scx_kfunc_ids_unlocked,
> +       .filter                 = scx_kfunc_ids_unlocked_filter,
>  };

why does sched-ext use so many id_set-s ?

        if ((ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
                                             &scx_kfunc_set_select_cpu)) ||
            (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,

&scx_kfunc_set_enqueue_dispatch)) ||
            (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
                                             &scx_kfunc_set_dispatch)) ||
            (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
                                             &scx_kfunc_set_cpu_release)) ||
            (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
                                             &scx_kfunc_set_unlocked)) ||

Can they all be rolled into one id_set then
the patches 2-6 will be collapsed into one patch and
one filter callback that will describe allowed hook/kfunc combinations?

Re: [RFC PATCH bpf-next 6/8] sched_ext: Add filter for scx_kfunc_ids_unlocked

[Andrea Righi]

On Fri, Feb 07, 2025 at 07:37:51PM -0800, Alexei Starovoitov wrote:
> On Wed, Feb 5, 2025 at 11:35 AM Juntong Deng <juntong.deng@outlook.com> wrote:
> >
> > This patch adds filter for scx_kfunc_ids_unlocked.
> >
> > The kfuncs in the scx_kfunc_ids_unlocked set can be used in init, exit,
> > cpu_online, cpu_offline, init_task, dump, cgroup_init, cgroup_exit,
> > cgroup_prep_move, cgroup_cancel_move, cgroup_move, cgroup_set_weight
> > operations.
> >
> > Signed-off-by: Juntong Deng <juntong.deng@outlook.com>
> > ---
> >  kernel/sched/ext.c | 30 ++++++++++++++++++++++++++++++
> >  1 file changed, 30 insertions(+)
> >
> > diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
> > index 7f039a32f137..955fb0f5fc5e 100644
> > --- a/kernel/sched/ext.c
> > +++ b/kernel/sched/ext.c
> > @@ -7079,9 +7079,39 @@ BTF_ID_FLAGS(func, scx_bpf_dispatch_from_dsq, KF_RCU)
> >  BTF_ID_FLAGS(func, scx_bpf_dispatch_vtime_from_dsq, KF_RCU)
> >  BTF_KFUNCS_END(scx_kfunc_ids_unlocked)
> >
> > +static int scx_kfunc_ids_unlocked_filter(const struct bpf_prog *prog, u32 kfunc_id)
> > +{
> > +       u32 moff;
> > +
> > +       if (!btf_id_set8_contains(&scx_kfunc_ids_unlocked, kfunc_id) ||
> > +           prog->aux->st_ops != &bpf_sched_ext_ops)
> > +               return 0;
> > +
> > +       moff = prog->aux->attach_st_ops_member_off;
> > +       if (moff == offsetof(struct sched_ext_ops, init) ||
> > +           moff == offsetof(struct sched_ext_ops, exit) ||
> > +           moff == offsetof(struct sched_ext_ops, cpu_online) ||
> > +           moff == offsetof(struct sched_ext_ops, cpu_offline) ||
> > +           moff == offsetof(struct sched_ext_ops, init_task) ||
> > +           moff == offsetof(struct sched_ext_ops, dump))
> > +               return 0;
> > +
> > +#ifdef CONFIG_EXT_GROUP_SCHED
> > +       if (moff == offsetof(struct sched_ext_ops, cgroup_init) ||
> > +           moff == offsetof(struct sched_ext_ops, cgroup_exit) ||
> > +           moff == offsetof(struct sched_ext_ops, cgroup_prep_move) ||
> > +           moff == offsetof(struct sched_ext_ops, cgroup_cancel_move) ||
> > +           moff == offsetof(struct sched_ext_ops, cgroup_move) ||
> > +           moff == offsetof(struct sched_ext_ops, cgroup_set_weight))
> > +               return 0;
> > +#endif
> > +       return -EACCES;
> > +}
> > +
> >  static const struct btf_kfunc_id_set scx_kfunc_set_unlocked = {
> >         .owner                  = THIS_MODULE,
> >         .set                    = &scx_kfunc_ids_unlocked,
> > +       .filter                 = scx_kfunc_ids_unlocked_filter,
> >  };
> 
> why does sched-ext use so many id_set-s ?
> 
>         if ((ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
>                                              &scx_kfunc_set_select_cpu)) ||
>             (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
> 
> &scx_kfunc_set_enqueue_dispatch)) ||
>             (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
>                                              &scx_kfunc_set_dispatch)) ||
>             (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
>                                              &scx_kfunc_set_cpu_release)) ||
>             (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
>                                              &scx_kfunc_set_unlocked)) ||
> 
> Can they all be rolled into one id_set then
> the patches 2-6 will be collapsed into one patch and
> one filter callback that will describe allowed hook/kfunc combinations?

I think the idea was to group them in different sets based on their context
usage, like scx_kfunc_set_select_cpu kfuncs can be used only from
ops.select_cpu(), scx_kfunc_set_dispatch kfuncs can be used only from
ops.dispatch(), etc.

However, since the actual context enforcement is done by scx_kf_allowed(),
it seems that we could have just 3 sets to classify the kfuncs based by
their prog type:
 1) BPF_PROG_TYPE_STRUCT_OPS
 2) BPF_PROG_TYPE_STRUCT_OPS + BPF_PROG_TYPE_SYSCALL
 3) BPF_PROG_TYPE_STRUCT_OPS + BPF_PROG_TYPE_SYSCALL + BPF_PROG_TYPE_TRACING

-Andrea

Re: [RFC PATCH bpf-next 6/8] sched_ext: Add filter for scx_kfunc_ids_unlocked

[Tejun Heo]

Hello,

On Fri, Feb 07, 2025 at 07:37:51PM -0800, Alexei Starovoitov wrote:
> Can they all be rolled into one id_set then
> the patches 2-6 will be collapsed into one patch and
> one filter callback that will describe allowed hook/kfunc combinations?

I thought the BPF side filtering may be declarative on the kfunc sets and
tried to group the sets by where they can be called from. As we're going
prodcedural, there's no point in separating out the sets and we can use
merged kfunc sets and do all the distinctions in the filter function.

Thanks.

-- 
tejun

Re: [RFC PATCH bpf-next 6/8] sched_ext: Add filter for scx_kfunc_ids_unlocked

[Juntong Deng]

On 2025/2/8 03:37, Alexei Starovoitov wrote:
> On Wed, Feb 5, 2025 at 11:35 AM Juntong Deng <juntong.deng@outlook.com> wrote:
>>
>> This patch adds filter for scx_kfunc_ids_unlocked.
>>
>> The kfuncs in the scx_kfunc_ids_unlocked set can be used in init, exit,
>> cpu_online, cpu_offline, init_task, dump, cgroup_init, cgroup_exit,
>> cgroup_prep_move, cgroup_cancel_move, cgroup_move, cgroup_set_weight
>> operations.
>>
>> Signed-off-by: Juntong Deng <juntong.deng@outlook.com>
>> ---
>>   kernel/sched/ext.c | 30 ++++++++++++++++++++++++++++++
>>   1 file changed, 30 insertions(+)
>>
>> diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
>> index 7f039a32f137..955fb0f5fc5e 100644
>> --- a/kernel/sched/ext.c
>> +++ b/kernel/sched/ext.c
>> @@ -7079,9 +7079,39 @@ BTF_ID_FLAGS(func, scx_bpf_dispatch_from_dsq, KF_RCU)
>>   BTF_ID_FLAGS(func, scx_bpf_dispatch_vtime_from_dsq, KF_RCU)
>>   BTF_KFUNCS_END(scx_kfunc_ids_unlocked)
>>
>> +static int scx_kfunc_ids_unlocked_filter(const struct bpf_prog *prog, u32 kfunc_id)
>> +{
>> +       u32 moff;
>> +
>> +       if (!btf_id_set8_contains(&scx_kfunc_ids_unlocked, kfunc_id) ||
>> +           prog->aux->st_ops != &bpf_sched_ext_ops)
>> +               return 0;
>> +
>> +       moff = prog->aux->attach_st_ops_member_off;
>> +       if (moff == offsetof(struct sched_ext_ops, init) ||
>> +           moff == offsetof(struct sched_ext_ops, exit) ||
>> +           moff == offsetof(struct sched_ext_ops, cpu_online) ||
>> +           moff == offsetof(struct sched_ext_ops, cpu_offline) ||
>> +           moff == offsetof(struct sched_ext_ops, init_task) ||
>> +           moff == offsetof(struct sched_ext_ops, dump))
>> +               return 0;
>> +
>> +#ifdef CONFIG_EXT_GROUP_SCHED
>> +       if (moff == offsetof(struct sched_ext_ops, cgroup_init) ||
>> +           moff == offsetof(struct sched_ext_ops, cgroup_exit) ||
>> +           moff == offsetof(struct sched_ext_ops, cgroup_prep_move) ||
>> +           moff == offsetof(struct sched_ext_ops, cgroup_cancel_move) ||
>> +           moff == offsetof(struct sched_ext_ops, cgroup_move) ||
>> +           moff == offsetof(struct sched_ext_ops, cgroup_set_weight))
>> +               return 0;
>> +#endif
>> +       return -EACCES;
>> +}
>> +
>>   static const struct btf_kfunc_id_set scx_kfunc_set_unlocked = {
>>          .owner                  = THIS_MODULE,
>>          .set                    = &scx_kfunc_ids_unlocked,
>> +       .filter                 = scx_kfunc_ids_unlocked_filter,
>>   };
> 
> why does sched-ext use so many id_set-s ?
> 
>          if ((ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
>                                               &scx_kfunc_set_select_cpu)) ||
>              (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
> 
> &scx_kfunc_set_enqueue_dispatch)) ||
>              (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
>                                               &scx_kfunc_set_dispatch)) ||
>              (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
>                                               &scx_kfunc_set_cpu_release)) ||
>              (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
>                                               &scx_kfunc_set_unlocked)) ||
> 
> Can they all be rolled into one id_set then
> the patches 2-6 will be collapsed into one patch and
> one filter callback that will describe allowed hook/kfunc combinations?

Yes, I agree that it would be ideal to put all kfuncs in the one id_set,
but I am not sure that this is better in implementation.

For filters, the only kfunc-related information that can be known is
the kfunc_id.

kfunc_id is not a stable value, for example, when we add a new kfunc to
the kernel, it may cause the kfunc_id of other kfuncs to change.

A simple experiment is to add a bpf_task_from_aaa kfunc, and then we
will find that the kfunc_id of bpf_task_from_pid has changed.

This means that it is simple for us to implement kfuncs grouping via
id_set because we only need to check if kfunc_id exists in a specific
id_set, we do not need to care about what kfunc_id is.

But if we implement grouping only in the filter, we may need to first
get the btf type of the corresponding kfunc based on the kfunc_id via
btf_type_by_id, and then further get the kfunc name, and then group
based on the kfunc name in the filter, which seems more complicated.

Re: [RFC PATCH bpf-next 6/8] sched_ext: Add filter for scx_kfunc_ids_unlocked

[Alexei Starovoitov]

On Mon, Feb 10, 2025 at 3:40 PM Juntong Deng <juntong.deng@outlook.com> wrote:
>
> On 2025/2/8 03:37, Alexei Starovoitov wrote:
> > On Wed, Feb 5, 2025 at 11:35 AM Juntong Deng <juntong.deng@outlook.com> wrote:
> >>
> >> This patch adds filter for scx_kfunc_ids_unlocked.
> >>
> >> The kfuncs in the scx_kfunc_ids_unlocked set can be used in init, exit,
> >> cpu_online, cpu_offline, init_task, dump, cgroup_init, cgroup_exit,
> >> cgroup_prep_move, cgroup_cancel_move, cgroup_move, cgroup_set_weight
> >> operations.
> >>
> >> Signed-off-by: Juntong Deng <juntong.deng@outlook.com>
> >> ---
> >>   kernel/sched/ext.c | 30 ++++++++++++++++++++++++++++++
> >>   1 file changed, 30 insertions(+)
> >>
> >> diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
> >> index 7f039a32f137..955fb0f5fc5e 100644
> >> --- a/kernel/sched/ext.c
> >> +++ b/kernel/sched/ext.c
> >> @@ -7079,9 +7079,39 @@ BTF_ID_FLAGS(func, scx_bpf_dispatch_from_dsq, KF_RCU)
> >>   BTF_ID_FLAGS(func, scx_bpf_dispatch_vtime_from_dsq, KF_RCU)
> >>   BTF_KFUNCS_END(scx_kfunc_ids_unlocked)
> >>
> >> +static int scx_kfunc_ids_unlocked_filter(const struct bpf_prog *prog, u32 kfunc_id)
> >> +{
> >> +       u32 moff;
> >> +
> >> +       if (!btf_id_set8_contains(&scx_kfunc_ids_unlocked, kfunc_id) ||
> >> +           prog->aux->st_ops != &bpf_sched_ext_ops)
> >> +               return 0;
> >> +
> >> +       moff = prog->aux->attach_st_ops_member_off;
> >> +       if (moff == offsetof(struct sched_ext_ops, init) ||
> >> +           moff == offsetof(struct sched_ext_ops, exit) ||
> >> +           moff == offsetof(struct sched_ext_ops, cpu_online) ||
> >> +           moff == offsetof(struct sched_ext_ops, cpu_offline) ||
> >> +           moff == offsetof(struct sched_ext_ops, init_task) ||
> >> +           moff == offsetof(struct sched_ext_ops, dump))
> >> +               return 0;
> >> +
> >> +#ifdef CONFIG_EXT_GROUP_SCHED
> >> +       if (moff == offsetof(struct sched_ext_ops, cgroup_init) ||
> >> +           moff == offsetof(struct sched_ext_ops, cgroup_exit) ||
> >> +           moff == offsetof(struct sched_ext_ops, cgroup_prep_move) ||
> >> +           moff == offsetof(struct sched_ext_ops, cgroup_cancel_move) ||
> >> +           moff == offsetof(struct sched_ext_ops, cgroup_move) ||
> >> +           moff == offsetof(struct sched_ext_ops, cgroup_set_weight))
> >> +               return 0;
> >> +#endif
> >> +       return -EACCES;
> >> +}
> >> +
> >>   static const struct btf_kfunc_id_set scx_kfunc_set_unlocked = {
> >>          .owner                  = THIS_MODULE,
> >>          .set                    = &scx_kfunc_ids_unlocked,
> >> +       .filter                 = scx_kfunc_ids_unlocked_filter,
> >>   };
> >
> > why does sched-ext use so many id_set-s ?
> >
> >          if ((ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
> >                                               &scx_kfunc_set_select_cpu)) ||
> >              (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
> >
> > &scx_kfunc_set_enqueue_dispatch)) ||
> >              (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
> >                                               &scx_kfunc_set_dispatch)) ||
> >              (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
> >                                               &scx_kfunc_set_cpu_release)) ||
> >              (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
> >                                               &scx_kfunc_set_unlocked)) ||
> >
> > Can they all be rolled into one id_set then
> > the patches 2-6 will be collapsed into one patch and
> > one filter callback that will describe allowed hook/kfunc combinations?
>
> Yes, I agree that it would be ideal to put all kfuncs in the one id_set,
> but I am not sure that this is better in implementation.
>
> For filters, the only kfunc-related information that can be known is
> the kfunc_id.
>
> kfunc_id is not a stable value, for example, when we add a new kfunc to
> the kernel, it may cause the kfunc_id of other kfuncs to change.
>
> A simple experiment is to add a bpf_task_from_aaa kfunc, and then we
> will find that the kfunc_id of bpf_task_from_pid has changed.
>
> This means that it is simple for us to implement kfuncs grouping via
> id_set because we only need to check if kfunc_id exists in a specific
> id_set, we do not need to care about what kfunc_id is.
>
> But if we implement grouping only in the filter, we may need to first
> get the btf type of the corresponding kfunc based on the kfunc_id via
> btf_type_by_id, and then further get the kfunc name, and then group
> based on the kfunc name in the filter, which seems more complicated.

I didn't mean to extract kfunc name as a string and do strcmp() on it.
That's a non-starter.
I imagined verifier-like approach of enum+set+list
where enum has all kfunc names,
set gives efficient btf_id_set8_contains() access,
and list[KF_bpf_foo] gives func_id to compare with.

But if the current break down of scx_kfunc_set_* fits well
with per struct_ops hook filtering then keep it.
But please think of a set approach for moff as well to avoid
+           moff == offsetof(struct sched_ext_ops, exit) ||
+           moff == offsetof(struct sched_ext_ops, cpu_online) ||
+           moff == offsetof(struct sched_ext_ops, cpu_offline) ||

Then it will be:
if (btf_id_set8_contains(&scx_kfunc_ids_unlocked, kfunc_id) ...
&& moff_set_containts(.._unlocked, moff)) // allow

There is SCX_OP_IDX(). Maybe it can be used to populate a set.

Something like this:
static const u32 ops_flags[] = {
  [SCX_OP_IDX(cpu_online)] = KF_UNLOCKED,
  ..
};

if (btf_id_set8_contains(&scx_kfunc_ids_unlocked, kfunc_id) &&
    (ops_flags[moff / sizeof(void (*)(void))] & KF_UNLOCKED)) // allow

Re: [RFC PATCH bpf-next 6/8] sched_ext: Add filter for scx_kfunc_ids_unlocked

[Juntong Deng]

On 2025/2/11 03:48, Alexei Starovoitov wrote:
> On Mon, Feb 10, 2025 at 3:40 PM Juntong Deng <juntong.deng@outlook.com> wrote:
>>
>> On 2025/2/8 03:37, Alexei Starovoitov wrote:
>>> On Wed, Feb 5, 2025 at 11:35 AM Juntong Deng <juntong.deng@outlook.com> wrote:
>>>>
>>>> This patch adds filter for scx_kfunc_ids_unlocked.
>>>>
>>>> The kfuncs in the scx_kfunc_ids_unlocked set can be used in init, exit,
>>>> cpu_online, cpu_offline, init_task, dump, cgroup_init, cgroup_exit,
>>>> cgroup_prep_move, cgroup_cancel_move, cgroup_move, cgroup_set_weight
>>>> operations.
>>>>
>>>> Signed-off-by: Juntong Deng <juntong.deng@outlook.com>
>>>> ---
>>>>    kernel/sched/ext.c | 30 ++++++++++++++++++++++++++++++
>>>>    1 file changed, 30 insertions(+)
>>>>
>>>> diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
>>>> index 7f039a32f137..955fb0f5fc5e 100644
>>>> --- a/kernel/sched/ext.c
>>>> +++ b/kernel/sched/ext.c
>>>> @@ -7079,9 +7079,39 @@ BTF_ID_FLAGS(func, scx_bpf_dispatch_from_dsq, KF_RCU)
>>>>    BTF_ID_FLAGS(func, scx_bpf_dispatch_vtime_from_dsq, KF_RCU)
>>>>    BTF_KFUNCS_END(scx_kfunc_ids_unlocked)
>>>>
>>>> +static int scx_kfunc_ids_unlocked_filter(const struct bpf_prog *prog, u32 kfunc_id)
>>>> +{
>>>> +       u32 moff;
>>>> +
>>>> +       if (!btf_id_set8_contains(&scx_kfunc_ids_unlocked, kfunc_id) ||
>>>> +           prog->aux->st_ops != &bpf_sched_ext_ops)
>>>> +               return 0;
>>>> +
>>>> +       moff = prog->aux->attach_st_ops_member_off;
>>>> +       if (moff == offsetof(struct sched_ext_ops, init) ||
>>>> +           moff == offsetof(struct sched_ext_ops, exit) ||
>>>> +           moff == offsetof(struct sched_ext_ops, cpu_online) ||
>>>> +           moff == offsetof(struct sched_ext_ops, cpu_offline) ||
>>>> +           moff == offsetof(struct sched_ext_ops, init_task) ||
>>>> +           moff == offsetof(struct sched_ext_ops, dump))
>>>> +               return 0;
>>>> +
>>>> +#ifdef CONFIG_EXT_GROUP_SCHED
>>>> +       if (moff == offsetof(struct sched_ext_ops, cgroup_init) ||
>>>> +           moff == offsetof(struct sched_ext_ops, cgroup_exit) ||
>>>> +           moff == offsetof(struct sched_ext_ops, cgroup_prep_move) ||
>>>> +           moff == offsetof(struct sched_ext_ops, cgroup_cancel_move) ||
>>>> +           moff == offsetof(struct sched_ext_ops, cgroup_move) ||
>>>> +           moff == offsetof(struct sched_ext_ops, cgroup_set_weight))
>>>> +               return 0;
>>>> +#endif
>>>> +       return -EACCES;
>>>> +}
>>>> +
>>>>    static const struct btf_kfunc_id_set scx_kfunc_set_unlocked = {
>>>>           .owner                  = THIS_MODULE,
>>>>           .set                    = &scx_kfunc_ids_unlocked,
>>>> +       .filter                 = scx_kfunc_ids_unlocked_filter,
>>>>    };
>>>
>>> why does sched-ext use so many id_set-s ?
>>>
>>>           if ((ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
>>>                                                &scx_kfunc_set_select_cpu)) ||
>>>               (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
>>>
>>> &scx_kfunc_set_enqueue_dispatch)) ||
>>>               (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
>>>                                                &scx_kfunc_set_dispatch)) ||
>>>               (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
>>>                                                &scx_kfunc_set_cpu_release)) ||
>>>               (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
>>>                                                &scx_kfunc_set_unlocked)) ||
>>>
>>> Can they all be rolled into one id_set then
>>> the patches 2-6 will be collapsed into one patch and
>>> one filter callback that will describe allowed hook/kfunc combinations?
>>
>> Yes, I agree that it would be ideal to put all kfuncs in the one id_set,
>> but I am not sure that this is better in implementation.
>>
>> For filters, the only kfunc-related information that can be known is
>> the kfunc_id.
>>
>> kfunc_id is not a stable value, for example, when we add a new kfunc to
>> the kernel, it may cause the kfunc_id of other kfuncs to change.
>>
>> A simple experiment is to add a bpf_task_from_aaa kfunc, and then we
>> will find that the kfunc_id of bpf_task_from_pid has changed.
>>
>> This means that it is simple for us to implement kfuncs grouping via
>> id_set because we only need to check if kfunc_id exists in a specific
>> id_set, we do not need to care about what kfunc_id is.
>>
>> But if we implement grouping only in the filter, we may need to first
>> get the btf type of the corresponding kfunc based on the kfunc_id via
>> btf_type_by_id, and then further get the kfunc name, and then group
>> based on the kfunc name in the filter, which seems more complicated.
> 
> I didn't mean to extract kfunc name as a string and do strcmp() on it.
> That's a non-starter.
> I imagined verifier-like approach of enum+set+list
> where enum has all kfunc names,
> set gives efficient btf_id_set8_contains() access,
> and list[KF_bpf_foo] gives func_id to compare with.
> 
> But if the current break down of scx_kfunc_set_* fits well
> with per struct_ops hook filtering then keep it.
> But please think of a set approach for moff as well to avoid
> +           moff == offsetof(struct sched_ext_ops, exit) ||
> +           moff == offsetof(struct sched_ext_ops, cpu_online) ||
> +           moff == offsetof(struct sched_ext_ops, cpu_offline) ||
> 
> Then it will be:
> if (btf_id_set8_contains(&scx_kfunc_ids_unlocked, kfunc_id) ...
> && moff_set_containts(.._unlocked, moff)) // allow
> 
> There is SCX_OP_IDX(). Maybe it can be used to populate a set.
> 
> Something like this:
> static const u32 ops_flags[] = {
>    [SCX_OP_IDX(cpu_online)] = KF_UNLOCKED,
>    ..
> };
> 
> if (btf_id_set8_contains(&scx_kfunc_ids_unlocked, kfunc_id) &&
>      (ops_flags[moff / sizeof(void (*)(void))] & KF_UNLOCKED)) // allow

Thanks for letting me know this method.

This is a good method.

I have used it in version 2 [0].

Also, I figured out a way to require only one filter.

[0]: 
https://lore.kernel.org/bpf/AM6PR03MB5080855B90C3FE9B6C4243B099FE2@AM6PR03MB5080.eurprd03.prod.outlook.com/T/#u

[RFC PATCH bpf-next 7/8] sched_ext: Removed mask-based runtime restrictions on calling kfuncs in different contexts

[Juntong Deng]

Currently, kfunc filters already support filtering based on struct_ops
context information.

The BPF verifier can check context-sensitive kfuncs before the SCX
program is run, avoiding runtime overhead.

Therefore we no longer need mask-based runtime restrictions.

This patch removes the mask-based runtime restrictions.

Signed-off-by: Juntong Deng <juntong.deng@outlook.com>
---
 include/linux/sched/ext.h |  24 ----
 kernel/sched/ext.c        | 232 ++++++++------------------------------
 2 files changed, 50 insertions(+), 206 deletions(-)

diff --git a/include/linux/sched/ext.h b/include/linux/sched/ext.h
index 1d70a9867fb1..867c24b88ace 100644
--- a/include/linux/sched/ext.h
+++ b/include/linux/sched/ext.h
@@ -96,29 +96,6 @@ enum scx_ent_dsq_flags {
 	SCX_TASK_DSQ_ON_PRIQ	= 1 << 0, /* task is queued on the priority queue of a dsq */
 };
 
-/*
- * Mask bits for scx_entity.kf_mask. Not all kfuncs can be called from
- * everywhere and the following bits track which kfunc sets are currently
- * allowed for %current. This simple per-task tracking works because SCX ops
- * nest in a limited way. BPF will likely implement a way to allow and disallow
- * kfuncs depending on the calling context which will replace this manual
- * mechanism. See scx_kf_allow().
- */
-enum scx_kf_mask {
-	SCX_KF_UNLOCKED		= 0,	  /* sleepable and not rq locked */
-	/* ENQUEUE and DISPATCH may be nested inside CPU_RELEASE */
-	SCX_KF_CPU_RELEASE	= 1 << 0, /* ops.cpu_release() */
-	/* ops.dequeue (in REST) may be nested inside DISPATCH */
-	SCX_KF_DISPATCH		= 1 << 1, /* ops.dispatch() */
-	SCX_KF_ENQUEUE		= 1 << 2, /* ops.enqueue() and ops.select_cpu() */
-	SCX_KF_SELECT_CPU	= 1 << 3, /* ops.select_cpu() */
-	SCX_KF_REST		= 1 << 4, /* other rq-locked operations */
-
-	__SCX_KF_RQ_LOCKED	= SCX_KF_CPU_RELEASE | SCX_KF_DISPATCH |
-				  SCX_KF_ENQUEUE | SCX_KF_SELECT_CPU | SCX_KF_REST,
-	__SCX_KF_TERMINAL	= SCX_KF_ENQUEUE | SCX_KF_SELECT_CPU | SCX_KF_REST,
-};
-
 enum scx_dsq_lnode_flags {
 	SCX_DSQ_LNODE_ITER_CURSOR = 1 << 0,
 
@@ -146,7 +123,6 @@ struct sched_ext_entity {
 	u32			weight;
 	s32			sticky_cpu;
 	s32			holding_cpu;
-	u32			kf_mask;	/* see scx_kf_mask above */
 	struct task_struct	*kf_tasks[2];	/* see SCX_CALL_OP_TASK() */
 	atomic_long_t		ops_state;
 
diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
index 955fb0f5fc5e..d2182f5e2b28 100644
--- a/kernel/sched/ext.c
+++ b/kernel/sched/ext.c
@@ -1043,19 +1043,6 @@ static long jiffies_delta_msecs(unsigned long at, unsigned long now)
 		return -(long)jiffies_to_msecs(now - at);
 }
 
-/* if the highest set bit is N, return a mask with bits [N+1, 31] set */
-static u32 higher_bits(u32 flags)
-{
-	return ~((1 << fls(flags)) - 1);
-}
-
-/* return the mask with only the highest bit set */
-static u32 highest_bit(u32 flags)
-{
-	int bit = fls(flags);
-	return ((u64)1 << bit) >> 1;
-}
-
 static bool u32_before(u32 a, u32 b)
 {
 	return (s32)(a - b) < 0;
@@ -1071,51 +1058,12 @@ static struct scx_dispatch_q *find_user_dsq(u64 dsq_id)
 	return rhashtable_lookup_fast(&dsq_hash, &dsq_id, dsq_hash_params);
 }
 
-/*
- * scx_kf_mask enforcement. Some kfuncs can only be called from specific SCX
- * ops. When invoking SCX ops, SCX_CALL_OP[_RET]() should be used to indicate
- * the allowed kfuncs and those kfuncs should use scx_kf_allowed() to check
- * whether it's running from an allowed context.
- *
- * @mask is constant, always inline to cull the mask calculations.
- */
-static __always_inline void scx_kf_allow(u32 mask)
-{
-	/* nesting is allowed only in increasing scx_kf_mask order */
-	WARN_ONCE((mask | higher_bits(mask)) & current->scx.kf_mask,
-		  "invalid nesting current->scx.kf_mask=0x%x mask=0x%x\n",
-		  current->scx.kf_mask, mask);
-	current->scx.kf_mask |= mask;
-	barrier();
-}
+#define SCX_CALL_OP(op, args...)	scx_ops.op(args)
 
-static void scx_kf_disallow(u32 mask)
-{
-	barrier();
-	current->scx.kf_mask &= ~mask;
-}
-
-#define SCX_CALL_OP(mask, op, args...)						\
-do {										\
-	if (mask) {								\
-		scx_kf_allow(mask);						\
-		scx_ops.op(args);						\
-		scx_kf_disallow(mask);						\
-	} else {								\
-		scx_ops.op(args);						\
-	}									\
-} while (0)
-
-#define SCX_CALL_OP_RET(mask, op, args...)					\
+#define SCX_CALL_OP_RET(op, args...)						\
 ({										\
 	__typeof__(scx_ops.op(args)) __ret;					\
-	if (mask) {								\
-		scx_kf_allow(mask);						\
-		__ret = scx_ops.op(args);					\
-		scx_kf_disallow(mask);						\
-	} else {								\
-		__ret = scx_ops.op(args);					\
-	}									\
+	__ret = scx_ops.op(args);						\
 	__ret;									\
 })
 
@@ -1130,74 +1078,36 @@ do {										\
  * scx_kf_allowed_on_arg_tasks() to test whether the invocation is allowed on
  * the specific task.
  */
-#define SCX_CALL_OP_TASK(mask, op, task, args...)				\
+#define SCX_CALL_OP_TASK(op, task, args...)					\
 do {										\
-	BUILD_BUG_ON((mask) & ~__SCX_KF_TERMINAL);				\
 	current->scx.kf_tasks[0] = task;					\
-	SCX_CALL_OP(mask, op, task, ##args);					\
+	SCX_CALL_OP(op, task, ##args);						\
 	current->scx.kf_tasks[0] = NULL;					\
 } while (0)
 
-#define SCX_CALL_OP_TASK_RET(mask, op, task, args...)				\
+#define SCX_CALL_OP_TASK_RET(op, task, args...)					\
 ({										\
 	__typeof__(scx_ops.op(task, ##args)) __ret;				\
-	BUILD_BUG_ON((mask) & ~__SCX_KF_TERMINAL);				\
 	current->scx.kf_tasks[0] = task;					\
-	__ret = SCX_CALL_OP_RET(mask, op, task, ##args);			\
+	__ret = SCX_CALL_OP_RET(op, task, ##args);				\
 	current->scx.kf_tasks[0] = NULL;					\
 	__ret;									\
 })
 
-#define SCX_CALL_OP_2TASKS_RET(mask, op, task0, task1, args...)			\
+#define SCX_CALL_OP_2TASKS_RET(op, task0, task1, args...)			\
 ({										\
 	__typeof__(scx_ops.op(task0, task1, ##args)) __ret;			\
-	BUILD_BUG_ON((mask) & ~__SCX_KF_TERMINAL);				\
 	current->scx.kf_tasks[0] = task0;					\
 	current->scx.kf_tasks[1] = task1;					\
-	__ret = SCX_CALL_OP_RET(mask, op, task0, task1, ##args);		\
+	__ret = SCX_CALL_OP_RET(op, task0, task1, ##args);			\
 	current->scx.kf_tasks[0] = NULL;					\
 	current->scx.kf_tasks[1] = NULL;					\
 	__ret;									\
 })
 
-/* @mask is constant, always inline to cull unnecessary branches */
-static __always_inline bool scx_kf_allowed(u32 mask)
-{
-	if (unlikely(!(current->scx.kf_mask & mask))) {
-		scx_ops_error("kfunc with mask 0x%x called from an operation only allowing 0x%x",
-			      mask, current->scx.kf_mask);
-		return false;
-	}
-
-	/*
-	 * Enforce nesting boundaries. e.g. A kfunc which can be called from
-	 * DISPATCH must not be called if we're running DEQUEUE which is nested
-	 * inside ops.dispatch(). We don't need to check boundaries for any
-	 * blocking kfuncs as the verifier ensures they're only called from
-	 * sleepable progs.
-	 */
-	if (unlikely(highest_bit(mask) == SCX_KF_CPU_RELEASE &&
-		     (current->scx.kf_mask & higher_bits(SCX_KF_CPU_RELEASE)))) {
-		scx_ops_error("cpu_release kfunc called from a nested operation");
-		return false;
-	}
-
-	if (unlikely(highest_bit(mask) == SCX_KF_DISPATCH &&
-		     (current->scx.kf_mask & higher_bits(SCX_KF_DISPATCH)))) {
-		scx_ops_error("dispatch kfunc called from a nested operation");
-		return false;
-	}
-
-	return true;
-}
-
 /* see SCX_CALL_OP_TASK() */
-static __always_inline bool scx_kf_allowed_on_arg_tasks(u32 mask,
-							struct task_struct *p)
+static __always_inline bool scx_kf_allowed_on_arg_tasks(struct task_struct *p)
 {
-	if (!scx_kf_allowed(mask))
-		return false;
-
 	if (unlikely((p != current->scx.kf_tasks[0] &&
 		      p != current->scx.kf_tasks[1]))) {
 		scx_ops_error("called on a task not being operated on");
@@ -1207,11 +1117,6 @@ static __always_inline bool scx_kf_allowed_on_arg_tasks(u32 mask,
 	return true;
 }
 
-static bool scx_kf_allowed_if_unlocked(void)
-{
-	return !current->scx.kf_mask;
-}
-
 /**
  * nldsq_next_task - Iterate to the next task in a non-local DSQ
  * @dsq: user dsq being interated
@@ -2027,7 +1932,7 @@ static void do_enqueue_task(struct rq *rq, struct task_struct *p, u64 enq_flags,
 	WARN_ON_ONCE(*ddsp_taskp);
 	*ddsp_taskp = p;
 
-	SCX_CALL_OP_TASK(SCX_KF_ENQUEUE, enqueue, p, enq_flags);
+	SCX_CALL_OP_TASK(enqueue, p, enq_flags);
 
 	*ddsp_taskp = NULL;
 	if (p->scx.ddsp_dsq_id != SCX_DSQ_INVALID)
@@ -2122,7 +2027,7 @@ static void enqueue_task_scx(struct rq *rq, struct task_struct *p, int enq_flags
 	add_nr_running(rq, 1);
 
 	if (SCX_HAS_OP(runnable) && !task_on_rq_migrating(p))
-		SCX_CALL_OP_TASK(SCX_KF_REST, runnable, p, enq_flags);
+		SCX_CALL_OP_TASK(runnable, p, enq_flags);
 
 	if (enq_flags & SCX_ENQ_WAKEUP)
 		touch_core_sched(rq, p);
@@ -2153,7 +2058,7 @@ static void ops_dequeue(struct task_struct *p, u64 deq_flags)
 		BUG();
 	case SCX_OPSS_QUEUED:
 		if (SCX_HAS_OP(dequeue))
-			SCX_CALL_OP_TASK(SCX_KF_REST, dequeue, p, deq_flags);
+			SCX_CALL_OP_TASK(dequeue, p, deq_flags);
 
 		if (atomic_long_try_cmpxchg(&p->scx.ops_state, &opss,
 					    SCX_OPSS_NONE))
@@ -2202,11 +2107,11 @@ static bool dequeue_task_scx(struct rq *rq, struct task_struct *p, int deq_flags
 	 */
 	if (SCX_HAS_OP(stopping) && task_current(rq, p)) {
 		update_curr_scx(rq);
-		SCX_CALL_OP_TASK(SCX_KF_REST, stopping, p, false);
+		SCX_CALL_OP_TASK(stopping, p, false);
 	}
 
 	if (SCX_HAS_OP(quiescent) && !task_on_rq_migrating(p))
-		SCX_CALL_OP_TASK(SCX_KF_REST, quiescent, p, deq_flags);
+		SCX_CALL_OP_TASK(quiescent, p, deq_flags);
 
 	if (deq_flags & SCX_DEQ_SLEEP)
 		p->scx.flags |= SCX_TASK_DEQD_FOR_SLEEP;
@@ -2226,7 +2131,7 @@ static void yield_task_scx(struct rq *rq)
 	struct task_struct *p = rq->curr;
 
 	if (SCX_HAS_OP(yield))
-		SCX_CALL_OP_2TASKS_RET(SCX_KF_REST, yield, p, NULL);
+		SCX_CALL_OP_2TASKS_RET(yield, p, NULL);
 	else
 		p->scx.slice = 0;
 }
@@ -2236,7 +2141,7 @@ static bool yield_to_task_scx(struct rq *rq, struct task_struct *to)
 	struct task_struct *from = rq->curr;
 
 	if (SCX_HAS_OP(yield))
-		return SCX_CALL_OP_2TASKS_RET(SCX_KF_REST, yield, from, to);
+		return SCX_CALL_OP_2TASKS_RET(yield, from, to);
 	else
 		return false;
 }
@@ -2763,7 +2668,7 @@ static int balance_one(struct rq *rq, struct task_struct *prev)
 		 * emitted in switch_class().
 		 */
 		if (SCX_HAS_OP(cpu_acquire))
-			SCX_CALL_OP(SCX_KF_REST, cpu_acquire, cpu_of(rq), NULL);
+			SCX_CALL_OP(cpu_acquire, cpu_of(rq), NULL);
 		rq->scx.cpu_released = false;
 	}
 
@@ -2808,8 +2713,7 @@ static int balance_one(struct rq *rq, struct task_struct *prev)
 	do {
 		dspc->nr_tasks = 0;
 
-		SCX_CALL_OP(SCX_KF_DISPATCH, dispatch, cpu_of(rq),
-			    prev_on_scx ? prev : NULL);
+		SCX_CALL_OP(dispatch, cpu_of(rq), prev_on_scx ? prev : NULL);
 
 		flush_dispatch_buf(rq);
 
@@ -2929,7 +2833,7 @@ static void set_next_task_scx(struct rq *rq, struct task_struct *p, bool first)
 
 	/* see dequeue_task_scx() on why we skip when !QUEUED */
 	if (SCX_HAS_OP(running) && (p->scx.flags & SCX_TASK_QUEUED))
-		SCX_CALL_OP_TASK(SCX_KF_REST, running, p);
+		SCX_CALL_OP_TASK(running, p);
 
 	clr_task_runnable(p, true);
 
@@ -3010,8 +2914,7 @@ static void switch_class(struct rq *rq, struct task_struct *next)
 				.task = next,
 			};
 
-			SCX_CALL_OP(SCX_KF_CPU_RELEASE,
-				    cpu_release, cpu_of(rq), &args);
+			SCX_CALL_OP(cpu_release, cpu_of(rq), &args);
 		}
 		rq->scx.cpu_released = true;
 	}
@@ -3024,7 +2927,7 @@ static void put_prev_task_scx(struct rq *rq, struct task_struct *p,
 
 	/* see dequeue_task_scx() on why we skip when !QUEUED */
 	if (SCX_HAS_OP(stopping) && (p->scx.flags & SCX_TASK_QUEUED))
-		SCX_CALL_OP_TASK(SCX_KF_REST, stopping, p, true);
+		SCX_CALL_OP_TASK(stopping, p, true);
 
 	if (p->scx.flags & SCX_TASK_QUEUED) {
 		set_task_runnable(rq, p);
@@ -3602,8 +3505,7 @@ static int select_task_rq_scx(struct task_struct *p, int prev_cpu, int wake_flag
 		WARN_ON_ONCE(*ddsp_taskp);
 		*ddsp_taskp = p;
 
-		cpu = SCX_CALL_OP_TASK_RET(SCX_KF_ENQUEUE | SCX_KF_SELECT_CPU,
-					   select_cpu, p, prev_cpu, wake_flags);
+		cpu = SCX_CALL_OP_TASK_RET(select_cpu, p, prev_cpu, wake_flags);
 		*ddsp_taskp = NULL;
 		if (ops_cpu_valid(cpu, "from ops.select_cpu()"))
 			return cpu;
@@ -3641,8 +3543,7 @@ static void set_cpus_allowed_scx(struct task_struct *p,
 	 * designation pointless. Cast it away when calling the operation.
 	 */
 	if (SCX_HAS_OP(set_cpumask))
-		SCX_CALL_OP_TASK(SCX_KF_REST, set_cpumask, p,
-				 (struct cpumask *)p->cpus_ptr);
+		SCX_CALL_OP_TASK(set_cpumask, p, (struct cpumask *)p->cpus_ptr);
 }
 
 static void reset_idle_masks(void)
@@ -3708,7 +3609,7 @@ void __scx_update_idle(struct rq *rq, bool idle, bool do_notify)
 	 * managed by put_prev_task_idle()/set_next_task_idle().
 	 */
 	if (SCX_HAS_OP(update_idle) && do_notify && !scx_rq_bypassing(rq))
-		SCX_CALL_OP(SCX_KF_REST, update_idle, cpu_of(rq), idle);
+		SCX_CALL_OP(update_idle, cpu_of(rq), idle);
 
 	/*
 	 * Update the idle masks:
@@ -3739,9 +3640,9 @@ static void handle_hotplug(struct rq *rq, bool online)
 		update_selcpu_topology();
 
 	if (online && SCX_HAS_OP(cpu_online))
-		SCX_CALL_OP(SCX_KF_UNLOCKED, cpu_online, cpu);
+		SCX_CALL_OP(cpu_online, cpu);
 	else if (!online && SCX_HAS_OP(cpu_offline))
-		SCX_CALL_OP(SCX_KF_UNLOCKED, cpu_offline, cpu);
+		SCX_CALL_OP(cpu_offline, cpu);
 	else
 		scx_ops_exit(SCX_ECODE_ACT_RESTART | SCX_ECODE_RSN_HOTPLUG,
 			     "cpu %d going %s, exiting scheduler", cpu,
@@ -3851,7 +3752,7 @@ static void task_tick_scx(struct rq *rq, struct task_struct *curr, int queued)
 		curr->scx.slice = 0;
 		touch_core_sched(rq, curr);
 	} else if (SCX_HAS_OP(tick)) {
-		SCX_CALL_OP(SCX_KF_REST, tick, curr);
+		SCX_CALL_OP(tick, curr);
 	}
 
 	if (!curr->scx.slice)
@@ -3928,7 +3829,7 @@ static int scx_ops_init_task(struct task_struct *p, struct task_group *tg, bool
 			.fork = fork,
 		};
 
-		ret = SCX_CALL_OP_RET(SCX_KF_UNLOCKED, init_task, p, &args);
+		ret = SCX_CALL_OP_RET(init_task, p, &args);
 		if (unlikely(ret)) {
 			ret = ops_sanitize_err("init_task", ret);
 			return ret;
@@ -3985,11 +3886,11 @@ static void scx_ops_enable_task(struct task_struct *p)
 	p->scx.weight = sched_weight_to_cgroup(weight);
 
 	if (SCX_HAS_OP(enable))
-		SCX_CALL_OP_TASK(SCX_KF_REST, enable, p);
+		SCX_CALL_OP_TASK(enable, p);
 	scx_set_task_state(p, SCX_TASK_ENABLED);
 
 	if (SCX_HAS_OP(set_weight))
-		SCX_CALL_OP_TASK(SCX_KF_REST, set_weight, p, p->scx.weight);
+		SCX_CALL_OP_TASK(set_weight, p, p->scx.weight);
 }
 
 static void scx_ops_disable_task(struct task_struct *p)
@@ -3998,7 +3899,7 @@ static void scx_ops_disable_task(struct task_struct *p)
 	WARN_ON_ONCE(scx_get_task_state(p) != SCX_TASK_ENABLED);
 
 	if (SCX_HAS_OP(disable))
-		SCX_CALL_OP(SCX_KF_REST, disable, p);
+		SCX_CALL_OP(disable, p);
 	scx_set_task_state(p, SCX_TASK_READY);
 }
 
@@ -4027,7 +3928,7 @@ static void scx_ops_exit_task(struct task_struct *p)
 	}
 
 	if (SCX_HAS_OP(exit_task))
-		SCX_CALL_OP(SCX_KF_REST, exit_task, p, &args);
+		SCX_CALL_OP(exit_task, p, &args);
 	scx_set_task_state(p, SCX_TASK_NONE);
 }
 
@@ -4136,7 +4037,7 @@ static void reweight_task_scx(struct rq *rq, struct task_struct *p,
 
 	p->scx.weight = sched_weight_to_cgroup(scale_load_down(lw->weight));
 	if (SCX_HAS_OP(set_weight))
-		SCX_CALL_OP_TASK(SCX_KF_REST, set_weight, p, p->scx.weight);
+		SCX_CALL_OP_TASK(set_weight, p, p->scx.weight);
 }
 
 static void prio_changed_scx(struct rq *rq, struct task_struct *p, int oldprio)
@@ -4152,8 +4053,7 @@ static void switching_to_scx(struct rq *rq, struct task_struct *p)
 	 * different scheduler class. Keep the BPF scheduler up-to-date.
 	 */
 	if (SCX_HAS_OP(set_cpumask))
-		SCX_CALL_OP_TASK(SCX_KF_REST, set_cpumask, p,
-				 (struct cpumask *)p->cpus_ptr);
+		SCX_CALL_OP_TASK(set_cpumask, p, (struct cpumask *)p->cpus_ptr);
 }
 
 static void switched_from_scx(struct rq *rq, struct task_struct *p)
@@ -4245,8 +4145,7 @@ int scx_tg_online(struct task_group *tg)
 			struct scx_cgroup_init_args args =
 				{ .weight = tg->scx_weight };
 
-			ret = SCX_CALL_OP_RET(SCX_KF_UNLOCKED, cgroup_init,
-					      tg->css.cgroup, &args);
+			ret = SCX_CALL_OP_RET(cgroup_init, tg->css.cgroup, &args);
 			if (ret)
 				ret = ops_sanitize_err("cgroup_init", ret);
 		}
@@ -4267,7 +4166,7 @@ void scx_tg_offline(struct task_group *tg)
 	percpu_down_read(&scx_cgroup_rwsem);
 
 	if (SCX_HAS_OP(cgroup_exit) && (tg->scx_flags & SCX_TG_INITED))
-		SCX_CALL_OP(SCX_KF_UNLOCKED, cgroup_exit, tg->css.cgroup);
+		SCX_CALL_OP(cgroup_exit, tg->css.cgroup);
 	tg->scx_flags &= ~(SCX_TG_ONLINE | SCX_TG_INITED);
 
 	percpu_up_read(&scx_cgroup_rwsem);
@@ -4300,8 +4199,7 @@ int scx_cgroup_can_attach(struct cgroup_taskset *tset)
 			continue;
 
 		if (SCX_HAS_OP(cgroup_prep_move)) {
-			ret = SCX_CALL_OP_RET(SCX_KF_UNLOCKED, cgroup_prep_move,
-					      p, from, css->cgroup);
+			ret = SCX_CALL_OP_RET(cgroup_prep_move, p, from, css->cgroup);
 			if (ret)
 				goto err;
 		}
@@ -4314,8 +4212,7 @@ int scx_cgroup_can_attach(struct cgroup_taskset *tset)
 err:
 	cgroup_taskset_for_each(p, css, tset) {
 		if (SCX_HAS_OP(cgroup_cancel_move) && p->scx.cgrp_moving_from)
-			SCX_CALL_OP(SCX_KF_UNLOCKED, cgroup_cancel_move, p,
-				    p->scx.cgrp_moving_from, css->cgroup);
+			SCX_CALL_OP(cgroup_cancel_move, p, p->scx.cgrp_moving_from, css->cgroup);
 		p->scx.cgrp_moving_from = NULL;
 	}
 
@@ -4346,8 +4243,7 @@ void scx_move_task(struct task_struct *p)
 	 * cgrp_moving_from set.
 	 */
 	if (SCX_HAS_OP(cgroup_move) && !WARN_ON_ONCE(!p->scx.cgrp_moving_from))
-		SCX_CALL_OP_TASK(SCX_KF_UNLOCKED, cgroup_move, p,
-			p->scx.cgrp_moving_from, tg_cgrp(task_group(p)));
+		SCX_CALL_OP_TASK(cgroup_move, p, p->scx.cgrp_moving_from, tg_cgrp(task_group(p)));
 	p->scx.cgrp_moving_from = NULL;
 }
 
@@ -4366,8 +4262,7 @@ void scx_cgroup_cancel_attach(struct cgroup_taskset *tset)
 
 	cgroup_taskset_for_each(p, css, tset) {
 		if (SCX_HAS_OP(cgroup_cancel_move) && p->scx.cgrp_moving_from)
-			SCX_CALL_OP(SCX_KF_UNLOCKED, cgroup_cancel_move, p,
-				    p->scx.cgrp_moving_from, css->cgroup);
+			SCX_CALL_OP(cgroup_cancel_move, p, p->scx.cgrp_moving_from, css->cgroup);
 		p->scx.cgrp_moving_from = NULL;
 	}
 out_unlock:
@@ -4380,8 +4275,7 @@ void scx_group_set_weight(struct task_group *tg, unsigned long weight)
 
 	if (scx_cgroup_enabled && tg->scx_weight != weight) {
 		if (SCX_HAS_OP(cgroup_set_weight))
-			SCX_CALL_OP(SCX_KF_UNLOCKED, cgroup_set_weight,
-				    tg_cgrp(tg), weight);
+			SCX_CALL_OP(cgroup_set_weight, tg_cgrp(tg), weight);
 		tg->scx_weight = weight;
 	}
 
@@ -4571,7 +4465,7 @@ static void scx_cgroup_exit(void)
 			continue;
 		rcu_read_unlock();
 
-		SCX_CALL_OP(SCX_KF_UNLOCKED, cgroup_exit, css->cgroup);
+		SCX_CALL_OP(cgroup_exit, css->cgroup);
 
 		rcu_read_lock();
 		css_put(css);
@@ -4614,8 +4508,7 @@ static int scx_cgroup_init(void)
 			continue;
 		rcu_read_unlock();
 
-		ret = SCX_CALL_OP_RET(SCX_KF_UNLOCKED, cgroup_init,
-				      css->cgroup, &args);
+		ret = SCX_CALL_OP_RET(cgroup_init, css->cgroup, &args);
 		if (ret) {
 			css_put(css);
 			scx_ops_error("ops.cgroup_init() failed (%d)", ret);
@@ -5078,7 +4971,7 @@ static void scx_ops_disable_workfn(struct kthread_work *work)
 	}
 
 	if (scx_ops.exit)
-		SCX_CALL_OP(SCX_KF_UNLOCKED, exit, ei);
+		SCX_CALL_OP(exit, ei);
 
 	cancel_delayed_work_sync(&scx_watchdog_work);
 
@@ -5284,7 +5177,7 @@ static void scx_dump_task(struct seq_buf *s, struct scx_dump_ctx *dctx,
 
 	if (SCX_HAS_OP(dump_task)) {
 		ops_dump_init(s, "    ");
-		SCX_CALL_OP(SCX_KF_REST, dump_task, dctx, p);
+		SCX_CALL_OP(dump_task, dctx, p);
 		ops_dump_exit();
 	}
 
@@ -5330,7 +5223,7 @@ static void scx_dump_state(struct scx_exit_info *ei, size_t dump_len)
 
 	if (SCX_HAS_OP(dump)) {
 		ops_dump_init(&s, "");
-		SCX_CALL_OP(SCX_KF_UNLOCKED, dump, &dctx);
+		SCX_CALL_OP(dump, &dctx);
 		ops_dump_exit();
 	}
 
@@ -5387,7 +5280,7 @@ static void scx_dump_state(struct scx_exit_info *ei, size_t dump_len)
 		used = seq_buf_used(&ns);
 		if (SCX_HAS_OP(dump_cpu)) {
 			ops_dump_init(&ns, "  ");
-			SCX_CALL_OP(SCX_KF_REST, dump_cpu, &dctx, cpu, idle);
+			SCX_CALL_OP(dump_cpu, &dctx, cpu, idle);
 			ops_dump_exit();
 		}
 
@@ -5607,7 +5500,7 @@ static int scx_ops_enable(struct sched_ext_ops *ops, struct bpf_link *link)
 	cpus_read_lock();
 
 	if (scx_ops.init) {
-		ret = SCX_CALL_OP_RET(SCX_KF_UNLOCKED, init);
+		ret = SCX_CALL_OP_RET(init);
 		if (ret) {
 			ret = ops_sanitize_err("init", ret);
 			cpus_read_unlock();
@@ -6383,9 +6276,6 @@ __bpf_kfunc s32 scx_bpf_select_cpu_dfl(struct task_struct *p, s32 prev_cpu,
 	if (!check_builtin_idle_enabled())
 		goto prev_cpu;
 
-	if (!scx_kf_allowed(SCX_KF_SELECT_CPU))
-		goto prev_cpu;
-
 #ifdef CONFIG_SMP
 	return scx_select_cpu_dfl(p, prev_cpu, wake_flags, is_idle);
 #endif
@@ -6450,9 +6340,6 @@ static const struct btf_kfunc_id_set scx_kfunc_set_select_cpu = {
 
 static bool scx_dsq_insert_preamble(struct task_struct *p, u64 enq_flags)
 {
-	if (!scx_kf_allowed(SCX_KF_ENQUEUE | SCX_KF_DISPATCH))
-		return false;
-
 	lockdep_assert_irqs_disabled();
 
 	if (unlikely(!p)) {
@@ -6637,9 +6524,6 @@ static bool scx_dsq_move(struct bpf_iter_scx_dsq_kern *kit,
 	bool in_balance;
 	unsigned long flags;
 
-	if (!scx_kf_allowed_if_unlocked() && !scx_kf_allowed(SCX_KF_DISPATCH))
-		return false;
-
 	/*
 	 * Can be called from either ops.dispatch() locking this_rq() or any
 	 * context where no rq lock is held. If latter, lock @p's task_rq which
@@ -6722,9 +6606,6 @@ __bpf_kfunc_start_defs();
  */
 __bpf_kfunc u32 scx_bpf_dispatch_nr_slots(void)
 {
-	if (!scx_kf_allowed(SCX_KF_DISPATCH))
-		return 0;
-
 	return scx_dsp_max_batch - __this_cpu_read(scx_dsp_ctx->cursor);
 }
 
@@ -6738,9 +6619,6 @@ __bpf_kfunc void scx_bpf_dispatch_cancel(void)
 {
 	struct scx_dsp_ctx *dspc = this_cpu_ptr(scx_dsp_ctx);
 
-	if (!scx_kf_allowed(SCX_KF_DISPATCH))
-		return;
-
 	if (dspc->cursor > 0)
 		dspc->cursor--;
 	else
@@ -6766,9 +6644,6 @@ __bpf_kfunc bool scx_bpf_dsq_move_to_local(u64 dsq_id)
 	struct scx_dsp_ctx *dspc = this_cpu_ptr(scx_dsp_ctx);
 	struct scx_dispatch_q *dsq;
 
-	if (!scx_kf_allowed(SCX_KF_DISPATCH))
-		return false;
-
 	flush_dispatch_buf(dspc->rq);
 
 	dsq = find_user_dsq(dsq_id);
@@ -6980,9 +6855,6 @@ __bpf_kfunc u32 scx_bpf_reenqueue_local(void)
 	struct rq *rq;
 	struct task_struct *p, *n;
 
-	if (!scx_kf_allowed(SCX_KF_CPU_RELEASE))
-		return 0;
-
 	rq = cpu_rq(smp_processor_id());
 	lockdep_assert_rq_held(rq);
 
@@ -7769,7 +7641,7 @@ __bpf_kfunc struct cgroup *scx_bpf_task_cgroup(struct task_struct *p)
 	struct task_group *tg = p->sched_task_group;
 	struct cgroup *cgrp = &cgrp_dfl_root.cgrp;
 
-	if (!scx_kf_allowed_on_arg_tasks(__SCX_KF_RQ_LOCKED, p))
+	if (!scx_kf_allowed_on_arg_tasks(p))
 		goto out;
 
 	cgrp = tg_cgrp(tg);
@@ -7891,10 +7763,6 @@ static int __init scx_init(void)
 	 *
 	 * Some kfuncs are context-sensitive and can only be called from
 	 * specific SCX ops. They are grouped into BTF sets accordingly.
-	 * Unfortunately, BPF currently doesn't have a way of enforcing such
-	 * restrictions. Eventually, the verifier should be able to enforce
-	 * them. For now, register them the same and make each kfunc explicitly
-	 * check using scx_kf_allowed().
 	 */
 	if ((ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
 					     &scx_kfunc_set_select_cpu)) ||
-- 
2.39.5

[RFC PATCH bpf-next 8/8] selftests/sched_ext: Update enq_select_cpu_fails to adapt to struct_ops context filter

[Juntong Deng]

After kfunc filters support struct_ops context information, SCX programs
that call incorrect context-sensitive kfuncs will fail to load.

This patch updates the enq_select_cpu_fails test case to adapt to the
failed load situation.

Signed-off-by: Juntong Deng <juntong.deng@outlook.com>
---
 .../sched_ext/enq_select_cpu_fails.c          | 35 +++----------------
 1 file changed, 4 insertions(+), 31 deletions(-)

diff --git a/tools/testing/selftests/sched_ext/enq_select_cpu_fails.c b/tools/testing/selftests/sched_ext/enq_select_cpu_fails.c
index dd1350e5f002..a04ad9a48a8f 100644
--- a/tools/testing/selftests/sched_ext/enq_select_cpu_fails.c
+++ b/tools/testing/selftests/sched_ext/enq_select_cpu_fails.c
@@ -11,51 +11,24 @@
 #include "enq_select_cpu_fails.bpf.skel.h"
 #include "scx_test.h"
 
-static enum scx_test_status setup(void **ctx)
+static enum scx_test_status run(void *ctx)
 {
 	struct enq_select_cpu_fails *skel;
 
 	skel = enq_select_cpu_fails__open_and_load();
-	if (!skel) {
-		SCX_ERR("Failed to open and load skel");
-		return SCX_TEST_FAIL;
-	}
-	*ctx = skel;
-
-	return SCX_TEST_PASS;
-}
-
-static enum scx_test_status run(void *ctx)
-{
-	struct enq_select_cpu_fails *skel = ctx;
-	struct bpf_link *link;
-
-	link = bpf_map__attach_struct_ops(skel->maps.enq_select_cpu_fails_ops);
-	if (!link) {
-		SCX_ERR("Failed to attach scheduler");
+	if (skel) {
+		enq_select_cpu_fails__destroy(skel);
+		SCX_ERR("This program should fail to load");
 		return SCX_TEST_FAIL;
 	}
 
-	sleep(1);
-
-	bpf_link__destroy(link);
-
 	return SCX_TEST_PASS;
 }
 
-static void cleanup(void *ctx)
-{
-	struct enq_select_cpu_fails *skel = ctx;
-
-	enq_select_cpu_fails__destroy(skel);
-}
-
 struct scx_test enq_select_cpu_fails = {
 	.name = "enq_select_cpu_fails",
 	.description = "Verify we fail to call scx_bpf_select_cpu_dfl() "
 		       "from ops.enqueue()",
-	.setup = setup,
 	.run = run,
-	.cleanup = cleanup,
 };
 REGISTER_SCX_TEST(&enq_select_cpu_fails)
-- 
2.39.5