[PATCHv2 bpf-next 9/9] selftests/bpf: Add uprobe session consumers test

[Jiri Olsa <jolsa@kernel.org>]

Adding test that attached/detaches multiple consumers on
single uprobe and verifies all were hit as expected.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 .../bpf/prog_tests/uprobe_multi_test.c        | 203 ++++++++++++++++++
 .../progs/uprobe_multi_session_consumers.c    |  53 +++++
 2 files changed, 256 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/progs/uprobe_multi_session_consumers.c

diff --git a/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c b/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
index b521590fdbb9..83eac954cf00 100644
--- a/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
+++ b/tools/testing/selftests/bpf/prog_tests/uprobe_multi_test.c
@@ -9,6 +9,7 @@
 #include "uprobe_multi_session.skel.h"
 #include "uprobe_multi_session_cookie.skel.h"
 #include "uprobe_multi_session_recursive.skel.h"
+#include "uprobe_multi_session_consumers.skel.h"
 #include "bpf/libbpf_internal.h"
 #include "testing_helpers.h"
 #include "../sdt.h"
@@ -739,6 +740,206 @@ static void test_session_recursive_skel_api(void)
 	uprobe_multi_session_recursive__destroy(skel);
 }
 
+static int uprobe_attach(struct uprobe_multi_session_consumers *skel, int bit)
+{
+	struct bpf_program **prog = &skel->progs.uprobe_0 + bit;
+	struct bpf_link **link = &skel->links.uprobe_0 + bit;
+	LIBBPF_OPTS(bpf_uprobe_multi_opts, opts);
+
+	/*
+	 * bit: 0,1 uprobe session
+	 * bit: 2,3 uprobe entry
+	 * bit: 4,5 uprobe return
+	 */
+	opts.session = bit < 2;
+	opts.retprobe = bit == 4 || bit == 5;
+
+	*link = bpf_program__attach_uprobe_multi(*prog, 0, "/proc/self/exe",
+						 "uprobe_session_consumer_test",
+						 &opts);
+	if (!ASSERT_OK_PTR(*link, "bpf_program__attach_uprobe_multi"))
+		return -1;
+	return 0;
+}
+
+static void uprobe_detach(struct uprobe_multi_session_consumers *skel, int bit)
+{
+	struct bpf_link **link = &skel->links.uprobe_0 + bit;
+
+	bpf_link__destroy(*link);
+	*link = NULL;
+}
+
+static bool test_bit(int bit, unsigned long val)
+{
+	return val & (1 << bit);
+}
+
+noinline int
+uprobe_session_consumer_test(struct uprobe_multi_session_consumers *skel,
+			     unsigned long before, unsigned long after)
+{
+	int bit;
+
+	/* detach uprobe for each unset bit in 'before' state ... */
+	for (bit = 0; bit < 6; bit++) {
+		if (test_bit(bit, before) && !test_bit(bit, after))
+			uprobe_detach(skel, bit);
+	}
+
+	/* ... and attach all new bits in 'after' state */
+	for (bit = 0; bit < 6; bit++) {
+		if (!test_bit(bit, before) && test_bit(bit, after)) {
+			if (!ASSERT_OK(uprobe_attach(skel, bit), "uprobe_attach_after"))
+				return -1;
+		}
+	}
+	return 0;
+}
+
+static void session_consumer_test(struct uprobe_multi_session_consumers *skel,
+				  unsigned long before, unsigned long after)
+{
+	int err, bit;
+
+	/* 'before' is each, we attach uprobe for every set bit */
+	for (bit = 0; bit < 6; bit++) {
+		if (test_bit(bit, before)) {
+			if (!ASSERT_OK(uprobe_attach(skel, bit), "uprobe_attach_before"))
+				goto cleanup;
+		}
+	}
+
+	err = uprobe_session_consumer_test(skel, before, after);
+	if (!ASSERT_EQ(err, 0, "uprobe_session_consumer_test"))
+		goto cleanup;
+
+	for (bit = 0; bit < 6; bit++) {
+		const char *fmt = "BUG";
+		__u64 val = 0;
+
+		if (bit == 0) {
+			/*
+			 * session with return
+			 *  +1 if defined in 'before'
+			 *  +1 if defined in 'after'
+			 */
+			if (test_bit(bit, before)) {
+				val++;
+				if (test_bit(bit, after))
+					val++;
+			}
+			fmt = "bit 0  : session with return";
+		} else if (bit == 1) {
+			/*
+			 * session without return
+			 *   +1 if defined in 'before'
+			 */
+			if (test_bit(bit, before))
+				val++;
+			fmt = "bit 1  : session with NO return";
+		} else if (bit < 4) {
+			/*
+			 * uprobe entry
+			 *   +1 if define in 'before'
+			 */
+			if (test_bit(bit, before))
+				val++;
+			fmt = "bit 3/4: uprobe";
+		} else {
+			/* uprobe return is tricky ;-)
+			 *
+			 * to trigger uretprobe consumer, the uretprobe needs to be installed,
+			 * which means one of the 'return' uprobes was alive when probe was hit:
+			 *
+			 *   bits: 0 (session with return) 4/5 uprobe return in 'installed' mask
+			 *
+			 * in addition if 'after' state removes everything that was installed in
+			 * 'before' state, then uprobe kernel object goes away and return uprobe
+			 * is not installed and we won't hit it even if it's in 'after' state.
+			 */
+			unsigned long installed = before & 0b110001; // is uretprobe installed
+			unsigned long exists    = before & after;    // did uprobe go away
+
+			if (installed && exists && test_bit(bit, after))
+				val++;
+			fmt = "bit 5/6: uretprobe";
+		}
+
+		ASSERT_EQ(skel->bss->uprobe_result[bit], val, fmt);
+		skel->bss->uprobe_result[bit] = 0;
+	}
+
+cleanup:
+	for (bit = 0; bit < 6; bit++) {
+		struct bpf_link **link = &skel->links.uprobe_0 + bit;
+
+		if (*link)
+			uprobe_detach(skel, bit);
+	}
+}
+
+static void test_session_consumers(void)
+{
+	struct uprobe_multi_session_consumers *skel;
+	int before, after;
+
+	skel = uprobe_multi_session_consumers__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "uprobe_multi_session_consumers__open_and_load"))
+		return;
+
+	/*
+	 * The idea of this test is to try all possible combinations of
+	 * uprobes consumers attached on single function.
+	 *
+	 *  - 1 uprobe session with return handler called
+	 *  - 1 uprobe session without return handler called
+	 *  - 2 uprobe entry consumer
+	 *  - 2 uprobe exit consumers
+	 *
+	 * The test uses 6 uprobes attached on single function, but that
+	 * translates into single uprobe with 6 consumers in kernel.
+	 *
+	 * The before/after values present the state of attached consumers
+	 * before and after the probed function:
+	 *
+	 *  bit 0   : uprobe session with return
+	 *  bit 1   : uprobe session with no return
+	 *  bit 2,3 : uprobe entry
+	 *  bit 4,5 : uprobe return
+	 *
+	 * For example for:
+	 *
+	 *   before = 0b10101
+	 *   after  = 0b00110
+	 *
+	 * it means that before we call 'uprobe_session_consumer_test' we
+	 * attach uprobes defined in 'before' value:
+	 *
+	 *   - bit 0: uprobe session with return
+	 *   - bit 2: uprobe entry
+	 *   - bit 4: uprobe return
+	 *
+	 * uprobe_session_consumer_test is called and inside it we attach
+	 * and detach * uprobes based on 'after' value:
+	 *
+	 *   - bit 0: uprobe session with return is detached
+	 *   - bit 1: uprobe session without return is attached
+	 *   - bit 2: stays untouched
+	 *   - bit 4: uprobe return is detached
+	 *
+	 * uprobe_session_consumer_test returs and we check counters values
+	 * increased by bpf programs on each uprobe to match the expected
+	 * count based on before/after bits.
+	 */
+	for (before = 0; before < 64; before++) {
+		for (after = 0; after < 64; after++)
+			session_consumer_test(skel, before, after);
+	}
+
+	uprobe_multi_session_consumers__destroy(skel);
+}
+
 static void test_bench_attach_uprobe(void)
 {
 	long attach_start_ns = 0, attach_end_ns = 0;
@@ -833,4 +1034,6 @@ void test_uprobe_multi_test(void)
 		test_session_cookie_skel_api();
 	if (test__start_subtest("session_cookie_recursive"))
 		test_session_recursive_skel_api();
+	if (test__start_subtest("session/consumers"))
+		test_session_consumers();
 }
diff --git a/tools/testing/selftests/bpf/progs/uprobe_multi_session_consumers.c b/tools/testing/selftests/bpf/progs/uprobe_multi_session_consumers.c
new file mode 100644
index 000000000000..035d31a0a7f8
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/uprobe_multi_session_consumers.c
@@ -0,0 +1,53 @@
+// SPDX-License-Identifier: GPL-2.0
+#include <linux/bpf.h>
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+#include <stdbool.h>
+#include "bpf_kfuncs.h"
+#include "bpf_misc.h"
+
+char _license[] SEC("license") = "GPL";
+
+__u64 uprobe_result[6];
+
+SEC("uprobe.session")
+int uprobe_0(struct pt_regs *ctx)
+{
+	uprobe_result[0]++;
+	return 0;
+}
+
+SEC("uprobe.session")
+int uprobe_1(struct pt_regs *ctx)
+{
+	uprobe_result[1]++;
+	return 1;
+}
+
+SEC("uprobe.multi")
+int uprobe_2(struct pt_regs *ctx)
+{
+	uprobe_result[2]++;
+	return 0;
+}
+
+SEC("uprobe.multi")
+int uprobe_3(struct pt_regs *ctx)
+{
+	uprobe_result[3]++;
+	return 0;
+}
+
+SEC("uprobe.multi")
+int uprobe_4(struct pt_regs *ctx)
+{
+	uprobe_result[4]++;
+	return 0;
+}
+
+SEC("uprobe.multi")
+int uprobe_5(struct pt_regs *ctx)
+{
+	uprobe_result[5]++;
+	return 0;
+}
-- 
2.45.2