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