1 /* SPDX-License-Identifier: GPL-2.0 */
2 /* Copyright (c) 2022 Meta Platforms, Inc. and affiliates. */
3
4 #ifndef _TASK_KFUNC_COMMON_H
5 #define _TASK_KFUNC_COMMON_H
6
7 #include <errno.h>
8 #include <vmlinux.h>
9 #include <bpf/bpf_helpers.h>
10 #include <bpf/bpf_tracing.h>
11
12 struct __tasks_kfunc_map_value {
13 struct task_struct __kptr * task;
14 };
15
16 struct {
17 __uint(type, BPF_MAP_TYPE_HASH);
18 __type(key, int);
19 __type(value, struct __tasks_kfunc_map_value);
20 __uint(max_entries, 1);
21 } __tasks_kfunc_map SEC(".maps");
22
23 struct task_kptr_lock_value {
24 struct bpf_spin_lock lock;
25 struct task_struct __kptr * task;
26 };
27
28 struct {
29 __uint(type, BPF_MAP_TYPE_ARRAY);
30 __type(key, int);
31 __type(value, struct task_kptr_lock_value);
32 __uint(max_entries, 1);
33 } task_kptr_lock_map SEC(".maps");
34
35 struct task_struct *bpf_task_acquire(struct task_struct *p) __ksym;
36 void bpf_task_release(struct task_struct *p) __ksym;
37 struct task_struct *bpf_task_from_pid(s32 pid) __ksym;
38 struct task_struct *bpf_task_from_vpid(s32 vpid) __ksym;
39 void bpf_rcu_read_lock(void) __ksym;
40 void bpf_rcu_read_unlock(void) __ksym;
41 void bpf_local_irq_save(unsigned long *flags) __weak __ksym;
42 void bpf_local_irq_restore(unsigned long *flags) __weak __ksym;
43
tasks_kfunc_map_value_lookup(struct task_struct * p)44 static inline struct __tasks_kfunc_map_value *tasks_kfunc_map_value_lookup(struct task_struct *p)
45 {
46 s32 pid;
47 long status;
48
49 status = bpf_probe_read_kernel(&pid, sizeof(pid), &p->pid);
50 if (status)
51 return NULL;
52
53 return bpf_map_lookup_elem(&__tasks_kfunc_map, &pid);
54 }
55
tasks_kfunc_map_insert(struct task_struct * p)56 static inline int tasks_kfunc_map_insert(struct task_struct *p)
57 {
58 struct __tasks_kfunc_map_value local, *v;
59 long status;
60 struct task_struct *acquired, *old;
61 s32 pid;
62
63 status = bpf_probe_read_kernel(&pid, sizeof(pid), &p->pid);
64 if (status)
65 return status;
66
67 local.task = NULL;
68 status = bpf_map_update_elem(&__tasks_kfunc_map, &pid, &local, BPF_NOEXIST);
69 if (status)
70 return status;
71
72 v = bpf_map_lookup_elem(&__tasks_kfunc_map, &pid);
73 if (!v) {
74 bpf_map_delete_elem(&__tasks_kfunc_map, &pid);
75 return -ENOENT;
76 }
77
78 acquired = bpf_task_acquire(p);
79 if (!acquired)
80 return -ENOENT;
81
82 old = bpf_kptr_xchg(&v->task, acquired);
83 if (old) {
84 bpf_task_release(old);
85 return -EEXIST;
86 }
87
88 return 0;
89 }
90
91 #endif /* _TASK_KFUNC_COMMON_H */
92