[PATCH bpf-next v4 16/16] selftests/bpf: Test an arena kfunc argument passed from a subprogram

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

Call a kfunc with an __arena argument from a global subprogram that does
nothing else with the arena, the address arrives as a number and no cast
is involved. The subprogram needs the arena base set up by its own
prologue for the JIT to rebase the argument, which the BPF CI review
found was not recorded for it.

Pass a NULL constant to a nullable __arena kfunc argument as well, from
a program that touches the arena nowhere else. The verifier takes its
early exit for a NULL argument, but the JIT rebases the call like any
other, behind the NULL test the function model asks for, so the arena
base has to be set up for that call alone.

Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 .../testing/selftests/bpf/progs/arena_kfunc.c | 51 +++++++++++++++++++
 1 file changed, 51 insertions(+)

diff --git a/tools/testing/selftests/bpf/progs/arena_kfunc.c b/tools/testing/selftests/bpf/progs/arena_kfunc.c
index a049d6371d09..5d1cb1a7e8ee 100644
--- a/tools/testing/selftests/bpf/progs/arena_kfunc.c
+++ b/tools/testing/selftests/bpf/progs/arena_kfunc.c
@@ -192,6 +192,57 @@ int arena_arg_unpopulated(void *ctx)
 	return 0;
 }
 
+/*
+ * A subprogram whose only arena use is the argument of a kfunc needs the
+ * arena base too, the JIT rebases the argument on it. The address arrives
+ * as a number, so no cast marks the subprogram either.
+ */
+__noinline int arena_arg_from_subprog(unsigned long addr)
+{
+	return bpf_kfunc_arena_arg_test((u64 *)addr) != 41;
+}
+
+/*
+ * A NULL constant for the __arena argument takes an early exit in the
+ * verifier; the call needs the arena base all the same, the JIT rebases
+ * the argument with a NULL test.
+ */
+SEC("syscall")
+__arch_x86_64
+__arch_arm64
+__success __retval(0)
+int arena_arg_null(void *ctx)
+{
+	return bpf_kfunc_arena_cap_nullable_test(NULL) != 0;
+}
+
+SEC("syscall")
+__arch_x86_64
+__arch_arm64
+__success __retval(0)
+int arena_arg_subprog(void *ctx)
+{
+#if defined(__BPF_FEATURE_ADDR_SPACE_CAST)
+	u64 __arena *val;
+	int ret;
+
+	val = bpf_arena_alloc_pages(&arena, NULL, 1, NUMA_NO_NODE, 0);
+	if (!val)
+		return 1;
+
+	*val = 41;
+	if (arena_arg_from_subprog((unsigned long)val))
+		ret = 2;
+	else
+		ret = *val != 42 ? 3 : 0;
+
+	bpf_arena_free_pages(&arena, (void __arena *)val, 1);
+	return ret;
+#else
+	return 0;
+#endif
+}
+
 SEC("syscall")
 __arch_x86_64
 __arch_arm64
-- 
2.53.0-Meta