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

static int page_size;
volatile bool wp_fired;

static void *fault_handler_thread(void *arg)
{
    const long uffd = (long) arg;
    struct pollfd pollfd = {
        .fd = uffd,
        .events = POLLIN,
    };
    int ret;

    while (true) {
        struct uffdio_writeprotect wp;
        struct uffd_msg msg;
        ssize_t nread;

        if (poll(&pollfd, 1, -1) == -1) {
            fprintf(stderr, "POLL failed: %s\n", strerror(errno));
            exit(-1);
        }
        if (read(uffd, &msg, sizeof(msg)) != sizeof(msg)) {
            fprintf(stderr, "READ failed\n");
            exit(-1);
        }
        if (msg.event == UFFD_EVENT_REMOVE) {
            continue;
        }
        if (msg.event != UFFD_EVENT_PAGEFAULT) {
            fprintf(stderr, "Not UFFD_EVENT_PAGEFAULT\n");
            exit(-1);
        }

        wp_fired = true;

        /* Simply unprotect and wake */
        wp.range.start = msg.arg.pagefault.address;
        wp.range.len = page_size;
        wp.mode = 0;
        if (ioctl(uffd, UFFDIO_WRITEPROTECT, &wp) == -1) {
            fprintf(stderr, "UFFDIO_WRITEPROTECT failed: %s\n", strerror(errno));
            exit(-1);
        }
    }
}

int main(void)
{
    struct uffdio_writeprotect wp;
    struct uffdio_zeropage zeropage;
    struct uffdio_register reg;
    struct uffdio_api api = {
        .api = UFFD_API,
        .features = UFFD_FEATURE_EVENT_REMOVE,
    };
    pthread_t fault, discard;
    long uffd;
    char *area;

    page_size = sysconf(_SC_PAGE_SIZE);

    uffd = syscall(__NR_userfaultfd, O_CLOEXEC | O_NONBLOCK);
    if (uffd == -1) {
        fprintf(stderr, "Could not create uffd: %s\n", strerror(errno));
        exit(-1);
    }
    if (ioctl(uffd, UFFDIO_API, &api) == -1) {
        fprintf(stderr, "UFFDIO_API failed: %s\n", strerror(errno));
        exit(-1);
    }

    area = mmap(NULL, page_size, PROT_READ | PROT_WRITE,
                MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    if (area == MAP_FAILED) {
        fprintf(stderr, "Could not allocate memory");
        exit(-1);
    }

    /* make sure the area is empty */
    if (madvise(area, page_size, MADV_DONTNEED)) {
        fprintf(stderr, "MADV_DONTNEED failed:%s\n", strerror(errno));
        exit(-1);
    }

    /* register a WP handler */
    reg.range.start = (uint64_t) area;
    reg.range.len = page_size,
    reg.mode = UFFDIO_REGISTER_MODE_WP;
    if (ioctl(uffd, UFFDIO_REGISTER, &reg) == -1) {
        fprintf(stderr, "UFFDIO_REGISTER failed: %s\n", strerror(errno));
        exit(-1);
    }

    /*
     * comment this out to make the wp actually fire
    zeropage.range = reg.range;
    zeropage.mode = UFFDIO_ZEROPAGE_MODE_DONTWAKE;
    if (ioctl(uffd, UFFDIO_ZEROPAGE, &zeropage) == -1) {
        fprintf(stderr, "UFFDIO_ZEROPAGE failed: %s\n", strerror(errno));
        exit(-1);
    }
    */

    /* protect the page */
    wp.range = reg.range;
    wp.mode = UFFDIO_WRITEPROTECT_MODE_WP;
    if (ioctl(uffd, UFFDIO_WRITEPROTECT, &wp) == -1) {
        fprintf(stderr, "UFFDIO_WRITEPROTECT failed: %s\n", strerror(errno));
        exit(-1);
    }

    if (pthread_create(&fault, NULL, fault_handler_thread,
                       (void *) uffd)) {
        fprintf(stderr, "Could not create fault handing thread");
        exit(-1);
    }

    *area = 0;
    if (wp_fired) {
        printf("WP fired\n");
        return 0;
    }
    printf("WP did not fire\n");
    return -1;
}