// SPDX-License-Identifier: GPL-2.0
/*
 * Trace raw_syscalls tracepoints to collect system call statistics.
 */

#include "vmlinux.h"
#include "syscall_summary.h"

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

/* This is to calculate a delta between sys-enter and sys-exit for each thread */
struct syscall_trace {
	int nr; /* syscall number is only available at sys-enter */
	int unused;
	u64 timestamp;
};

#define MAX_ENTRIES	(128 * 1024)

struct syscall_trace_map {
	__uint(type, BPF_MAP_TYPE_HASH);
	__type(key, int); /* tid */
	__type(value, struct syscall_trace);
	__uint(max_entries, MAX_ENTRIES);
} syscall_trace_map SEC(".maps");

struct syscall_stats_map {
	__uint(type, BPF_MAP_TYPE_HASH);
	__type(key, struct syscall_key);
	__type(value, struct syscall_stats);
	__uint(max_entries, MAX_ENTRIES);
} syscall_stats_map SEC(".maps");

int enabled; /* controlled from userspace */

const volatile enum syscall_aggr_mode aggr_mode;
const volatile int use_cgroup_v2;

int perf_subsys_id = -1;

static inline __u64 get_current_cgroup_id(void)
{
	struct task_struct *task;
	struct cgroup *cgrp;

	if (use_cgroup_v2)
		return bpf_get_current_cgroup_id();

	task = bpf_get_current_task_btf();

	if (perf_subsys_id == -1) {
#if __has_builtin(__builtin_preserve_enum_value)
		perf_subsys_id = bpf_core_enum_value(enum cgroup_subsys_id,
						     perf_event_cgrp_id);
#else
		perf_subsys_id = perf_event_cgrp_id;
#endif
	}

	cgrp = BPF_CORE_READ(task, cgroups, subsys[perf_subsys_id], cgroup);
	return BPF_CORE_READ(cgrp, kn, id);
}

static void update_stats(int cpu_or_tid, u64 cgroup_id, int nr, s64 duration,
			 long ret)
{
	struct syscall_key key = {
		.cpu_or_tid = cpu_or_tid,
		.cgroup = cgroup_id,
		.nr = nr,
	};
	struct syscall_stats *stats;

	stats = bpf_map_lookup_elem(&syscall_stats_map, &key);
	if (stats == NULL) {
		struct syscall_stats zero = {};

		bpf_map_update_elem(&syscall_stats_map, &key, &zero, BPF_NOEXIST);
		stats = bpf_map_lookup_elem(&syscall_stats_map, &key);
		if (stats == NULL)
			return;
	}

	__sync_fetch_and_add(&stats->count, 1);
	if (ret < 0)
		__sync_fetch_and_add(&stats->error, 1);

	if (duration > 0) {
		__sync_fetch_and_add(&stats->total_time, duration);
		__sync_fetch_and_add(&stats->squared_sum, duration * duration);
		if (stats->max_time < duration)
			stats->max_time = duration;
		if (stats->min_time > duration || stats->min_time == 0)
			stats->min_time = duration;
	}

	return;
}

SEC("tp_btf/sys_enter")
int sys_enter(u64 *ctx)
{
	int tid;
	struct syscall_trace st;

	if (!enabled)
		return 0;

	st.nr = ctx[1]; /* syscall number */
	st.unused = 0;
	st.timestamp = bpf_ktime_get_ns();

	tid = bpf_get_current_pid_tgid();
	bpf_map_update_elem(&syscall_trace_map, &tid, &st, BPF_ANY);

	return 0;
}

SEC("tp_btf/sys_exit")
int sys_exit(u64 *ctx)
{
	int tid;
	int key = 0;
	u64 cgroup = 0;
	long ret = ctx[1]; /* return value of the syscall */
	struct syscall_trace *st;
	s64 delta;

	if (!enabled)
		return 0;

	tid = bpf_get_current_pid_tgid();
	st = bpf_map_lookup_elem(&syscall_trace_map, &tid);
	if (st == NULL)
		return 0;

	if (aggr_mode == SYSCALL_AGGR_THREAD)
		key = tid;
	else if (aggr_mode == SYSCALL_AGGR_CGROUP)
		cgroup = get_current_cgroup_id();
	else
		key = bpf_get_smp_processor_id();

	delta = bpf_ktime_get_ns() - st->timestamp;
	update_stats(key, cgroup, st->nr, delta, ret);

	bpf_map_delete_elem(&syscall_trace_map, &tid);
	return 0;
}

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