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