tools/testing/selftests/bpf/progs/task_kfunc_common.h

   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 hash_map {
  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_struct *bpf_task_acquire(struct task_struct *p) __ksym;
  24 void bpf_task_release(struct task_struct *p) __ksym;
  25 struct task_struct *bpf_task_from_pid(s32 pid) __ksym;
  26 void bpf_rcu_read_lock(void) __ksym;
  27 void bpf_rcu_read_unlock(void) __ksym;
  28
  29 static inline struct __tasks_kfunc_map_value *tasks_kfunc_map_value_lookup(struct task_struct *p)
  30 {
  31         s32 pid;
  32         long status;
  33
  34         status = bpf_probe_read_kernel(&pid, sizeof(pid), &p->pid);
  35         if (status)
  36                 return NULL;
  37
  38         return bpf_map_lookup_elem(&__tasks_kfunc_map, &pid);
  39 }
  40
  41 static inline int tasks_kfunc_map_insert(struct task_struct *p)
  42 {
  43         struct __tasks_kfunc_map_value local, *v;
  44         long status;
  45         struct task_struct *acquired, *old;
  46         s32 pid;
  47
  48         status = bpf_probe_read_kernel(&pid, sizeof(pid), &p->pid);
  49         if (status)
  50                 return status;
  51
  52         local.task = NULL;
  53         status = bpf_map_update_elem(&__tasks_kfunc_map, &pid, &local, BPF_NOEXIST);
  54         if (status)
  55                 return status;
  56
  57         v = bpf_map_lookup_elem(&__tasks_kfunc_map, &pid);
  58         if (!v) {
  59                 bpf_map_delete_elem(&__tasks_kfunc_map, &pid);
  60                 return -ENOENT;
  61         }
  62
  63         acquired = bpf_task_acquire(p);
  64         if (!acquired)
  65                 return -ENOENT;
  66
  67         old = bpf_kptr_xchg(&v->task, acquired);
  68         if (old) {
  69                 bpf_task_release(old);
  70                 return -EEXIST;
  71         }
  72
  73         return 0;
  74 }
  75
  76 #endif /* _TASK_KFUNC_COMMON_H */