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}