// SPDX-License-Identifier: GPL-2.0
// Copyright (C) 2024 KylinSoft Co., Ltd. All rights reserved.
//
// Base on capable.py - Brendan Gregg

#include "vmlinux.h"
#include "capable.h"
#include <bpf/bpf_core_read.h>
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_tracing.h>
#include "compat.bpf.h"

#define MAX_ENTRIES	10240

const volatile enum uniqueness unique_type = UNQ_OFF;
const volatile bool kernel_stack = false;
const volatile bool user_stack = false;
const volatile bool filter_cg = false;
const volatile pid_t target_pid = -1;
const volatile bool use_pidns = false;
const volatile __u64 pidns_dev = 0;
const volatile __u64 pidns_ino = 0;
/*
 * Set by userspace based on BTF detection.
 * True if cap_out parameter is a bitmask (kernel >= 5.1),
 * false if cap_out is a simple audit flag (kernel < 5.1)
 */
const volatile bool has_cap_flags = false;

struct args_t {
	int cap;
	int cap_out;
};

struct unique_key {
	int cap;
	u32 tgid;
	u64 cgroupid;
};

struct {
	__uint(type, BPF_MAP_TYPE_HASH);
	__uint(max_entries, MAX_ENTRIES);
	__type(key, u64);
	__type(value, struct args_t);
} start SEC(".maps");

struct {
	__uint(type, BPF_MAP_TYPE_CGROUP_ARRAY);
	__uint(max_entries, 1);
	__type(key, u32);
	__type(value, u32);
} cgroup_map SEC(".maps");

struct {
	__uint(type, BPF_MAP_TYPE_STACK_TRACE);
	__type(key, u32);
} stackmap SEC(".maps");

struct {
	__uint(type, BPF_MAP_TYPE_HASH);
	__uint(max_entries, MAX_ENTRIES);
	__type(key, struct key_t);
	__type(value, struct cap_event);
} info SEC(".maps");

struct {
	__uint(type, BPF_MAP_TYPE_HASH);
	__uint(max_entries, MAX_ENTRIES);
	__type(key, struct unique_key);
	__type(value, u64);
} seen SEC(".maps");

SEC("kprobe/cap_capable")
int BPF_KPROBE(kprobe__cap_capable_entry, const struct cred *cred,
	       struct user_namespace *target_ns, int cap, int cap_out)
{
	__u64 pid_tgid;
	struct bpf_pidns_info nsdata;

	if (filter_cg && !bpf_current_task_under_cgroup(&cgroup_map, 0))
		return 0;

	if (use_pidns) {
		if (bpf_get_ns_current_pid_tgid(pidns_dev, pidns_ino, &nsdata,
						sizeof(struct bpf_pidns_info)))
			return 0;
		pid_tgid = (__u64)nsdata.tgid << 32 | nsdata.pid;
	} else {
		pid_tgid = bpf_get_current_pid_tgid();
	}

	if (target_pid != -1 && target_pid != (u32)pid_tgid)
		return 0;

	struct args_t args = {};
	args.cap = cap;
	args.cap_out = cap_out;

	bpf_map_update_elem(&start, &pid_tgid, &args, BPF_ANY);

	return 0;
}

SEC("kretprobe/cap_capable")
int BPF_KRETPROBE(kretprobe__cap_capable_exit)
{
	__u64 pid_tgid;
	struct args_t *argsp;
	struct key_t i_key;
	struct bpf_pidns_info nsdata;

	if (use_pidns) {
		if (bpf_get_ns_current_pid_tgid(pidns_dev, pidns_ino, &nsdata,
						sizeof(struct bpf_pidns_info)))
			return 0;
		pid_tgid = (__u64)nsdata.tgid << 32 | nsdata.pid;
	} else {
		pid_tgid = bpf_get_current_pid_tgid();
	}

	argsp = bpf_map_lookup_elem(&start, &pid_tgid);
	if (!argsp)
		return 0;

	bpf_map_delete_elem(&start, &pid_tgid);

	struct cap_event *event = reserve_buf(sizeof(*event));
	if (!event)
		return 0;

	event->pid = pid_tgid >> 32;
	event->tgid = pid_tgid;
	event->cap = argsp->cap;
	event->uid = bpf_get_current_uid_gid();
	bpf_get_current_comm(&event->task, sizeof(event->task));
	event->ret = PT_REGS_RC(ctx);

	if (has_cap_flags) {
		/* @opts: bitmask of options defined in include/linux/security.h */
		event->audit = (argsp->cap_out & 0b10) == 0;
		event->insetid = (argsp->cap_out & 0b100) != 0;
	} else {
		event->audit = argsp->cap_out;
		event->insetid = -1;
	}

	if (unique_type) {
		struct unique_key key = { .cap = argsp->cap };

		if (unique_type == UNQ_CGROUP)
			key.cgroupid = bpf_get_current_cgroup_id();
		else
			key.tgid = pid_tgid;

		if (bpf_map_lookup_elem(&seen, &key))
			return 0;

		u64 zero = 0;
		bpf_map_update_elem(&seen, &key, &zero, BPF_ANY);
	}

	if (kernel_stack || user_stack) {
		i_key.pid = pid_tgid >> 32;
		i_key.tgid = pid_tgid;

		i_key.kernel_stack_id = i_key.user_stack_id = -1;
		if (user_stack)
			i_key.user_stack_id = bpf_get_stackid(ctx, &stackmap, BPF_F_USER_STACK);
		if (kernel_stack)
			i_key.kernel_stack_id = bpf_get_stackid(ctx, &stackmap, 0);

		struct cap_event cap_event = {};
		cap_event.pid = event->pid;
		cap_event.tgid = event->tgid;
		cap_event.cap = event->cap;
		cap_event.uid = event->uid;
		__builtin_memcpy(&cap_event.task, &event->task, sizeof(cap_event.task));
		cap_event.ret = event->ret;
		cap_event.audit = event->audit;
		cap_event.insetid = event->insetid;

		bpf_map_update_elem(&info, &i_key, &cap_event, BPF_NOEXIST);
	}

	/* Submit the event data */
	submit_buf(ctx, event, sizeof(*event));

	return 0;
}

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