141d76c72SYonghong Song // SPDX-License-Identifier: GPL-2.0
241d76c72SYonghong Song /* Copyright (c) 2022 Meta Platforms, Inc. and affiliates. */
341d76c72SYonghong Song 
441d76c72SYonghong Song #include "bpf_iter.h"
541d76c72SYonghong Song #include <bpf/bpf_helpers.h>
641d76c72SYonghong Song #include <bpf/bpf_tracing.h>
741d76c72SYonghong Song #include "bpf_misc.h"
841d76c72SYonghong Song 
941d76c72SYonghong Song char _license[] SEC("license") = "GPL";
1041d76c72SYonghong Song 
1141d76c72SYonghong Song struct {
1241d76c72SYonghong Song 	__uint(type, BPF_MAP_TYPE_CGRP_STORAGE);
1341d76c72SYonghong Song 	__uint(map_flags, BPF_F_NO_PREALLOC);
1441d76c72SYonghong Song 	__type(key, int);
1541d76c72SYonghong Song 	__type(value, long);
1641d76c72SYonghong Song } map_a SEC(".maps");
1741d76c72SYonghong Song 
1841d76c72SYonghong Song __u32 target_pid;
1941d76c72SYonghong Song __u64 cgroup_id;
2041d76c72SYonghong Song 
2141d76c72SYonghong Song void bpf_rcu_read_lock(void) __ksym;
2241d76c72SYonghong Song void bpf_rcu_read_unlock(void) __ksym;
2341d76c72SYonghong Song 
2441d76c72SYonghong Song SEC("?iter.s/cgroup")
cgroup_iter(struct bpf_iter__cgroup * ctx)2541d76c72SYonghong Song int cgroup_iter(struct bpf_iter__cgroup *ctx)
2641d76c72SYonghong Song {
2741d76c72SYonghong Song 	struct cgroup *cgrp = ctx->cgroup;
2841d76c72SYonghong Song 	long *ptr;
2941d76c72SYonghong Song 
3041d76c72SYonghong Song 	if (cgrp == NULL)
3141d76c72SYonghong Song 		return 0;
3241d76c72SYonghong Song 
3341d76c72SYonghong Song 	ptr = bpf_cgrp_storage_get(&map_a, cgrp, 0,
3441d76c72SYonghong Song 				   BPF_LOCAL_STORAGE_GET_F_CREATE);
3541d76c72SYonghong Song 	if (ptr)
3641d76c72SYonghong Song 		cgroup_id = cgrp->kn->id;
3741d76c72SYonghong Song 	return 0;
3841d76c72SYonghong Song }
3941d76c72SYonghong Song 
4041d76c72SYonghong Song SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
no_rcu_lock(void * ctx)4141d76c72SYonghong Song int no_rcu_lock(void *ctx)
4241d76c72SYonghong Song {
4341d76c72SYonghong Song 	struct task_struct *task;
4441d76c72SYonghong Song 	struct cgroup *cgrp;
4541d76c72SYonghong Song 	long *ptr;
4641d76c72SYonghong Song 
4741d76c72SYonghong Song 	task = bpf_get_current_task_btf();
4841d76c72SYonghong Song 	if (task->pid != target_pid)
4941d76c72SYonghong Song 		return 0;
5041d76c72SYonghong Song 
51*6fcd486bSAlexei Starovoitov 	/* task->cgroups is untrusted in sleepable prog outside of RCU CS */
5241d76c72SYonghong Song 	cgrp = task->cgroups->dfl_cgrp;
5341d76c72SYonghong Song 	ptr = bpf_cgrp_storage_get(&map_a, cgrp, 0,
5441d76c72SYonghong Song 				   BPF_LOCAL_STORAGE_GET_F_CREATE);
5541d76c72SYonghong Song 	if (ptr)
5641d76c72SYonghong Song 		cgroup_id = cgrp->kn->id;
5741d76c72SYonghong Song 	return 0;
5841d76c72SYonghong Song }
5941d76c72SYonghong Song 
6041d76c72SYonghong Song SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
yes_rcu_lock(void * ctx)6141d76c72SYonghong Song int yes_rcu_lock(void *ctx)
6241d76c72SYonghong Song {
6341d76c72SYonghong Song 	struct task_struct *task;
6441d76c72SYonghong Song 	struct cgroup *cgrp;
6541d76c72SYonghong Song 	long *ptr;
6641d76c72SYonghong Song 
6741d76c72SYonghong Song 	task = bpf_get_current_task_btf();
6841d76c72SYonghong Song 	if (task->pid != target_pid)
6941d76c72SYonghong Song 		return 0;
7041d76c72SYonghong Song 
7141d76c72SYonghong Song 	bpf_rcu_read_lock();
7241d76c72SYonghong Song 	cgrp = task->cgroups->dfl_cgrp;
73*6fcd486bSAlexei Starovoitov 	/* cgrp is trusted under RCU CS */
7441d76c72SYonghong Song 	ptr = bpf_cgrp_storage_get(&map_a, cgrp, 0, BPF_LOCAL_STORAGE_GET_F_CREATE);
7541d76c72SYonghong Song 	if (ptr)
7641d76c72SYonghong Song 		cgroup_id = cgrp->kn->id;
7741d76c72SYonghong Song 	bpf_rcu_read_unlock();
7841d76c72SYonghong Song 	return 0;
7941d76c72SYonghong Song }
80