[PATCH bpf-next 5/5] selftests/bpf: add test cases for multiple attach of tracing program

[Menglong Dong <dongmenglong.8@bytedance.com>]

In this commit, we add the testcases for multiple attaching of tracing,
include FENTRY, FEXIT, MODIFY_RETURN.

Signed-off-by: Menglong Dong <dongmenglong.8@bytedance.com>
---
 .../selftests/bpf/bpf_testmod/bpf_testmod.c   |  49 ++++++
 .../bpf/prog_tests/tracing_multi_attach.c     | 155 ++++++++++++++++++
 .../selftests/bpf/progs/tracing_multi_test.c  |  72 ++++++++
 3 files changed, 276 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/prog_tests/tracing_multi_attach.c
 create mode 100644 tools/testing/selftests/bpf/progs/tracing_multi_test.c

diff --git a/tools/testing/selftests/bpf/bpf_testmod/bpf_testmod.c b/tools/testing/selftests/bpf/bpf_testmod/bpf_testmod.c
index 66787e99ba1b..237eeb7daa07 100644
--- a/tools/testing/selftests/bpf/bpf_testmod/bpf_testmod.c
+++ b/tools/testing/selftests/bpf/bpf_testmod/bpf_testmod.c
@@ -98,12 +98,61 @@ bpf_testmod_test_struct_arg_8(u64 a, void *b, short c, int d, void *e,
 	return bpf_testmod_test_struct_arg_result;
 }
 
+noinline int
+bpf_testmod_test_struct_arg_9(struct bpf_testmod_struct_arg_2 a,
+			      struct bpf_testmod_struct_arg_1 b) {
+	bpf_testmod_test_struct_arg_result = a.a + a.b  + b.a;
+	return bpf_testmod_test_struct_arg_result;
+}
+
+noinline int
+bpf_testmod_test_struct_arg_10(int a, struct bpf_testmod_struct_arg_2 b) {
+	bpf_testmod_test_struct_arg_result = a + b.a + b.b;
+	return bpf_testmod_test_struct_arg_result;
+}
+
+noinline struct bpf_testmod_struct_arg_2 *
+bpf_testmod_test_struct_arg_11(int a, struct bpf_testmod_struct_arg_2 b, int c) {
+	bpf_testmod_test_struct_arg_result = a + b.a + b.b + c;
+	return (void *)bpf_testmod_test_struct_arg_result;
+}
+
+noinline int
+bpf_testmod_test_struct_arg_12(int a, struct bpf_testmod_struct_arg_2 b, int *c) {
+	bpf_testmod_test_struct_arg_result = a + b.a + b.b + *c;
+	return bpf_testmod_test_struct_arg_result;
+}
+
 noinline int
 bpf_testmod_test_arg_ptr_to_struct(struct bpf_testmod_struct_arg_1 *a) {
 	bpf_testmod_test_struct_arg_result = a->a;
 	return bpf_testmod_test_struct_arg_result;
 }
 
+noinline int
+bpf_testmod_test_arg_ptr_1(struct bpf_testmod_struct_arg_1 *a) {
+	bpf_testmod_test_struct_arg_result = a->a;
+	return bpf_testmod_test_struct_arg_result;
+}
+
+noinline int
+bpf_testmod_test_arg_ptr_2(struct bpf_testmod_struct_arg_2 *a) {
+	bpf_testmod_test_struct_arg_result = a->a + a->b;
+	return bpf_testmod_test_struct_arg_result;
+}
+
+noinline int
+bpf_testmod_test_arg_ptr_3(int a, struct bpf_testmod_struct_arg_2 *b) {
+	bpf_testmod_test_struct_arg_result = a + b->a + b->b;
+	return bpf_testmod_test_struct_arg_result;
+}
+
+noinline int
+bpf_testmod_test_arg_ptr_4(struct bpf_testmod_struct_arg_2 *a, int b) {
+	bpf_testmod_test_struct_arg_result = a->a + a->b + b;
+	return bpf_testmod_test_struct_arg_result;
+}
+
 __bpf_kfunc void
 bpf_testmod_test_mod_kfunc(int i)
 {
diff --git a/tools/testing/selftests/bpf/prog_tests/tracing_multi_attach.c b/tools/testing/selftests/bpf/prog_tests/tracing_multi_attach.c
new file mode 100644
index 000000000000..6162d41cca9e
--- /dev/null
+++ b/tools/testing/selftests/bpf/prog_tests/tracing_multi_attach.c
@@ -0,0 +1,155 @@
+// SPDX-License-Identifier: GPL-2.0
+/* Copyright (c) 2024 Bytedance. */
+
+#include <test_progs.h>
+#include "tracing_multi_test.skel.h"
+
+struct test_item {
+	char *prog;
+	char *target;
+	int attach_type;
+	bool success;
+	int link_fd;
+};
+
+static struct test_item test_items[] = {
+	{
+		.prog = "fentry_test1", .target = "bpf_testmod_test_struct_arg_9",
+		.attach_type = BPF_TRACE_FENTRY, .success = true,
+	},
+	{
+		.prog = "fentry_test1", .target = "bpf_testmod_test_struct_arg_1",
+		.attach_type = BPF_TRACE_FENTRY, .success = false,
+	},
+	{
+		.prog = "fentry_test1", .target = "bpf_testmod_test_struct_arg_2",
+		.attach_type = BPF_TRACE_FENTRY, .success = false,
+	},
+	{
+		.prog = "fentry_test1", .target = "bpf_testmod_test_arg_ptr_2",
+		.attach_type = BPF_TRACE_FENTRY, .success = false,
+	},
+	{
+		.prog = "fentry_test2", .target = "bpf_testmod_test_struct_arg_2",
+		.attach_type = BPF_TRACE_FENTRY, .success = false,
+	},
+	{
+		.prog = "fentry_test2", .target = "bpf_testmod_test_struct_arg_10",
+		.attach_type = BPF_TRACE_FENTRY, .success = true,
+	},
+	{
+		.prog = "fentry_test2", .target = "bpf_testmod_test_struct_arg_9",
+		.attach_type = BPF_TRACE_FENTRY, .success = false,
+	},
+	{
+		.prog = "fentry_test2", .target = "bpf_testmod_test_arg_ptr_3",
+		.attach_type = BPF_TRACE_FENTRY, .success = false,
+	},
+	{
+		.prog = "fentry_test3", .target = "bpf_testmod_test_arg_ptr_3",
+		.attach_type = BPF_TRACE_FENTRY, .success = false,
+	},
+	{
+		.prog = "fentry_test3", .target = "bpf_testmod_test_arg_ptr_4",
+		.attach_type = BPF_TRACE_FENTRY, .success = true,
+	},
+	{
+		.prog = "fentry_test4", .target = "bpf_testmod_test_struct_arg_4",
+		.attach_type = BPF_TRACE_FENTRY, .success = true,
+	},
+	{
+		.prog = "fentry_test4", .target = "bpf_testmod_test_struct_arg_2",
+		.attach_type = BPF_TRACE_FENTRY, .success = true,
+	},
+	{
+		.prog = "fentry_test4", .target = "bpf_testmod_test_struct_arg_12",
+		.attach_type = BPF_TRACE_FENTRY, .success = false,
+	},
+	{
+		.prog = "fexit_test1", .target = "bpf_testmod_test_struct_arg_2",
+		.attach_type = BPF_TRACE_FEXIT, .success = true,
+	},
+	{
+		.prog = "fexit_test1", .target = "bpf_testmod_test_struct_arg_3",
+		.attach_type = BPF_TRACE_FEXIT, .success = true,
+	},
+	{
+		.prog = "fexit_test1", .target = "bpf_testmod_test_struct_arg_4",
+		.attach_type = BPF_TRACE_FEXIT, .success = false,
+	},
+	{
+		.prog = "fexit_test2", .target = "bpf_testmod_test_struct_arg_10",
+		.attach_type = BPF_TRACE_FEXIT, .success = false,
+	},
+	{
+		.prog = "fexit_test2", .target = "bpf_testmod_test_struct_arg_11",
+		.attach_type = BPF_TRACE_FEXIT, .success = false,
+	},
+	{
+		.prog = "fexit_test2", .target = "bpf_testmod_test_struct_arg_12",
+		.attach_type = BPF_TRACE_FEXIT, .success = true,
+	},
+	{
+		.prog = "fmod_ret_test1", .target = "bpf_modify_return_test2",
+		.attach_type = BPF_MODIFY_RETURN, .success = true,
+	},
+};
+
+static int do_test_item(struct tracing_multi_test *skel, struct test_item *item)
+{
+	LIBBPF_OPTS(bpf_link_create_opts, link_opts);
+	struct bpf_program *prog;
+	int err, btf_fd = 0, btf_type_id;
+
+	err = libbpf_find_kernel_btf_id(item->target, item->attach_type,
+					&btf_fd, &btf_type_id);
+	if (!ASSERT_OK(err, "find_vmlinux_btf_id"))
+		return -1;
+
+	link_opts.target_btf_id = btf_type_id;
+	prog = bpf_object__find_program_by_name(skel->obj, item->prog);
+	if (!ASSERT_OK_PTR(prog, "find_program_by_name"))
+		return -1;
+
+	err = bpf_link_create(bpf_program__fd(prog), btf_fd, item->attach_type,
+			      &link_opts);
+	item->link_fd = err;
+	if (item->success) {
+		if (!ASSERT_GE(err, 0, "link_create"))
+			return -1;
+	} else {
+		if (!ASSERT_LT(err, 0, "link_create"))
+			return -1;
+	}
+
+	return 0;
+}
+
+void test_tracing_multi_attach(void)
+{
+	struct tracing_multi_test *skel;
+	int i = 0, err, fd;
+
+	skel = tracing_multi_test__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "tracing_multi_test__open_and_load"))
+		return;
+
+	err = tracing_multi_test__attach(skel);
+	if (!ASSERT_OK(err, "tracing_multi_test__attach"))
+		goto destroy_skel;
+
+	for (; i < ARRAY_SIZE(test_items); i++) {
+		if (do_test_item(skel, &test_items[i]))
+			break;
+	}
+
+	for (i = 0; i < ARRAY_SIZE(test_items); i++) {
+		fd = test_items[i].link_fd;
+		if (fd >= 0)
+			close(fd);
+	}
+
+	tracing_multi_test__detach(skel);
+destroy_skel:
+	tracing_multi_test__destroy(skel);
+}
diff --git a/tools/testing/selftests/bpf/progs/tracing_multi_test.c b/tools/testing/selftests/bpf/progs/tracing_multi_test.c
new file mode 100644
index 000000000000..f1ca8b64ed16
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/tracing_multi_test.c
@@ -0,0 +1,72 @@
+// SPDX-License-Identifier: GPL-2.0
+/* Copyright (c) 2024 ByteDance */
+#include <linux/bpf.h>
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+
+char _license[] SEC("license") = "GPL";
+
+struct bpf_testmod_struct_arg_1 {
+	int a;
+};
+struct bpf_testmod_struct_arg_2 {
+	long a;
+	long b;
+};
+
+__u64 fentry_test1_result = 0;
+SEC("fentry/bpf_testmod_test_struct_arg_1")
+int BPF_PROG2(fentry_test1, struct bpf_testmod_struct_arg_2, a)
+{
+	fentry_test1_result = a.a + a.b;
+	return 0;
+}
+
+__u64 fentry_test2_result = 0;
+SEC("fentry/bpf_testmod_test_struct_arg_2")
+int BPF_PROG2(fentry_test2, int, a, struct bpf_testmod_struct_arg_2, b)
+{
+	fentry_test2_result = a + b.a + b.b;
+	return 0;
+}
+
+__u64 fentry_test3_result = 0;
+SEC("fentry/bpf_testmod_test_arg_ptr_2")
+int BPF_PROG(fentry_test3, struct bpf_testmod_struct_arg_2 *a)
+{
+	fentry_test3_result = a->a + a->b;
+	return 0;
+}
+
+__u64 fentry_test4_result = 0;
+SEC("fentry/bpf_testmod_test_struct_arg_1")
+int BPF_PROG2(fentry_test4, struct bpf_testmod_struct_arg_2, a, int, b,
+	      int, c)
+{
+	fentry_test3_result = c;
+	return 0;
+}
+
+__u64 fexit_test1_result = 0;
+SEC("fexit/bpf_testmod_test_struct_arg_1")
+int BPF_PROG2(fexit_test1, struct bpf_testmod_struct_arg_2, a, int, b,
+	      int, c, int, retval)
+{
+	fexit_test1_result = retval;
+	return 0;
+}
+
+__u64 fexit_test2_result = 0;
+SEC("fexit/bpf_testmod_test_struct_arg_2")
+int BPF_PROG2(fexit_test2, int, a, struct bpf_testmod_struct_arg_2, b,
+	      int, c, int, retval)
+{
+	fexit_test2_result = a + b.a + b.b + retval;
+	return 0;
+}
+
+SEC("fmod_ret/bpf_modify_return_test")
+int BPF_PROG(fmod_ret_test1, int a, int *b)
+{
+	return 0;
+}
-- 
2.39.2