[PATCH bpf-next v2 11/21] bpf: Add tracing_multi link support for bpf progs

[Leon Hwang <leon.hwang@linux.dev>]

Similar to the tracing_multi link support for kernel functions, add
support for bpf progs by reusing the 'ids' fields as BTF ids, and adding
fds array to indicate the target bpf prog. Then, for each pair of fd+id,
they will be the attach targets of tracing_multi link.

When attaching bpf prog, it attaches to the targets via
bpf_trampoline_multi_prog_attach(), that would work on all
trampoline-supported archs.

Since tracing_multi prog cannot recognize its targets at load time, reject
attaching to any other tracing prog to avoid two-levels nested tracing
attachment: e.g. fentry <- fentry.multi <- another fentry.

Assisted-by: Codex:gpt-5.6-sol
Signed-off-by: Leon Hwang <leon.hwang@linux.dev>
---
 include/uapi/linux/bpf.h       |   2 +
 kernel/trace/bpf_trace.c       | 177 ++++++++++++++++++++++++++++-----
 tools/include/uapi/linux/bpf.h |   2 +
 3 files changed, 156 insertions(+), 25 deletions(-)

diff --git a/include/uapi/linux/bpf.h b/include/uapi/linux/bpf.h
index 6f70db7515dd..94302a3b680c 100644
--- a/include/uapi/linux/bpf.h
+++ b/include/uapi/linux/bpf.h
@@ -1927,6 +1927,8 @@ union bpf_attr {
 				__aligned_u64	ids;
 				__aligned_u64	cookies;
 				__u32		cnt;
+				__u32		:32;
+				__aligned_u64	fds;
 			} tracing_multi;
 		};
 	} link_create;
diff --git a/kernel/trace/bpf_trace.c b/kernel/trace/bpf_trace.c
index 69559158ab0b..d4af42e25fe5 100644
--- a/kernel/trace/bpf_trace.c
+++ b/kernel/trace/bpf_trace.c
@@ -3679,12 +3679,19 @@ __bpf_kfunc int bpf_copy_from_user_task_str_dynptr(const struct bpf_dynptr *dptr
 
 __bpf_kfunc_end_defs();
 
+static void bpf_put_progs(struct bpf_prog **progs, int cnt);
+
 static void bpf_tracing_multi_link_release(struct bpf_link *link)
 {
 	struct bpf_tracing_multi_link *tr_link =
 		container_of(link, struct bpf_tracing_multi_link, link);
 
-	bpf_trampoline_multi_detach(link->prog, tr_link);
+	if (tr_link->tgt_progs) {
+		bpf_trampoline_multi_prog_detach(tr_link);
+		bpf_put_progs(tr_link->tgt_progs, tr_link->nodes_cnt);
+	} else {
+		bpf_trampoline_multi_detach(link->prog, tr_link);
+	}
 }
 
 static void bpf_tracing_multi_link_dealloc(struct bpf_link *link)
@@ -3692,6 +3699,8 @@ static void bpf_tracing_multi_link_dealloc(struct bpf_link *link)
 	struct bpf_tracing_multi_link *tr_link =
 		container_of(link, struct bpf_tracing_multi_link, link);
 
+	kvfree(tr_link->pokes);
+	kvfree(tr_link->tgt_progs);
 	kvfree(tr_link->fexits);
 	kvfree(tr_link->cookies);
 	kvfree(tr_link);
@@ -3787,64 +3796,129 @@ static const struct bpf_link_ops bpf_tracing_multi_link_lops = {
 #endif
 };
 
-static int ids_cmp_r(const void *pa, const void *pb, const void *priv __maybe_unused)
+struct tracing_multi_sort_data {
+	u64 *keys;
+	struct bpf_prog **progs;
+	u32 *ids;
+	u64 *cookies;
+};
+
+static int keys_cmp_r(const void *pa, const void *pb, const void *priv __maybe_unused)
 {
-	u32 a = *(u32 *) pa;
-	u32 b = *(u32 *) pb;
+	const u64 *key_a = pa, *key_b = pb;
 
-	return (a > b) - (a < b);
+	return (*key_a > *key_b) - (*key_a < *key_b);
 }
 
-static void ids_swap_r(void *a, void *b, int size __maybe_unused,
-		       const void *priv __maybe_unused)
+static void keys_swap_r(void *a, void *b, int size __maybe_unused,
+			const void *priv __maybe_unused)
 {
+	const struct tracing_multi_sort_data *data = priv;
+	struct bpf_prog **prog_a, **prog_b;
 	u64 *cookie_a, *cookie_b, *cookies;
-	u32 *id_a = a, *id_b = b, *ids;
-	void **data = (void **) priv;
+	u64 *key_a = a, *key_b = b, *keys;
+	u32 *id_a, *id_b, *ids;
+
+	keys    = data->keys;
+	ids     = data->ids;
+	cookies = data->cookies;
+	id_a    = ids + (key_a - keys);
+	id_b    = ids + (key_b - keys);
 
-	ids     = data[0];
-	cookies = data[1];
+	if (data->progs) {
+		prog_a = data->progs + (key_a - keys);
+		prog_b = data->progs + (key_b - keys);
+		swap(*prog_a, *prog_b);
+	}
 
 	if (cookies) {
 		cookie_a = cookies + (id_a - ids);
 		cookie_b = cookies + (id_b - ids);
 		swap(*cookie_a, *cookie_b);
 	}
+
 	swap(*id_a, *id_b);
+	swap(*key_a, *key_b);
 }
 
-static int check_dup_ids(u32 *ids, u64 *cookies, u32 cnt)
+static int check_dup_keys(u64 *keys, struct bpf_prog **progs, u32 *ids,
+			  u64 *cookies, u32 cnt)
 {
-	void *data[2] = { ids, cookies };
-	int err = 0;
+	struct tracing_multi_sort_data data = {
+		.keys = keys,
+		.progs = progs,
+		.ids = ids,
+		.cookies = cookies,
+	};
 
 	/*
-	 * Sort ids array (together with cookies array if defined)
-	 * and check it for duplicates. The ids and cookies arrays
-	 * are left sorted.
+	 * Sort trampoline keys together with target programs, ids, and cookies,
+	 * then check for duplicates. The arrays are left sorted.
 	 */
-	sort_r_nonatomic(ids, cnt, sizeof(ids[0]), ids_cmp_r, ids_swap_r, data);
+	sort_r_nonatomic(keys, cnt, sizeof(keys[0]), keys_cmp_r, keys_swap_r, &data);
+
+	for (int i = 1; i < cnt; i++)
+		if (keys[i] == keys[i - 1])
+			return -EINVAL;
 
-	for (int i = 1; i < cnt; i++) {
-		if (ids[i] == ids[i - 1]) {
+	return 0;
+}
+
+static void bpf_put_progs(struct bpf_prog **progs, int cnt)
+{
+	int i;
+
+	if (!progs)
+		return;
+
+	for (i = 0; i < cnt; i++)
+		bpf_prog_put(progs[i]);
+}
+
+static int bpf_get_progs(struct bpf_prog **progs, int *fds, int cnt)
+{
+	int err, i, put_cnt = 0;
+	struct bpf_prog *prog;
+
+	for (i = 0; i < cnt; i++) {
+		prog = bpf_prog_get(fds[i]);
+		if (IS_ERR(prog)) {
+			err = PTR_ERR(prog);
+			goto error;
+		}
+
+		progs[i] = prog;
+		put_cnt++;
+
+		if (prog->type == BPF_PROG_TYPE_TRACING) {
 			err = -EINVAL;
-			break;
+			goto error;
 		}
 	}
+	return 0;
+
+error:
+	bpf_put_progs(progs, put_cnt);
 	return err;
 }
 
 int bpf_tracing_multi_attach(struct bpf_prog *prog, const union bpf_attr *attr)
 {
 	struct bpf_tracing_multi_link *link = NULL;
+	struct btf *btf = prog->aux->attach_btf;
 	struct bpf_tramp_node *fexits = NULL;
 	struct bpf_link_primer link_primer;
+	struct bpf_text_poke **pokes = NULL;
+	u64 *cookies = NULL, *keys = NULL;
+	struct bpf_prog **progs = NULL;
 	u32 cnt, *ids = NULL;
 	u64 __user *ucookies;
-	u64 *cookies = NULL;
 	u32 __user *uids;
+	int __user *ufds;
+	int *fds = NULL;
 	int err;
 
+	ufds = u64_to_user_ptr(attr->link_create.tracing_multi.fds);
 	uids = u64_to_user_ptr(attr->link_create.tracing_multi.ids);
 	cnt = attr->link_create.tracing_multi.cnt;
 
@@ -3855,7 +3929,7 @@ int bpf_tracing_multi_attach(struct bpf_prog *prog, const union bpf_attr *attr)
 	if (attr->link_create.flags || attr->link_create.target_fd)
 		return -EINVAL;
 
-	if (!IS_ENABLED(CONFIG_HAVE_SINGLE_FTRACE_DIRECT_OPS))
+	if (!ufds && !IS_ENABLED(CONFIG_HAVE_SINGLE_FTRACE_DIRECT_OPS))
 		return -EOPNOTSUPP;
 
 	ids = kvmalloc_objs(*ids, cnt);
@@ -3867,6 +3941,37 @@ int bpf_tracing_multi_attach(struct bpf_prog *prog, const union bpf_attr *attr)
 		goto error;
 	}
 
+	if (ufds) {
+		fds = kvmalloc_objs(*fds, cnt);
+		if (!fds) {
+			err = -ENOMEM;
+			goto error;
+		}
+
+		if (copy_from_user(fds, ufds, cnt * sizeof(*fds))) {
+			err = -EFAULT;
+			goto error;
+		}
+
+		progs = kvmalloc_objs(*progs, cnt);
+		if (!progs) {
+			err = -ENOMEM;
+			goto error;
+		}
+
+		err = bpf_get_progs(progs, fds, cnt);
+		if (err) {
+			cnt = 0;
+			goto error;
+		}
+
+		pokes = kvmalloc_objs(*pokes, cnt);
+		if (!pokes) {
+			err = -ENOMEM;
+			goto error;
+		}
+	}
+
 	ucookies = u64_to_user_ptr(attr->link_create.tracing_multi.cookies);
 	if (ucookies) {
 		cookies = kvmalloc_objs(*cookies, cnt);
@@ -3880,7 +3985,16 @@ int bpf_tracing_multi_attach(struct bpf_prog *prog, const union bpf_attr *attr)
 		}
 	}
 
-	err = check_dup_ids(ids, cookies, cnt);
+	keys = kvmalloc_objs(*keys, cnt);
+	if (!keys) {
+		err = -ENOMEM;
+		goto error;
+	}
+
+	for (int i = 0; i < cnt; i++)
+		keys[i] = bpf_trampoline_compute_key(progs ? progs[i] : NULL, btf, ids[i]);
+
+	err = check_dup_keys(keys, progs, ids, cookies, cnt);
 	if (err)
 		goto error;
 
@@ -3908,19 +4022,32 @@ int bpf_tracing_multi_attach(struct bpf_prog *prog, const union bpf_attr *attr)
 	link->nodes_cnt = cnt;
 	link->cookies = cookies;
 	link->fexits = fexits;
+	link->tgt_progs = progs;
+	link->pokes = pokes;
 
-	err = bpf_trampoline_multi_attach(prog, ids, link);
+	if (progs)
+		err = bpf_trampoline_multi_prog_attach(prog, ids, keys, link);
+	else
+		err = bpf_trampoline_multi_attach(prog, ids, link);
 	kvfree(ids);
+	kvfree(keys);
+	kvfree(fds);
 	if (err) {
+		bpf_put_progs(progs, cnt);
 		bpf_link_cleanup(&link_primer);
 		return err;
 	}
 	return bpf_link_settle(&link_primer);
 
 error:
+	bpf_put_progs(progs, cnt);
 	kvfree(fexits);
 	kvfree(cookies);
 	kvfree(ids);
+	kvfree(keys);
+	kvfree(fds);
+	kvfree(progs);
+	kvfree(pokes);
 	kvfree(link);
 	return err;
 }
diff --git a/tools/include/uapi/linux/bpf.h b/tools/include/uapi/linux/bpf.h
index 6f70db7515dd..94302a3b680c 100644
--- a/tools/include/uapi/linux/bpf.h
+++ b/tools/include/uapi/linux/bpf.h
@@ -1927,6 +1927,8 @@ union bpf_attr {
 				__aligned_u64	ids;
 				__aligned_u64	cookies;
 				__u32		cnt;
+				__u32		:32;
+				__aligned_u64	fds;
 			} tracing_multi;
 		};
 	} link_create;
-- 
2.55.0