tools/testing/selftests/bpf/progs/cgrp_ls_recursion.c

   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_helpers.h>
   6 #include <bpf/bpf_tracing.h>
   7
   8 char _license[] SEC("license") = "GPL";
   9
  10 struct {
  11         __uint(type, BPF_MAP_TYPE_CGRP_STORAGE);
  12         __uint(map_flags, BPF_F_NO_PREALLOC);
  13         __type(key, int);
  14         __type(value, long);
  15 } map_a SEC(".maps");
  16
  17 struct {
  18         __uint(type, BPF_MAP_TYPE_CGRP_STORAGE);
  19         __uint(map_flags, BPF_F_NO_PREALLOC);
  20         __type(key, int);
  21         __type(value, long);
  22 } map_b SEC(".maps");
  23
  24 int target_hid = 0;
  25 bool is_cgroup1 = 0;
  26
  27 struct cgroup *bpf_task_get_cgroup1(struct task_struct *task, int hierarchy_id) __ksym;
  28 void bpf_cgroup_release(struct cgroup *cgrp) __ksym;
  29
  30 static void __on_update(struct cgroup *cgrp)
  31 {
  32         long *ptr;
  33
  34         ptr = bpf_cgrp_storage_get(&map_a, cgrp, 0, BPF_LOCAL_STORAGE_GET_F_CREATE);
  35         if (ptr)
  36                 *ptr += 1;
  37
  38         ptr = bpf_cgrp_storage_get(&map_b, cgrp, 0, BPF_LOCAL_STORAGE_GET_F_CREATE);
  39         if (ptr)
  40                 *ptr += 1;
  41 }
  42
  43 SEC("fentry/bpf_local_storage_update")
  44 int BPF_PROG(on_update)
  45 {
  46         struct task_struct *task = bpf_get_current_task_btf();
  47         struct cgroup *cgrp;
  48
  49         if (is_cgroup1) {
  50                 cgrp = bpf_task_get_cgroup1(task, target_hid);
  51                 if (!cgrp)
  52                         return 0;
  53
  54                 __on_update(cgrp);
  55                 bpf_cgroup_release(cgrp);
  56                 return 0;
  57         }
  58
  59         __on_update(task->cgroups->dfl_cgrp);
  60         return 0;
  61 }
  62
  63 static void __on_enter(struct pt_regs *regs, long id, struct cgroup *cgrp)
  64 {
  65         long *ptr;
  66
  67         ptr = bpf_cgrp_storage_get(&map_a, cgrp, 0, BPF_LOCAL_STORAGE_GET_F_CREATE);
  68         if (ptr)
  69                 *ptr = 200;
  70
  71         ptr = bpf_cgrp_storage_get(&map_b, cgrp, 0, BPF_LOCAL_STORAGE_GET_F_CREATE);
  72         if (ptr)
  73                 *ptr = 100;
  74 }
  75
  76 SEC("tp_btf/sys_enter")
  77 int BPF_PROG(on_enter, struct pt_regs *regs, long id)
  78 {
  79         struct task_struct *task = bpf_get_current_task_btf();
  80         struct cgroup *cgrp;
  81
  82         if (is_cgroup1) {
  83                 cgrp = bpf_task_get_cgroup1(task, target_hid);
  84                 if (!cgrp)
  85                         return 0;
  86
  87                 __on_enter(regs, id, cgrp);
  88                 bpf_cgroup_release(cgrp);
  89                 return 0;
  90         }
  91
  92         __on_enter(regs, id, task->cgroups->dfl_cgrp);
  93         return 0;
  94 }