[RFC PATCH bpf-next 05/17] bpf: Add bpf_tramp_id object

[Jiri Olsa <jolsa@kernel.org>]

Adding bpf_tramp_id object that allows to store multiple BTF
function ids together with their resolved addresses.

It will be used in following changes to identify and attach
multiple functions to trampolines.

The bpf_tramp_id object will be shared between trampoline and
link in following changes, so it keeps refcount for that.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 include/linux/bpf.h     | 12 ++++++++++++
 kernel/bpf/trampoline.c | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+)

diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 32168ea92551..a5738d57f6bd 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -27,6 +27,7 @@
 #include <linux/bpfptr.h>
 #include <linux/btf.h>
 #include <linux/rcupdate_trace.h>
+#include <linux/refcount.h>
 
 struct bpf_verifier_env;
 struct bpf_verifier_log;
@@ -846,6 +847,15 @@ struct bpf_tramp_image {
 	};
 };
 
+struct bpf_tramp_id {
+	u32 max;
+	u32 cnt;
+	u32 obj_id;
+	u32 *id;
+	void **addr;
+	refcount_t refcnt;
+};
+
 struct bpf_shim_tramp_link;
 
 struct bpf_trampoline {
@@ -917,6 +927,8 @@ int bpf_trampoline_unlink_prog(struct bpf_tramp_prog *tp, struct bpf_trampoline
 struct bpf_trampoline *bpf_trampoline_get(u64 key,
 					  struct bpf_attach_target_info *tgt_info);
 void bpf_trampoline_put(struct bpf_trampoline *tr);
+struct bpf_tramp_id *bpf_tramp_id_alloc(u32 cnt);
+void bpf_tramp_id_put(struct bpf_tramp_id *id);
 int arch_prepare_bpf_dispatcher(void *image, s64 *funcs, int num_funcs);
 #define BPF_DISPATCHER_INIT(_name) {				\
 	.mutex = __MUTEX_INITIALIZER(_name.mutex),		\
diff --git a/kernel/bpf/trampoline.c b/kernel/bpf/trampoline.c
index 56899d63c08c..c0983ff5aa3a 100644
--- a/kernel/bpf/trampoline.c
+++ b/kernel/bpf/trampoline.c
@@ -149,6 +149,44 @@ void bpf_image_ksym_del(struct bpf_ksym *ksym)
 			   PAGE_SIZE, true, ksym->name);
 }
 
+struct bpf_tramp_id *bpf_tramp_id_alloc(u32 max)
+{
+	struct bpf_tramp_id *id;
+
+	id = kzalloc(sizeof(*id), GFP_KERNEL);
+	if (id) {
+		id->id = kcalloc(max, sizeof(u32), GFP_KERNEL);
+		id->addr = kcalloc(max, sizeof(*id->addr), GFP_KERNEL);
+		if (!id->id || !id->addr) {
+			kfree(id->id);
+			kfree(id->addr);
+			kfree(id);
+			return NULL;
+		}
+		id->max = max;
+		refcount_set(&id->refcnt, 1);
+	}
+	return id;
+}
+
+__maybe_unused
+static struct bpf_tramp_id *bpf_tramp_id_get(struct bpf_tramp_id *id)
+{
+	refcount_inc(&id->refcnt);
+	return id;
+}
+
+void bpf_tramp_id_put(struct bpf_tramp_id *id)
+{
+	if (!id)
+		return;
+	if (!refcount_dec_and_test(&id->refcnt))
+		return;
+	kfree(id->addr);
+	kfree(id->id);
+	kfree(id);
+}
+
 static struct bpf_trampoline *bpf_trampoline_lookup(u64 key)
 {
 	struct bpf_trampoline *tr;
-- 
2.37.1