1// SPDX-License-Identifier: GPL-2.0
  2/* Copyright (c) 2022 Meta Platforms, Inc. and affiliates. */
  3#include <vmlinux.h>
  4#include <bpf/bpf_helpers.h>
  5#include <bpf/bpf_tracing.h>
  6#include "bpf_misc.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
 17__s32 target_pid;
 18__u64 cgroup_id;
 19int target_hid;
 20bool is_cgroup1;
 21
 22struct cgroup *bpf_task_get_cgroup1(struct task_struct *task, int hierarchy_id) __ksym;
 23void bpf_cgroup_release(struct cgroup *cgrp) __ksym;
 24void bpf_rcu_read_lock(void) __ksym;
 25void bpf_rcu_read_unlock(void) __ksym;
 26
 27SEC("?iter.s/cgroup")
 28int cgroup_iter(struct bpf_iter__cgroup *ctx)
 29{
 30	struct cgroup *cgrp = ctx->cgroup;
 31	long *ptr;
 32
 33	if (cgrp == NULL)
 34		return 0;
 35
 36	ptr = bpf_cgrp_storage_get(&map_a, cgrp, 0,
 37				   BPF_LOCAL_STORAGE_GET_F_CREATE);
 38	if (ptr)
 39		cgroup_id = cgrp->kn->id;
 40	return 0;
 41}
 42
 43static void __no_rcu_lock(struct cgroup *cgrp)
 44{
 45	long *ptr;
 46
 47	/* Note that trace rcu is held in sleepable prog, so we can use
 48	 * bpf_cgrp_storage_get() in sleepable prog.
 49	 */
 50	ptr = bpf_cgrp_storage_get(&map_a, cgrp, 0,
 51				   BPF_LOCAL_STORAGE_GET_F_CREATE);
 52	if (ptr)
 53		cgroup_id = cgrp->kn->id;
 54}
 55
 56SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
 57int cgrp1_no_rcu_lock(void *ctx)
 58{
 59	struct task_struct *task;
 60	struct cgroup *cgrp;
 61
 62	task = bpf_get_current_task_btf();
 63	if (task->pid != target_pid)
 64		return 0;
 65
 66	/* bpf_task_get_cgroup1 can work in sleepable prog */
 67	cgrp = bpf_task_get_cgroup1(task, target_hid);
 68	if (!cgrp)
 69		return 0;
 70
 71	__no_rcu_lock(cgrp);
 72	bpf_cgroup_release(cgrp);
 73	return 0;
 74}
 75
 76SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
 77int no_rcu_lock(void *ctx)
 78{
 79	struct task_struct *task;
 80
 81	task = bpf_get_current_task_btf();
 82	if (task->pid != target_pid)
 83		return 0;
 84
 85	/* task->cgroups is untrusted in sleepable prog outside of RCU CS */
 86	__no_rcu_lock(task->cgroups->dfl_cgrp);
 87	return 0;
 88}
 89
 90SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
 91int yes_rcu_lock(void *ctx)
 92{
 93	struct task_struct *task;
 94	struct cgroup *cgrp;
 95	long *ptr;
 96
 97	task = bpf_get_current_task_btf();
 98	if (task->pid != target_pid)
 99		return 0;
100
101	if (is_cgroup1) {
102		bpf_rcu_read_lock();
103		cgrp = bpf_task_get_cgroup1(task, target_hid);
104		if (!cgrp) {
105			bpf_rcu_read_unlock();
106			return 0;
107		}
108
109		ptr = bpf_cgrp_storage_get(&map_a, cgrp, 0, BPF_LOCAL_STORAGE_GET_F_CREATE);
110		if (ptr)
111			cgroup_id = cgrp->kn->id;
112		bpf_cgroup_release(cgrp);
113		bpf_rcu_read_unlock();
114		return 0;
115	}
116
117	bpf_rcu_read_lock();
118	cgrp = task->cgroups->dfl_cgrp;
119	/* cgrp is trusted under RCU CS */
120	ptr = bpf_cgrp_storage_get(&map_a, cgrp, 0, BPF_LOCAL_STORAGE_GET_F_CREATE);
121	if (ptr)
122		cgroup_id = cgrp->kn->id;
123	bpf_rcu_read_unlock();
124	return 0;
125}