#include "bpf_experimental.h"

struct val_t {
        long b, c, d;
};

struct elem {
        long sum;
        struct val_t __percpu_kptr *pc;
};

struct {
        __uint(type, BPF_MAP_TYPE_ARRAY);
        __uint(max_entries, 1);
        __type(key, int);
        __type(value, struct elem);
} array SEC(".maps");

void bpf_rcu_read_lock(void) __ksym;
void bpf_rcu_read_unlock(void) __ksym;

const volatile int nr_cpus;

/* Initialize the percpu object */
SEC("?fentry/bpf_fentry_test1")
int BPF_PROG(test_array_map_1)
{
        struct val_t __percpu_kptr *p;
        struct elem *e;
        int index = 0;

        e = bpf_map_lookup_elem(&array, &index);
        if (!e)
                return 0;

        p = bpf_percpu_obj_new(struct val_t);
        if (!p)
                return 0;

        p = bpf_kptr_xchg(&e->pc, p);
        if (p)
                bpf_percpu_obj_drop(p);

        return 0;
}

/* Update percpu data */
SEC("?fentry/bpf_fentry_test2")
int BPF_PROG(test_array_map_2)
{
        struct val_t __percpu_kptr *p;
        struct val_t *v;
        struct elem *e;
        int index = 0;

        e = bpf_map_lookup_elem(&array, &index);
        if (!e)
                return 0;

        p = e->pc;
        if (!p)
                return 0;

        v = bpf_per_cpu_ptr(p, 0);
        if (!v)
                return 0;
        v->c = 1;
        v->d = 2;

        return 0;
}

int cpu0_field_d, sum_field_c;
int my_pid;

/* Summarize percpu data */
SEC("?fentry/bpf_fentry_test3")
int BPF_PROG(test_array_map_3)
{
        struct val_t __percpu_kptr *p;
        int i, index = 0;
        struct val_t *v;
        struct elem *e;

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

        e = bpf_map_lookup_elem(&array, &index);
        if (!e)
                return 0;

        p = e->pc;
        if (!p)
                return 0;

        bpf_for(i, 0, nr_cpus) {
                v = bpf_per_cpu_ptr(p, i);
                if (v) {
                        if (i == 0)
                                cpu0_field_d = v->d;
                        sum_field_c += v->c;
                }
        }

        return 0;
}

/* Explicitly free allocated percpu data */
SEC("?fentry/bpf_fentry_test4")
int BPF_PROG(test_array_map_4)
{
        struct val_t __percpu_kptr *p;
        struct elem *e;
        int index = 0;

        e = bpf_map_lookup_elem(&array, &index);
        if (!e)
                return 0;

        /* delete */
        p = bpf_kptr_xchg(&e->pc, NULL);
        if (p) {
                bpf_percpu_obj_drop(p);
        }

        return 0;
}

SEC("?fentry.s/bpf_fentry_test1")
int BPF_PROG(test_array_map_10)
{
        struct val_t __percpu_kptr *p, *p1;
        int i, index = 0;
        struct val_t *v;
        struct elem *e;

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

        e = bpf_map_lookup_elem(&array, &index);
        if (!e)
                return 0;

        bpf_rcu_read_lock();
        p = e->pc;
        if (!p) {
                p = bpf_percpu_obj_new(struct val_t);
                if (!p)
                        goto out;

                p1 = bpf_kptr_xchg(&e->pc, p);
                if (p1) {
                        /* race condition */
                        bpf_percpu_obj_drop(p1);
                }
        }

        v = bpf_this_cpu_ptr(p);
        v->c = 3;
        v = bpf_this_cpu_ptr(p);
        v->c = 0;

        v = bpf_per_cpu_ptr(p, 0);
        if (!v)
                goto out;
        v->c = 1;
        v->d = 2;

        /* delete */
        p1 = bpf_kptr_xchg(&e->pc, NULL);
        if (!p1)
                goto out;

        bpf_for(i, 0, nr_cpus) {
                v = bpf_per_cpu_ptr(p, i);
                if (v) {
                        if (i == 0)
                                cpu0_field_d = v->d;
                        sum_field_c += v->c;
                }
        }

        /* finally release p */
        bpf_percpu_obj_drop(p1);
out:
        bpf_rcu_read_unlock();
        return 0;
}

struct {
        __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
        __uint(max_entries, 2);
        __type(key, int);
        __type(value, u32);
} percpu SEC(".maps");

SEC("?fentry/bpf_fentry_test1")
int BPF_PROG(test_percpu_array, int x)
{
        u64 value = 0xDEADC0DE;
        int key = 0;

        bpf_map_update_elem(&percpu, &key, &value, BPF_ANY);
        return 0;
}

struct {
        __uint(type, BPF_MAP_TYPE_PERCPU_CGROUP_STORAGE);
        __type(key, struct bpf_cgroup_storage_key);
        __type(value, u32);
} percpu_cgroup_storage SEC(".maps");

SEC("cgroup_skb/egress")
int cgroup_egress(struct __sk_buff *skb)
{
        u32 *val = bpf_get_local_storage(&percpu_cgroup_storage, 0);

        *val = 1;
        return 1;
}

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