[PATCH bpf-next v2 17/21] selftests/bpf: Add tracing_multi bpf prog cookie test

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

Verify that the cookie also works for bpf prog targets of tracing_multi
link.

Assisted-by: Codex:gpt-5.5
Signed-off-by: Leon Hwang <leon.hwang@linux.dev>
---
 .../selftests/bpf/prog_tests/tracing_multi.c  | 21 +++++++++++++++++--
 .../selftests/bpf/progs/tracing_multi_bpf.c   | 13 ++++++++++--
 2 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/tools/testing/selftests/bpf/prog_tests/tracing_multi.c b/tools/testing/selftests/bpf/prog_tests/tracing_multi.c
index d428b5e2600c..bcf933760bc5 100644
--- a/tools/testing/selftests/bpf/prog_tests/tracing_multi.c
+++ b/tools/testing/selftests/bpf/prog_tests/tracing_multi.c
@@ -1143,9 +1143,12 @@ static int run_bpf_target(struct bpf_program *prog, __u32 retval, const char *na
 	return 0;
 }
 
-static void test_link_api_bpf_prog(void)
+static void test_link_api_bpf_prog(bool test_cookies)
 {
 	const char *funcs[] = { "target_1", "target_2" };
+	__u64 fentry_cookies[] = { 0xfeed01, 0xfeed02 };
+	__u64 fexit_cookies[] = { 0xfeed11, 0xfeed12 };
+	__u64 fsession_cookies[] = { 0xfeed21, 0xfeed22 };
 	LIBBPF_OPTS(bpf_tracing_multi_opts, opts);
 	struct tracing_multi_bpf *skel = NULL;
 	int fds[ARRAY_SIZE(funcs)];
@@ -1158,27 +1161,39 @@ static void test_link_api_bpf_prog(void)
 	fds[0] = bpf_program__fd(skel->progs.target_1);
 	fds[1] = bpf_program__fd(skel->progs.target_2);
 
+	skel->bss->test_cookies = test_cookies;
+
 	opts.fds = fds;
 	opts.funcs = funcs;
 	opts.cnt = ARRAY_SIZE(fds);
 
+	opts.cookies = test_cookies ? fentry_cookies : NULL;
 	link = bpf_program__attach_tracing_multi(skel->progs.test_fentry, NULL, &opts);
 	if (!ASSERT_OK_PTR(link, "attach_fentry"))
 		goto cleanup;
 	skel->links.test_fentry = link;
 
+	opts.cookies = test_cookies ? fexit_cookies : NULL;
 	link = bpf_program__attach_tracing_multi(skel->progs.test_fexit, NULL, &opts);
 	if (!ASSERT_OK_PTR(link, "attach_fexit"))
 		goto cleanup;
 	skel->links.test_fexit = link;
 
+	opts.cookies = test_cookies ? fsession_cookies : NULL;
 	link = bpf_program__attach_tracing_multi(skel->progs.test_fsession, NULL, &opts);
 	if (!ASSERT_OK_PTR(link, "attach_fsession"))
 		goto cleanup;
 	skel->links.test_fsession = link;
 
+	skel->bss->fentry_cookie = fentry_cookies[0];
+	skel->bss->fexit_cookie = fexit_cookies[0];
+	skel->bss->fsession_cookie = fsession_cookies[0];
 	if (run_bpf_target(skel->progs.target_1, 1, "target_1"))
 		goto cleanup;
+
+	skel->bss->fentry_cookie = fentry_cookies[1];
+	skel->bss->fexit_cookie = fexit_cookies[1];
+	skel->bss->fsession_cookie = fsession_cookies[1];
 	if (run_bpf_target(skel->progs.target_2, 2, "target_2"))
 		goto cleanup;
 
@@ -1296,7 +1311,9 @@ static void test_attach_api_bpf_prog_fails(void)
 void test_tracing_multi_prog_test(void)
 {
 	if (test__start_subtest("link_api"))
-		test_link_api_bpf_prog();
+		test_link_api_bpf_prog(false);
 	if (test__start_subtest("attach_api_fails"))
 		test_attach_api_bpf_prog_fails();
+	if (test__start_subtest("cookies"))
+		test_link_api_bpf_prog(true);
 }
diff --git a/tools/testing/selftests/bpf/progs/tracing_multi_bpf.c b/tools/testing/selftests/bpf/progs/tracing_multi_bpf.c
index 3e02602558ad..651b5d86203f 100644
--- a/tools/testing/selftests/bpf/progs/tracing_multi_bpf.c
+++ b/tools/testing/selftests/bpf/progs/tracing_multi_bpf.c
@@ -5,6 +5,10 @@
 
 char _license[] SEC("license") = "GPL";
 
+bool test_cookies;
+__u64 fentry_cookie;
+__u64 fexit_cookie;
+__u64 fsession_cookie;
 __u64 test_result_fentry;
 __u64 test_result_fexit;
 __u64 test_result_fsession_entry;
@@ -29,20 +33,25 @@ int target_2(void *ctx)
 SEC("fentry.multi")
 int BPF_PROG(test_fentry)
 {
-	test_result_fentry++;
+	if (!test_cookies || bpf_get_attach_cookie(ctx) == fentry_cookie)
+		test_result_fentry++;
 	return 0;
 }
 
 SEC("fexit.multi")
 int BPF_PROG(test_fexit)
 {
-	test_result_fexit++;
+	if (!test_cookies || bpf_get_attach_cookie(ctx) == fexit_cookie)
+		test_result_fexit++;
 	return 0;
 }
 
 SEC("fsession.multi")
 int BPF_PROG(test_fsession)
 {
+	if (test_cookies && bpf_get_attach_cookie(ctx) != fsession_cookie)
+		return 0;
+
 	if (bpf_session_is_return(ctx))
 		test_result_fsession_exit++;
 	else
-- 
2.55.0