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
__on_lookup(struct cgroup * cgrp)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")
BPF_PROG(on_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
__on_update(struct cgroup * cgrp)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")
BPF_PROG(on_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
__on_enter(struct pt_regs * regs,long id,struct cgroup * cgrp)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")
BPF_PROG(on_enter,struct pt_regs * regs,long id)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