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_tracing.h>
6 #include <bpf/bpf_helpers.h>
7 
8 #include "task_kfunc_common.h"
9 
10 char _license[] SEC("license") = "GPL";
11 
12 int err, pid;
13 
14 /* Prototype for all of the program trace events below:
15  *
16  * TRACE_EVENT(task_newtask,
17  *         TP_PROTO(struct task_struct *p, u64 clone_flags)
18  */
19 
20 static bool is_test_kfunc_task(void)
21 {
22 	int cur_pid = bpf_get_current_pid_tgid() >> 32;
23 
24 	return pid == cur_pid;
25 }
26 
27 static int test_acquire_release(struct task_struct *task)
28 {
29 	struct task_struct *acquired;
30 
31 	acquired = bpf_task_acquire(task);
32 	bpf_task_release(acquired);
33 
34 	return 0;
35 }
36 
37 SEC("tp_btf/task_newtask")
38 int BPF_PROG(test_task_acquire_release_argument, struct task_struct *task, u64 clone_flags)
39 {
40 	if (!is_test_kfunc_task())
41 		return 0;
42 
43 	return test_acquire_release(task);
44 }
45 
46 SEC("tp_btf/task_newtask")
47 int BPF_PROG(test_task_acquire_release_current, struct task_struct *task, u64 clone_flags)
48 {
49 	if (!is_test_kfunc_task())
50 		return 0;
51 
52 	return test_acquire_release(bpf_get_current_task_btf());
53 }
54 
55 SEC("tp_btf/task_newtask")
56 int BPF_PROG(test_task_acquire_leave_in_map, struct task_struct *task, u64 clone_flags)
57 {
58 	long status;
59 
60 	if (!is_test_kfunc_task())
61 		return 0;
62 
63 	status = tasks_kfunc_map_insert(task);
64 	if (status)
65 		err = 1;
66 
67 	return 0;
68 }
69 
70 SEC("tp_btf/task_newtask")
71 int BPF_PROG(test_task_xchg_release, struct task_struct *task, u64 clone_flags)
72 {
73 	struct task_struct *kptr;
74 	struct __tasks_kfunc_map_value *v;
75 	long status;
76 
77 	if (!is_test_kfunc_task())
78 		return 0;
79 
80 	status = tasks_kfunc_map_insert(task);
81 	if (status) {
82 		err = 1;
83 		return 0;
84 	}
85 
86 	v = tasks_kfunc_map_value_lookup(task);
87 	if (!v) {
88 		err = 2;
89 		return 0;
90 	}
91 
92 	kptr = bpf_kptr_xchg(&v->task, NULL);
93 	if (!kptr) {
94 		err = 3;
95 		return 0;
96 	}
97 
98 	bpf_task_release(kptr);
99 
100 	return 0;
101 }
102 
103 SEC("tp_btf/task_newtask")
104 int BPF_PROG(test_task_get_release, struct task_struct *task, u64 clone_flags)
105 {
106 	struct task_struct *kptr;
107 	struct __tasks_kfunc_map_value *v;
108 	long status;
109 
110 	if (!is_test_kfunc_task())
111 		return 0;
112 
113 	status = tasks_kfunc_map_insert(task);
114 	if (status) {
115 		err = 1;
116 		return 0;
117 	}
118 
119 	v = tasks_kfunc_map_value_lookup(task);
120 	if (!v) {
121 		err = 2;
122 		return 0;
123 	}
124 
125 	kptr = bpf_task_kptr_get(&v->task);
126 	if (!kptr) {
127 		err = 3;
128 		return 0;
129 	}
130 
131 	bpf_task_release(kptr);
132 
133 	return 0;
134 }
135 
136 SEC("tp_btf/task_newtask")
137 int BPF_PROG(test_task_current_acquire_release, struct task_struct *task, u64 clone_flags)
138 {
139 	struct task_struct *current, *acquired;
140 
141 	if (!is_test_kfunc_task())
142 		return 0;
143 
144 	current = bpf_get_current_task_btf();
145 	acquired = bpf_task_acquire(current);
146 	bpf_task_release(acquired);
147 
148 	return 0;
149 }
150 
151 static void lookup_compare_pid(const struct task_struct *p)
152 {
153 	struct task_struct *acquired;
154 
155 	acquired = bpf_task_from_pid(p->pid);
156 	if (!acquired) {
157 		err = 1;
158 		return;
159 	}
160 
161 	if (acquired->pid != p->pid)
162 		err = 2;
163 	bpf_task_release(acquired);
164 }
165 
166 SEC("tp_btf/task_newtask")
167 int BPF_PROG(test_task_from_pid_arg, struct task_struct *task, u64 clone_flags)
168 {
169 	struct task_struct *acquired;
170 
171 	if (!is_test_kfunc_task())
172 		return 0;
173 
174 	lookup_compare_pid(task);
175 	return 0;
176 }
177 
178 SEC("tp_btf/task_newtask")
179 int BPF_PROG(test_task_from_pid_current, struct task_struct *task, u64 clone_flags)
180 {
181 	struct task_struct *current, *acquired;
182 
183 	if (!is_test_kfunc_task())
184 		return 0;
185 
186 	lookup_compare_pid(bpf_get_current_task_btf());
187 	return 0;
188 }
189 
190 static int is_pid_lookup_valid(s32 pid)
191 {
192 	struct task_struct *acquired;
193 
194 	acquired = bpf_task_from_pid(pid);
195 	if (acquired) {
196 		bpf_task_release(acquired);
197 		return 1;
198 	}
199 
200 	return 0;
201 }
202 
203 SEC("tp_btf/task_newtask")
204 int BPF_PROG(test_task_from_pid_invalid, struct task_struct *task, u64 clone_flags)
205 {
206 	struct task_struct *acquired;
207 
208 	if (!is_test_kfunc_task())
209 		return 0;
210 
211 	if (is_pid_lookup_valid(-1)) {
212 		err = 1;
213 		return 0;
214 	}
215 
216 	if (is_pid_lookup_valid(0xcafef00d)) {
217 		err = 2;
218 		return 0;
219 	}
220 
221 	return 0;
222 }
223