| 
 | 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 | +}  | 
0 commit comments