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

[David Vernet <void@manifault.com>]

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