// compile with "gcc -o khugepaged-vs-uffd khugepaged-vs-uffd.c -pthread"
#define _GNU_SOURCE
#include <pthread.h>
#include <err.h>
#include <sched.h>
#include <stdio.h>
#include <string.h>
#include <fcntl.h>
#include <signal.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/syscall.h>
#include <sys/stat.h>
#include <sys/prctl.h>
#include <sys/mount.h>
#include <sys/mman.h>
#include <sys/ioctl.h>
#include <linux/userfaultfd.h>

#ifndef MADV_COLLAPSE
#define MADV_COLLAPSE 25
#endif

#ifndef UFFD_USER_MODE_ONLY
#define UFFD_USER_MODE_ONLY 1
#endif

#define SYSCHK(x) ({          \
  typeof(x) __res = (x);      \
  if (__res == (typeof(x))-1) \
    err(1, "SYSCHK(" #x ")"); \
  __res;                      \
})

static void write_file(char *name, char *buf) {
  int fd = SYSCHK(open(name, O_WRONLY));
  if (write(fd, buf, strlen(buf)) != strlen(buf))
    err(1, "write %s", name);
  close(fd);
}

static void write_map(char *name, int outer_id) {
  char buf[100];
  sprintf(buf, "0 %d 1", outer_id);
  write_file(name, buf);
}

static void *thread_fn(void *dummy) {
  system("head -n50 /proc/$PPID/smaps;echo;echo");
  SYSCHK(prctl(PR_SET_NAME, "DELAYME"));
  SYSCHK(madvise((void*)0x200000UL, 0x200000, MADV_COLLAPSE));
  SYSCHK(prctl(PR_SET_NAME, "thread"));
  system("head -n50 /proc/$PPID/smaps");
  return NULL;
}

int main(void) {
  int outer_uid = getuid();
  int outer_gid = getgid();
  SYSCHK(unshare(CLONE_NEWNS|CLONE_NEWUSER));
  SYSCHK(mount(NULL, "/", NULL, MS_PRIVATE|MS_REC, NULL));
  write_file("/proc/self/setgroups", "deny");
  write_map("/proc/self/uid_map", outer_uid);
  write_map("/proc/self/gid_map", outer_gid);

  SYSCHK(mount("none", "/tmp", "tmpfs", MS_NOSUID|MS_NODEV, "huge=always"));
  int fd = SYSCHK(open("/tmp/a", O_RDWR|O_CREAT, 0600));
  SYSCHK(ftruncate(fd, 0x200000));
  void *ptr = SYSCHK(mmap((void*)0x200000UL, 0x100000, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_FIXED_NOREPLACE, fd, 0));
  *(volatile char *)ptr;
  SYSCHK(mmap((void*)0x300000UL, 0x100000, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_FIXED_NOREPLACE, fd, 0x100000));
  for (int i=0; i<512; i++)
    *(volatile char *)(0x200000UL + 0x1000 * i);

  int uffd = SYSCHK(syscall(__NR_userfaultfd, UFFD_USER_MODE_ONLY));

  struct uffdio_api api = { .api = UFFD_API, .features = 0 };
  SYSCHK(ioctl(uffd, UFFDIO_API, &api));

  struct uffdio_register reg = {
    .range = { .start = 0x200000, .len = 0x200000 },
    .mode = UFFDIO_REGISTER_MODE_MISSING
  };
  SYSCHK(ioctl(uffd, UFFDIO_REGISTER, &reg));

  pthread_t thread;
  if (pthread_create(&thread, NULL, thread_fn, NULL))
    errx(1, "pthread_create");

  sleep(1);

  unsigned char dummy_page[0x1000] = {1};
  struct uffdio_copy copy = {
    .dst = 0x201000,
    .src = (unsigned long)dummy_page,
    .len = 0x1000,
    .mode = 0
  };
  SYSCHK(ioctl(uffd, UFFDIO_COPY, &copy));

  if (pthread_join(thread, NULL))
    errx(1, "pthread_join");

  //system("cat /proc/$PPID/smaps");
}