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