1*12bb6ca4SYonghong Song // SPDX-License-Identifier: GPL-2.0
2*12bb6ca4SYonghong Song /* Copyright (c) 2022 Meta Platforms, Inc. and affiliates. */
3*12bb6ca4SYonghong Song 
4*12bb6ca4SYonghong Song #include "vmlinux.h"
5*12bb6ca4SYonghong Song #include <bpf/bpf_helpers.h>
6*12bb6ca4SYonghong Song #include <bpf/bpf_tracing.h>
7*12bb6ca4SYonghong Song 
8*12bb6ca4SYonghong Song char _license[] SEC("license") = "GPL";
9*12bb6ca4SYonghong Song 
10*12bb6ca4SYonghong Song struct {
11*12bb6ca4SYonghong Song 	__uint(type, BPF_MAP_TYPE_CGRP_STORAGE);
12*12bb6ca4SYonghong Song 	__uint(map_flags, BPF_F_NO_PREALLOC);
13*12bb6ca4SYonghong Song 	__type(key, int);
14*12bb6ca4SYonghong Song 	__type(value, long);
15*12bb6ca4SYonghong Song } map_a SEC(".maps");
16*12bb6ca4SYonghong Song 
17*12bb6ca4SYonghong Song struct {
18*12bb6ca4SYonghong Song 	__uint(type, BPF_MAP_TYPE_CGRP_STORAGE);
19*12bb6ca4SYonghong Song 	__uint(map_flags, BPF_F_NO_PREALLOC);
20*12bb6ca4SYonghong Song 	__type(key, int);
21*12bb6ca4SYonghong Song 	__type(value, long);
22*12bb6ca4SYonghong Song } map_b SEC(".maps");
23*12bb6ca4SYonghong Song 
24*12bb6ca4SYonghong Song #define MAGIC_VALUE 0xabcd1234
25*12bb6ca4SYonghong Song 
26*12bb6ca4SYonghong Song pid_t target_pid = 0;
27*12bb6ca4SYonghong Song int mismatch_cnt = 0;
28*12bb6ca4SYonghong Song int enter_cnt = 0;
29*12bb6ca4SYonghong Song int exit_cnt = 0;
30*12bb6ca4SYonghong Song 
31*12bb6ca4SYonghong Song SEC("tp_btf/sys_enter")
BPF_PROG(on_enter,struct pt_regs * regs,long id)32*12bb6ca4SYonghong Song int BPF_PROG(on_enter, struct pt_regs *regs, long id)
33*12bb6ca4SYonghong Song {
34*12bb6ca4SYonghong Song 	struct task_struct *task;
35*12bb6ca4SYonghong Song 	long *ptr;
36*12bb6ca4SYonghong Song 	int err;
37*12bb6ca4SYonghong Song 
38*12bb6ca4SYonghong Song 	task = bpf_get_current_task_btf();
39*12bb6ca4SYonghong Song 	if (task->pid != target_pid)
40*12bb6ca4SYonghong Song 		return 0;
41*12bb6ca4SYonghong Song 
42*12bb6ca4SYonghong Song 	/* populate value 0 */
43*12bb6ca4SYonghong Song 	ptr = bpf_cgrp_storage_get(&map_a, task->cgroups->dfl_cgrp, 0,
44*12bb6ca4SYonghong Song 				   BPF_LOCAL_STORAGE_GET_F_CREATE);
45*12bb6ca4SYonghong Song 	if (!ptr)
46*12bb6ca4SYonghong Song 		return 0;
47*12bb6ca4SYonghong Song 
48*12bb6ca4SYonghong Song 	/* delete value 0 */
49*12bb6ca4SYonghong Song 	err = bpf_cgrp_storage_delete(&map_a, task->cgroups->dfl_cgrp);
50*12bb6ca4SYonghong Song 	if (err)
51*12bb6ca4SYonghong Song 		return 0;
52*12bb6ca4SYonghong Song 
53*12bb6ca4SYonghong Song 	/* value is not available */
54*12bb6ca4SYonghong Song 	ptr = bpf_cgrp_storage_get(&map_a, task->cgroups->dfl_cgrp, 0, 0);
55*12bb6ca4SYonghong Song 	if (ptr)
56*12bb6ca4SYonghong Song 		return 0;
57*12bb6ca4SYonghong Song 
58*12bb6ca4SYonghong Song 	/* re-populate the value */
59*12bb6ca4SYonghong Song 	ptr = bpf_cgrp_storage_get(&map_a, task->cgroups->dfl_cgrp, 0,
60*12bb6ca4SYonghong Song 				   BPF_LOCAL_STORAGE_GET_F_CREATE);
61*12bb6ca4SYonghong Song 	if (!ptr)
62*12bb6ca4SYonghong Song 		return 0;
63*12bb6ca4SYonghong Song 	__sync_fetch_and_add(&enter_cnt, 1);
64*12bb6ca4SYonghong Song 	*ptr = MAGIC_VALUE + enter_cnt;
65*12bb6ca4SYonghong Song 
66*12bb6ca4SYonghong Song 	return 0;
67*12bb6ca4SYonghong Song }
68*12bb6ca4SYonghong Song 
69*12bb6ca4SYonghong Song SEC("tp_btf/sys_exit")
BPF_PROG(on_exit,struct pt_regs * regs,long id)70*12bb6ca4SYonghong Song int BPF_PROG(on_exit, struct pt_regs *regs, long id)
71*12bb6ca4SYonghong Song {
72*12bb6ca4SYonghong Song 	struct task_struct *task;
73*12bb6ca4SYonghong Song 	long *ptr;
74*12bb6ca4SYonghong Song 
75*12bb6ca4SYonghong Song 	task = bpf_get_current_task_btf();
76*12bb6ca4SYonghong Song 	if (task->pid != target_pid)
77*12bb6ca4SYonghong Song 		return 0;
78*12bb6ca4SYonghong Song 
79*12bb6ca4SYonghong Song 	ptr = bpf_cgrp_storage_get(&map_a, task->cgroups->dfl_cgrp, 0,
80*12bb6ca4SYonghong Song 				   BPF_LOCAL_STORAGE_GET_F_CREATE);
81*12bb6ca4SYonghong Song 	if (!ptr)
82*12bb6ca4SYonghong Song 		return 0;
83*12bb6ca4SYonghong Song 
84*12bb6ca4SYonghong Song 	__sync_fetch_and_add(&exit_cnt, 1);
85*12bb6ca4SYonghong Song 	if (*ptr != MAGIC_VALUE + exit_cnt)
86*12bb6ca4SYonghong Song 		__sync_fetch_and_add(&mismatch_cnt, 1);
87*12bb6ca4SYonghong Song 	return 0;
88*12bb6ca4SYonghong Song }
89