// SPDX-License-Identifier: GPL-2.0
/* Copyright (c) 2023 Bytedance */

#include <vmlinux.h>
#include <bpf/bpf_tracing.h>
#include <bpf/bpf_helpers.h>

#include "bpf_misc.h"

struct cgroup *bpf_cgroup_from_id(u64 cgid) __ksym;
long bpf_task_under_cgroup(struct task_struct *task, struct cgroup *ancestor) __ksym;
void bpf_cgroup_release(struct cgroup *p) __ksym;
struct task_struct *bpf_task_acquire(struct task_struct *p) __ksym;
void bpf_task_release(struct task_struct *p) __ksym;

const volatile int local_pid;
const volatile __u64 cgid;
int remote_pid;

SEC("tp_btf/task_newtask")
int BPF_PROG(tp_btf_run, struct task_struct *task, u64 clone_flags)
{
	struct cgroup *cgrp = NULL;
	struct task_struct *acquired;

	if (local_pid != (bpf_get_current_pid_tgid() >> 32))
		return 0;

	acquired = bpf_task_acquire(task);
	if (!acquired)
		return 0;

	if (local_pid == acquired->tgid)
		goto out;

	cgrp = bpf_cgroup_from_id(cgid);
	if (!cgrp)
		goto out;

	if (bpf_task_under_cgroup(acquired, cgrp))
		remote_pid = acquired->tgid;

out:
	if (cgrp)
		bpf_cgroup_release(cgrp);
	bpf_task_release(acquired);

	return 0;
}

SEC("lsm.s/bpf")
int BPF_PROG(lsm_run, int cmd, union bpf_attr *attr, unsigned int size)
{
	struct cgroup *cgrp = NULL;
	struct task_struct *task;
	int ret = 0;

	task = bpf_get_current_task_btf();
	if (local_pid != task->pid)
		return 0;

	if (cmd != BPF_LINK_CREATE)
		return 0;

	/* 1 is the root cgroup */
	cgrp = bpf_cgroup_from_id(1);
	if (!cgrp)
		goto out;
	if (!bpf_task_under_cgroup(task, cgrp))
		ret = -1;
	bpf_cgroup_release(cgrp);

out:
	return ret;
}

char _license[] SEC("license") = "GPL";