[PATCH bpf-next 14/29] bpf: Add support to store multiple ids in bpf_tramp_id object

[Jiri Olsa <jolsa@xxxxxxxxxx>]

Adding support to store multiple ids in bpf_tramp_id object,
to have id for trampolines with multiple functions assigned.

Extra array of u32 values is allocated within bpf_tramp_id
object allocation.

Signed-off-by: Jiri Olsa <jolsa@xxxxxxxxxx>
---
 include/linux/bpf.h     |  6 ++++--
 kernel/bpf/syscall.c    |  6 +++---
 kernel/bpf/trampoline.c | 39 +++++++++++++++++++++++++++++++--------
 kernel/bpf/verifier.c   |  2 +-
 4 files changed, 39 insertions(+), 14 deletions(-)

diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 2dbc00904a84..47e25d8be600 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -672,8 +672,10 @@ struct bpf_tramp_image {
 };
 
 struct bpf_tramp_id {
+	u32 max;
+	u32 cnt;
 	u32 obj_id;
-	u32 btf_id;
+	u32 *id;
 	void *addr;
 };
 
@@ -749,7 +751,7 @@ static __always_inline __nocfi unsigned int bpf_dispatcher_nop_func(
 	return bpf_func(ctx, insnsi);
 }
 #ifdef CONFIG_BPF_JIT
-struct bpf_tramp_id *bpf_tramp_id_alloc(void);
+struct bpf_tramp_id *bpf_tramp_id_alloc(u32 cnt);
 void bpf_tramp_id_free(struct bpf_tramp_id *id);
 bool bpf_tramp_id_is_empty(struct bpf_tramp_id *id);
 int bpf_tramp_id_is_equal(struct bpf_tramp_id *a, struct bpf_tramp_id *b);
diff --git a/kernel/bpf/syscall.c b/kernel/bpf/syscall.c
index a65c1862ab68..216fcce07326 100644
--- a/kernel/bpf/syscall.c
+++ b/kernel/bpf/syscall.c
@@ -2704,7 +2704,7 @@ static int bpf_tracing_link_fill_link_info(const struct bpf_link *link,
 
 	info->tracing.attach_type = tr_link->attach_type;
 	info->tracing.target_obj_id = attach->id->obj_id;
-	info->tracing.target_btf_id = attach->id->btf_id;
+	info->tracing.target_btf_id = attach->id->id[0];
 
 	return 0;
 }
@@ -2766,7 +2766,7 @@ static int bpf_tracing_prog_attach(struct bpf_prog *prog,
 			goto out_put_prog;
 		}
 
-		id = bpf_tramp_id_alloc();
+		id = bpf_tramp_id_alloc(1);
 		if (!id) {
 			err = -ENOMEM;
 			goto out_put_prog;
@@ -2829,7 +2829,7 @@ static int bpf_tracing_prog_attach(struct bpf_prog *prog,
 			goto out_unlock;
 		}
 
-		id = bpf_tramp_id_alloc();
+		id = bpf_tramp_id_alloc(1);
 		if (!id) {
 			err = -ENOMEM;
 			goto out_unlock;
diff --git a/kernel/bpf/trampoline.c b/kernel/bpf/trampoline.c
index 16fc4c14319b..d65f463c532d 100644
--- a/kernel/bpf/trampoline.c
+++ b/kernel/bpf/trampoline.c
@@ -60,27 +60,45 @@ void bpf_image_ksym_del(struct bpf_ksym *ksym)
 			   PAGE_SIZE, true, ksym->name);
 }
 
+static bool bpf_tramp_id_is_multi(struct bpf_tramp_id *id)
+{
+	return id->cnt > 1;
+}
+
 static u64 bpf_tramp_id_key(struct bpf_tramp_id *id)
 {
-	return ((u64) id->obj_id << 32) | id->btf_id;
+	if (bpf_tramp_id_is_multi(id))
+		return (u64) &id;
+	else
+		return ((u64) id->obj_id << 32) | id->id[0];
 }
 
 bool bpf_tramp_id_is_empty(struct bpf_tramp_id *id)
 {
-	return !id || (!id->obj_id && !id->btf_id);
+	return !id || id->cnt == 0;
 }
 
 int bpf_tramp_id_is_equal(struct bpf_tramp_id *a,
 			  struct bpf_tramp_id *b)
 {
-	return !memcmp(a, b, sizeof(*a));
+	return a->obj_id == b->obj_id && a->cnt == b->cnt &&
+	       !memcmp(a->id, b->id, a->cnt * sizeof(*a->id));
 }
 
-struct bpf_tramp_id *bpf_tramp_id_alloc(void)
+struct bpf_tramp_id *bpf_tramp_id_alloc(u32 max)
 {
 	struct bpf_tramp_id *id;
 
-	return kzalloc(sizeof(*id), GFP_KERNEL);
+	id = kzalloc(sizeof(*id), GFP_KERNEL);
+	if (id) {
+		id->id = kzalloc(sizeof(u32) * max, GFP_KERNEL);
+		if (!id->id) {
+			kfree(id);
+			return NULL;
+		}
+		id->max = max;
+	}
+	return id;
 }
 
 void bpf_tramp_id_init(struct bpf_tramp_id *id,
@@ -91,11 +109,15 @@ void bpf_tramp_id_init(struct bpf_tramp_id *id,
 		id->obj_id = tgt_prog->aux->id;
 	else
 		id->obj_id = btf_obj_id(btf);
-	id->btf_id = btf_id;
+	id->id[0] = btf_id;
+	id->cnt = 1;
 }
 
 void bpf_tramp_id_free(struct bpf_tramp_id *id)
 {
+	if (!id)
+		return;
+	kfree(id->id);
 	kfree(id);
 }
 
@@ -362,7 +384,8 @@ bpf_tramp_image_alloc(struct bpf_tramp_id *id, u32 idx)
 	ksym = &im->ksym;
 	INIT_LIST_HEAD_RCU(&ksym->lnode);
 	key = bpf_tramp_id_key(id);
-	snprintf(ksym->name, KSYM_NAME_LEN, "bpf_trampoline_%llu_%u", key, idx);
+	snprintf(ksym->name, KSYM_NAME_LEN, "bpf_trampoline_%llu_%u%s", key, idx,
+		 bpf_tramp_id_is_multi(id) ? "_multi" : "");
 	bpf_image_ksym_add(image, ksym);
 	return im;
 
@@ -597,7 +620,7 @@ struct bpf_tramp_attach *bpf_tramp_attach(struct bpf_tramp_id *id,
 	if (!node)
 		goto out;
 
-	err = bpf_check_attach_model(prog, tgt_prog, id->btf_id, &tr->func.model);
+	err = bpf_check_attach_model(prog, tgt_prog, id->id[0], &tr->func.model);
 	if (err)
 		goto out;
 
diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
index e05f39fd2708..1903d5d256b6 100644
--- a/kernel/bpf/verifier.c
+++ b/kernel/bpf/verifier.c
@@ -13995,7 +13995,7 @@ static int check_attach_btf_id(struct bpf_verifier_env *env)
 		return -EINVAL;
 	}
 
-	id = bpf_tramp_id_alloc();
+	id = bpf_tramp_id_alloc(1);
 	if (!id)
 		return -ENOMEM;
 
-- 
2.31.1