1// SPDX-License-Identifier: GPL-2.0
2/* Copyright (c) 2022 Meta Platforms, Inc. and affiliates. */
3
4#include "vmlinux.h"
5#include <bpf/bpf_helpers.h>
6#include <bpf/bpf_tracing.h>
7#include "bpf_tracing_net.h"
8#include "bpf_misc.h"
9
10char _license[] SEC("license") = "GPL";
11
12struct {
13	__uint(type, BPF_MAP_TYPE_TASK_STORAGE);
14	__uint(map_flags, BPF_F_NO_PREALLOC);
15	__type(key, int);
16	__type(value, long);
17} map_a SEC(".maps");
18
19__u32 user_data, key_serial, target_pid;
20__u64 flags, task_storage_val, cgroup_id;
21
22struct bpf_key *bpf_lookup_user_key(__u32 serial, __u64 flags) __ksym;
23void bpf_key_put(struct bpf_key *key) __ksym;
24void bpf_rcu_read_lock(void) __ksym;
25void bpf_rcu_read_unlock(void) __ksym;
26struct task_struct *bpf_task_acquire(struct task_struct *p) __ksym;
27void bpf_task_release(struct task_struct *p) __ksym;
28
29SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
30int get_cgroup_id(void *ctx)
31{
32	struct task_struct *task;
33	struct css_set *cgroups;
34
35	task = bpf_get_current_task_btf();
36	if (task->pid != target_pid)
37		return 0;
38
39	/* simulate bpf_get_current_cgroup_id() helper */
40	bpf_rcu_read_lock();
41	cgroups = task->cgroups;
42	if (!cgroups)
43		goto unlock;
44	cgroup_id = cgroups->dfl_cgrp->kn->id;
45unlock:
46	bpf_rcu_read_unlock();
47	return 0;
48}
49
50SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
51int task_succ(void *ctx)
52{
53	struct task_struct *task, *real_parent;
54	long init_val = 2;
55	long *ptr;
56
57	task = bpf_get_current_task_btf();
58	if (task->pid != target_pid)
59		return 0;
60
61	bpf_rcu_read_lock();
62	/* region including helper using rcu ptr real_parent */
63	real_parent = task->real_parent;
64	if (!real_parent)
65		goto out;
66	ptr = bpf_task_storage_get(&map_a, real_parent, &init_val,
67				   BPF_LOCAL_STORAGE_GET_F_CREATE);
68	if (!ptr)
69		goto out;
70	ptr = bpf_task_storage_get(&map_a, real_parent, 0, 0);
71	if (!ptr)
72		goto out;
73	task_storage_val = *ptr;
74out:
75	bpf_rcu_read_unlock();
76	return 0;
77}
78
79SEC("?fentry.s/" SYS_PREFIX "sys_nanosleep")
80int no_lock(void *ctx)
81{
82	struct task_struct *task, *real_parent;
83
84	/* old style ptr_to_btf_id is not allowed in sleepable */
85	task = bpf_get_current_task_btf();
86	real_parent = task->real_parent;
87	(void)bpf_task_storage_get(&map_a, real_parent, 0, 0);
88	return 0;
89}
90
91SEC("?fentry.s/" SYS_PREFIX "sys_nanosleep")
92int two_regions(void *ctx)
93{
94	struct task_struct *task, *real_parent;
95
96	/* two regions */
97	task = bpf_get_current_task_btf();
98	bpf_rcu_read_lock();
99	bpf_rcu_read_unlock();
100	bpf_rcu_read_lock();
101	real_parent = task->real_parent;
102	if (!real_parent)
103		goto out;
104	(void)bpf_task_storage_get(&map_a, real_parent, 0, 0);
105out:
106	bpf_rcu_read_unlock();
107	return 0;
108}
109
110SEC("?fentry/" SYS_PREFIX "sys_getpgid")
111int non_sleepable_1(void *ctx)
112{
113	struct task_struct *task, *real_parent;
114
115	task = bpf_get_current_task_btf();
116	bpf_rcu_read_lock();
117	real_parent = task->real_parent;
118	if (!real_parent)
119		goto out;
120	(void)bpf_task_storage_get(&map_a, real_parent, 0, 0);
121out:
122	bpf_rcu_read_unlock();
123	return 0;
124}
125
126SEC("?fentry/" SYS_PREFIX "sys_getpgid")
127int non_sleepable_2(void *ctx)
128{
129	struct task_struct *task, *real_parent;
130
131	bpf_rcu_read_lock();
132	task = bpf_get_current_task_btf();
133	bpf_rcu_read_unlock();
134
135	bpf_rcu_read_lock();
136	real_parent = task->real_parent;
137	if (!real_parent)
138		goto out;
139	(void)bpf_task_storage_get(&map_a, real_parent, 0, 0);
140out:
141	bpf_rcu_read_unlock();
142	return 0;
143}
144
145SEC("?fentry.s/" SYS_PREFIX "sys_nanosleep")
146int task_acquire(void *ctx)
147{
148	struct task_struct *task, *real_parent, *gparent;
149
150	task = bpf_get_current_task_btf();
151	bpf_rcu_read_lock();
152	real_parent = task->real_parent;
153	if (!real_parent)
154		goto out;
155
156	/* rcu_ptr->rcu_field */
157	gparent = real_parent->real_parent;
158	if (!gparent)
159		goto out;
160
161	/* acquire a reference which can be used outside rcu read lock region */
162	gparent = bpf_task_acquire(gparent);
163	if (!gparent)
164		goto out;
165
166	(void)bpf_task_storage_get(&map_a, gparent, 0, 0);
167	bpf_task_release(gparent);
168out:
169	bpf_rcu_read_unlock();
170	return 0;
171}
172
173SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
174int miss_lock(void *ctx)
175{
176	struct task_struct *task;
177
178	/* missing bpf_rcu_read_lock() */
179	task = bpf_get_current_task_btf();
180	bpf_rcu_read_lock();
181	(void)bpf_task_storage_get(&map_a, task, 0, 0);
182	bpf_rcu_read_unlock();
183	bpf_rcu_read_unlock();
184	return 0;
185}
186
187SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
188int miss_unlock(void *ctx)
189{
190	struct task_struct *task;
191
192	/* missing bpf_rcu_read_unlock() */
193	task = bpf_get_current_task_btf();
194	bpf_rcu_read_lock();
195	(void)bpf_task_storage_get(&map_a, task, 0, 0);
196	return 0;
197}
198
199SEC("?fentry/" SYS_PREFIX "sys_getpgid")
200int non_sleepable_rcu_mismatch(void *ctx)
201{
202	struct task_struct *task, *real_parent;
203
204	task = bpf_get_current_task_btf();
205	/* non-sleepable: missing bpf_rcu_read_unlock() in one path */
206	bpf_rcu_read_lock();
207	real_parent = task->real_parent;
208	if (!real_parent)
209		goto out;
210	(void)bpf_task_storage_get(&map_a, real_parent, 0, 0);
211	if (real_parent)
212		bpf_rcu_read_unlock();
213out:
214	return 0;
215}
216
217SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
218int inproper_sleepable_helper(void *ctx)
219{
220	struct task_struct *task, *real_parent;
221	struct pt_regs *regs;
222	__u32 value = 0;
223	void *ptr;
224
225	task = bpf_get_current_task_btf();
226	/* sleepable helper in rcu read lock region */
227	bpf_rcu_read_lock();
228	real_parent = task->real_parent;
229	if (!real_parent)
230		goto out;
231	regs = (struct pt_regs *)bpf_task_pt_regs(real_parent);
232	if (!regs)
233		goto out;
234
235	ptr = (void *)PT_REGS_IP(regs);
236	(void)bpf_copy_from_user_task(&value, sizeof(uint32_t), ptr, task, 0);
237	user_data = value;
238	(void)bpf_task_storage_get(&map_a, real_parent, 0, 0);
239out:
240	bpf_rcu_read_unlock();
241	return 0;
242}
243
244SEC("?lsm.s/bpf")
245int BPF_PROG(inproper_sleepable_kfunc, int cmd, union bpf_attr *attr, unsigned int size)
246{
247	struct bpf_key *bkey;
248
249	/* sleepable kfunc in rcu read lock region */
250	bpf_rcu_read_lock();
251	bkey = bpf_lookup_user_key(key_serial, flags);
252	bpf_rcu_read_unlock();
253	if (!bkey)
254		return -1;
255	bpf_key_put(bkey);
256
257	return 0;
258}
259
260SEC("?fentry.s/" SYS_PREFIX "sys_nanosleep")
261int nested_rcu_region(void *ctx)
262{
263	struct task_struct *task, *real_parent;
264
265	/* nested rcu read lock regions */
266	task = bpf_get_current_task_btf();
267	bpf_rcu_read_lock();
268	bpf_rcu_read_lock();
269	real_parent = task->real_parent;
270	if (!real_parent)
271		goto out;
272	(void)bpf_task_storage_get(&map_a, real_parent, 0, 0);
273out:
274	bpf_rcu_read_unlock();
275	bpf_rcu_read_unlock();
276	return 0;
277}
278
279SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
280int task_trusted_non_rcuptr(void *ctx)
281{
282	struct task_struct *task, *group_leader;
283
284	task = bpf_get_current_task_btf();
285	bpf_rcu_read_lock();
286	/* the pointer group_leader is explicitly marked as trusted */
287	group_leader = task->real_parent->group_leader;
288	(void)bpf_task_storage_get(&map_a, group_leader, 0, 0);
289	bpf_rcu_read_unlock();
290	return 0;
291}
292
293SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
294int task_untrusted_rcuptr(void *ctx)
295{
296	struct task_struct *task, *real_parent;
297
298	task = bpf_get_current_task_btf();
299	bpf_rcu_read_lock();
300	real_parent = task->real_parent;
301	bpf_rcu_read_unlock();
302	/* helper use of rcu ptr outside the rcu read lock region */
303	(void)bpf_task_storage_get(&map_a, real_parent, 0, 0);
304	return 0;
305}
306
307SEC("?fentry.s/" SYS_PREFIX "sys_nanosleep")
308int cross_rcu_region(void *ctx)
309{
310	struct task_struct *task, *real_parent;
311
312	/* rcu ptr define/use in different regions */
313	task = bpf_get_current_task_btf();
314	bpf_rcu_read_lock();
315	real_parent = task->real_parent;
316	bpf_rcu_read_unlock();
317	bpf_rcu_read_lock();
318	(void)bpf_task_storage_get(&map_a, real_parent, 0, 0);
319	bpf_rcu_read_unlock();
320	return 0;
321}
322
323__noinline
324static int static_subprog(void *ctx)
325{
326	volatile int ret = 0;
327
328	if (bpf_get_prandom_u32())
329		return ret + 42;
330	return ret + bpf_get_prandom_u32();
331}
332
333__noinline
334int global_subprog(u64 a)
335{
336	volatile int ret = a;
337
338	return ret + static_subprog(NULL);
339}
340
341__noinline
342static int static_subprog_lock(void *ctx)
343{
344	volatile int ret = 0;
345
346	bpf_rcu_read_lock();
347	if (bpf_get_prandom_u32())
348		return ret + 42;
349	return ret + bpf_get_prandom_u32();
350}
351
352__noinline
353int global_subprog_lock(u64 a)
354{
355	volatile int ret = a;
356
357	return ret + static_subprog_lock(NULL);
358}
359
360__noinline
361static int static_subprog_unlock(void *ctx)
362{
363	volatile int ret = 0;
364
365	bpf_rcu_read_unlock();
366	if (bpf_get_prandom_u32())
367		return ret + 42;
368	return ret + bpf_get_prandom_u32();
369}
370
371__noinline
372int global_subprog_unlock(u64 a)
373{
374	volatile int ret = a;
375
376	return ret + static_subprog_unlock(NULL);
377}
378
379SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
380int rcu_read_lock_subprog(void *ctx)
381{
382	volatile int ret = 0;
383
384	bpf_rcu_read_lock();
385	if (bpf_get_prandom_u32())
386		ret += static_subprog(ctx);
387	bpf_rcu_read_unlock();
388	return 0;
389}
390
391SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
392int rcu_read_lock_global_subprog(void *ctx)
393{
394	volatile int ret = 0;
395
396	bpf_rcu_read_lock();
397	if (bpf_get_prandom_u32())
398		ret += global_subprog(ret);
399	bpf_rcu_read_unlock();
400	return 0;
401}
402
403SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
404int rcu_read_lock_subprog_lock(void *ctx)
405{
406	volatile int ret = 0;
407
408	ret += static_subprog_lock(ctx);
409	bpf_rcu_read_unlock();
410	return 0;
411}
412
413SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
414int rcu_read_lock_global_subprog_lock(void *ctx)
415{
416	volatile int ret = 0;
417
418	ret += global_subprog_lock(ret);
419	bpf_rcu_read_unlock();
420	return 0;
421}
422
423SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
424int rcu_read_lock_subprog_unlock(void *ctx)
425{
426	volatile int ret = 0;
427
428	bpf_rcu_read_lock();
429	ret += static_subprog_unlock(ctx);
430	return 0;
431}
432
433SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
434int rcu_read_lock_global_subprog_unlock(void *ctx)
435{
436	volatile int ret = 0;
437
438	bpf_rcu_read_lock();
439	ret += global_subprog_unlock(ret);
440	return 0;
441}
442