95 lines · c
1// SPDX-License-Identifier: GPL-2.02/* 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 24int target_hid = 0;25bool is_cgroup1 = 0;26 27struct cgroup *bpf_task_get_cgroup1(struct task_struct *task, int hierarchy_id) __ksym;28void bpf_cgroup_release(struct cgroup *cgrp) __ksym;29 30static void __on_update(struct cgroup *cgrp)31{32 long *ptr;33 34 ptr = bpf_cgrp_storage_get(&map_a, cgrp, 0, BPF_LOCAL_STORAGE_GET_F_CREATE);35 if (ptr)36 *ptr += 1;37 38 ptr = bpf_cgrp_storage_get(&map_b, cgrp, 0, BPF_LOCAL_STORAGE_GET_F_CREATE);39 if (ptr)40 *ptr += 1;41}42 43SEC("fentry/bpf_local_storage_update")44int BPF_PROG(on_update)45{46 struct task_struct *task = bpf_get_current_task_btf();47 struct cgroup *cgrp;48 49 if (is_cgroup1) {50 cgrp = bpf_task_get_cgroup1(task, target_hid);51 if (!cgrp)52 return 0;53 54 __on_update(cgrp);55 bpf_cgroup_release(cgrp);56 return 0;57 }58 59 __on_update(task->cgroups->dfl_cgrp);60 return 0;61}62 63static void __on_enter(struct pt_regs *regs, long id, struct cgroup *cgrp)64{65 long *ptr;66 67 ptr = bpf_cgrp_storage_get(&map_a, cgrp, 0, BPF_LOCAL_STORAGE_GET_F_CREATE);68 if (ptr)69 *ptr = 200;70 71 ptr = bpf_cgrp_storage_get(&map_b, cgrp, 0, BPF_LOCAL_STORAGE_GET_F_CREATE);72 if (ptr)73 *ptr = 100;74}75 76SEC("tp_btf/sys_enter")77int BPF_PROG(on_enter, struct pt_regs *regs, long id)78{79 struct task_struct *task = bpf_get_current_task_btf();80 struct cgroup *cgrp;81 82 if (is_cgroup1) {83 cgrp = bpf_task_get_cgroup1(task, target_hid);84 if (!cgrp)85 return 0;86 87 __on_enter(regs, id, cgrp);88 bpf_cgroup_release(cgrp);89 return 0;90 }91 92 __on_enter(regs, id, task->cgroups->dfl_cgrp);93 return 0;94}95