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
8char _license[] SEC("license") = "GPL";
9
10struct {
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
17struct {
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#define MAGIC_VALUE 0xabcd1234
25
26pid_t target_pid = 0;
27int mismatch_cnt = 0;
28int enter_cnt = 0;
29int exit_cnt = 0;
30int target_hid = 0;
31bool is_cgroup1 = 0;
32
33struct cgroup *bpf_task_get_cgroup1(struct task_struct *task, int hierarchy_id) __ksym;
34void bpf_cgroup_release(struct cgroup *cgrp) __ksym;
35
36static void __on_enter(struct pt_regs *regs, long id, struct cgroup *cgrp)
37{
38	long *ptr;
39	int err;
40
41	/* populate value 0 */
42	ptr = bpf_cgrp_storage_get(&map_a, cgrp, 0,
43				   BPF_LOCAL_STORAGE_GET_F_CREATE);
44	if (!ptr)
45		return;
46
47	/* delete value 0 */
48	err = bpf_cgrp_storage_delete(&map_a, cgrp);
49	if (err)
50		return;
51
52	/* value is not available */
53	ptr = bpf_cgrp_storage_get(&map_a, cgrp, 0, 0);
54	if (ptr)
55		return;
56
57	/* re-populate the value */
58	ptr = bpf_cgrp_storage_get(&map_a, cgrp, 0,
59				   BPF_LOCAL_STORAGE_GET_F_CREATE);
60	if (!ptr)
61		return;
62	__sync_fetch_and_add(&enter_cnt, 1);
63	*ptr = MAGIC_VALUE + enter_cnt;
64}
65
66SEC("tp_btf/sys_enter")
67int BPF_PROG(on_enter, struct pt_regs *regs, long id)
68{
69	struct task_struct *task;
70	struct cgroup *cgrp;
71
72	task = bpf_get_current_task_btf();
73	if (task->pid != target_pid)
74		return 0;
75
76	if (is_cgroup1) {
77		cgrp = bpf_task_get_cgroup1(task, target_hid);
78		if (!cgrp)
79			return 0;
80
81		__on_enter(regs, id, cgrp);
82		bpf_cgroup_release(cgrp);
83		return 0;
84	}
85
86	__on_enter(regs, id, task->cgroups->dfl_cgrp);
87	return 0;
88}
89
90static void __on_exit(struct pt_regs *regs, long id, struct cgroup *cgrp)
91{
92	long *ptr;
93
94	ptr = bpf_cgrp_storage_get(&map_a, cgrp, 0,
95				   BPF_LOCAL_STORAGE_GET_F_CREATE);
96	if (!ptr)
97		return;
98
99	__sync_fetch_and_add(&exit_cnt, 1);
100	if (*ptr != MAGIC_VALUE + exit_cnt)
101		__sync_fetch_and_add(&mismatch_cnt, 1);
102}
103
104SEC("tp_btf/sys_exit")
105int BPF_PROG(on_exit, struct pt_regs *regs, long id)
106{
107	struct task_struct *task;
108	struct cgroup *cgrp;
109
110	task = bpf_get_current_task_btf();
111	if (task->pid != target_pid)
112		return 0;
113
114	if (is_cgroup1) {
115		cgrp = bpf_task_get_cgroup1(task, target_hid);
116		if (!cgrp)
117			return 0;
118
119		__on_exit(regs, id, cgrp);
120		bpf_cgroup_release(cgrp);
121		return 0;
122	}
123
124	__on_exit(regs, id, task->cgroups->dfl_cgrp);
125	return 0;
126}
127