Re: [PATCH v3 bpf-next 3/3] selftests/bpf: add custom SEC() handling selftest

[Alan Maguire <alan.maguire@oracle.com>]

On Mon, 14 Feb 2022, Alexei Starovoitov wrote:

> Alan,
> can you demo your "okprobe" feature based on this api?
> Any rough patches would do.

sure; see below.  Requires Andrii's v3 patches to be applied first,
and demonstrates okprobe handling for a kprobe function that exists
and one that doesn't - the important thing is skeleton attach
can succeed even when a function is missing (as it would be if
the associated module wasn't loaded).

> The "o" handling will be done in which callback?
> 

We set program type at init and do custom attach using the function
name (specified in the program section after the "okprobe" prefix).  
However we make sure to catch -ENOENT attach failures and return 0
with a NULL link so skeleton attach can proceed.

From 9bbd615b71f8f59ff743608bc86d7a2a346da2a8 Mon Sep 17 00:00:00 2001
From: Alan Maguire <alan.maguire@oracle.com>
Date: Mon, 14 Feb 2022 22:57:56 +0000
Subject: [PATCH bpf-next] selftests/bpf: demonstrate further use of custom
 SEC() handling

Register and use SEC() handling for "okprobe/" kprobe programs
(Optional kprobe) which should be attached as kprobes but
critically should not stop skeleton loading if attach fails
due to non-existence of the to-be-probed function.  This mode
of SEC() handling is useful for tracing module functions
where the module might not be loaded.

Note - this patch is based on the v3 of Andrii's section
handling patches [1] and these need to be applied for it to
apply cleanly.

[1] https://lore.kernel.org/bpf/20220211211450.2224877-1-andrii@kernel.org/

Signed-off-by: Alan Maguire <alan.maguire@oracle.com>
---
 .../selftests/bpf/prog_tests/custom_sec_handlers.c | 33 ++++++++++++++++++++++
 .../selftests/bpf/progs/test_custom_sec_handlers.c | 17 +++++++++++
 2 files changed, 50 insertions(+)

diff --git a/tools/testing/selftests/bpf/prog_tests/custom_sec_handlers.c b/tools/testing/selftests/bpf/prog_tests/custom_sec_handlers.c
index 2826452..5da1375 100644
--- a/tools/testing/selftests/bpf/prog_tests/custom_sec_handlers.c
+++ b/tools/testing/selftests/bpf/prog_tests/custom_sec_handlers.c
@@ -9,11 +9,14 @@
 #define COOKIE_CUSTOM 3
 #define COOKIE_FALLBACK 4
 #define COOKIE_KPROBE 5
+#define COOKIE_OKPROBE 6
 
 static int custom_init_prog(struct bpf_program *prog, long cookie)
 {
 	if (cookie == COOKIE_ABC1)
 		bpf_program__set_autoload(prog, false);
+	else if (cookie == COOKIE_OKPROBE)
+		bpf_program__set_type(prog, BPF_PROG_TYPE_KPROBE);
 
 	return 0;
 }
@@ -32,6 +35,8 @@ static int custom_preload_prog(struct bpf_program *prog,
 static int custom_attach_prog(const struct bpf_program *prog, long cookie,
 			      struct bpf_link **link)
 {
+	const char *func_name;
+
 	switch (cookie) {
 	case COOKIE_ABC2:
 		*link = bpf_program__attach_raw_tracepoint(prog, "sys_enter");
@@ -39,6 +44,15 @@ static int custom_attach_prog(const struct bpf_program *prog, long cookie,
 	case COOKIE_CUSTOM:
 		*link = bpf_program__attach_tracepoint(prog, "syscalls", "sys_enter_nanosleep");
 		return libbpf_get_error(*link);
+	case COOKIE_OKPROBE:
+		func_name = bpf_program__section_name(prog) + strlen("okprobe/");
+		*link = bpf_program__attach_kprobe(prog, false, func_name);
+		/* it's ok if func doesn't exist. */
+		if (libbpf_get_error(*link) == -ENOENT) {
+			*link = NULL;
+			return 0;
+		}
+		return libbpf_get_error(*link);
 	case COOKIE_KPROBE:
 	case COOKIE_FALLBACK:
 		/* no auto-attach for SEC("xyz") and SEC("kprobe") */
@@ -55,6 +69,7 @@ static int custom_attach_prog(const struct bpf_program *prog, long cookie,
 static int custom_id;
 static int fallback_id;
 static int kprobe_id;
+static int okprobe_id;
 
 __attribute__((constructor))
 static void register_sec_handlers(void)
@@ -77,10 +92,18 @@ static void register_sec_handlers(void)
 		.preload_fn = NULL,
 		.attach_fn = custom_attach_prog,
 	);
+	LIBBPF_OPTS(libbpf_prog_handler_opts, okprobe_opts,
+		.cookie = COOKIE_OKPROBE,
+		.init_fn = custom_init_prog,
+		.preload_fn = NULL,
+		.attach_fn = custom_attach_prog,
+	);
 
+	
 	abc1_id = libbpf_register_prog_handler("abc", BPF_PROG_TYPE_RAW_TRACEPOINT, 0, &abc1_opts);
 	abc2_id = libbpf_register_prog_handler("abc/", BPF_PROG_TYPE_RAW_TRACEPOINT, 0, &abc2_opts);
 	custom_id = libbpf_register_prog_handler("custom+", BPF_PROG_TYPE_TRACEPOINT, 0, &custom_opts);
+	okprobe_id = libbpf_register_prog_handler("okprobe/", BPF_PROG_TYPE_KPROBE, 0, &okprobe_opts);
 }
 
 __attribute__((destructor))
@@ -89,6 +112,7 @@ static void unregister_sec_handlers(void)
 	libbpf_unregister_prog_handler(abc1_id);
 	libbpf_unregister_prog_handler(abc2_id);
 	libbpf_unregister_prog_handler(custom_id);
+	libbpf_unregister_prog_handler(okprobe_id);
 }
 
 void test_custom_sec_handlers(void)
@@ -104,6 +128,7 @@ void test_custom_sec_handlers(void)
 	ASSERT_GT(abc1_id, 0, "abc1_id");
 	ASSERT_GT(abc2_id, 0, "abc2_id");
 	ASSERT_GT(custom_id, 0, "custom_id");
+	ASSERT_GT(okprobe_id, 0, "okprobe_id");
 
 	/* override libbpf's handle of SEC("kprobe/...") but also allow pure
 	 * SEC("kprobe") due to "kprobe+" specifier. Register it as
@@ -138,6 +163,8 @@ void test_custom_sec_handlers(void)
 	ASSERT_EQ(bpf_program__type(skel->progs.custom2), BPF_PROG_TYPE_TRACEPOINT, "custom2_type");
 	ASSERT_EQ(bpf_program__type(skel->progs.kprobe1), BPF_PROG_TYPE_TRACEPOINT, "kprobe1_type");
 	ASSERT_EQ(bpf_program__type(skel->progs.xyz), BPF_PROG_TYPE_SYSCALL, "xyz_type");
+	ASSERT_EQ(bpf_program__type(skel->progs.kprobe2), BPF_PROG_TYPE_KPROBE, "kprobe2_type");
+	ASSERT_EQ(bpf_program__type(skel->progs.kprobe3), BPF_PROG_TYPE_KPROBE, "kprobe3_type");
 
 	skel->rodata->my_pid = getpid();
 
@@ -167,6 +194,12 @@ void test_custom_sec_handlers(void)
 	ASSERT_FALSE(skel->bss->kprobe1_called, "kprobe1_called");
 	/* SEC("xyz") shouldn't be auto-attached */
 	ASSERT_FALSE(skel->bss->xyz_called, "xyz_called");
+	/* SEC("okprobe/sys_nanosleep") should be auto-attached */
+	ASSERT_TRUE(skel->bss->kprobe2_called, "kprobe2_called");
+	/* SEC("okprobe/nonexistent_function") shouldn't attach, but
+	 * this shouldn't prevent overall skeleton attach.
+	 */
+	ASSERT_FALSE(skel->bss->kprobe3_called, "kprobe3_called");
 
 cleanup:
 	test_custom_sec_handlers__destroy(skel);
diff --git a/tools/testing/selftests/bpf/progs/test_custom_sec_handlers.c b/tools/testing/selftests/bpf/progs/test_custom_sec_handlers.c
index 4061f70..6e9e051f 100644
--- a/tools/testing/selftests/bpf/progs/test_custom_sec_handlers.c
+++ b/tools/testing/selftests/bpf/progs/test_custom_sec_handlers.c
@@ -4,6 +4,7 @@
 #include "vmlinux.h"
 #include <bpf/bpf_helpers.h>
 #include <bpf/bpf_tracing.h>
+#include "bpf_misc.h"
 
 const volatile int my_pid;
 
@@ -12,6 +13,8 @@
 bool custom1_called;
 bool custom2_called;
 bool kprobe1_called;
+bool kprobe2_called;
+bool kprobe3_called;
 bool xyz_called;
 
 SEC("abc")
@@ -60,4 +63,18 @@ int xyz(void *ctx)
 	return 0;
 }
 
+SEC("okprobe/" SYS_PREFIX "sys_nanosleep")
+int kprobe2(void *ctx)
+{
+	kprobe2_called = true;
+	return 0;
+}
+
+SEC("okprobe/nonexistent_function")
+int kprobe3(void *ctx)
+{
+	kprobe3_called = true;
+	return 0;
+}
+
 char _license[] SEC("license") = "GPL";
-- 
1.8.3.1