[PATCH bpf-next 18/29] bpf: Add refcount_t to struct bpf_tramp_id

[Jiri Olsa <jolsa@redhat.com>]

Adding refcount_t to struct bpf_tramp_id so we can
track its allocation and safely use one object on
more places in following changes.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 include/linux/bpf.h     |  3 +++
 kernel/bpf/syscall.c    |  2 +-
 kernel/bpf/trampoline.c | 16 +++++++++++++---
 kernel/bpf/verifier.c   |  2 +-
 4 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index dda24339e4b1..04ada1d2495e 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -23,6 +23,7 @@
 #include <linux/slab.h>
 #include <linux/percpu-refcount.h>
 #include <linux/bpfptr.h>
+#include <linux/refcount.h>
 
 struct bpf_verifier_env;
 struct bpf_verifier_log;
@@ -677,6 +678,7 @@ struct bpf_tramp_id {
 	u32 obj_id;
 	u32 *id;
 	void **addr;
+	refcount_t refcnt;
 };
 
 struct bpf_tramp_node {
@@ -753,6 +755,7 @@ static __always_inline __nocfi unsigned int bpf_dispatcher_nop_func(
 #ifdef CONFIG_BPF_JIT
 struct bpf_tramp_id *bpf_tramp_id_alloc(u32 cnt);
 void bpf_tramp_id_free(struct bpf_tramp_id *id);
+void bpf_tramp_id_put(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);
 struct bpf_tramp_id *bpf_tramp_id_single(const struct bpf_prog *tgt_prog,
diff --git a/kernel/bpf/syscall.c b/kernel/bpf/syscall.c
index 0acf6cb0fdc7..bfbd81869818 100644
--- a/kernel/bpf/syscall.c
+++ b/kernel/bpf/syscall.c
@@ -2901,7 +2901,7 @@ static int bpf_tracing_prog_attach(struct bpf_prog *prog,
 out_put_prog:
 	if (tgt_prog_fd && tgt_prog)
 		bpf_prog_put(tgt_prog);
-	bpf_tramp_id_free(id);
+	bpf_tramp_id_put(id);
 	return err;
 }
 
diff --git a/kernel/bpf/trampoline.c b/kernel/bpf/trampoline.c
index e6a73088ecee..39600fb78c9e 100644
--- a/kernel/bpf/trampoline.c
+++ b/kernel/bpf/trampoline.c
@@ -100,6 +100,7 @@ struct bpf_tramp_id *bpf_tramp_id_alloc(u32 max)
 			return NULL;
 		}
 		id->max = max;
+		refcount_set(&id->refcnt, 1);
 	}
 	return id;
 }
@@ -133,10 +134,18 @@ struct bpf_tramp_id *bpf_tramp_id_single(const struct bpf_prog *tgt_prog,
 	return id;
 }
 
-void bpf_tramp_id_free(struct bpf_tramp_id *id)
+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);
@@ -162,7 +171,7 @@ static struct bpf_trampoline *bpf_trampoline_get(struct bpf_tramp_id *id)
 	if (!tr)
 		goto out;
 
-	tr->id = id;
+	tr->id = bpf_tramp_id_get(id);
 	INIT_HLIST_NODE(&tr->hlist);
 	hlist_add_head(&tr->hlist, head);
 	refcount_set(&tr->refcnt, 1);
@@ -592,6 +601,7 @@ void bpf_trampoline_put(struct bpf_trampoline *tr)
 	 * multiple rcu callbacks.
 	 */
 	hlist_del(&tr->hlist);
+	bpf_tramp_id_put(tr->id);
 	kfree(tr);
 out:
 	mutex_unlock(&trampoline_mutex);
@@ -663,7 +673,7 @@ void bpf_tramp_detach(struct bpf_tramp_attach *attach)
 	hlist_for_each_entry_safe(node, n, &attach->nodes, hlist_attach)
 		node_free(node);
 
-	bpf_tramp_id_free(attach->id);
+	bpf_tramp_id_put(attach->id);
 	kfree(attach);
 }
 
diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
index 8d56d43489aa..6a87180ac2bb 100644
--- a/kernel/bpf/verifier.c
+++ b/kernel/bpf/verifier.c
@@ -14001,7 +14001,7 @@ static int check_attach_btf_id(struct bpf_verifier_env *env)
 
 	attach = bpf_tramp_attach(id, tgt_prog, prog);
 	if (IS_ERR(attach)) {
-		bpf_tramp_id_free(id);
+		bpf_tramp_id_put(id);
 		return PTR_ERR(attach);
 	}
 
-- 
2.31.1