[PATCH bpf-next 2/2] selftests/bpf: Test bpf_for_each_map_elem on BPF_MAP_TYPE_HASH_OF_MAPS

[David Vernet]

Now that we support using the bpf_for_each_map_elem() on the
BPF_MAP_TYPE_HASH_OF_MAPS map type, we should add selftests that
verify expected behavior.

This patch addds those selftests. Included in this is a new
test_btf_map_in_map_failure.c BPF prog that can be used to verify
expected map-in-map failures in the verifier.

Signed-off-by: David Vernet <void@xxxxxxxxxxxxx>
---
 .../selftests/bpf/prog_tests/btf_map_in_map.c | 44 +++++++++++
 .../selftests/bpf/progs/test_btf_map_in_map.c | 73 ++++++++++++++++++-
 .../bpf/progs/test_btf_map_in_map_failure.c   | 39 ++++++++++
 3 files changed, 155 insertions(+), 1 deletion(-)
 create mode 100644 tools/testing/selftests/bpf/progs/test_btf_map_in_map_failure.c

diff --git a/tools/testing/selftests/bpf/prog_tests/btf_map_in_map.c b/tools/testing/selftests/bpf/prog_tests/btf_map_in_map.c
index a8b53b8736f0..257b5b5a08ba 100644
--- a/tools/testing/selftests/bpf/prog_tests/btf_map_in_map.c
+++ b/tools/testing/selftests/bpf/prog_tests/btf_map_in_map.c
@@ -4,6 +4,7 @@
 #include <test_progs.h>
 
 #include "test_btf_map_in_map.skel.h"
+#include "test_btf_map_in_map_failure.skel.h"
 
 static int duration;
 
@@ -154,6 +155,43 @@ static void test_diff_size(void)
 	test_btf_map_in_map__destroy(skel);
 }
 
+static void test_hash_iter(const char *prog_name)
+{
+	struct test_btf_map_in_map *skel;
+	struct bpf_program *prog;
+	struct bpf_link *link = NULL;
+	int err, child_pid, status;
+
+	skel = test_btf_map_in_map__open();
+	if (!ASSERT_OK_PTR(skel, "test_btf_map_in_map__open\n"))
+		return;
+
+	skel->bss->pid = getpid();
+	err = test_btf_map_in_map__load(skel);
+	if (!ASSERT_OK(err, "test_btf_map_in_map__load"))
+		goto cleanup;
+
+	prog = bpf_object__find_program_by_name(skel->obj, prog_name);
+	if (!ASSERT_OK_PTR(prog, "bpf_object__find_program_by_name"))
+		goto cleanup;
+
+	link = bpf_program__attach(prog);
+	if (!ASSERT_OK_PTR(link, "bpf_program__attach"))
+		goto cleanup;
+
+	child_pid = fork();
+	if (!ASSERT_GT(child_pid, -1, "child_pid"))
+		goto cleanup;
+	if (child_pid == 0)
+		_exit(0);
+	waitpid(child_pid, &status, 0);
+	ASSERT_OK(skel->bss->err, "post_wait_err");
+
+cleanup:
+	bpf_link__destroy(link);
+	test_btf_map_in_map__destroy(skel);
+}
+
 void test_btf_map_in_map(void)
 {
 	if (test__start_subtest("lookup_update"))
@@ -161,4 +199,10 @@ void test_btf_map_in_map(void)
 
 	if (test__start_subtest("diff_size"))
 		test_diff_size();
+
+	if (test__start_subtest("hash_iter"))
+		test_hash_iter("test_iter_hash_of_maps");
+
+	RUN_TESTS(test_btf_map_in_map);
+	RUN_TESTS(test_btf_map_in_map_failure);
 }
diff --git a/tools/testing/selftests/bpf/progs/test_btf_map_in_map.c b/tools/testing/selftests/bpf/progs/test_btf_map_in_map.c
index c218cf8989a9..905392de6a3b 100644
--- a/tools/testing/selftests/bpf/progs/test_btf_map_in_map.c
+++ b/tools/testing/selftests/bpf/progs/test_btf_map_in_map.c
@@ -1,7 +1,13 @@
 /* SPDX-License-Identifier: GPL-2.0 */
 /* Copyright (c) 2020 Facebook */
-#include <linux/bpf.h>
+#include <linux/types.h>
 #include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+#include <vmlinux.h>
+
+#include "bpf_misc.h"
+
+int err, pid;
 
 struct inner_map {
 	__uint(type, BPF_MAP_TYPE_ARRAY);
@@ -120,6 +126,13 @@ struct outer_sockarr_sz1 {
 
 int input = 0;
 
+static bool is_test_task(void)
+{
+	int cur_pid = bpf_get_current_pid_tgid() >> 32;
+
+	return pid == cur_pid;
+}
+
 SEC("raw_tp/sys_enter")
 int handle__sys_enter(void *ctx)
 {
@@ -147,4 +160,62 @@ int handle__sys_enter(void *ctx)
 	return 0;
 }
 
+struct callback_ctx {
+	bool invoked;
+	bool failed;
+};
+
+static __u64 set_invoked(struct bpf_map *map, __u64 *key, __u64 *val, struct callback_ctx *ctx)
+{
+	struct bpf_map *inner_map;
+
+	ctx->invoked = true;
+	inner_map = bpf_map_lookup_elem(map, key);
+	if (!inner_map) {
+		ctx->failed = true;
+		return 1;
+	}
+
+	return 0;
+}
+
+SEC("tp_btf/task_newtask")
+int BPF_PROG(test_iter_hash_of_maps, struct task_struct *task, u64 clone_flags)
+{
+	long ret;
+	struct callback_ctx callback_ctx = {
+		.invoked = false,
+		.failed = false,
+	};
+
+	if (!is_test_task())
+		return 0;
+
+	ret = bpf_for_each_map_elem(&outer_hash, set_invoked, &callback_ctx, 0);
+	if (ret < 1)
+		err = 1;
+
+	if (!callback_ctx.invoked)
+		err = 2;
+
+	if (callback_ctx.failed)
+		err = 3;
+
+	return 0;
+}
+
+static __u64 empty_cb(struct bpf_map *map, __u64 *key, __u64 *val, void *ctx)
+{
+	return 0;
+}
+
+SEC("tp_btf/task_newtask")
+__success
+int BPF_PROG(test_iter_hash_of_maps_no_ctx, struct task_struct *task, u64 clone_flags)
+{
+	/* Should be able to iterate with no context as well. */
+	bpf_for_each_map_elem(&outer_hash, empty_cb, NULL, 0);
+	return 0;
+}
+
 char _license[] SEC("license") = "GPL";
diff --git a/tools/testing/selftests/bpf/progs/test_btf_map_in_map_failure.c b/tools/testing/selftests/bpf/progs/test_btf_map_in_map_failure.c
new file mode 100644
index 000000000000..6b7a94fe15c7
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/test_btf_map_in_map_failure.c
@@ -0,0 +1,39 @@
+// SPDX-License-Identifier: GPL-2.0
+/* Copyright (c) 2023 Meta Platforms, Inc. and affiliates. */
+
+#include <linux/types.h>
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+#include <vmlinux.h>
+
+#include "bpf_misc.h"
+
+struct inner_map {
+	__uint(type, BPF_MAP_TYPE_ARRAY);
+	__uint(max_entries, 1);
+	__type(key, int);
+	__type(value, int);
+} inner_map1 SEC(".maps"),
+  inner_map2 SEC(".maps");
+
+struct outer_hash {
+	__uint(type, BPF_MAP_TYPE_HASH_OF_MAPS);
+	__uint(max_entries, 2);
+	__type(key, int);
+	__array(values, struct inner_map);
+} outer_hash SEC(".maps") = {
+	.values = {
+		[0] = &inner_map2,
+		[1] = &inner_map1,
+	},
+};
+
+SEC("tp_btf/task_newtask")
+__failure
+__msg("R2 type=scalar expected=func")
+int BPF_PROG(test_iter_hash_of_maps_null_cb, struct task_struct *task, u64 clone_flags)
+{
+	/* Can't iterate over a NULL callback. */
+	bpf_for_each_map_elem(&outer_hash, NULL, NULL, 0);
+	return 0;
+}
-- 
2.40.1