Re: [PATCH 01/10] sched_ext: Drop TRACING access to select_cpu kfuncs

[Andrea Righi <arighi@nvidia.com>]

Hi Tejun,

On Thu, Apr 09, 2026 at 08:30:37PM -1000, Tejun Heo wrote:
> The select_cpu kfuncs - scx_bpf_select_cpu_dfl(), scx_bpf_select_cpu_and()
> and __scx_bpf_select_cpu_and() - take task_rq_lock() internally. Exposing
> them via scx_kfunc_set_idle to BPF_PROG_TYPE_TRACING is unsafe: arbitrary
> tracing contexts (kprobes, tracepoints, fentry, LSM) may run with @p's
> pi_lock state unknown.
> 
> Move them out of scx_kfunc_ids_idle into a new scx_kfunc_ids_select_cpu
> set registered only for STRUCT_OPS and SYSCALL.

In addition of being unsafe it also makes sense from a logical perspective to
not "allocate" idle CPUs from a BPF_PROG_TYPE_TRACING context.

> 
> Extracted from a larger verifier-time kfunc context filter patch
> originally written by Juntong Deng.
> 
> Original-patch-by: Juntong Deng <juntong.deng@outlook.com>
> Cc: Cheng-Yang Chou <yphbchou0911@gmail.com>
> Signed-off-by: Tejun Heo <tj@kernel.org>

Reviewed-by: Andrea Righi <arighi@nvidia.com>

Thanks,
-Andrea

> ---
>  kernel/sched/ext_idle.c | 25 +++++++++++++++++++++----
>  1 file changed, 21 insertions(+), 4 deletions(-)
> 
> diff --git a/kernel/sched/ext_idle.c b/kernel/sched/ext_idle.c
> index ecf7e09b54ae..cd88aee47bd8 100644
> --- a/kernel/sched/ext_idle.c
> +++ b/kernel/sched/ext_idle.c
> @@ -1469,9 +1469,6 @@ BTF_ID_FLAGS(func, scx_bpf_pick_idle_cpu_node, KF_IMPLICIT_ARGS | KF_RCU)
>  BTF_ID_FLAGS(func, scx_bpf_pick_idle_cpu, KF_IMPLICIT_ARGS | KF_RCU)
>  BTF_ID_FLAGS(func, scx_bpf_pick_any_cpu_node, KF_IMPLICIT_ARGS | KF_RCU)
>  BTF_ID_FLAGS(func, scx_bpf_pick_any_cpu, KF_IMPLICIT_ARGS | KF_RCU)
> -BTF_ID_FLAGS(func, __scx_bpf_select_cpu_and, KF_IMPLICIT_ARGS | KF_RCU)
> -BTF_ID_FLAGS(func, scx_bpf_select_cpu_and, KF_RCU)
> -BTF_ID_FLAGS(func, scx_bpf_select_cpu_dfl, KF_IMPLICIT_ARGS | KF_RCU)
>  BTF_KFUNCS_END(scx_kfunc_ids_idle)
>  
>  static const struct btf_kfunc_id_set scx_kfunc_set_idle = {
> @@ -1479,13 +1476,33 @@ static const struct btf_kfunc_id_set scx_kfunc_set_idle = {
>  	.set			= &scx_kfunc_ids_idle,
>  };
>  
> +/*
> + * The select_cpu kfuncs internally call task_rq_lock() when invoked from an
> + * rq-unlocked context, and thus cannot be safely called from arbitrary tracing
> + * contexts where @p's pi_lock state is unknown. Keep them out of
> + * BPF_PROG_TYPE_TRACING by registering them in their own set which is exposed
> + * only to STRUCT_OPS and SYSCALL programs.
> + */
> +BTF_KFUNCS_START(scx_kfunc_ids_select_cpu)
> +BTF_ID_FLAGS(func, __scx_bpf_select_cpu_and, KF_IMPLICIT_ARGS | KF_RCU)
> +BTF_ID_FLAGS(func, scx_bpf_select_cpu_and, KF_RCU)
> +BTF_ID_FLAGS(func, scx_bpf_select_cpu_dfl, KF_IMPLICIT_ARGS | KF_RCU)
> +BTF_KFUNCS_END(scx_kfunc_ids_select_cpu)
> +
> +static const struct btf_kfunc_id_set scx_kfunc_set_select_cpu = {
> +	.owner			= THIS_MODULE,
> +	.set			= &scx_kfunc_ids_select_cpu,
> +};
> +
>  int scx_idle_init(void)
>  {
>  	int ret;
>  
>  	ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS, &scx_kfunc_set_idle) ||
>  	      register_btf_kfunc_id_set(BPF_PROG_TYPE_TRACING, &scx_kfunc_set_idle) ||
> -	      register_btf_kfunc_id_set(BPF_PROG_TYPE_SYSCALL, &scx_kfunc_set_idle);
> +	      register_btf_kfunc_id_set(BPF_PROG_TYPE_SYSCALL, &scx_kfunc_set_idle) ||
> +	      register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS, &scx_kfunc_set_select_cpu) ||
> +	      register_btf_kfunc_id_set(BPF_PROG_TYPE_SYSCALL, &scx_kfunc_set_select_cpu);
>  
>  	return ret;
>  }
> -- 
> 2.53.0
>