[RFC PATCH bpf-next v1 15/16] selftests/bpf: Test typed arena page allocation and release

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

Allocate a page of typed arena objects, write one, reach it again through
an opaque copy of its address and its neighbour by arithmetic on that copy,
and check that the kfunc receives the registered typed arena in place of
the type ID. Take the chunk a value names and see the same chunk refused
twice, once for the chunk and once for a page inside it. Ask for one page
of an object that spans several and see the request rounded up to the
object's chunk, with the granted count written back. Release a chunk and
see it keep its objects and stay taken within the same invocation, since
the release is deferred behind a grace period.

Reject a use of the returned pointer without a NULL check, a struct without
special fields, and a map that is not the program's arena. Declare the two
kfuncs in bpf_experimental.h.

Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 .../testing/selftests/bpf/bpf_experimental.h  |  12 ++
 .../bpf/progs/verifier_typed_arena.c          | 142 ++++++++++++++++++
 2 files changed, 154 insertions(+)

diff --git a/tools/testing/selftests/bpf/bpf_experimental.h b/tools/testing/selftests/bpf/bpf_experimental.h
index 128654997328..cc1537a5e07a 100644
--- a/tools/testing/selftests/bpf/bpf_experimental.h
+++ b/tools/testing/selftests/bpf/bpf_experimental.h
@@ -15,6 +15,18 @@
  */
 #define __typed_arena_size(sz) __attribute__((btf_decl_tag("typed_arena_size:" #sz)))
 
+/*
+ * Back the chunks covering *page_cnt pages of the typed arena of the struct
+ * whose local BTF type ID is type_id with zeroed objects, at addr, a typed
+ * pointer naming the first chunk, or anywhere for NULL; round the count up
+ * to whole chunks and write it back; return a pointer to the first object,
+ * or NULL. Release such a range after a grace period, after which its
+ * objects read as the dummy object. Both are safe under a spin lock.
+ */
+extern void *bpf_typed_arena_alloc_pages(void *map, __u64 type_id, void *addr, __u32 *page_cnt,
+					 int node_id) __ksym;
+extern void bpf_typed_arena_free_pages(void *map, __u64 type_id, void *ptr, __u32 page_cnt) __ksym;
+
 /* Convenience macro to wrap over bpf_obj_new */
 #define bpf_obj_new(type) ((type *)bpf_obj_new(bpf_core_type_id_local(type)))
 
diff --git a/tools/testing/selftests/bpf/progs/verifier_typed_arena.c b/tools/testing/selftests/bpf/progs/verifier_typed_arena.c
index 808a65611d4e..8ec75fd1f97b 100644
--- a/tools/testing/selftests/bpf/progs/verifier_typed_arena.c
+++ b/tools/testing/selftests/bpf/progs/verifier_typed_arena.c
@@ -1157,6 +1157,148 @@ int ptr_field_sanitize_on_one_path(void *ctx)
 	return 0;
 }
 
+#define TYPE_ID(T) bpf_core_type_id_local(T)
+
+SEC("syscall")
+__description("allocated pages hold real objects, reachable through any value that lands in them")
+__success __retval(0)
+__xlated("r2 = 0x{{[0-9a-f]+[0-9a-f][0-9a-f][0-9a-f][0-9a-f][0-9a-f][0-9a-f][0-9a-f][0-9a-f]}}")
+__xlated("call kernel-function")
+int pages_alloc(void *ctx)
+{
+	struct typed_obj *obj, *again, *next;
+	void *opaque;
+	__u32 cnt = 1;
+
+	obj = bpf_typed_arena_alloc_pages(&arena, TYPE_ID(struct typed_obj), NULL, &cnt,
+					  NUMA_NO_NODE);
+	if (!obj)
+		return 1;
+	if (cnt != 1)
+		return 2;
+	obj->value = 5;
+	/* The object through an opaque value, and its neighbor by arithmetic on that value */
+	opaque = obj;
+	again = opaque;
+	if (again != obj || again->value != 5)
+		return 3;
+	next = opaque + sizeof(*obj);
+	next->value = 6;
+	if (next == obj || next->value != 6 || obj->value != 5)
+		return 4;
+	return 0;
+}
+
+SEC("syscall")
+__description("a fixed request takes its chunk once")
+__success __retval(0)
+int pages_alloc_fixed(void *ctx)
+{
+	struct typed_obj *obj, *hint;
+	__u32 cnt;
+
+	/* The chunk user space names, the first one here */
+	hint = ptr;
+	cnt = 2;
+	obj = bpf_typed_arena_alloc_pages(&arena, TYPE_ID(struct typed_obj), hint, &cnt,
+					  NUMA_NO_NODE);
+	if (!obj)
+		return 1;
+	if (obj != hint || cnt != 2)
+		return 2;
+	cnt = 1;
+	if (bpf_typed_arena_alloc_pages(&arena, TYPE_ID(struct typed_obj), hint, &cnt,
+					NUMA_NO_NODE))
+		return 3;
+	/* The page after it, part of the first request */
+	hint = (void *)hint + __PAGE_SIZE;
+	cnt = 1;
+	if (bpf_typed_arena_alloc_pages(&arena, TYPE_ID(struct typed_obj), hint, &cnt,
+					NUMA_NO_NODE))
+		return 4;
+	return 0;
+}
+
+SEC("syscall")
+__description("a request is rounded up to whole chunks and the granted count written back")
+__success __retval(0)
+int pages_alloc_granted_count(void *ctx)
+{
+	__u32 cnt = 1, chunk_pages = 16384 > __PAGE_SIZE ? 16384 / __PAGE_SIZE : 1;
+	struct big_obj *obj;
+
+	obj = bpf_typed_arena_alloc_pages(&arena, TYPE_ID(struct big_obj), NULL, &cnt,
+					  NUMA_NO_NODE);
+	if (!obj)
+		return 1;
+	if (cnt != chunk_pages)
+		return 2;
+	obj->pad[8191] = 1;
+	return obj->pad[8191] - 1;
+}
+
+SEC("syscall")
+__description("a released chunk keeps its objects and stays taken until the grace period has passed")
+__success __retval(0)
+int pages_free(void *ctx)
+{
+	struct typed_obj *obj;
+	__u32 cnt = 1;
+
+	obj = bpf_typed_arena_alloc_pages(&arena, TYPE_ID(struct typed_obj), NULL, &cnt,
+					  NUMA_NO_NODE);
+	if (!obj)
+		return 1;
+	obj->value = 7;
+	bpf_typed_arena_free_pages(&arena, TYPE_ID(struct typed_obj), obj, 1);
+	if (obj->value != 7)
+		return 2;
+	cnt = 1;
+	if (bpf_typed_arena_alloc_pages(&arena, TYPE_ID(struct typed_obj), obj, &cnt, NUMA_NO_NODE))
+		return 3;
+	return 0;
+}
+
+SEC("syscall")
+__description("an allocation must be checked for NULL")
+__failure __msg("invalid mem access 'typed_arena_ptr_or_null_'")
+int pages_alloc_null_check(void *ctx)
+{
+	struct typed_obj *obj;
+	__u32 cnt = 1;
+
+	obj = bpf_typed_arena_alloc_pages(&arena, TYPE_ID(struct typed_obj), NULL, &cnt,
+					  NUMA_NO_NODE);
+	obj->value = 1;
+	return 0;
+}
+
+SEC("syscall")
+__description("the page kfuncs register the type like a cast: a struct without special fields belongs in the raw arena")
+__failure __msg("struct plain_obj has no special fields and needs no typed arena")
+int pages_alloc_plain_struct(void *ctx)
+{
+	struct plain_obj *obj;
+	__u32 cnt = 1;
+
+	obj = bpf_typed_arena_alloc_pages(&arena, TYPE_ID(struct plain_obj), NULL, &cnt,
+					  NUMA_NO_NODE);
+	return obj == NULL;
+}
+
+SEC("syscall")
+__description("the page kfuncs need the program's arena")
+__failure __msg("can only be used in a program that has an associated arena")
+int pages_alloc_needs_arena(void *ctx)
+{
+	struct typed_obj *obj;
+	__u32 cnt = 1;
+
+	obj = bpf_typed_arena_alloc_pages(&not_an_arena, TYPE_ID(struct typed_obj), NULL, &cnt,
+					  NUMA_NO_NODE);
+	return obj == NULL;
+}
+
 #endif /* __BPF_FEATURE_TYPED_ARENA_CAST */
 
 char _license[] SEC("license") = "GPL";
-- 
2.53.0