selftests/bpf: Add custom SEC() handling selftest

Add a selftest validating various aspects of libbpf's handling of custom
SEC() handlers. It also demonstrates how libraries can ensure very early
callbacks registration and unregistration using
__attribute__((constructor))/__attribute__((destructor)) functions.

Signed-off-by: Andrii Nakryiko <andrii@kernel.org>
Signed-off-by: Alexei Starovoitov <ast@kernel.org>
Tested-by: Alan Maguire <alan.maguire@oracle.com>
Reviewed-by: Alan Maguire <alan.maguire@oracle.com>
Link: https://lore.kernel.org/bpf/20220305010129.1549719-4-andrii@kernel.org

diff --git a/tools/testing/selftests/bpf/prog_tests/custom_sec_handlers.c b/tools/testing/selftests/bpf/prog_tests/custom_sec_handlers.c
new file mode 100644 (file)
index 0000000..b2dfc59
--- /dev/null
+++ b/tools/testing/selftests/bpf/prog_tests/custom_sec_handlers.c
@@ -0,0 +1,176 @@
+// SPDX-License-Identifier: GPL-2.0
+/* Copyright (c) 2022 Facebook */
+
+#include <test_progs.h>
+#include "test_custom_sec_handlers.skel.h"
+
+#define COOKIE_ABC1 1
+#define COOKIE_ABC2 2
+#define COOKIE_CUSTOM 3
+#define COOKIE_FALLBACK 4
+#define COOKIE_KPROBE 5
+
+static int custom_setup_prog(struct bpf_program *prog, long cookie)
+{
+       if (cookie == COOKIE_ABC1)
+               bpf_program__set_autoload(prog, false);
+
+       return 0;
+}
+
+static int custom_prepare_load_prog(struct bpf_program *prog,
+                                   struct bpf_prog_load_opts *opts, long cookie)
+{
+       if (cookie == COOKIE_FALLBACK)
+               opts->prog_flags |= BPF_F_SLEEPABLE;
+       else if (cookie == COOKIE_ABC1)
+               ASSERT_FALSE(true, "unexpected preload for abc");
+
+       return 0;
+}
+
+static int custom_attach_prog(const struct bpf_program *prog, long cookie,
+                             struct bpf_link **link)
+{
+       switch (cookie) {
+       case COOKIE_ABC2:
+               *link = bpf_program__attach_raw_tracepoint(prog, "sys_enter");
+               return libbpf_get_error(*link);
+       case COOKIE_CUSTOM:
+               *link = bpf_program__attach_tracepoint(prog, "syscalls", "sys_enter_nanosleep");
+               return libbpf_get_error(*link);
+       case COOKIE_KPROBE:
+       case COOKIE_FALLBACK:
+               /* no auto-attach for SEC("xyz") and SEC("kprobe") */
+               *link = NULL;
+               return 0;
+       default:
+               ASSERT_FALSE(true, "unexpected cookie");
+               return -EINVAL;
+       }
+}
+
+static int abc1_id;
+static int abc2_id;
+static int custom_id;
+static int fallback_id;
+static int kprobe_id;
+
+__attribute__((constructor))
+static void register_sec_handlers(void)
+{
+       LIBBPF_OPTS(libbpf_prog_handler_opts, abc1_opts,
+               .cookie = COOKIE_ABC1,
+               .prog_setup_fn = custom_setup_prog,
+               .prog_prepare_load_fn = custom_prepare_load_prog,
+               .prog_attach_fn = NULL,
+       );
+       LIBBPF_OPTS(libbpf_prog_handler_opts, abc2_opts,
+               .cookie = COOKIE_ABC2,
+               .prog_setup_fn = custom_setup_prog,
+               .prog_prepare_load_fn = custom_prepare_load_prog,
+               .prog_attach_fn = custom_attach_prog,
+       );
+       LIBBPF_OPTS(libbpf_prog_handler_opts, custom_opts,
+               .cookie = COOKIE_CUSTOM,
+               .prog_setup_fn = NULL,
+               .prog_prepare_load_fn = NULL,
+               .prog_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);
+}
+
+__attribute__((destructor))
+static void unregister_sec_handlers(void)
+{
+       libbpf_unregister_prog_handler(abc1_id);
+       libbpf_unregister_prog_handler(abc2_id);
+       libbpf_unregister_prog_handler(custom_id);
+}
+
+void test_custom_sec_handlers(void)
+{
+       LIBBPF_OPTS(libbpf_prog_handler_opts, opts,
+               .prog_setup_fn = custom_setup_prog,
+               .prog_prepare_load_fn = custom_prepare_load_prog,
+               .prog_attach_fn = custom_attach_prog,
+       );
+       struct test_custom_sec_handlers* skel;
+       int err;
+
+       ASSERT_GT(abc1_id, 0, "abc1_id");
+       ASSERT_GT(abc2_id, 0, "abc2_id");
+       ASSERT_GT(custom_id, 0, "custom_id");
+
+       /* override libbpf's handle of SEC("kprobe/...") but also allow pure
+        * SEC("kprobe") due to "kprobe+" specifier. Register it as
+        * TRACEPOINT, just for fun.
+        */
+       opts.cookie = COOKIE_KPROBE;
+       kprobe_id = libbpf_register_prog_handler("kprobe+", BPF_PROG_TYPE_TRACEPOINT, 0, &opts);
+       /* fallback treats everything as BPF_PROG_TYPE_SYSCALL program to test
+        * setting custom BPF_F_SLEEPABLE bit in preload handler
+        */
+       opts.cookie = COOKIE_FALLBACK;
+       fallback_id = libbpf_register_prog_handler(NULL, BPF_PROG_TYPE_SYSCALL, 0, &opts);
+
+       if (!ASSERT_GT(fallback_id, 0, "fallback_id") /* || !ASSERT_GT(kprobe_id, 0, "kprobe_id")*/) {
+               if (fallback_id > 0)
+                       libbpf_unregister_prog_handler(fallback_id);
+               if (kprobe_id > 0)
+                       libbpf_unregister_prog_handler(kprobe_id);
+               return;
+       }
+
+       /* open skeleton and validate assumptions */
+       skel = test_custom_sec_handlers__open();
+       if (!ASSERT_OK_PTR(skel, "skel_open"))
+               goto cleanup;
+
+       ASSERT_EQ(bpf_program__type(skel->progs.abc1), BPF_PROG_TYPE_RAW_TRACEPOINT, "abc1_type");
+       ASSERT_FALSE(bpf_program__autoload(skel->progs.abc1), "abc1_autoload");
+
+       ASSERT_EQ(bpf_program__type(skel->progs.abc2), BPF_PROG_TYPE_RAW_TRACEPOINT, "abc2_type");
+       ASSERT_EQ(bpf_program__type(skel->progs.custom1), BPF_PROG_TYPE_TRACEPOINT, "custom1_type");
+       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");
+
+       skel->rodata->my_pid = getpid();
+
+       /* now attempt to load everything */
+       err = test_custom_sec_handlers__load(skel);
+       if (!ASSERT_OK(err, "skel_load"))
+               goto cleanup;
+
+       /* now try to auto-attach everything */
+       err = test_custom_sec_handlers__attach(skel);
+       if (!ASSERT_OK(err, "skel_attach"))
+               goto cleanup;
+
+       skel->links.xyz = bpf_program__attach(skel->progs.kprobe1);
+       ASSERT_EQ(errno, EOPNOTSUPP, "xyz_attach_err");
+       ASSERT_ERR_PTR(skel->links.xyz, "xyz_attach");
+
+       /* trigger programs */
+       usleep(1);
+
+       /* SEC("abc") is set to not auto-loaded */
+       ASSERT_FALSE(skel->bss->abc1_called, "abc1_called");
+       ASSERT_TRUE(skel->bss->abc2_called, "abc2_called");
+       ASSERT_TRUE(skel->bss->custom1_called, "custom1_called");
+       ASSERT_TRUE(skel->bss->custom2_called, "custom2_called");
+       /* SEC("kprobe") shouldn't be auto-attached */
+       ASSERT_FALSE(skel->bss->kprobe1_called, "kprobe1_called");
+       /* SEC("xyz") shouldn't be auto-attached */
+       ASSERT_FALSE(skel->bss->xyz_called, "xyz_called");
+
+cleanup:
+       test_custom_sec_handlers__destroy(skel);
+
+       ASSERT_OK(libbpf_unregister_prog_handler(fallback_id), "unregister_fallback");
+       ASSERT_OK(libbpf_unregister_prog_handler(kprobe_id), "unregister_kprobe");
+}

diff --git a/tools/testing/selftests/bpf/progs/test_custom_sec_handlers.c b/tools/testing/selftests/bpf/progs/test_custom_sec_handlers.c
new file mode 100644 (file)
index 0000000..4061f70
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/test_custom_sec_handlers.c
@@ -0,0 +1,63 @@
+// SPDX-License-Identifier: GPL-2.0
+/* Copyright (c) 2022 Facebook */
+
+#include "vmlinux.h"
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+
+const volatile int my_pid;
+
+bool abc1_called;
+bool abc2_called;
+bool custom1_called;
+bool custom2_called;
+bool kprobe1_called;
+bool xyz_called;
+
+SEC("abc")
+int abc1(void *ctx)
+{
+       abc1_called = true;
+       return 0;
+}
+
+SEC("abc/whatever")
+int abc2(void *ctx)
+{
+       abc2_called = true;
+       return 0;
+}
+
+SEC("custom")
+int custom1(void *ctx)
+{
+       custom1_called = true;
+       return 0;
+}
+
+SEC("custom/something")
+int custom2(void *ctx)
+{
+       custom2_called = true;
+       return 0;
+}
+
+SEC("kprobe")
+int kprobe1(void *ctx)
+{
+       kprobe1_called = true;
+       return 0;
+}
+
+SEC("xyz/blah")
+int xyz(void *ctx)
+{
+       int whatever;
+
+       /* use sleepable helper, custom handler should set sleepable flag */
+       bpf_copy_from_user(&whatever, sizeof(whatever), NULL);
+       xyz_called = true;
+       return 0;
+}
+
+char _license[] SEC("license") = "GPL";