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 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 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