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_lookup(struct cgroup *cgrp) 31 { 32 bpf_cgrp_storage_delete(&map_a, cgrp); 33 bpf_cgrp_storage_delete(&map_b, cgrp); 34 } 35 36 SEC("fentry/bpf_local_storage_lookup") 37 int BPF_PROG(on_lookup) 38 { 39 struct task_struct *task = bpf_get_current_task_btf(); 40 struct cgroup *cgrp; 41 42 if (is_cgroup1) { 43 cgrp = bpf_task_get_cgroup1(task, target_hid); 44 if (!cgrp) 45 return 0; 46 47 __on_lookup(cgrp); 48 bpf_cgroup_release(cgrp); 49 return 0; 50 } 51 52 __on_lookup(task->cgroups->dfl_cgrp); 53 return 0; 54 } 55 56 static void __on_update(struct cgroup *cgrp) 57 { 58 long *ptr; 59 60 ptr = bpf_cgrp_storage_get(&map_a, cgrp, 0, BPF_LOCAL_STORAGE_GET_F_CREATE); 61 if (ptr) 62 *ptr += 1; 63 64 ptr = bpf_cgrp_storage_get(&map_b, cgrp, 0, BPF_LOCAL_STORAGE_GET_F_CREATE); 65 if (ptr) 66 *ptr += 1; 67 } 68 69 SEC("fentry/bpf_local_storage_update") 70 int BPF_PROG(on_update) 71 { 72 struct task_struct *task = bpf_get_current_task_btf(); 73 struct cgroup *cgrp; 74 75 if (is_cgroup1) { 76 cgrp = bpf_task_get_cgroup1(task, target_hid); 77 if (!cgrp) 78 return 0; 79 80 __on_update(cgrp); 81 bpf_cgroup_release(cgrp); 82 return 0; 83 } 84 85 __on_update(task->cgroups->dfl_cgrp); 86 return 0; 87 } 88 89 static void __on_enter(struct pt_regs *regs, long id, struct cgroup *cgrp) 90 { 91 long *ptr; 92 93 ptr = bpf_cgrp_storage_get(&map_a, cgrp, 0, BPF_LOCAL_STORAGE_GET_F_CREATE); 94 if (ptr) 95 *ptr = 200; 96 97 ptr = bpf_cgrp_storage_get(&map_b, cgrp, 0, BPF_LOCAL_STORAGE_GET_F_CREATE); 98 if (ptr) 99 *ptr = 100; 100 } 101 102 SEC("tp_btf/sys_enter") 103 int BPF_PROG(on_enter, struct pt_regs *regs, long id) 104 { 105 struct task_struct *task = bpf_get_current_task_btf(); 106 struct cgroup *cgrp; 107 108 if (is_cgroup1) { 109 cgrp = bpf_task_get_cgroup1(task, target_hid); 110 if (!cgrp) 111 return 0; 112 113 __on_enter(regs, id, cgrp); 114 bpf_cgroup_release(cgrp); 115 return 0; 116 } 117 118 __on_enter(regs, id, task->cgroups->dfl_cgrp); 119 return 0; 120 } 121