[PATCH bpf-next v5 1/2] bpf: Add sleepable version of bpf_get_stack in raw_tp_prog_func_proto and tp_prog_func_proto

[Feng Yang <yangfeng59949@163.com>]

From: Feng Yang <yangfeng@kylinos.cn>

bpf_get_stack() has two implementations.  bpf_get_stack() passes
may_fault == false and resolves build IDs with build_id_parse_nofault(),
so an IP whose backing file page is not resident falls back to
BPF_STACK_BUILD_ID_IP.  bpf_get_stack_sleepable() passes may_fault == true
and uses build_id_parse_file(), which faults the page in and returns a real
build ID.

Only kprobe/uprobe programs are routed to the sleepable implementation
today, so the other program types that can be sleepable keep losing
build IDs for non-resident pages:

 - tracing_prog_func_proto() serves BPF_TRACE_ITER, fentry/fexit/fmod_ret
   and BPF_TRACE_RAW_TP programs, and falls back to
   raw_tp_prog_func_proto() for BPF_FUNC_get_stack, which is also the
   get_func_proto() of BPF_PROG_TYPE_RAW_TRACEPOINT and is reached by
   BPF_PROG_TYPE_LSM and BPF_PROG_TYPE_SYSCALL through the default cases
   of bpf_lsm_func_proto() and syscall_prog_func_proto().

 - tp_prog_func_proto() serves BPF_PROG_TYPE_TRACEPOINT programs.

Route them to bpf_get_stack_sleepable() as well and mark the new
prototypes might_sleep, as
commit 9d02927fdf4e ("bpf: Mark faultable stack helpers as sleepable")
did for the faultable stack helper prototypes.

The verifier then rejects bpf_get_stack() from these program types
marked sleepable whenever in_sleepable_context() is false:
inside bpf_rcu_read_lock(), bpf_preempt_disable(), and so on. This also
covers calls that do not pass BPF_F_USER_BUILD_ID and would not sleep;
such programs load today. This is consistent with the existing constraints.

Signed-off-by: Feng Yang <yangfeng@kylinos.cn>
---
v5: selftest: Remove fork-related code, use getpgid. Thanks, Andrii
Nakryiko.

v4: Update commit messages. Test modification: limit capture to one hit, add s390 flags.
https://lore.kernel.org/all/20261105093212.826479-1-yangfeng59949@163.com/

v3: Add selftest. Update commit messages.
https://lore.kernel.org/all/20261104083833.971911-1-yangfeng59949@163.com/

v2: Move regs onto the stack.
https://lore.kernel.org/all/20261102073714.1103088-1-yangfeng59949@163.com/

v1: https://lore.kernel.org/all/20261028093548.790977-1-yangfeng59949@163.com/
---
 kernel/trace/bpf_trace.c | 54 ++++++++++++++++++++++++++++++++++++++--
 1 file changed, 52 insertions(+), 2 deletions(-)

diff --git a/kernel/trace/bpf_trace.c b/kernel/trace/bpf_trace.c
index 2aec56d7f3d3..1229ee06dd71 100644
--- a/kernel/trace/bpf_trace.c
+++ b/kernel/trace/bpf_trace.c
@@ -84,6 +84,7 @@ static struct bpf_raw_event_map *bpf_get_raw_tracepoint_module(const char *name)
 
 u64 bpf_get_stackid(u64 r1, u64 r2, u64 r3, u64 r4, u64 r5);
 u64 bpf_get_stack(u64 r1, u64 r2, u64 r3, u64 r4, u64 r5);
+u64 bpf_get_stack_sleepable(u64 r1, u64 r2, u64 r3, u64 r4, u64 r5);
 
 static int bpf_btf_printf_prepare(struct btf_ptr *ptr, u32 btf_ptr_size,
 				  u64 flags, const struct btf **btf,
@@ -1477,6 +1478,26 @@ static const struct bpf_func_proto bpf_get_stack_proto_tp = {
 	.arg4_type	= ARG_ANYTHING,
 };
 
+BPF_CALL_4(bpf_get_stack_tp_sleepable, void *, tp_buff, void *, buf, u32, size,
+	   u64, flags)
+{
+	struct pt_regs *regs = *(struct pt_regs **)tp_buff;
+
+	return bpf_get_stack_sleepable((unsigned long) regs, (unsigned long) buf,
+				       (unsigned long) size, flags, 0);
+}
+
+static const struct bpf_func_proto bpf_get_stack_proto_tp_sleepable = {
+	.func		= bpf_get_stack_tp_sleepable,
+	.gpl_only	= true,
+	.might_sleep	= true,
+	.ret_type	= RET_INTEGER,
+	.arg1_type	= ARG_PTR_TO_CTX,
+	.arg2_type	= ARG_PTR_TO_UNINIT_MEM,
+	.arg3_type	= ARG_MEM_SIZE_OR_ZERO,
+	.arg4_type	= ARG_ANYTHING,
+};
+
 static const struct bpf_func_proto *
 tp_prog_func_proto(enum bpf_func_id func_id, const struct bpf_prog *prog)
 {
@@ -1486,7 +1507,8 @@ tp_prog_func_proto(enum bpf_func_id func_id, const struct bpf_prog *prog)
 	case BPF_FUNC_get_stackid:
 		return &bpf_get_stackid_proto_tp;
 	case BPF_FUNC_get_stack:
-		return &bpf_get_stack_proto_tp;
+		return prog->sleepable ? &bpf_get_stack_proto_tp_sleepable
+				       : &bpf_get_stack_proto_tp;
 	case BPF_FUNC_get_attach_cookie:
 		return &bpf_get_attach_cookie_proto_trace;
 	default:
@@ -1716,6 +1738,33 @@ static const struct bpf_func_proto bpf_get_stack_proto_raw_tp = {
 	.arg4_type	= ARG_ANYTHING,
 };
 
+BPF_CALL_4(bpf_get_stack_raw_tp_sleepable, struct bpf_raw_tracepoint_args *, args,
+	   void *, buf, u32, size, u64, flags)
+{
+	/*
+	 * Don't take a per-CPU bpf_raw_tp_regs slot here: those slots are a
+	 * LIFO nest counter and are vulnerable to slot overwrite and
+	 * corruption across sleep.  Sleepable programs never run in
+	 * IRQ/NMI context, so a local pt_regs is sufficient.
+	 */
+	struct pt_regs regs = {};
+
+	perf_fetch_caller_regs(&regs);
+	return bpf_get_stack_sleepable((unsigned long) &regs, (unsigned long) buf,
+				       (unsigned long) size, flags, 0);
+}
+
+static const struct bpf_func_proto bpf_get_stack_proto_raw_tp_sleepable = {
+	.func		= bpf_get_stack_raw_tp_sleepable,
+	.gpl_only	= true,
+	.might_sleep	= true,
+	.ret_type	= RET_INTEGER,
+	.arg1_type	= ARG_PTR_TO_CTX,
+	.arg2_type	= ARG_PTR_TO_UNINIT_MEM,
+	.arg3_type	= ARG_MEM_SIZE_OR_ZERO,
+	.arg4_type	= ARG_ANYTHING,
+};
+
 static const struct bpf_func_proto *
 raw_tp_prog_func_proto(enum bpf_func_id func_id, const struct bpf_prog *prog)
 {
@@ -1725,7 +1774,8 @@ raw_tp_prog_func_proto(enum bpf_func_id func_id, const struct bpf_prog *prog)
 	case BPF_FUNC_get_stackid:
 		return &bpf_get_stackid_proto_raw_tp;
 	case BPF_FUNC_get_stack:
-		return &bpf_get_stack_proto_raw_tp;
+		return prog->sleepable ? &bpf_get_stack_proto_raw_tp_sleepable
+				       : &bpf_get_stack_proto_raw_tp;
 	case BPF_FUNC_get_attach_cookie:
 		return &bpf_get_attach_cookie_proto_tracing;
 	default:
-- 
2.25.1