[PATCHv2 bpf-next 05/11] bpf: Disable preemption in bpf_get_stackid

[Jiri Olsa <jolsa@kernel.org>]

The get_perf_callchain call needs disabled preemption plus we need
it disabled as long as we access its returned trace entries buffer.

Note the bpf_get_stackid_pe function is executed already with
preemption disabled.

Cc: stable@vger.kernel.org
Fixes: d5a3b1f69186 ("bpf: introduce BPF_MAP_TYPE_STACK_TRACE")
Reported-by: Tao Chen <chen.dylane@linux.dev>
Closes: https://lore.kernel.org/bpf/20260206090653.1336687-2-chen.dylane@linux.dev/
Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 kernel/bpf/stackmap.c | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/kernel/bpf/stackmap.c b/kernel/bpf/stackmap.c
index 43b1e5436047..e48ddb4edd47 100644
--- a/kernel/bpf/stackmap.c
+++ b/kernel/bpf/stackmap.c
@@ -624,30 +624,37 @@ BPF_CALL_3(bpf_get_stackid, struct pt_regs *, regs, struct bpf_map *, map,
 	struct perf_callchain_entry *trace;
 	struct stackid stackid;
 	bool kernel = !user;
+	int err = -EFAULT;
 	u32 max_depth;
-	int err;
 
 	if (unlikely(flags & ~(BPF_F_SKIP_FIELD_MASK | BPF_F_USER_STACK |
 			       BPF_F_FAST_STACK_CMP | BPF_F_REUSE_STACKID)))
 		return -EINVAL;
 
 	max_depth = stack_map_calculate_max_depth(map->value_size, elem_size, flags);
+
+	preempt_disable();
 	trace = get_perf_callchain(regs, kernel, user, max_depth,
 				   false, false, 0);
 
 	if (unlikely(!trace))
 		/* couldn't fetch the stack trace */
-		return -EFAULT;
+		goto out;
 
 	err = stackid_fastpath(&stackid, map, trace, flags);
 	if (err != -ENOENT)
-		return err;
+		goto out;
 
 	new_bucket = stackid_new_bucket(&stackid, map);
-	if (!new_bucket)
-		return -ENOMEM;
+	if (new_bucket) {
+		preempt_enable();
+		return stackid_install(&stackid, map, new_bucket, flags);
+	}
+	err = -ENOMEM;
 
-	return stackid_install(&stackid, map, new_bucket, flags);
+out:
+	preempt_enable();
+	return err;
 }
 
 const struct bpf_func_proto bpf_get_stackid_proto = {
-- 
2.54.0