package ebpf

/*
#include <stdlib.h>
#include <stdint.h>
#include <unistd.h>
#include <linux/unistd.h>
#include <linux/bpf.h>

// See the bpf man page for details on the following functions

static int bpf(enum bpf_cmd cmd, union bpf_attr *attr, unsigned int size)
{
    return syscall(__NR_bpf, cmd, attr, size);
}

static int bpf_create_map(
	 enum bpf_map_type map_type, unsigned int key_size, unsigned int value_size,
	 unsigned int max_entries)
{
	union bpf_attr attr = {
		.map_type    = map_type,
		.key_size    = key_size,
		.value_size  = value_size,
		.max_entries = max_entries};

	return bpf(BPF_MAP_CREATE, &attr, sizeof(attr));
}

int bpf_get_next_key(int fd, const void *key, void *next_key)
{
	union bpf_attr attr = {
		.map_fd   = fd,
		.key      = (__u64) (unsigned long) key,
		.next_key = (__u64) (unsigned long) next_key};

	return bpf(BPF_MAP_GET_NEXT_KEY, &attr, sizeof(attr));
}


int bpf_lookup_elem(int fd, const void *key, void *value)
{
	union bpf_attr attr = {
		.map_fd = fd,
		.key    = (__u64) (unsigned long) key,
		.value  = (__u64) (unsigned long) value,
	};

	return bpf(BPF_MAP_LOOKUP_ELEM, &attr, sizeof(attr));
}

int bpf_update_elem(int fd, const void *key, const void *value, uint64_t flags)
{
	union bpf_attr attr = {
		.map_fd = fd,
		.key    = (__u64) (unsigned long) key,
		.value  = (__u64) (unsigned long) value,
		.flags  = flags,
	};

	return bpf(BPF_MAP_UPDATE_ELEM, &attr, sizeof(attr));
}

*/
import "C"
import (
	"fmt"
	"syscall"
	"unsafe"
)

// MapFD is a file descriptor representing a eBPF map
type MapFD int

// CreateMap creates an eBPF map from int->uint64. The file descriptor of the
// created map is returned.
func CreateMap() (MapFD, error) {

	r, errno := C.bpf_create_map(C.BPF_MAP_TYPE_HASH, 4, 8, 64)

	if r == -1 {
		return MapFD(r), fmt.Errorf("bpf_create_map with errno %d", errno)
	}

	return MapFD(r), nil
}

// GetMap gets the key/value pairs from the specified eBPF map as a Go map
func (mfd MapFD) GetMap() (map[int]uint64, error) {
	retMap := make(map[int]uint64)

	mapFDC := C.int(mfd)
	keyC := C.int(0)
	nextKeyC := C.int(0)
	valC := C.uint64_t(0)

	for {
		r, errno := C.bpf_get_next_key(mapFDC, unsafe.Pointer(&keyC), unsafe.Pointer(&nextKeyC))
		if r == -1 {
			if errno == syscall.ENOENT {
				// The provided key was the last element. We're done iterating.
				return retMap, nil
			}
			return nil, fmt.Errorf("bpf_get_next_key failed with errno %d", errno)
		}

		r = C.bpf_lookup_elem(mapFDC, unsafe.Pointer(&nextKeyC), unsafe.Pointer(&valC))
		if r == -1 {
			return nil, fmt.Errorf("bpf_lookup_elem failed")
		}
		retMap[int(nextKeyC)] = uint64(valC)

		keyC = nextKeyC
	}
}

// Add puts the (key, value) into the eBPF map, only if the key does not exist
// yet in the map. It returns an error otherwise.
func (mfd MapFD) Add(key int, value uint64) error {
	return mfd.updateElement(key, value, C.BPF_NOEXIST)
}

// Change changes the value to an existing key in the eBPF map.  It returns an
// error otherwise.
func (mfd MapFD) Change(key int, value uint64) error {
	return mfd.updateElement(key, value, C.BPF_EXIST)
}

// Set puts the (key, value) into the eBPF map.  It will create or overwrite an
// existing entry for that key.
func (mfd MapFD) Set(key int, value uint64) error {
	return mfd.updateElement(key, value, C.BPF_ANY)
}

// updateElement is the low level wrapper to bpf_update_elem, used from Add(),
// Set() and Change().
func (mfd MapFD) updateElement(key int, value uint64, flag C.uint64_t) error {

	r, errno := C.bpf_update_elem(C.int(mfd),
		unsafe.Pointer(&key),
		unsafe.Pointer(&value),
		flag)

	if r == -1 {
		return fmt.Errorf("bpf_update_elem failed with errno %d", errno)
	}

	return nil
}