#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdbool.h>
#include <inttypes.h>
#include <string.h>
#include <fcntl.h>
#include <unistd.h>
#include <errno.h>
#include <poll.h>
#include <pthread.h>
#include <sys/mman.h>
#include <sys/syscall.h>
#include <sys/ioctl.h>
#include <linux/memfd.h>
#include <linux/userfaultfd.h>

static size_t pagesize;
static int uffd;
static int pagemap_fd;

#define SIZE (128*1024*1024ull)
#define barrier() __asm__ __volatile__("": : :"memory")

volatile bool uffd_triggered;

static void uffd_wp_range(char *start, size_t size, bool wp)
{
	struct uffdio_writeprotect uffd_writeprotect;

	uffd_writeprotect.range.start = (unsigned long) start;
	uffd_writeprotect.range.len = size;
	if (wp) {
		uffd_writeprotect.mode = UFFDIO_WRITEPROTECT_MODE_WP;
	} else {
		uffd_writeprotect.mode = 0;
	}
	if (ioctl(uffd, UFFDIO_WRITEPROTECT, &uffd_writeprotect)) {
		fprintf(stderr, "UFFDIO_WRITEPROTECT failed: %d\n", errno);
		exit(1);
	}
}

static void *uffd_thread_fn(void *arg)
{
	static struct uffd_msg msg;
	ssize_t nread;

	while (1) {
		struct pollfd pollfd;
		int nready;

		pollfd.fd = uffd;
		pollfd.events = POLLIN;
		nready = poll(&pollfd, 1, -1);
		if (nready == -1) {
			fprintf(stderr, "poll() failed: %d\n", errno);
			exit(1);
		}

		nread = read(uffd, &msg, sizeof(msg));
		if (nread <= 0)
			continue;

		if (msg.event != UFFD_EVENT_PAGEFAULT ||
		    !(msg.arg.pagefault.flags & UFFD_PAGEFAULT_FLAG_WP)) {
			printf("FAIL: wrong uffd-wp event fired\n");
			exit(1);
		}

		/* un-protect */
		uffd_triggered = true;
		uffd_wp_range((char *)(uintptr_t)msg.arg.pagefault.address, pagesize, false);
	}
	return arg;
}

static int setup_uffd(char *map, size_t size)
{
	struct uffdio_api uffdio_api;
	struct uffdio_register uffdio_register;
	pthread_t thread;

	uffd = syscall(__NR_userfaultfd,
		       O_CLOEXEC | O_NONBLOCK | UFFD_USER_MODE_ONLY);
	if (uffd < 0) {
		fprintf(stderr, "syscall() failed: %d\n", errno);
		return -errno;
	}

	uffdio_api.api = UFFD_API;
	uffdio_api.features = UFFD_FEATURE_PAGEFAULT_FLAG_WP;
	if (ioctl(uffd, UFFDIO_API, &uffdio_api) < 0) {
		fprintf(stderr, "UFFDIO_API failed: %d\n", errno);
		return -errno;
	}

	if (!(uffdio_api.features & UFFD_FEATURE_PAGEFAULT_FLAG_WP)) {
		fprintf(stderr, "UFFD_FEATURE_WRITEPROTECT missing\n");
		return -ENOSYS;
	}

	uffdio_register.range.start = (unsigned long) map;
	uffdio_register.range.len = size;
	uffdio_register.mode = UFFDIO_REGISTER_MODE_WP;
	if (ioctl(uffd, UFFDIO_REGISTER, &uffdio_register) < 0) {
		fprintf(stderr, "UFFDIO_REGISTER failed: %d\n", errno);
		return -errno;
	}

	pthread_create(&thread, NULL, uffd_thread_fn, NULL);

	return 0;
}

int main(void)
{
	const size_t size = SIZE;
	char *map, *cur;

	pagesize = getpagesize();
	pagemap_fd = open("/proc/self/pagemap", O_RDONLY);
	if (pagemap_fd < 0) {
		fprintf(stderr, "Opening /proc/self/pagemap failed\n");
		return 1;
	}

	map = mmap(NULL, size, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANON, -1, 0);
	if (map == MAP_FAILED) {
		fprintf(stderr, "mmap() failed\n");
		return -errno;
	}

	if (madvise(map, size, MADV_HUGEPAGE)) {
		fprintf(stderr, "MADV_HUGEPAGE failed\n");
		return -errno;
	}

	if (setup_uffd(map, size))
		return 1;

	/* Populate the shared zeropage, hopefullt also the huge one.*/
	madvise(map, size, MADV_POPULATE_READ);

	/* ... and write-protect all pages using uffd-wp. */
	uffd_wp_range(map, size, true);

	/*
	 * Discard every second odd page, this will split any huge zero
	 * THP and will hopefully keep the PTE protected using uffd-wp.
	 *
	 * Any mechanism to PTE-map the THP would do.
	 */
	for (cur = map + pagesize; cur < map + size; cur += 2 * pagesize)
		madvise(cur, pagesize, MADV_DONTNEED);

	/*
	 * Test every second even page (-> all remaining ones) if we get our
	 * uffd-wp events.
	 */
	for (cur = map; cur < map + size; cur += 2 * pagesize) {
		uffd_triggered = false;

		barrier();
		/* Trigger a write fault. */
		*cur = 1;
		barrier();

		if (!uffd_triggered) {
			printf("FAIL: uffd-wp did not trigger\n");
			return 1;
		}
	}

	printf("PASS: uffd-wp triggered\n");
	return 0;
}