1130861Strhodes// SPDX-License-Identifier: GPL-2.0
2227281Smarius/* Copyright (c) 2022 Meta Platforms, Inc. and affiliates. */
3130861Strhodes
4130861Strhodes#include <vmlinux.h>
5130861Strhodes#include <bpf/bpf_tracing.h>
6130861Strhodes#include <bpf/bpf_helpers.h>
7130861Strhodes
8130861Strhodes#include "task_kfunc_common.h"
9130861Strhodes
10130861Strhodeschar _license[] SEC("license") = "GPL";
11130861Strhodes
12130861Strhodesint err, pid;
13130861Strhodes
14130861Strhodes/* Prototype for all of the program trace events below:
15130861Strhodes *
16130861Strhodes * TRACE_EVENT(task_newtask,
17130861Strhodes *         TP_PROTO(struct task_struct *p, u64 clone_flags)
18130861Strhodes */
19130861Strhodes
20130861Strhodesstruct task_struct *bpf_task_acquire(struct task_struct *p) __ksym __weak;
21130861Strhodes
22130861Strhodesstruct task_struct *bpf_task_acquire___one(struct task_struct *task) __ksym __weak;
23130861Strhodes/* The two-param bpf_task_acquire doesn't exist */
24130861Strhodesstruct task_struct *bpf_task_acquire___two(struct task_struct *p, void *ctx) __ksym __weak;
25130861Strhodes/* Incorrect type for first param */
26130861Strhodesstruct task_struct *bpf_task_acquire___three(void *ctx) __ksym __weak;
27130861Strhodes
28227281Smariusvoid invalid_kfunc(void) __ksym __weak;
29130861Strhodesvoid bpf_testmod_test_mod_kfunc(int i) __ksym __weak;
30130861Strhodes
31130861Strhodesstatic bool is_test_kfunc_task(void)
32130861Strhodes{
33227281Smarius	int cur_pid = bpf_get_current_pid_tgid() >> 32;
34130861Strhodes
35227281Smarius	return pid == cur_pid;
36159719Sbrueffer}
37159719Sbrueffer
38159719Sbruefferstatic int test_acquire_release(struct task_struct *task)
39159719Sbrueffer{
40159719Sbrueffer	struct task_struct *acquired = NULL;
41227281Smarius
42227281Smarius	if (!bpf_ksym_exists(bpf_task_acquire)) {
43227281Smarius		err = 3;
44227281Smarius		return 0;
45227281Smarius	}
46227281Smarius	if (!bpf_ksym_exists(bpf_testmod_test_mod_kfunc)) {
47227281Smarius		err = 4;
48227281Smarius		return 0;
49130861Strhodes	}
50130861Strhodes	if (bpf_ksym_exists(invalid_kfunc)) {
51130861Strhodes		/* the verifier's dead code elimination should remove this */
52227281Smarius		err = 5;
53227281Smarius		asm volatile ("goto -1"); /* for (;;); */
54227281Smarius	}
55227281Smarius
56227281Smarius	acquired = bpf_task_acquire(task);
57227281Smarius	if (acquired)
58227281Smarius		bpf_task_release(acquired);
59227281Smarius	else
60130861Strhodes		err = 6;
61130861Strhodes
62227281Smarius	return 0;
63227281Smarius}
64130861Strhodes
65227281SmariusSEC("tp_btf/task_newtask")
66130861Strhodesint BPF_PROG(test_task_kfunc_flavor_relo, struct task_struct *task, u64 clone_flags)
67227281Smarius{
68227281Smarius	struct task_struct *acquired = NULL;
69227281Smarius	int fake_ctx = 42;
70227281Smarius
71227281Smarius	if (bpf_ksym_exists(bpf_task_acquire___one)) {
72227281Smarius		acquired = bpf_task_acquire___one(task);
73227281Smarius	} else if (bpf_ksym_exists(bpf_task_acquire___two)) {
74227281Smarius		/* Here, bpf_object__resolve_ksym_func_btf_id's find_ksym_btf_id
75227281Smarius		 * call will find vmlinux's bpf_task_acquire, but subsequent
76227281Smarius		 * bpf_core_types_are_compat will fail
77227281Smarius		 */
78130861Strhodes		acquired = bpf_task_acquire___two(task, &fake_ctx);
79130861Strhodes		err = 3;
80227281Smarius		return 0;
81227281Smarius	} else if (bpf_ksym_exists(bpf_task_acquire___three)) {
82227281Smarius		/* bpf_core_types_are_compat will fail similarly to above case */
83227281Smarius		acquired = bpf_task_acquire___three(&fake_ctx);
84227281Smarius		err = 4;
85227281Smarius		return 0;
86227281Smarius	}
87227282Smarius
88130861Strhodes	if (acquired)
89130861Strhodes		bpf_task_release(acquired);
90130861Strhodes	else
91130861Strhodes		err = 5;
92227281Smarius	return 0;
93227281Smarius}
94227281Smarius
95227281SmariusSEC("tp_btf/task_newtask")
96227281Smariusint BPF_PROG(test_task_kfunc_flavor_relo_not_found, struct task_struct *task, u64 clone_flags)
97227281Smarius{
98227281Smarius	/* Neither symbol should successfully resolve.
99227281Smarius	 * Success or failure of one ___flavor should not affect others
100130861Strhodes	 */
101130861Strhodes	if (bpf_ksym_exists(bpf_task_acquire___two))
102227281Smarius		err = 1;
103227281Smarius	else if (bpf_ksym_exists(bpf_task_acquire___three))
104130861Strhodes		err = 2;
105267938Sbapt
106227281Smarius	return 0;
107267938Sbapt}
108130861Strhodes
109130861StrhodesSEC("tp_btf/task_newtask")
110130861Strhodesint BPF_PROG(test_task_acquire_release_argument, struct task_struct *task, u64 clone_flags)
111227281Smarius{
112227281Smarius	if (!is_test_kfunc_task())
113227281Smarius		return 0;
114
115	return test_acquire_release(task);
116}
117
118SEC("tp_btf/task_newtask")
119int BPF_PROG(test_task_acquire_release_current, struct task_struct *task, u64 clone_flags)
120{
121	if (!is_test_kfunc_task())
122		return 0;
123
124	return test_acquire_release(bpf_get_current_task_btf());
125}
126
127SEC("tp_btf/task_newtask")
128int BPF_PROG(test_task_acquire_leave_in_map, struct task_struct *task, u64 clone_flags)
129{
130	long status;
131
132	if (!is_test_kfunc_task())
133		return 0;
134
135	status = tasks_kfunc_map_insert(task);
136	if (status)
137		err = 1;
138
139	return 0;
140}
141
142SEC("tp_btf/task_newtask")
143int BPF_PROG(test_task_xchg_release, struct task_struct *task, u64 clone_flags)
144{
145	struct task_struct *kptr;
146	struct __tasks_kfunc_map_value *v;
147	long status;
148
149	if (!is_test_kfunc_task())
150		return 0;
151
152	status = tasks_kfunc_map_insert(task);
153	if (status) {
154		err = 1;
155		return 0;
156	}
157
158	v = tasks_kfunc_map_value_lookup(task);
159	if (!v) {
160		err = 2;
161		return 0;
162	}
163
164	kptr = bpf_kptr_xchg(&v->task, NULL);
165	if (!kptr) {
166		err = 3;
167		return 0;
168	}
169
170	bpf_task_release(kptr);
171
172	return 0;
173}
174
175SEC("tp_btf/task_newtask")
176int BPF_PROG(test_task_map_acquire_release, struct task_struct *task, u64 clone_flags)
177{
178	struct task_struct *kptr;
179	struct __tasks_kfunc_map_value *v;
180	long status;
181
182	if (!is_test_kfunc_task())
183		return 0;
184
185	status = tasks_kfunc_map_insert(task);
186	if (status) {
187		err = 1;
188		return 0;
189	}
190
191	v = tasks_kfunc_map_value_lookup(task);
192	if (!v) {
193		err = 2;
194		return 0;
195	}
196
197	bpf_rcu_read_lock();
198	kptr = v->task;
199	if (!kptr) {
200		err = 3;
201	} else {
202		kptr = bpf_task_acquire(kptr);
203		if (!kptr)
204			err = 4;
205		else
206			bpf_task_release(kptr);
207	}
208	bpf_rcu_read_unlock();
209
210	return 0;
211}
212
213SEC("tp_btf/task_newtask")
214int BPF_PROG(test_task_current_acquire_release, struct task_struct *task, u64 clone_flags)
215{
216	struct task_struct *current, *acquired;
217
218	if (!is_test_kfunc_task())
219		return 0;
220
221	current = bpf_get_current_task_btf();
222	acquired = bpf_task_acquire(current);
223	if (acquired)
224		bpf_task_release(acquired);
225	else
226		err = 1;
227
228	return 0;
229}
230
231static void lookup_compare_pid(const struct task_struct *p)
232{
233	struct task_struct *acquired;
234
235	acquired = bpf_task_from_pid(p->pid);
236	if (!acquired) {
237		err = 1;
238		return;
239	}
240
241	if (acquired->pid != p->pid)
242		err = 2;
243	bpf_task_release(acquired);
244}
245
246SEC("tp_btf/task_newtask")
247int BPF_PROG(test_task_from_pid_arg, struct task_struct *task, u64 clone_flags)
248{
249	if (!is_test_kfunc_task())
250		return 0;
251
252	lookup_compare_pid(task);
253	return 0;
254}
255
256SEC("tp_btf/task_newtask")
257int BPF_PROG(test_task_from_pid_current, struct task_struct *task, u64 clone_flags)
258{
259	if (!is_test_kfunc_task())
260		return 0;
261
262	lookup_compare_pid(bpf_get_current_task_btf());
263	return 0;
264}
265
266static int is_pid_lookup_valid(s32 pid)
267{
268	struct task_struct *acquired;
269
270	acquired = bpf_task_from_pid(pid);
271	if (acquired) {
272		bpf_task_release(acquired);
273		return 1;
274	}
275
276	return 0;
277}
278
279SEC("tp_btf/task_newtask")
280int BPF_PROG(test_task_from_pid_invalid, struct task_struct *task, u64 clone_flags)
281{
282	if (!is_test_kfunc_task())
283		return 0;
284
285	bpf_strncmp(task->comm, 12, "foo");
286	bpf_strncmp(task->comm, 16, "foo");
287	bpf_strncmp(&task->comm[8], 4, "foo");
288
289	if (is_pid_lookup_valid(-1)) {
290		err = 1;
291		return 0;
292	}
293
294	if (is_pid_lookup_valid(0xcafef00d)) {
295		err = 2;
296		return 0;
297	}
298
299	return 0;
300}
301
302SEC("tp_btf/task_newtask")
303int BPF_PROG(task_kfunc_acquire_trusted_walked, struct task_struct *task, u64 clone_flags)
304{
305	struct task_struct *acquired;
306
307	/* task->group_leader is listed as a trusted, non-NULL field of task struct. */
308	acquired = bpf_task_acquire(task->group_leader);
309	if (acquired)
310		bpf_task_release(acquired);
311	else
312		err = 1;
313
314
315	return 0;
316}
317