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 SEC("fentry/bpf_local_storage_lookup")
BPF_PROG(on_lookup)25*12bb6ca4SYonghong Song int BPF_PROG(on_lookup)
26*12bb6ca4SYonghong Song {
27*12bb6ca4SYonghong Song 	struct task_struct *task = bpf_get_current_task_btf();
28*12bb6ca4SYonghong Song 
29*12bb6ca4SYonghong Song 	bpf_cgrp_storage_delete(&map_a, task->cgroups->dfl_cgrp);
30*12bb6ca4SYonghong Song 	bpf_cgrp_storage_delete(&map_b, task->cgroups->dfl_cgrp);
31*12bb6ca4SYonghong Song 	return 0;
32*12bb6ca4SYonghong Song }
33*12bb6ca4SYonghong Song 
34*12bb6ca4SYonghong Song SEC("fentry/bpf_local_storage_update")
BPF_PROG(on_update)35*12bb6ca4SYonghong Song int BPF_PROG(on_update)
36*12bb6ca4SYonghong Song {
37*12bb6ca4SYonghong Song 	struct task_struct *task = bpf_get_current_task_btf();
38*12bb6ca4SYonghong Song 	long *ptr;
39*12bb6ca4SYonghong Song 
40*12bb6ca4SYonghong Song 	ptr = bpf_cgrp_storage_get(&map_a, task->cgroups->dfl_cgrp, 0,
41*12bb6ca4SYonghong Song 				   BPF_LOCAL_STORAGE_GET_F_CREATE);
42*12bb6ca4SYonghong Song 	if (ptr)
43*12bb6ca4SYonghong Song 		*ptr += 1;
44*12bb6ca4SYonghong Song 
45*12bb6ca4SYonghong Song 	ptr = bpf_cgrp_storage_get(&map_b, task->cgroups->dfl_cgrp, 0,
46*12bb6ca4SYonghong Song 				   BPF_LOCAL_STORAGE_GET_F_CREATE);
47*12bb6ca4SYonghong Song 	if (ptr)
48*12bb6ca4SYonghong Song 		*ptr += 1;
49*12bb6ca4SYonghong Song 
50*12bb6ca4SYonghong Song 	return 0;
51*12bb6ca4SYonghong Song }
52*12bb6ca4SYonghong Song 
53*12bb6ca4SYonghong Song SEC("tp_btf/sys_enter")
BPF_PROG(on_enter,struct pt_regs * regs,long id)54*12bb6ca4SYonghong Song int BPF_PROG(on_enter, struct pt_regs *regs, long id)
55*12bb6ca4SYonghong Song {
56*12bb6ca4SYonghong Song 	struct task_struct *task;
57*12bb6ca4SYonghong Song 	long *ptr;
58*12bb6ca4SYonghong Song 
59*12bb6ca4SYonghong Song 	task = bpf_get_current_task_btf();
60*12bb6ca4SYonghong Song 	ptr = bpf_cgrp_storage_get(&map_a, task->cgroups->dfl_cgrp, 0,
61*12bb6ca4SYonghong Song 				   BPF_LOCAL_STORAGE_GET_F_CREATE);
62*12bb6ca4SYonghong Song 	if (ptr)
63*12bb6ca4SYonghong Song 		*ptr = 200;
64*12bb6ca4SYonghong Song 
65*12bb6ca4SYonghong Song 	ptr = bpf_cgrp_storage_get(&map_b, task->cgroups->dfl_cgrp, 0,
66*12bb6ca4SYonghong Song 				   BPF_LOCAL_STORAGE_GET_F_CREATE);
67*12bb6ca4SYonghong Song 	if (ptr)
68*12bb6ca4SYonghong Song 		*ptr = 100;
69*12bb6ca4SYonghong Song 	return 0;
70*12bb6ca4SYonghong Song }
71