[PATCH bpf-next v1 5/6] bpf: Scope the bpf_for_each_map_elem() array key to the callback frame

[Ihor Solodrai <ihor.solodrai@linux.dev>]

bpf_for_each_array_elem() passes the callback the address of a u32 held
in its own stack frame:

	u32 i, key, num_elems = 0;
	...
		key = i;
		ret = callback_fn((u64)(long)map, (u64)(long)&key, ...);

The callback can store that PTR_TO_MAP_KEY into callback_ctx and the
program can dereference it after the iteration finishes. Loads through
PTR_TO_MAP_KEY are not fault-protected and array key_size is fixed at 4,
so it is a four-byte read-only leak of kernel stack.

Declare the key frame-scoped for the array map ops rather than in the
shared map_set_for_each_callback_args(): of the map_for_each_callback
implementations, only array and percpu-array pass a key from their own
frame. The hash family passes elem->key, which stays valid for as long
as the program runs.

map_key_from_value() does the same for array maps, for the timer, wq and
task_work callbacks. Those are left alone: their signature is
(map, key, value) with R4 and R5 uninitialised, so there is no
callback_ctx to park a typed pointer in, and a pointer written into map
memory loses its type.

The element value is unaffected in every case: it lives until map
teardown.

The frame-owned reference consumes an id, so leak_prog and nested_cb now
report id 5 rather than id 4; update the expected messages.

Reported-by: Nicholas Carlini <npc@anthropic.com>
Signed-off-by: Ihor Solodrai <ihor.solodrai@linux.dev>
---
 kernel/bpf/arraymap.c                          | 18 ++++++++++++++++--
 .../testing/selftests/bpf/prog_tests/cb_refs.c |  4 ++--
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/kernel/bpf/arraymap.c b/kernel/bpf/arraymap.c
index 0ce26b538075..44bd229873ca 100644
--- a/kernel/bpf/arraymap.c
+++ b/kernel/bpf/arraymap.c
@@ -855,6 +855,20 @@ static u64 array_map_mem_usage(const struct bpf_map *map)
 	return usage;
 }
 
+static int array_map_set_for_each_callback_args(struct bpf_verifier_env *env,
+						struct bpf_func_state *caller,
+						struct bpf_func_state *callee)
+{
+	int err;
+
+	err = map_set_for_each_callback_args(env, caller, callee);
+	if (err)
+		return err;
+
+	mark_frame_scoped_arg(callee, BPF_REG_2);
+	return 0;
+}
+
 BTF_ID_LIST_SINGLE(array_map_btf_ids, struct, bpf_array)
 const struct bpf_map_ops array_map_ops = {
 	.map_meta_equal = array_map_meta_equal,
@@ -875,7 +889,7 @@ const struct bpf_map_ops array_map_ops = {
 	.map_check_btf = array_map_check_btf,
 	.map_lookup_batch = generic_map_lookup_batch,
 	.map_update_batch = generic_map_update_batch,
-	.map_set_for_each_callback_args = map_set_for_each_callback_args,
+	.map_set_for_each_callback_args = array_map_set_for_each_callback_args,
 	.map_for_each_callback = bpf_for_each_array_elem,
 	.map_mem_usage = array_map_mem_usage,
 	.map_btf_id = &array_map_btf_ids[0],
@@ -900,7 +914,7 @@ const struct bpf_map_ops percpu_array_map_ops = {
 	.map_check_btf = array_map_check_btf,
 	.map_lookup_batch = generic_map_lookup_batch,
 	.map_update_batch = generic_map_update_batch,
-	.map_set_for_each_callback_args = map_set_for_each_callback_args,
+	.map_set_for_each_callback_args = array_map_set_for_each_callback_args,
 	.map_for_each_callback = bpf_for_each_array_elem,
 	.map_mem_usage = array_map_mem_usage,
 	.map_btf_id = &array_map_btf_ids[0],
diff --git a/tools/testing/selftests/bpf/prog_tests/cb_refs.c b/tools/testing/selftests/bpf/prog_tests/cb_refs.c
index c32c6dab49bc..645f065c21f5 100644
--- a/tools/testing/selftests/bpf/prog_tests/cb_refs.c
+++ b/tools/testing/selftests/bpf/prog_tests/cb_refs.c
@@ -12,8 +12,8 @@ struct {
 	const char *err_msg;
 } cb_refs_tests[] = {
 	{ "underflow_prog", "R1 type=scalar expected=ptr_, trusted_ptr_, rcu_ptr_" },
-	{ "leak_prog", "Unreleased reference id=4 alloc_insn=3" }, /* alloc_insn=3{2,3} */
-	{ "nested_cb", "Unreleased reference id=4 alloc_insn=2" }, /* alloc_insn=2{4,5} */
+	{ "leak_prog", "Unreleased reference id=5 alloc_insn=3" }, /* alloc_insn=3{2,3} */
+	{ "nested_cb", "Unreleased reference id=5 alloc_insn=2" }, /* alloc_insn=2{4,5} */
 	{ "non_cb_transfer_ref", "Unreleased reference id=4 alloc_insn=1" }, /* alloc_insn=1{1,2} */
 };
 
-- 
2.55.0