[RFC PATCH bpf-next 1/2] libbpf: support selective kernel module BTF loading via .kmod_btfs section

[Fuyu Zhao]

Add support for a new ELF section ".kmod_btfs" that allows BPF programs
to declare which kernel modules need BTF loading. This avoids loading
all module BTFs and speeds up program load when only a subset of modules
is needed.

Specifically, this patch introduces two internal functions:
  - bpf_object__collect_kmod_btf_names(): parses the ".kmod_btfs"
    section and collects the declared module names.
  - is_kmod_btf_needed(): determines whether a given module's BTF
    should be loaded, allowing libbpf to skip unneeded modules.

Signed-off-by: Fuyu Zhao <zhaofuyu@xxxxxxxx>
---
 tools/lib/bpf/bpf_helpers.h |  14 +++++
 tools/lib/bpf/libbpf.c      | 112 ++++++++++++++++++++++++++++++++++++
 2 files changed, 126 insertions(+)

diff --git a/tools/lib/bpf/bpf_helpers.h b/tools/lib/bpf/bpf_helpers.h
index 9d160b5b9c0e..171ea055cd32 100644
--- a/tools/lib/bpf/bpf_helpers.h
+++ b/tools/lib/bpf/bpf_helpers.h
@@ -188,6 +188,20 @@ enum libbpf_tristate {
 	TRI_MODULE = 2,
 };
 
+/* Helper typedef for declaring kernel module names that need BTF loading.
+ *
+ * Usage: define an array in the ".kmod_btfs" ELF section to specify
+ * which modules need BTF loading:
+ *
+ *   DEFINE_KMOD_BTFS(_needed_kmods)= { "module1", "module2", ... };
+ *
+ * This avoids unnecessary BTF loading and speeds up the BPF program
+ * load process.
+ */
+#define KMOD_NAME_LEN 64
+#define DEFINE_KMOD_BTFS(name) \
+	SEC(".kmod_btfs") char name[][KMOD_NAME_LEN]
+
 #define __kconfig __attribute__((section(".kconfig")))
 #define __ksym __attribute__((section(".ksyms")))
 #define __kptr_untrusted __attribute__((btf_type_tag("kptr_untrusted")))
diff --git a/tools/lib/bpf/libbpf.c b/tools/lib/bpf/libbpf.c
index 514e4e9daa82..5cefe82e4a67 100644
--- a/tools/lib/bpf/libbpf.c
+++ b/tools/lib/bpf/libbpf.c
@@ -548,6 +548,7 @@ struct bpf_struct_ops {
 #define STRUCT_OPS_SEC ".struct_ops"
 #define STRUCT_OPS_LINK_SEC ".struct_ops.link"
 #define ARENA_SEC ".addr_space.1"
+#define KMODS_BTFS_SEC ".kmod_btfs"
 
 enum libbpf_map_type {
 	LIBBPF_MAP_UNSPEC,
@@ -703,6 +704,12 @@ enum bpf_object_state {
 	OBJ_LOADED,
 };
 
+/* Should match typedef in bpf_helpers.h */
+#define KMOD_NAME_LEN 64
+
+#define KMODS_BTF_UNLOADED 0UL
+#define KMODS_BTF_LOADED   1UL
+
 struct bpf_object {
 	char name[BPF_OBJ_NAME_LEN];
 	char license[64];
@@ -779,6 +786,14 @@ struct bpf_object {
 	char *token_path;
 	int token_fd;
 
+	/* kernel module BTFs to load, declared in ".kmod_btfs" ELF section */
+	struct {
+		char (*data)[KMOD_NAME_LEN];
+		size_t nr_names;
+		size_t nr_loaded;
+		struct hashmap *hashmap;
+	} *kmod_btfs;
+
 	char path[];
 };
 
@@ -901,6 +916,68 @@ bpf_object__init_prog(struct bpf_object *obj, struct bpf_program *prog,
 	return -ENOMEM;
 }
 
+static size_t mod_name_hash_fn(long key, void *ctx)
+{
+	return str_hash((char *)key);
+}
+
+static bool mod_name_equal_fn(long key1, long key2, void *ctx)
+{
+	return strcmp((char *)key1, (char *)key2) == 0;
+}
+
+static int
+bpf_object__collect_kmod_btf_names(struct bpf_object *obj, Elf_Data *sec_data,
+				   const char *sec_name)
+{
+	int module_cnt, i, err = 0;
+
+	if (sec_data->d_size % KMOD_NAME_LEN != 0) {
+		pr_warn("sec '%s': size %zu should be multiple of %d\n",
+			sec_name, sec_data->d_size, KMOD_NAME_LEN);
+		return -EINVAL;
+	}
+
+	module_cnt = sec_data->d_size / KMOD_NAME_LEN;
+	obj->kmod_btfs = calloc(1, sizeof(*obj->kmod_btfs));
+	if (!obj->kmod_btfs)
+		return -ENOMEM;
+
+	obj->kmod_btfs->data = calloc(module_cnt, KMOD_NAME_LEN);
+	if (!obj->kmod_btfs->data)
+		goto err_out;
+	memcpy(obj->kmod_btfs->data, sec_data->d_buf, sec_data->d_size);
+
+	obj->kmod_btfs->hashmap = hashmap__new(mod_name_hash_fn,
+					       mod_name_equal_fn, NULL);
+	if (IS_ERR(obj->kmod_btfs->hashmap)) {
+		err = PTR_ERR(obj->kmod_btfs->hashmap);
+		goto err_out;
+	}
+
+	for (i = 0; i < module_cnt; i++) {
+		obj->kmod_btfs->data[i][KMOD_NAME_LEN - 1] = '\0';
+		if (hashmap__find(obj->kmod_btfs->hashmap,
+				  obj->kmod_btfs->data[i], NULL)) {
+			pr_warn("sec '%s': ignored duplicate module '%s'\n",
+				sec_name, obj->kmod_btfs->data[i]);
+			continue;
+		}
+		err = hashmap__set(obj->kmod_btfs->hashmap, obj->kmod_btfs->data[i],
+				   KMODS_BTF_UNLOADED, NULL, NULL);
+		if (err)
+			goto err_out;
+		obj->kmod_btfs->nr_names++;
+	}
+	return 0;
+
+err_out:
+	hashmap__free(obj->kmod_btfs->hashmap);
+	zfree(&obj->kmod_btfs->data);
+	zfree(&obj->kmod_btfs);
+	return err;
+}
+
 static int
 bpf_object__add_programs(struct bpf_object *obj, Elf_Data *sec_data,
 			 const char *sec_name, int sec_idx)
@@ -4034,6 +4111,10 @@ static int bpf_object__elf_collect(struct bpf_object *obj)
 				memcpy(obj->jumptables_data, data->d_buf, data->d_size);
 				obj->jumptables_data_sz = data->d_size;
 				obj->efile.jumptables_data_shndx = idx;
+			} else if (strcmp(name, KMODS_BTFS_SEC) == 0) {
+				err = bpf_object__collect_kmod_btf_names(obj, data, name);
+				if (err)
+					return err;
 			} else {
 				pr_info("elf: skipping unrecognized data section(%d) %s\n",
 					idx, name);
@@ -5803,6 +5884,21 @@ int bpf_core_add_cands(struct bpf_core_cand *local_cand,
 	return 0;
 }
 
+static bool is_kmod_btf_needed(struct bpf_object *obj, const char *name)
+{
+	uintptr_t val;
+
+	if (!hashmap__find(obj->kmod_btfs->hashmap, name, &val))
+		return false;
+
+	if (val == KMODS_BTF_LOADED)
+		return false;
+
+	hashmap__set(obj->kmod_btfs->hashmap, name, KMODS_BTF_LOADED, NULL, NULL);
+	obj->kmod_btfs->nr_loaded++;
+	return true;
+}
+
 static int load_module_btfs(struct bpf_object *obj)
 {
 	struct bpf_btf_info info;
@@ -5867,6 +5963,12 @@ static int load_module_btfs(struct bpf_object *obj)
 			continue;
 		}
 
+		if (obj->kmod_btfs && obj->kmod_btfs->hashmap &&
+		    !is_kmod_btf_needed(obj, name)) {
+			close(fd);
+			continue;
+		}
+
 		btf = btf_get_from_fd(fd, obj->btf_vmlinux);
 		err = libbpf_get_error(btf);
 		if (err) {
@@ -5891,6 +5993,10 @@ static int load_module_btfs(struct bpf_object *obj)
 			break;
 		}
 		obj->btf_module_cnt++;
+
+		if (obj->kmod_btfs &&
+		    obj->kmod_btfs->nr_names == obj->kmod_btfs->nr_loaded)
+			break;
 	}
 
 	if (err) {
@@ -9030,6 +9136,12 @@ static void bpf_object_cleanup_btf(struct bpf_object *obj)
 	/* clean up vmlinux BTF */
 	btf__free(obj->btf_vmlinux);
 	obj->btf_vmlinux = NULL;
+
+	if (obj->kmod_btfs) {
+		hashmap__free(obj->kmod_btfs->hashmap);
+		zfree(&obj->kmod_btfs->data);
+		zfree(&obj->kmod_btfs);
+	}
 }
 
 static void bpf_object_post_load_cleanup(struct bpf_object *obj)
-- 
2.34.1