tools/testing/selftests/bpf/prog_tests/custom_sec_handlers.c

   1 // SPDX-License-Identifier: GPL-2.0
   2 /* Copyright (c) 2022 Facebook */
   3
   4 #include <test_progs.h>
   5 #include "test_custom_sec_handlers.skel.h"
   6
   7 #define COOKIE_ABC1 1
   8 #define COOKIE_ABC2 2
   9 #define COOKIE_CUSTOM 3
  10 #define COOKIE_FALLBACK 4
  11 #define COOKIE_KPROBE 5
  12
  13 static int custom_setup_prog(struct bpf_program *prog, long cookie)
  14 {
  15         if (cookie == COOKIE_ABC1)
  16                 bpf_program__set_autoload(prog, false);
  17
  18         return 0;
  19 }
  20
  21 static int custom_prepare_load_prog(struct bpf_program *prog,
  22                                     struct bpf_prog_load_opts *opts, long cookie)
  23 {
  24         if (cookie == COOKIE_FALLBACK)
  25                 opts->prog_flags |= BPF_F_SLEEPABLE;
  26         else if (cookie == COOKIE_ABC1)
  27                 ASSERT_FALSE(true, "unexpected preload for abc");
  28
  29         return 0;
  30 }
  31
  32 static int custom_attach_prog(const struct bpf_program *prog, long cookie,
  33                               struct bpf_link **link)
  34 {
  35         switch (cookie) {
  36         case COOKIE_ABC2:
  37                 *link = bpf_program__attach_raw_tracepoint(prog, "sys_enter");
  38                 return libbpf_get_error(*link);
  39         case COOKIE_CUSTOM:
  40                 *link = bpf_program__attach_tracepoint(prog, "syscalls", "sys_enter_nanosleep");
  41                 return libbpf_get_error(*link);
  42         case COOKIE_KPROBE:
  43         case COOKIE_FALLBACK:
  44                 /* no auto-attach for SEC("xyz") and SEC("kprobe") */
  45                 *link = NULL;
  46                 return 0;
  47         default:
  48                 ASSERT_FALSE(true, "unexpected cookie");
  49                 return -EINVAL;
  50         }
  51 }
  52
  53 static int abc1_id;
  54 static int abc2_id;
  55 static int custom_id;
  56 static int fallback_id;
  57 static int kprobe_id;
  58
  59 __attribute__((constructor))
  60 static void register_sec_handlers(void)
  61 {
  62         LIBBPF_OPTS(libbpf_prog_handler_opts, abc1_opts,
  63                 .cookie = COOKIE_ABC1,
  64                 .prog_setup_fn = custom_setup_prog,
  65                 .prog_prepare_load_fn = custom_prepare_load_prog,
  66                 .prog_attach_fn = NULL,
  67         );
  68         LIBBPF_OPTS(libbpf_prog_handler_opts, abc2_opts,
  69                 .cookie = COOKIE_ABC2,
  70                 .prog_setup_fn = custom_setup_prog,
  71                 .prog_prepare_load_fn = custom_prepare_load_prog,
  72                 .prog_attach_fn = custom_attach_prog,
  73         );
  74         LIBBPF_OPTS(libbpf_prog_handler_opts, custom_opts,
  75                 .cookie = COOKIE_CUSTOM,
  76                 .prog_setup_fn = NULL,
  77                 .prog_prepare_load_fn = NULL,
  78                 .prog_attach_fn = custom_attach_prog,
  79         );
  80
  81         abc1_id = libbpf_register_prog_handler("abc", BPF_PROG_TYPE_RAW_TRACEPOINT, 0, &abc1_opts);
  82         abc2_id = libbpf_register_prog_handler("abc/", BPF_PROG_TYPE_RAW_TRACEPOINT, 0, &abc2_opts);
  83         custom_id = libbpf_register_prog_handler("custom+", BPF_PROG_TYPE_TRACEPOINT, 0, &custom_opts);
  84 }
  85
  86 __attribute__((destructor))
  87 static void unregister_sec_handlers(void)
  88 {
  89         libbpf_unregister_prog_handler(abc1_id);
  90         libbpf_unregister_prog_handler(abc2_id);
  91         libbpf_unregister_prog_handler(custom_id);
  92 }
  93
  94 void test_custom_sec_handlers(void)
  95 {
  96         LIBBPF_OPTS(libbpf_prog_handler_opts, opts,
  97                 .prog_setup_fn = custom_setup_prog,
  98                 .prog_prepare_load_fn = custom_prepare_load_prog,
  99                 .prog_attach_fn = custom_attach_prog,
 100         );
 101         struct test_custom_sec_handlers* skel;
 102         int err;
 103
 104         ASSERT_GT(abc1_id, 0, "abc1_id");
 105         ASSERT_GT(abc2_id, 0, "abc2_id");
 106         ASSERT_GT(custom_id, 0, "custom_id");
 107
 108         /* override libbpf's handle of SEC("kprobe/...") but also allow pure
 109          * SEC("kprobe") due to "kprobe+" specifier. Register it as
 110          * TRACEPOINT, just for fun.
 111          */
 112         opts.cookie = COOKIE_KPROBE;
 113         kprobe_id = libbpf_register_prog_handler("kprobe+", BPF_PROG_TYPE_TRACEPOINT, 0, &opts);
 114         /* fallback treats everything as BPF_PROG_TYPE_SYSCALL program to test
 115          * setting custom BPF_F_SLEEPABLE bit in preload handler
 116          */
 117         opts.cookie = COOKIE_FALLBACK;
 118         fallback_id = libbpf_register_prog_handler(NULL, BPF_PROG_TYPE_SYSCALL, 0, &opts);
 119
 120         if (!ASSERT_GT(fallback_id, 0, "fallback_id") /* || !ASSERT_GT(kprobe_id, 0, "kprobe_id")*/) {
 121                 if (fallback_id > 0)
 122                         libbpf_unregister_prog_handler(fallback_id);
 123                 if (kprobe_id > 0)
 124                         libbpf_unregister_prog_handler(kprobe_id);
 125                 return;
 126         }
 127
 128         /* open skeleton and validate assumptions */
 129         skel = test_custom_sec_handlers__open();
 130         if (!ASSERT_OK_PTR(skel, "skel_open"))
 131                 goto cleanup;
 132
 133         ASSERT_EQ(bpf_program__type(skel->progs.abc1), BPF_PROG_TYPE_RAW_TRACEPOINT, "abc1_type");
 134         ASSERT_FALSE(bpf_program__autoload(skel->progs.abc1), "abc1_autoload");
 135
 136         ASSERT_EQ(bpf_program__type(skel->progs.abc2), BPF_PROG_TYPE_RAW_TRACEPOINT, "abc2_type");
 137         ASSERT_EQ(bpf_program__type(skel->progs.custom1), BPF_PROG_TYPE_TRACEPOINT, "custom1_type");
 138         ASSERT_EQ(bpf_program__type(skel->progs.custom2), BPF_PROG_TYPE_TRACEPOINT, "custom2_type");
 139         ASSERT_EQ(bpf_program__type(skel->progs.kprobe1), BPF_PROG_TYPE_TRACEPOINT, "kprobe1_type");
 140         ASSERT_EQ(bpf_program__type(skel->progs.xyz), BPF_PROG_TYPE_SYSCALL, "xyz_type");
 141
 142         skel->rodata->my_pid = getpid();
 143
 144         /* now attempt to load everything */
 145         err = test_custom_sec_handlers__load(skel);
 146         if (!ASSERT_OK(err, "skel_load"))
 147                 goto cleanup;
 148
 149         /* now try to auto-attach everything */
 150         err = test_custom_sec_handlers__attach(skel);
 151         if (!ASSERT_OK(err, "skel_attach"))
 152                 goto cleanup;
 153
 154         skel->links.xyz = bpf_program__attach(skel->progs.kprobe1);
 155         ASSERT_EQ(errno, EOPNOTSUPP, "xyz_attach_err");
 156         ASSERT_ERR_PTR(skel->links.xyz, "xyz_attach");
 157
 158         /* trigger programs */
 159         usleep(1);
 160
 161         /* SEC("abc") is set to not auto-loaded */
 162         ASSERT_FALSE(skel->bss->abc1_called, "abc1_called");
 163         ASSERT_TRUE(skel->bss->abc2_called, "abc2_called");
 164         ASSERT_TRUE(skel->bss->custom1_called, "custom1_called");
 165         ASSERT_TRUE(skel->bss->custom2_called, "custom2_called");
 166         /* SEC("kprobe") shouldn't be auto-attached */
 167         ASSERT_FALSE(skel->bss->kprobe1_called, "kprobe1_called");
 168         /* SEC("xyz") shouldn't be auto-attached */
 169         ASSERT_FALSE(skel->bss->xyz_called, "xyz_called");
 170
 171 cleanup:
 172         test_custom_sec_handlers__destroy(skel);
 173
 174         ASSERT_OK(libbpf_unregister_prog_handler(fallback_id), "unregister_fallback");
 175         ASSERT_OK(libbpf_unregister_prog_handler(kprobe_id), "unregister_kprobe");
 176 }