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