// SPDX-License-Identifier: GPL-2.0
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include <errno.h>
#include <poll.h>
#include <unistd.h>
#include <linux/perf_event.h>
#include <sys/mman.h>
#include "trace_helpers.h"

#define MAX_SYMS 300000
static struct ksym syms[MAX_SYMS];
static int sym_cnt;

static int ksym_cmp(const void *p1, const void *p2)
{
        return ((struct ksym *)p1)->addr - ((struct ksym *)p2)->addr;
}

int load_kallsyms(void)
{
        FILE *f = fopen("/proc/kallsyms", "r");
        char func[256], buf[256];
        char symbol;
        void *addr;
        int i = 0;

        if (!f)
                return -ENOENT;

        while (!feof(f)) {
                if (!fgets(buf, sizeof(buf), f))
                        break;
                if (sscanf(buf, "%p %c %s", &addr, &symbol, func) != 3)
                        break;
                if (!addr)
                        continue;
                syms[i].addr = (long) addr;
                syms[i].name = strdup(func);
                i++;
        }
        sym_cnt = i;
        qsort(syms, sym_cnt, sizeof(struct ksym), ksym_cmp);
        return 0;
}

struct ksym *ksym_search(long key)
{
        int start = 0, end = sym_cnt;
        int result;

        while (start < end) {
                size_t mid = start + (end - start) / 2;

                result = key - syms[mid].addr;
                if (result < 0)
                        end = mid;
                else if (result > 0)
                        start = mid + 1;
                else
                        return &syms[mid];
        }

        if (start >= 1 && syms[start - 1].addr < key &&
            key < syms[start].addr)
                /* valid ksym */
                return &syms[start - 1];

        /* out of range. return _stext */
        return &syms[0];
}

long ksym_get_addr(const char *name)
{
        int i;

        for (i = 0; i < sym_cnt; i++) {
                if (strcmp(syms[i].name, name) == 0)
                        return syms[i].addr;
        }

        return 0;
}

static int page_size;
static int page_cnt = 8;
static struct perf_event_mmap_page *header;

int perf_event_mmap_header(int fd, struct perf_event_mmap_page **header)
{
        void *base;
        int mmap_size;

        page_size = getpagesize();
        mmap_size = page_size * (page_cnt + 1);

        base = mmap(NULL, mmap_size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
        if (base == MAP_FAILED) {
                printf("mmap err\n");
                return -1;
        }

        *header = base;
        return 0;
}

int perf_event_mmap(int fd)
{
        return perf_event_mmap_header(fd, &header);
}

static int perf_event_poll(int fd)
{
        struct pollfd pfd = { .fd = fd, .events = POLLIN };

        return poll(&pfd, 1, 1000);
}

struct perf_event_sample {
        struct perf_event_header header;
        __u32 size;
        char data[];
};

static enum bpf_perf_event_ret bpf_perf_event_print(void *event, void *priv)
{
        struct perf_event_sample *e = event;
        perf_event_print_fn fn = priv;
        int ret;

        if (e->header.type == PERF_RECORD_SAMPLE) {
                ret = fn(e->data, e->size);
                if (ret != LIBBPF_PERF_EVENT_CONT)
                        return ret;
        } else if (e->header.type == PERF_RECORD_LOST) {
                struct {
                        struct perf_event_header header;
                        __u64 id;
                        __u64 lost;
                } *lost = (void *) e;
                printf("lost %lld events\n", lost->lost);
        } else {
                printf("unknown event type=%d size=%d\n",
                       e->header.type, e->header.size);
        }

        return LIBBPF_PERF_EVENT_CONT;
}

int perf_event_poller(int fd, perf_event_print_fn output_fn)
{
        enum bpf_perf_event_ret ret;
        void *buf = NULL;
        size_t len = 0;

        for (;;) {
                perf_event_poll(fd);
                ret = bpf_perf_event_read_simple(header, page_cnt * page_size,
                                                 page_size, &buf, &len,
                                                 bpf_perf_event_print,
                                                 output_fn);
                if (ret != LIBBPF_PERF_EVENT_CONT)
                        break;
        }
        free(buf);

        return ret;
}

int perf_event_poller_multi(int *fds, struct perf_event_mmap_page **headers,
                            int num_fds, perf_event_print_fn output_fn)
{
        enum bpf_perf_event_ret ret;
        struct pollfd *pfds;
        void *buf = NULL;
        size_t len = 0;
        int i;

        pfds = calloc(num_fds, sizeof(*pfds));
        if (!pfds)
                return LIBBPF_PERF_EVENT_ERROR;

        for (i = 0; i < num_fds; i++) {
                pfds[i].fd = fds[i];
                pfds[i].events = POLLIN;
        }

        for (;;) {
                poll(pfds, num_fds, 1000);
                for (i = 0; i < num_fds; i++) {
                        if (!pfds[i].revents)
                                continue;

                        ret = bpf_perf_event_read_simple(headers[i],
                                                         page_cnt * page_size,
                                                         page_size, &buf, &len,
                                                         bpf_perf_event_print,
                                                         output_fn);
                        if (ret != LIBBPF_PERF_EVENT_CONT)
                                break;
                }
        }
        free(buf);
        free(pfds);

        return ret;
}