// 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";