#define _GNU_SOURCE 

#include <dirent.h>
#include <endian.h>
#include <errno.h>
#include <fcntl.h>
#include <pthread.h>
#include <sched.h>
#include <signal.h>
#include <stdarg.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mount.h>
#include <sys/prctl.h>
#include <sys/resource.h>
#include <sys/stat.h>
#include <sys/syscall.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <time.h>
#include <unistd.h>

#include <linux/capability.h>
#include <linux/futex.h>
#include <linux/sched.h>
#include <linux/kvm.h>

#ifndef __NR_clone3
#define __NR_clone3 435
#endif

int count = 0;
static void sleep_ms(uint64_t ms)
{
	usleep(ms * 1000);
}

static uint64_t current_time_ms(void)
{
	struct timespec ts;
	if (clock_gettime(CLOCK_MONOTONIC, &ts))
		exit(1);
	return (uint64_t)ts.tv_sec * 1000 + (uint64_t)ts.tv_nsec / 1000000;
}

static void thread_start(void* (*fn)(void*), void* arg)
{
	pthread_t th;
	pthread_attr_t attr;
	pthread_attr_init(&attr);
	pthread_attr_setstacksize(&attr, 128 << 10);
	int i = 0;
	for (; i < 100; i++) {
		if (pthread_create(&th, &attr, fn, arg) == 0) {
			pthread_attr_destroy(&attr);
			return;
		}
		if (errno == EAGAIN) {
			usleep(50);
			continue;
		}
		break;
	}
	exit(1);
}

typedef struct {
	int state;
} event_t;

static void event_init(event_t* ev)
{
	ev->state = 0;
}

static void event_reset(event_t* ev)
{
	ev->state = 0;
}

static void event_set(event_t* ev)
{
	if (ev->state)
		exit(1);
	__atomic_store_n(&ev->state, 1, __ATOMIC_RELEASE);
	syscall(SYS_futex, &ev->state, FUTEX_WAKE | FUTEX_PRIVATE_FLAG, 1000000);
}

static void event_wait(event_t* ev)
{
	while (!__atomic_load_n(&ev->state, __ATOMIC_ACQUIRE))
		syscall(SYS_futex, &ev->state, FUTEX_WAIT | FUTEX_PRIVATE_FLAG, 0, 0);
}

static int event_isset(event_t* ev)
{
	return __atomic_load_n(&ev->state, __ATOMIC_ACQUIRE);
}

static int event_timedwait(event_t* ev, uint64_t timeout)
{
	uint64_t start = current_time_ms();
	uint64_t now = start;
	for (;;) {
		uint64_t remain = timeout - (now - start);
		struct timespec ts;
		ts.tv_sec = remain / 1000;
		ts.tv_nsec = (remain % 1000) * 1000 * 1000;
		syscall(SYS_futex, &ev->state, FUTEX_WAIT | FUTEX_PRIVATE_FLAG, 0, &ts);
		if (__atomic_load_n(&ev->state, __ATOMIC_ACQUIRE))
			return 1;
		now = current_time_ms();
		if (now - start > timeout)
			return 0;
	}
}

#define MAX_FDS 30

static void close_fds()
{
	for (int fd = 3; fd < MAX_FDS; fd++)
		close(fd);
}

struct thread_t {
	int created, call;
	event_t ready, done;
};

static struct thread_t threads[16];
static void execute_call(int call);
static int running;

static void* thr(void* arg)
{
	struct thread_t* th = (struct thread_t*)arg;
	for (;;) {
		event_wait(&th->ready);
		event_reset(&th->ready);
		execute_call(th->call);
		event_set(&th->done);
	}
	return 0;
}

static void execute_one(void)
{
	int i, call, thread;
	for (call = 0; call < 8; call++) {
		for (thread = 0; thread < (int)(sizeof(threads) / sizeof(threads[0])); thread++) {
			struct thread_t* th = &threads[thread];
			if (!th->created) {
				th->created = 1;
				event_init(&th->ready);
				event_init(&th->done);
				event_set(&th->done);
				thread_start(thr, th);
			}
			if (!event_isset(&th->done))
				continue;
			event_reset(&th->done);
			th->call = call;
			__atomic_fetch_add(&running, 1, __ATOMIC_RELAXED);
			event_set(&th->ready);
			if (call == 6)
				break;
			event_timedwait(&th->done, 50);
			break;
		}
	}
	for (i = 0; i < 100 && __atomic_load_n(&running, __ATOMIC_RELAXED); i++)
		sleep_ms(1);
	close_fds();
}

static void execute_one(void);

#define WAIT_FLAGS __WALL

static void loop(void)
{
	int iter = 0;
	for (;; iter++) {
		int pid = fork();
		if (pid < 0)
			exit(1);
		if (pid == 0) {
			execute_one();
			exit(0);
		}
		int status = 0;
		uint64_t start = current_time_ms();
		for (;;) {
			if (waitpid(-1, &status, WNOHANG | WAIT_FLAGS) == pid)
				break;
			sleep_ms(1);
			if (current_time_ms() - start < 5000)
				continue;
			kill(-pid, SIGKILL);
			kill(pid, SIGKILL);
			break;
		}
	}
}

uint64_t r[3] = {0xffffffffffffffff, 0xffffffffffffffff, 0xffffffffffffffff};

void execute_call(int call)
{
		intptr_t res = 0;
	switch (call) {
	case 0:
		printf("%d: open_kvm\n", getpid());
		res = syscall(__NR_openat, 0xffffffffffffff9cul, "/dev/kvm", 0ul, 0ul);
		if (res != -1)
				r[0] = res;
		break;
	case 1:
		printf("%d: KVM_CREATE_VM\n", getpid());
		res = syscall(__NR_ioctl, r[0], 0xae01, 0ul);
		if (res != -1)
				r[1] = res;
		break;
	case 2:
		printf("%d: KVM_ENABLE_CAP\n", getpid());
		*(uint32_t*)0x20000640 = 0xc0;
		*(uint32_t*)0x20000644 = 0;
		*(uint64_t*)0x20000648 = 0x8000;
		syscall(__NR_ioctl, r[1], 0x4068aea3, 0x20000640ul);
		break;
	case 3:
		printf("%d: KVM_SET_USER_MEMORY_REGION\n", getpid());
		*(uint32_t*)0x20000040 = 0;
		*(uint32_t*)0x20000044 = 1;
		*(uint64_t*)0x20000048 = 0;
		*(uint64_t*)0x20000050 = 0x2000;
		*(uint64_t*)0x20000058 = 0x20ffe000;
		syscall(__NR_ioctl, r[1], 0x4020ae46, 0x20000040ul);
		break;
	case 4:
		printf("%d: KVM_CREATE_VCPU\n", getpid());
		res = syscall(__NR_ioctl, r[1], 0xae41, 0ul);
		if (res != -1)
				r[2] = res;
		break;
	case 5:
		printf("%d: KVM_SET_MSRS\n", getpid());
		memcpy((void*)0x200003c0, "\x01\x00\x00\x00\xfa\xff\xff\xff\x03\x4d\x56\x4b\x00\x00\x00\x00\x01", 17);
		syscall(__NR_ioctl, r[2], KVM_SET_MSRS, 0x200003c0ul); 
		break;
	case 6:
		printf("%d: KVM_RUN\n", getpid());
		syscall(__NR_ioctl, r[2], 0xae80, 0ul);
		break;
	case 7:
		printf("%d: clone\n", getpid());
		*(uint64_t*)0x20000400 = 0x802000;
		*(uint64_t*)0x20000408 = 0;
		*(uint64_t*)0x20000410 = 0;
		*(uint64_t*)0x20000418 = 0;
		*(uint32_t*)0x20000420 = 0x2e;
		*(uint64_t*)0x20000428 = 0;
		*(uint64_t*)0x20000430 = 0;
		*(uint64_t*)0x20000438 = 0;
		*(uint64_t*)0x20000440 = 0x20000340;
		*(uint32_t*)0x20000340 = 0;
		*(uint32_t*)0x20000344 = -1;
		*(uint32_t*)0x20000348 = -1;
		*(uint32_t*)0x2000034c = -1;
		*(uint32_t*)0x20000350 = -1;
		*(uint32_t*)0x20000354 = 0;
		*(uint32_t*)0x20000358 = 0;
		*(uint32_t*)0x2000035c = -1;
		*(uint32_t*)0x20000360 = 0;
		*(uint32_t*)0x20000364 = -1;
		*(uint64_t*)0x20000448 = 0xa;
		*(uint32_t*)0x20000450 = -1;
		syscall(__NR_clone3, (void*)0x20000400, 0x58);
		break;
	}

}
int main(void)
{
	syscall(__NR_mmap, 0x1ffff000ul, 0x1000ul, 0ul, 0x32ul, -1, 0ul);
	syscall(__NR_mmap, 0x20000000ul, 0x1000000ul, 7ul, 0x32ul, -1, 0ul);
	syscall(__NR_mmap, 0x21000000ul, 0x1000ul, 0ul, 0x32ul, -1, 0ul);
	loop();
	return 0;
}