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(cgrp: 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: regs, id: id, cgrp); |
88 | bpf_cgroup_release(cgrp); |
89 | return 0; |
90 | } |
91 | |
92 | __on_enter(regs: regs, id: id, cgrp: task->cgroups->dfl_cgrp); |
93 | return 0; |
94 | } |
95 | |