[PATCH bpf-next v1 11/14] bpf: Replace arena kfunc argument flags with suffixes

[Kumar Kartikeya Dwivedi <memxor@gmail.com>]

The arena allocation kfuncs still identify pointer arguments with
KF_ARENA_ARG2. These flags cover only the first two parameters and
duplicate the __arena suffix mechanism used by other kfuncs.

Annotate the optional allocation address with __arena__nullable. Mark the
free and reserve addresses with __arena so a valid address whose low 32
bits are zero is rebased unconditionally instead of becoming NULL.

The JIT now passes kernel arena addresses to these kfuncs. Translate them
back to the lower-32-bit user addresses expected by the existing arena
helpers by subtracting kern_vm_start. This preserves allocation-anywhere,
freeing the first page of a 4 GiB arena, and reservation at address zero.

Drop KF_ARENA_ARG1 and KF_ARENA_ARG2 from the kernel interface and remove
the flags from the arena kfunc sets. KF_ARENA_RET remains responsible for
annotating the allocation return value.

Keep the affected selftests synchronized with the conversion. Associate
an arena before the iterator map-pointer failures so they still reach the
intended diagnostics, account for the extra nullable branch in JIT labels,
and treat 1ULL << 32 as the same allocation-anywhere request as NULL after
the required 32-bit truncation.

Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 include/linux/btf.h                           |  2 -
 kernel/bpf/arena.c                            | 40 ++++++++++++++-----
 .../selftests/bpf/progs/arena_kfunc_jit.c     | 16 ++++----
 .../selftests/bpf/progs/verifier_arena.c      |  6 +++
 .../bpf/progs/verifier_arena_large.c          |  4 +-
 5 files changed, 46 insertions(+), 22 deletions(-)

diff --git a/include/linux/btf.h b/include/linux/btf.h
index 89d5a5c4f117..65e5f11dc27e 100644
--- a/include/linux/btf.h
+++ b/include/linux/btf.h
@@ -76,8 +76,6 @@
 #define KF_RCU_PROTECTED (1 << 11) /* kfunc should be protected by rcu cs when they are invoked */
 #define KF_FASTCALL     (1 << 12) /* kfunc supports bpf_fastcall protocol */
 #define KF_ARENA_RET    (1 << 13) /* kfunc returns an arena pointer */
-#define KF_ARENA_ARG1   (1 << 14) /* kfunc takes an arena pointer as its first argument */
-#define KF_ARENA_ARG2   (1 << 15) /* kfunc takes an arena pointer as its second argument */
 #define KF_IMPLICIT_ARGS (1 << 16) /* kfunc has implicit arguments supplied by the verifier */
 #define KF_SPINLOCK_SAFE (1 << 17) /* kfunc is allowed inside bpf_spin_lock-ed region */
 
diff --git a/kernel/bpf/arena.c b/kernel/bpf/arena.c
index 7b6847200b43..6c34a0d34b3f 100644
--- a/kernel/bpf/arena.c
+++ b/kernel/bpf/arena.c
@@ -1044,18 +1044,28 @@ static void arena_free_irq(struct irq_work *iw)
 	schedule_work(&arena->free_work);
 }
 
+static long arena_kaddr_to_uaddr(struct bpf_arena *arena, const void *addr)
+{
+	if (!addr)
+		return 0;
+
+	return (long)addr - bpf_arena_get_kern_vm_start(arena);
+}
+
 __bpf_kfunc_start_defs();
 
-__bpf_kfunc void *bpf_arena_alloc_pages(void *p__map, void *addr__ign, u32 page_cnt,
-					int node_id, u64 flags)
+__bpf_kfunc void *bpf_arena_alloc_pages(void *p__map, void *addr__arena__nullable,
+					u32 page_cnt, int node_id, u64 flags)
 {
 	struct bpf_map *map = p__map;
 	struct bpf_arena *arena = container_of(map, struct bpf_arena, map);
+	long addr;
 
 	if (map->map_type != BPF_MAP_TYPE_ARENA || flags || !page_cnt)
 		return NULL;
 
-	return (void *)arena_alloc_pages(arena, (long)addr__ign, page_cnt, node_id, true);
+	addr = arena_kaddr_to_uaddr(arena, addr__arena__nullable);
+	return (void *)arena_alloc_pages(arena, addr, page_cnt, node_id, true);
 }
 
 void *bpf_arena_alloc_pages_non_sleepable(void *p__map, void *addr__ign, u32 page_cnt,
@@ -1082,14 +1092,20 @@ void *bpf_arena_alloc_pages_sleepable(void *p__map, void *addr__ign, u32 page_cn
 	return (void *)arena_alloc_pages(arena, (long)addr__ign, page_cnt, node_id, true);
 }
 
-__bpf_kfunc void bpf_arena_free_pages(void *p__map, void *ptr__ign, u32 page_cnt)
+/*
+ * A valid arena address can have zero low 32 bits, so ptr must be rebased
+ * unconditionally instead of being treated as nullable.
+ */
+__bpf_kfunc void bpf_arena_free_pages(void *p__map, void *ptr__arena, u32 page_cnt)
 {
 	struct bpf_map *map = p__map;
 	struct bpf_arena *arena = container_of(map, struct bpf_arena, map);
+	long ptr;
 
-	if (map->map_type != BPF_MAP_TYPE_ARENA || !page_cnt || !ptr__ign)
+	if (map->map_type != BPF_MAP_TYPE_ARENA || !page_cnt)
 		return;
-	arena_free_pages(arena, (long)ptr__ign, page_cnt, true);
+	ptr = arena_kaddr_to_uaddr(arena, ptr__arena);
+	arena_free_pages(arena, ptr, page_cnt, true);
 }
 
 void bpf_arena_free_pages_non_sleepable(void *p__map, void *ptr__ign, u32 page_cnt)
@@ -1102,10 +1118,11 @@ void bpf_arena_free_pages_non_sleepable(void *p__map, void *ptr__ign, u32 page_c
 	arena_free_pages(arena, (long)ptr__ign, page_cnt, false);
 }
 
-__bpf_kfunc int bpf_arena_reserve_pages(void *p__map, void *ptr__ign, u32 page_cnt)
+__bpf_kfunc int bpf_arena_reserve_pages(void *p__map, void *ptr__arena, u32 page_cnt)
 {
 	struct bpf_map *map = p__map;
 	struct bpf_arena *arena = container_of(map, struct bpf_arena, map);
+	long ptr;
 
 	if (map->map_type != BPF_MAP_TYPE_ARENA)
 		return -EINVAL;
@@ -1113,14 +1130,15 @@ __bpf_kfunc int bpf_arena_reserve_pages(void *p__map, void *ptr__ign, u32 page_c
 	if (!page_cnt)
 		return 0;
 
-	return arena_reserve_pages(arena, (long)ptr__ign, page_cnt);
+	ptr = arena_kaddr_to_uaddr(arena, ptr__arena);
+	return arena_reserve_pages(arena, ptr, page_cnt);
 }
 __bpf_kfunc_end_defs();
 
 BTF_KFUNCS_START(arena_kfuncs)
-BTF_ID_FLAGS(func, bpf_arena_alloc_pages, KF_ARENA_RET | KF_ARENA_ARG2 | KF_SPINLOCK_SAFE)
-BTF_ID_FLAGS(func, bpf_arena_free_pages, KF_ARENA_ARG2 | KF_SPINLOCK_SAFE)
-BTF_ID_FLAGS(func, bpf_arena_reserve_pages, KF_ARENA_ARG2 | KF_SPINLOCK_SAFE)
+BTF_ID_FLAGS(func, bpf_arena_alloc_pages, KF_ARENA_RET | KF_SPINLOCK_SAFE)
+BTF_ID_FLAGS(func, bpf_arena_free_pages, KF_SPINLOCK_SAFE)
+BTF_ID_FLAGS(func, bpf_arena_reserve_pages, KF_SPINLOCK_SAFE)
 BTF_KFUNCS_END(arena_kfuncs)
 
 static const struct btf_kfunc_id_set common_kfunc_set = {
diff --git a/tools/testing/selftests/bpf/progs/arena_kfunc_jit.c b/tools/testing/selftests/bpf/progs/arena_kfunc_jit.c
index b5a01cbc33a7..c9af35c683b3 100644
--- a/tools/testing/selftests/bpf/progs/arena_kfunc_jit.c
+++ b/tools/testing/selftests/bpf/progs/arena_kfunc_jit.c
@@ -49,15 +49,15 @@ __arch_x86_64
 __jited("...")
 __jited("	movl	%edi, %edi")
 __jited("	testl	%edi, %edi")
-__jited("	je	L0")
+__jited("	je	L1")
 __jited("	addq	%r12, %rdi")
-__jited("L0:	callq	{{.*}}")
+__jited("L1:	callq	{{.*}}")
 __arch_arm64
 __jited("...")
 __jited("	mov	w0, w0")
-__jited("	cbz	w0, L0")
+__jited("	cbz	w0, L1")
 __jited("	add	x0, x28, w0, uxtw")
-__jited("L0:	{{.*}}")
+__jited("L1:	{{.*}}")
 __success
 int arena_arg_jit_nullable(void *ctx)
 {
@@ -79,9 +79,9 @@ __jited("	movl	%ecx, %ecx")
 __jited("	addq	%r12, %rcx")
 __jited("	movl	%r8d, %r8d")
 __jited("	testl	%r8d, %r8d")
-__jited("	je	L0")
+__jited("	je	L1")
 __jited("	addq	%r12, %r8")
-__jited("L0:	callq	{{.*}}")
+__jited("L1:	callq	{{.*}}")
 __arch_arm64
 __jited("...")
 __jited("	add	x0, x28, w0, uxtw")
@@ -89,9 +89,9 @@ __jited("	add	x1, x28, w1, uxtw")
 __jited("	add	x2, x28, w2, uxtw")
 __jited("	add	x3, x28, w3, uxtw")
 __jited("	mov	w4, w4")
-__jited("	cbz	w4, L0")
+__jited("	cbz	w4, L1")
 __jited("	add	x4, x28, w4, uxtw")
-__jited("L0:	{{.*}}")
+__jited("L1:	{{.*}}")
 __success
 int arena_arg_jit_args5(void *ctx)
 {
diff --git a/tools/testing/selftests/bpf/progs/verifier_arena.c b/tools/testing/selftests/bpf/progs/verifier_arena.c
index 815f342eb4b0..d76490e059f9 100644
--- a/tools/testing/selftests/bpf/progs/verifier_arena.c
+++ b/tools/testing/selftests/bpf/progs/verifier_arena.c
@@ -445,6 +445,8 @@ int iter_maps1(struct bpf_iter__bpf_map *ctx)
 
 	if (!map)
 		return 0;
+	/* Associate an arena before testing the generic map-pointer path. */
+	bpf_arena_reserve_pages(&arena, NULL, 0);
 	bpf_arena_alloc_pages(map, NULL, map->max_entries, 0, 0);
 	return 0;
 }
@@ -455,6 +457,8 @@ int iter_maps2(struct bpf_iter__bpf_map *ctx)
 {
 	struct seq_file *seq = ctx->meta->seq;
 
+	/* Associate an arena before testing the generic map-pointer path. */
+	bpf_arena_reserve_pages(&arena, NULL, 0);
 	bpf_arena_alloc_pages((void *)seq, NULL, 1, 0, 0);
 	return 0;
 }
@@ -467,6 +471,8 @@ int iter_maps3(struct bpf_iter__bpf_map *ctx)
 
 	if (!map)
 		return 0;
+	/* Associate an arena before testing the generic map-pointer path. */
+	bpf_arena_reserve_pages(&arena, NULL, 0);
 	bpf_arena_alloc_pages(map->inner_map_meta, NULL, map->max_entries, 0, 0);
 	return 0;
 }
diff --git a/tools/testing/selftests/bpf/progs/verifier_arena_large.c b/tools/testing/selftests/bpf/progs/verifier_arena_large.c
index 6ab8730d4878..f6515e0e9b17 100644
--- a/tools/testing/selftests/bpf/progs/verifier_arena_large.c
+++ b/tools/testing/selftests/bpf/progs/verifier_arena_large.c
@@ -49,8 +49,10 @@ int big_alloc1(void *ctx)
 
 	no_page = bpf_arena_alloc_pages(&arena, (void __arena *)ARENA_SIZE,
 					1, NUMA_NO_NODE, 0);
-	if (no_page)
+	/* Only the low 32 bits contribute, so this is equivalent to NULL. */
+	if (!no_page)
 		return 3;
+	bpf_arena_free_pages(&arena, (void __arena *)no_page, 1);
 	if (*page1 != 1)
 		return 4;
 	if (*page2 != 2)
-- 
2.53.0