148671232SYonghong Song // SPDX-License-Identifier: GPL-2.0
248671232SYonghong Song /* Copyright (c) 2022 Meta Platforms, Inc. and affiliates. */
348671232SYonghong Song 
448671232SYonghong Song #include "vmlinux.h"
548671232SYonghong Song #include <bpf/bpf_helpers.h>
648671232SYonghong Song #include <bpf/bpf_tracing.h>
748671232SYonghong Song #include "bpf_tracing_net.h"
848671232SYonghong Song #include "bpf_misc.h"
948671232SYonghong Song 
1048671232SYonghong Song char _license[] SEC("license") = "GPL";
1148671232SYonghong Song 
1248671232SYonghong Song struct {
1348671232SYonghong Song 	__uint(type, BPF_MAP_TYPE_TASK_STORAGE);
1448671232SYonghong Song 	__uint(map_flags, BPF_F_NO_PREALLOC);
1548671232SYonghong Song 	__type(key, int);
1648671232SYonghong Song 	__type(value, long);
1748671232SYonghong Song } map_a SEC(".maps");
1848671232SYonghong Song 
1948671232SYonghong Song __u32 user_data, key_serial, target_pid;
2048671232SYonghong Song __u64 flags, task_storage_val, cgroup_id;
2148671232SYonghong Song 
2248671232SYonghong Song struct bpf_key *bpf_lookup_user_key(__u32 serial, __u64 flags) __ksym;
2348671232SYonghong Song void bpf_key_put(struct bpf_key *key) __ksym;
2448671232SYonghong Song void bpf_rcu_read_lock(void) __ksym;
2548671232SYonghong Song void bpf_rcu_read_unlock(void) __ksym;
26*f85671c6SDavid Vernet struct task_struct *bpf_task_acquire(struct task_struct *p) __ksym;
2748671232SYonghong Song void bpf_task_release(struct task_struct *p) __ksym;
2848671232SYonghong Song 
2948671232SYonghong Song SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
get_cgroup_id(void * ctx)3048671232SYonghong Song int get_cgroup_id(void *ctx)
3148671232SYonghong Song {
3248671232SYonghong Song 	struct task_struct *task;
338723ec22SYonghong Song 	struct css_set *cgroups;
3448671232SYonghong Song 
3548671232SYonghong Song 	task = bpf_get_current_task_btf();
3648671232SYonghong Song 	if (task->pid != target_pid)
3748671232SYonghong Song 		return 0;
3848671232SYonghong Song 
3948671232SYonghong Song 	/* simulate bpf_get_current_cgroup_id() helper */
4048671232SYonghong Song 	bpf_rcu_read_lock();
418723ec22SYonghong Song 	cgroups = task->cgroups;
428723ec22SYonghong Song 	if (!cgroups)
438723ec22SYonghong Song 		goto unlock;
448723ec22SYonghong Song 	cgroup_id = cgroups->dfl_cgrp->kn->id;
458723ec22SYonghong Song unlock:
4648671232SYonghong Song 	bpf_rcu_read_unlock();
4748671232SYonghong Song 	return 0;
4848671232SYonghong Song }
4948671232SYonghong Song 
5048671232SYonghong Song SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
task_succ(void * ctx)5148671232SYonghong Song int task_succ(void *ctx)
5248671232SYonghong Song {
5348671232SYonghong Song 	struct task_struct *task, *real_parent;
5448671232SYonghong Song 	long init_val = 2;
5548671232SYonghong Song 	long *ptr;
5648671232SYonghong Song 
5748671232SYonghong Song 	task = bpf_get_current_task_btf();
5848671232SYonghong Song 	if (task->pid != target_pid)
5948671232SYonghong Song 		return 0;
6048671232SYonghong Song 
6148671232SYonghong Song 	bpf_rcu_read_lock();
6248671232SYonghong Song 	/* region including helper using rcu ptr real_parent */
6348671232SYonghong Song 	real_parent = task->real_parent;
648723ec22SYonghong Song 	if (!real_parent)
658723ec22SYonghong Song 		goto out;
6648671232SYonghong Song 	ptr = bpf_task_storage_get(&map_a, real_parent, &init_val,
6748671232SYonghong Song 				   BPF_LOCAL_STORAGE_GET_F_CREATE);
6848671232SYonghong Song 	if (!ptr)
6948671232SYonghong Song 		goto out;
7048671232SYonghong Song 	ptr = bpf_task_storage_get(&map_a, real_parent, 0, 0);
7148671232SYonghong Song 	if (!ptr)
7248671232SYonghong Song 		goto out;
7348671232SYonghong Song 	task_storage_val = *ptr;
7448671232SYonghong Song out:
7548671232SYonghong Song 	bpf_rcu_read_unlock();
7648671232SYonghong Song 	return 0;
7748671232SYonghong Song }
7848671232SYonghong Song 
7948671232SYonghong Song SEC("?fentry.s/" SYS_PREFIX "sys_nanosleep")
no_lock(void * ctx)8048671232SYonghong Song int no_lock(void *ctx)
8148671232SYonghong Song {
8248671232SYonghong Song 	struct task_struct *task, *real_parent;
8348671232SYonghong Song 
846fcd486bSAlexei Starovoitov 	/* old style ptr_to_btf_id is not allowed in sleepable */
8548671232SYonghong Song 	task = bpf_get_current_task_btf();
8648671232SYonghong Song 	real_parent = task->real_parent;
8748671232SYonghong Song 	(void)bpf_task_storage_get(&map_a, real_parent, 0, 0);
8848671232SYonghong Song 	return 0;
8948671232SYonghong Song }
9048671232SYonghong Song 
9148671232SYonghong Song SEC("?fentry.s/" SYS_PREFIX "sys_nanosleep")
two_regions(void * ctx)9248671232SYonghong Song int two_regions(void *ctx)
9348671232SYonghong Song {
9448671232SYonghong Song 	struct task_struct *task, *real_parent;
9548671232SYonghong Song 
9648671232SYonghong Song 	/* two regions */
9748671232SYonghong Song 	task = bpf_get_current_task_btf();
9848671232SYonghong Song 	bpf_rcu_read_lock();
9948671232SYonghong Song 	bpf_rcu_read_unlock();
10048671232SYonghong Song 	bpf_rcu_read_lock();
10148671232SYonghong Song 	real_parent = task->real_parent;
1028723ec22SYonghong Song 	if (!real_parent)
1038723ec22SYonghong Song 		goto out;
10448671232SYonghong Song 	(void)bpf_task_storage_get(&map_a, real_parent, 0, 0);
1058723ec22SYonghong Song out:
10648671232SYonghong Song 	bpf_rcu_read_unlock();
10748671232SYonghong Song 	return 0;
10848671232SYonghong Song }
10948671232SYonghong Song 
11048671232SYonghong Song SEC("?fentry/" SYS_PREFIX "sys_getpgid")
non_sleepable_1(void * ctx)11148671232SYonghong Song int non_sleepable_1(void *ctx)
11248671232SYonghong Song {
11348671232SYonghong Song 	struct task_struct *task, *real_parent;
11448671232SYonghong Song 
11548671232SYonghong Song 	task = bpf_get_current_task_btf();
11648671232SYonghong Song 	bpf_rcu_read_lock();
11748671232SYonghong Song 	real_parent = task->real_parent;
1188723ec22SYonghong Song 	if (!real_parent)
1198723ec22SYonghong Song 		goto out;
12048671232SYonghong Song 	(void)bpf_task_storage_get(&map_a, real_parent, 0, 0);
1218723ec22SYonghong Song out:
12248671232SYonghong Song 	bpf_rcu_read_unlock();
12348671232SYonghong Song 	return 0;
12448671232SYonghong Song }
12548671232SYonghong Song 
12648671232SYonghong Song SEC("?fentry/" SYS_PREFIX "sys_getpgid")
non_sleepable_2(void * ctx)12748671232SYonghong Song int non_sleepable_2(void *ctx)
12848671232SYonghong Song {
12948671232SYonghong Song 	struct task_struct *task, *real_parent;
13048671232SYonghong Song 
13148671232SYonghong Song 	bpf_rcu_read_lock();
13248671232SYonghong Song 	task = bpf_get_current_task_btf();
13348671232SYonghong Song 	bpf_rcu_read_unlock();
13448671232SYonghong Song 
13548671232SYonghong Song 	bpf_rcu_read_lock();
13648671232SYonghong Song 	real_parent = task->real_parent;
1378723ec22SYonghong Song 	if (!real_parent)
1388723ec22SYonghong Song 		goto out;
13948671232SYonghong Song 	(void)bpf_task_storage_get(&map_a, real_parent, 0, 0);
1408723ec22SYonghong Song out:
14148671232SYonghong Song 	bpf_rcu_read_unlock();
14248671232SYonghong Song 	return 0;
14348671232SYonghong Song }
14448671232SYonghong Song 
14548671232SYonghong Song SEC("?fentry.s/" SYS_PREFIX "sys_nanosleep")
task_acquire(void * ctx)14648671232SYonghong Song int task_acquire(void *ctx)
14748671232SYonghong Song {
1488723ec22SYonghong Song 	struct task_struct *task, *real_parent, *gparent;
14948671232SYonghong Song 
15048671232SYonghong Song 	task = bpf_get_current_task_btf();
15148671232SYonghong Song 	bpf_rcu_read_lock();
15248671232SYonghong Song 	real_parent = task->real_parent;
1538723ec22SYonghong Song 	if (!real_parent)
1548723ec22SYonghong Song 		goto out;
1558723ec22SYonghong Song 
1568723ec22SYonghong Song 	/* rcu_ptr->rcu_field */
1578723ec22SYonghong Song 	gparent = real_parent->real_parent;
1588723ec22SYonghong Song 	if (!gparent)
1598723ec22SYonghong Song 		goto out;
1608723ec22SYonghong Song 
16148671232SYonghong Song 	/* acquire a reference which can be used outside rcu read lock region */
162*f85671c6SDavid Vernet 	gparent = bpf_task_acquire(gparent);
1638723ec22SYonghong Song 	if (!gparent)
1648723ec22SYonghong Song 		goto out;
1658723ec22SYonghong Song 
1668723ec22SYonghong Song 	(void)bpf_task_storage_get(&map_a, gparent, 0, 0);
1678723ec22SYonghong Song 	bpf_task_release(gparent);
1688723ec22SYonghong Song out:
16948671232SYonghong Song 	bpf_rcu_read_unlock();
17048671232SYonghong Song 	return 0;
17148671232SYonghong Song }
17248671232SYonghong Song 
17348671232SYonghong Song SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
miss_lock(void * ctx)17448671232SYonghong Song int miss_lock(void *ctx)
17548671232SYonghong Song {
17648671232SYonghong Song 	struct task_struct *task;
17748671232SYonghong Song 
17848671232SYonghong Song 	/* missing bpf_rcu_read_lock() */
17948671232SYonghong Song 	task = bpf_get_current_task_btf();
18048671232SYonghong Song 	bpf_rcu_read_lock();
18148671232SYonghong Song 	(void)bpf_task_storage_get(&map_a, task, 0, 0);
18248671232SYonghong Song 	bpf_rcu_read_unlock();
18348671232SYonghong Song 	bpf_rcu_read_unlock();
18448671232SYonghong Song 	return 0;
18548671232SYonghong Song }
18648671232SYonghong Song 
18748671232SYonghong Song SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
miss_unlock(void * ctx)18848671232SYonghong Song int miss_unlock(void *ctx)
18948671232SYonghong Song {
19048671232SYonghong Song 	struct task_struct *task;
19148671232SYonghong Song 
19248671232SYonghong Song 	/* missing bpf_rcu_read_unlock() */
19348671232SYonghong Song 	task = bpf_get_current_task_btf();
19448671232SYonghong Song 	bpf_rcu_read_lock();
19548671232SYonghong Song 	(void)bpf_task_storage_get(&map_a, task, 0, 0);
19648671232SYonghong Song 	return 0;
19748671232SYonghong Song }
19848671232SYonghong Song 
19948671232SYonghong Song SEC("?fentry/" SYS_PREFIX "sys_getpgid")
non_sleepable_rcu_mismatch(void * ctx)20048671232SYonghong Song int non_sleepable_rcu_mismatch(void *ctx)
20148671232SYonghong Song {
20248671232SYonghong Song 	struct task_struct *task, *real_parent;
20348671232SYonghong Song 
20448671232SYonghong Song 	task = bpf_get_current_task_btf();
20548671232SYonghong Song 	/* non-sleepable: missing bpf_rcu_read_unlock() in one path */
20648671232SYonghong Song 	bpf_rcu_read_lock();
20748671232SYonghong Song 	real_parent = task->real_parent;
2088723ec22SYonghong Song 	if (!real_parent)
2098723ec22SYonghong Song 		goto out;
21048671232SYonghong Song 	(void)bpf_task_storage_get(&map_a, real_parent, 0, 0);
21148671232SYonghong Song 	if (real_parent)
21248671232SYonghong Song 		bpf_rcu_read_unlock();
2138723ec22SYonghong Song out:
21448671232SYonghong Song 	return 0;
21548671232SYonghong Song }
21648671232SYonghong Song 
21748671232SYonghong Song SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
inproper_sleepable_helper(void * ctx)21848671232SYonghong Song int inproper_sleepable_helper(void *ctx)
21948671232SYonghong Song {
22048671232SYonghong Song 	struct task_struct *task, *real_parent;
22148671232SYonghong Song 	struct pt_regs *regs;
22248671232SYonghong Song 	__u32 value = 0;
22348671232SYonghong Song 	void *ptr;
22448671232SYonghong Song 
22548671232SYonghong Song 	task = bpf_get_current_task_btf();
22648671232SYonghong Song 	/* sleepable helper in rcu read lock region */
22748671232SYonghong Song 	bpf_rcu_read_lock();
22848671232SYonghong Song 	real_parent = task->real_parent;
2298723ec22SYonghong Song 	if (!real_parent)
2308723ec22SYonghong Song 		goto out;
23148671232SYonghong Song 	regs = (struct pt_regs *)bpf_task_pt_regs(real_parent);
2328723ec22SYonghong Song 	if (!regs)
2338723ec22SYonghong Song 		goto out;
23448671232SYonghong Song 
23548671232SYonghong Song 	ptr = (void *)PT_REGS_IP(regs);
23648671232SYonghong Song 	(void)bpf_copy_from_user_task(&value, sizeof(uint32_t), ptr, task, 0);
23748671232SYonghong Song 	user_data = value;
23848671232SYonghong Song 	(void)bpf_task_storage_get(&map_a, real_parent, 0, 0);
2398723ec22SYonghong Song out:
24048671232SYonghong Song 	bpf_rcu_read_unlock();
24148671232SYonghong Song 	return 0;
24248671232SYonghong Song }
24348671232SYonghong Song 
24448671232SYonghong Song SEC("?lsm.s/bpf")
BPF_PROG(inproper_sleepable_kfunc,int cmd,union bpf_attr * attr,unsigned int size)24548671232SYonghong Song int BPF_PROG(inproper_sleepable_kfunc, int cmd, union bpf_attr *attr, unsigned int size)
24648671232SYonghong Song {
24748671232SYonghong Song 	struct bpf_key *bkey;
24848671232SYonghong Song 
24948671232SYonghong Song 	/* sleepable kfunc in rcu read lock region */
25048671232SYonghong Song 	bpf_rcu_read_lock();
25148671232SYonghong Song 	bkey = bpf_lookup_user_key(key_serial, flags);
25248671232SYonghong Song 	bpf_rcu_read_unlock();
25348671232SYonghong Song 	if (!bkey)
25448671232SYonghong Song 		return -1;
25548671232SYonghong Song 	bpf_key_put(bkey);
25648671232SYonghong Song 
25748671232SYonghong Song 	return 0;
25848671232SYonghong Song }
25948671232SYonghong Song 
26048671232SYonghong Song SEC("?fentry.s/" SYS_PREFIX "sys_nanosleep")
nested_rcu_region(void * ctx)26148671232SYonghong Song int nested_rcu_region(void *ctx)
26248671232SYonghong Song {
26348671232SYonghong Song 	struct task_struct *task, *real_parent;
26448671232SYonghong Song 
26548671232SYonghong Song 	/* nested rcu read lock regions */
26648671232SYonghong Song 	task = bpf_get_current_task_btf();
26748671232SYonghong Song 	bpf_rcu_read_lock();
26848671232SYonghong Song 	bpf_rcu_read_lock();
26948671232SYonghong Song 	real_parent = task->real_parent;
2708723ec22SYonghong Song 	if (!real_parent)
2718723ec22SYonghong Song 		goto out;
27248671232SYonghong Song 	(void)bpf_task_storage_get(&map_a, real_parent, 0, 0);
2738723ec22SYonghong Song out:
27448671232SYonghong Song 	bpf_rcu_read_unlock();
27548671232SYonghong Song 	bpf_rcu_read_unlock();
27648671232SYonghong Song 	return 0;
27748671232SYonghong Song }
27848671232SYonghong Song 
27948671232SYonghong Song SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
task_trusted_non_rcuptr(void * ctx)2806fcd486bSAlexei Starovoitov int task_trusted_non_rcuptr(void *ctx)
28148671232SYonghong Song {
282ec9230b1SYonghong Song 	struct task_struct *task, *group_leader;
28348671232SYonghong Song 
28448671232SYonghong Song 	task = bpf_get_current_task_btf();
28548671232SYonghong Song 	bpf_rcu_read_lock();
2866fcd486bSAlexei Starovoitov 	/* the pointer group_leader is explicitly marked as trusted */
287ec9230b1SYonghong Song 	group_leader = task->real_parent->group_leader;
288ec9230b1SYonghong Song 	(void)bpf_task_storage_get(&map_a, group_leader, 0, 0);
28948671232SYonghong Song 	bpf_rcu_read_unlock();
29048671232SYonghong Song 	return 0;
29148671232SYonghong Song }
29248671232SYonghong Song 
29348671232SYonghong Song SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
task_untrusted_rcuptr(void * ctx)29448671232SYonghong Song int task_untrusted_rcuptr(void *ctx)
29548671232SYonghong Song {
29648671232SYonghong Song 	struct task_struct *task, *real_parent;
29748671232SYonghong Song 
29848671232SYonghong Song 	task = bpf_get_current_task_btf();
29948671232SYonghong Song 	bpf_rcu_read_lock();
30048671232SYonghong Song 	real_parent = task->real_parent;
30148671232SYonghong Song 	bpf_rcu_read_unlock();
30248671232SYonghong Song 	/* helper use of rcu ptr outside the rcu read lock region */
30348671232SYonghong Song 	(void)bpf_task_storage_get(&map_a, real_parent, 0, 0);
30448671232SYonghong Song 	return 0;
30548671232SYonghong Song }
30648671232SYonghong Song 
30748671232SYonghong Song SEC("?fentry.s/" SYS_PREFIX "sys_nanosleep")
cross_rcu_region(void * ctx)30848671232SYonghong Song int cross_rcu_region(void *ctx)
30948671232SYonghong Song {
31048671232SYonghong Song 	struct task_struct *task, *real_parent;
31148671232SYonghong Song 
31248671232SYonghong Song 	/* rcu ptr define/use in different regions */
31348671232SYonghong Song 	task = bpf_get_current_task_btf();
31448671232SYonghong Song 	bpf_rcu_read_lock();
31548671232SYonghong Song 	real_parent = task->real_parent;
31648671232SYonghong Song 	bpf_rcu_read_unlock();
31748671232SYonghong Song 	bpf_rcu_read_lock();
31848671232SYonghong Song 	(void)bpf_task_storage_get(&map_a, real_parent, 0, 0);
31948671232SYonghong Song 	bpf_rcu_read_unlock();
32048671232SYonghong Song 	return 0;
32148671232SYonghong Song }
322