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
 13static 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
 21static 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
 32static 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
 53static int abc1_id;
 54static int abc2_id;
 55static int custom_id;
 56static int fallback_id;
 57static int kprobe_id;
 58
 59__attribute__((constructor))
 60static 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))
 87static 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
 94void 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
171cleanup:
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}