[PATCH bpf-next 1/2] bpf: Support bpf_for_each_map_elem() for BPF_MAP_TYPE_HASH_OF_MAPS maps

[PATCH bpf-next 1/2] bpf: Support bpf_for_each_map_elem() for BPF_MAP_TYPE_HASH_OF_MAPS maps

[David Vernet]

The bpf_for_each_map_elem() helper can be used to iterate over all of
the elements of a map. The helper is not supported for all maps, which
currently includes the BPF_MAP_TYPE_HASH_OF_MAPS map type. Other hash
map types, such as BPF_MAP_TYPE_HASH, BPF_MAP_TYPE_PERCPU_HASH, etc, do
support this helper, so adding support for a hash of maps map type
doesn't require much code change.

The current use case for this is sched_ext, where we want to populate a
hashmap of cgroup id -> array of integers before a scheduler is
attached, so the scheduler can iterate over the map and assign the array
of CPUs to each cgroup as a cpumask.

This patch therefore adds support for using bpf_for_each_map_elem() for
BPF_MAP_TYPE_HASH_OF_MAPS. A subsequent patch will add selftests which
validate its behavior.

Signed-off-by: David Vernet <void@manifault.com>
---
 include/uapi/linux/bpf.h | 3 ++-
 kernel/bpf/hashtab.c     | 2 ++
 kernel/bpf/verifier.c    | 6 +++++-
 3 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/include/uapi/linux/bpf.h b/include/uapi/linux/bpf.h
index a7b5e91dd768..fc5f027c4837 100644
--- a/include/uapi/linux/bpf.h
+++ b/include/uapi/linux/bpf.h
@@ -4909,7 +4909,8 @@ union bpf_attr {
  *
  *		BPF_MAP_TYPE_HASH, BPF_MAP_TYPE_PERCPU_HASH,
  *		BPF_MAP_TYPE_LRU_HASH, BPF_MAP_TYPE_LRU_PERCPU_HASH,
- *		BPF_MAP_TYPE_ARRAY, BPF_MAP_TYPE_PERCPU_ARRAY
+ *		BPF_MAP_TYPE_HASH_OF_MAPS, BPF_MAP_TYPE_ARRAY,
+ *		BPF_MAP_TYPE_PERCPU_ARRAY
  *
  *		long (\*callback_fn)(struct bpf_map \*map, const void \*key, void \*value, void \*ctx);
  *
diff --git a/kernel/bpf/hashtab.c b/kernel/bpf/hashtab.c
index 9901efee4339..fef7b94ed389 100644
--- a/kernel/bpf/hashtab.c
+++ b/kernel/bpf/hashtab.c
@@ -2580,6 +2580,8 @@ const struct bpf_map_ops htab_of_maps_map_ops = {
 	.map_fd_sys_lookup_elem = bpf_map_fd_sys_lookup_elem,
 	.map_gen_lookup = htab_of_map_gen_lookup,
 	.map_check_btf = map_check_no_btf,
+	.map_set_for_each_callback_args = map_set_for_each_callback_args,
+	.map_for_each_callback = bpf_for_each_hash_elem,
 	.map_mem_usage = htab_map_mem_usage,
 	BATCH_OPS(htab),
 	.map_btf_id = &htab_map_btf_ids[0],
diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
index 086b2a14905b..ba74b3b8f181 100644
--- a/kernel/bpf/verifier.c
+++ b/kernel/bpf/verifier.c
@@ -8174,10 +8174,14 @@ static int check_map_func_compatibility(struct bpf_verifier_env *env,
 			goto error;
 		break;
 	case BPF_MAP_TYPE_ARRAY_OF_MAPS:
-	case BPF_MAP_TYPE_HASH_OF_MAPS:
 		if (func_id != BPF_FUNC_map_lookup_elem)
 			goto error;
 		break;
+	case BPF_MAP_TYPE_HASH_OF_MAPS:
+		if (func_id != BPF_FUNC_map_lookup_elem &&
+		    func_id != BPF_FUNC_for_each_map_elem)
+			goto error;
+		break;
 	case BPF_MAP_TYPE_SOCKMAP:
 		if (func_id != BPF_FUNC_sk_redirect_map &&
 		    func_id != BPF_FUNC_sock_map_update &&
-- 
2.40.1

[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@manifault.com>
---
 .../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

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

[Alexei Starovoitov]

On Mon, Jun 5, 2023 at 1:05 PM David Vernet <void@manifault.com> wrote:
>
> +
> +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;
> +       }

This doesn't look right. 'val' is not 'u64 *'.
It probably should be 'struct bpf_map *',
so lookup shouldn't be necessary.

map_set_for_each_callback_args() just sets it to PTR_TO_MAP_VALUE.
It probably should be CONST_PTR_TO_MAP.
Also for normal hash map for_each_elem 'val' is writeable.
It shouldn't be in this case.
So to cleanly support iterating hash_of_maps patch 1 needs more work.

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

[David Vernet]

On Mon, Jun 05, 2023 at 03:07:21PM -0700, Alexei Starovoitov wrote:
> On Mon, Jun 5, 2023 at 1:05 PM David Vernet <void@manifault.com> wrote:
> >
> > +
> > +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;
> > +       }
> 
> This doesn't look right. 'val' is not 'u64 *'.
> It probably should be 'struct bpf_map *',
> so lookup shouldn't be necessary.
> 
> map_set_for_each_callback_args() just sets it to PTR_TO_MAP_VALUE.
> It probably should be CONST_PTR_TO_MAP.
> Also for normal hash map for_each_elem 'val' is writeable.
> It shouldn't be in this case.
> So to cleanly support iterating hash_of_maps patch 1 needs more work.

Thanks for the review and pointing out these issues. This all makes
sense, I'll incorporate it into v2.